attention-viz-demo
This is a demo that allows a user to visualize the top-k attention scores from OpenFold for a particular protein. The demo plots the top attention scores from one residue to another on the predicted 3D structure. Additionally, the demo produces Arc Diagrams in 2D of the top attention scores from one residue to another.
Science Score: 44.0%
This score indicates how likely this project is to be science-related based on various indicators:
-
✓CITATION.cff file
Found CITATION.cff file -
✓codemeta.json file
Found codemeta.json file -
✓.zenodo.json file
Found .zenodo.json file -
○DOI references
-
○Academic publication links
-
○Academic email domains
-
○Institutional organization owner
-
○JOSS paper metadata
-
○Scientific vocabulary similarity
Low similarity (10.5%) to scientific vocabulary
Repository
This is a demo that allows a user to visualize the top-k attention scores from OpenFold for a particular protein. The demo plots the top attention scores from one residue to another on the predicted 3D structure. Additionally, the demo produces Arc Diagrams in 2D of the top attention scores from one residue to another.
Basic Info
- Host: GitHub
- Owner: vizfold
- License: apache-2.0
- Language: Python
- Default Branch: main
- Size: 8.79 MB
Statistics
- Stars: 0
- Watchers: 2
- Forks: 1
- Open Issues: 0
- Releases: 0
Metadata Files
README.md
OpenFold-based Attention Visualization Demo
This is a lightweight extension of OpenFold that enables interactive visualization of attention mechanisms in protein structure prediction. It provides tools to render MSA row and Triangle attention scores as:
- Arc diagrams (sequence space)
- 3D PyMOL overlays (structure space)
Key Features
- Compatible with OpenFold outputs (
.pdb, attention text dumps) - Support for layer- and head-specific visualizations
- Integrated residue highlighting
- Notebook-friendly and HPC-friendly workflow
Installation
This repo assumes you have already installed OpenFold and its dependencies, or you are using CyberShuttle (see cybershuttle.yml)
You will also need:
- PyMOL (open-source version is sufficient)
- matplotlib, numpy, scipy, pandas
- biopython (for sequence parsing)
Interactive Demo: viz_attention_demo.ipynb
The notebook viz_attention_demo.ipynb demonstrates the full visualization pipeline using OpenFold.
It performs the following steps:
- Runs inference using OpenFold with precomputed alignments
- Extracts top-k residue–residue attention scores from each layer and head
- Saves these scores to text files
- Visualizes attention:
- As arc diagrams (residue–residue attention on the sequence)
- As 3D PyMOL overlays (on the predicted structure)
We focus on two attention types: - MSA Row Attention - Triangle Start Attention
The thickness of the lines (in both arc diagrams and 3D renderings) indicates the strength of the attention score.
MSA Row Attention (Layer 47, Protein 6KWC)
- Shows pairwise attention between residues as inferred from multiple sequence alignments
- Visualized across all attention heads at a selected model layer
Arc diagram (Head 2):

All heads subplot:

Triangle Start Attention (Layer 47, Residue 18)
- Focuses on attention from a single residue to others, as part of triangle-based geometric reasoning
- The selected residue is highlighted:
- In arc diagrams: using a blue label
- In 3D visualizations: using a sphere
Arc diagram (Head 0):

All heads subplot:

