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.

https://github.com/vizfold/attention-viz-demo

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
Last synced: 10 months ago · JSON representation ·

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
Created about 1 year ago · Last pushed about 1 year ago
Metadata Files
Readme License Citation

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:

  1. Runs inference using OpenFold with precomputed alignments
  2. Extracts top-k residue–residue attention scores from each layer and head
  3. Saves these scores to text files
  4. 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):

msa_row_arc

All heads subplot:

msa_row_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):

triangle_start_arc

All heads subplot:

triangle_start_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

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