Acknowledgements
This project is based on OpenFold, an open-source reimplementation of AlphaFold, distributed under the Apache License 2.0.
We have extended OpenFold with: - Custom visualization tools for attention maps (3D + arc diagrams) - Demo scripts and configuration for interactive analysis - Modifications to the inference pipeline for simplified usage
This repository includes source code originally developed by the OpenFold contributors. All original rights and attributions are retained in accordance with the Apache 2.0 License.
License
This project is licensed under the Apache License 2.0.
See the LICENSE file for details.
Owner
- Name: vizfold
- Login: vizfold
- Kind: organization
- Repositories: 1
- Profile: https://github.com/vizfold
Citation (CITATION.cff)
cff-version: 1.2.0
preferred-citation:
authors:
- family-names: "Ahdritz"
given-names: "Gustaf"
orcid: https://orcid.org/0000-0001-8283-5324
- family-names: "Bouatta"
given-names: "Nazim"
orcid: https://orcid.org/0000-0002-6524-874X
- family-names: "Kadyan"
given-names: "Sachin"
orcid: https://orcid.org/0000-0002-6079-7627
- family-names: "Xia"
given-names: "Qinghui"
- family-names: "Gerecke"
given-names: "William"
orcid: https://orcid.org/0000-0002-9777-6192
- family-names: "O'Donnell"
given-names: "Timothy J"
orcid: https://orcid.org/0000-0002-9949-069X
- family-names: "Berenberg"
given-names: "Daniel"
orcid: https://orcid.org/0000-0003-4631-0947
- family-names: "Fisk"
given-names: "Ian"
- family-names: "Zanichelli"
given-names: "Niccolò"
orcid: https://orcid.org/0000-0002-3093-3587
- family-names: "Zhang"
given-names: "Bo"
orcid: https://orcid.org/0000-0002-9714-2827
- family-names: "Nowaczynski"
given-names: "Arkadiusz"
orcid: https://orcid.org/0000-0002-3351-9584
- family-names: "Wang"
given-names: "Bei"
orcid: https://orcid.org/0000-0003-4942-9652
- family-names: "Stepniewska-Dziubinska"
given-names: "Marta M"
orcid: https://orcid.org/0000-0003-4942-9652
- family-names: "Zhang"
given-names: "Shang"
orcid: https://orcid.org/0000-0003-0759-2080
- family-names: "Ojewole"
given-names: "Adegoke"
orcid: https://orcid.org/0000-0003-2661-4388
- family-names: "Guney"
given-names: "Murat Efe"
- family-names: "Biderman"
given-names: "Stella"
orcid: https://orcid.org/0000-0001-8228-1042
- family-names: "Watkins"
given-names: "Andrew M"
orcid: https://orcid.org/0000-0003-1617-1720
- family-names: "Ra"
given-names: "Stephen"
orcid: https://orcid.org/0000-0002-2820-0050
- family-names: "Lorenzo"
given-names: "Pablo Ribalta"
orcid: https://orcid.org/0000-0002-3657-8053
- family-names: "Nivon"
given-names: "Lucas"
- family-names: "Weitzner"
given-names: "Brian"
orcid: https://orcid.org/0000-0002-1909-0961
- family-names: "Ban"
given-names: "Yih-En"
orcid: https://orcid.org/0000-0003-3698-3574
- family-names: "Ban"
given-names: "Yih-En Andrew"
orcid: https://orcid.org/0000-0003-3698-3574
- family-names: "Sorger"
given-names: "Peter K"
orcid: https://orcid.org/0000-0002-3364-1838
- family-names: "Mostaque"
given-names: "Emad"
- family-names: "Zhang"
given-names: "Zhao"
orcid: https://orcid.org/0000-0001-5921-0035
- family-names: "Bonneau"
given-names: "Richard"
orcid: https://orcid.org/0000-0003-4354-7906
- family-names: "AlQuraishi"
given-names: "Mohammed"
orcid: https://orcid.org/0000-0001-6817-1322
title: "OpenFold: Retraining AlphaFold2 yields new insights into its learning mechanisms and capacity for generalization"
type: article
doi: 10.1101/2022.11.20.517210
doi: 10.1101/2022.11.20.517210
date-released: 2021-11-12
url: "https://doi.org/10.1101/2022.11.20.517210"
GitHub Events
Total
- Issue comment event: 1
- Push event: 12
- Public event: 1
- Pull request event: 2
- Fork event: 1
- Create event: 1
Last Year
- Issue comment event: 1
- Push event: 12
- Public event: 1
- Pull request event: 2
- Fork event: 1
- Create event: 1