BERTology Meets Biology: Interpreting Attention in Protein Language Models

TL;DR

The paper investigates whether attention in protein-focused Transformer models captures biologically meaningful structure and function. By analyzing five pretrained architectures across multiple protein datasets, it demonstrates that attention aligns with contact maps in deep layers, targets binding sites (and to a lesser extent PTMs), and encodes higher-level biophysical properties as depth increases. The authors introduce cross-layer probing and 3D visualizations to contextualize attention within protein structure, and show that attention correlates with traditional substitution matrices, suggesting biochemical relevance beyond spurious correlations. These findings support the use of attention-based interpretability as a tool for scientific discovery in proteomics and motivate further cross-domain visualization and probing approaches.

Abstract

Transformer architectures have proven to learn useful representations for protein classification and generation tasks. However, these representations present challenges in interpretability. In this work, we demonstrate a set of methods for analyzing protein Transformer models through the lens of attention. We show that attention: (1) captures the folding structure of proteins, connecting amino acids that are far apart in the underlying sequence, but spatially close in the three-dimensional structure, (2) targets binding sites, a key functional component of proteins, and (3) focuses on progressively more complex biophysical properties with increasing layer depth. We find this behavior to be consistent across three Transformer architectures (BERT, ALBERT, XLNet) and two distinct protein datasets. We also present a three-dimensional visualization of the interaction between attention and protein structure. Code for visualization and analysis is available at https://github.com/salesforce/provis.

Figures (19)

• Figure 1: Examples of how specialized attention heads in a Transformer recover protein structure and function, based solely on language model pre-training. Orange lines depict attention between amino acids (line width proportional to attention weight; values below 0.1 hidden). Heads were selected based on correlation with ground-truth annotations of contact maps and binding sites. Visualizations based on the NGL Viewer rose2018nglrose2015ngl10.1093/bioinformatics/btx789.
• Figure 2: Agreement between attention and contact maps across five pretrained Transformer models from TAPE (a) and ProtTrans (b--e). The heatmaps show the proportion of high-confidence attention weights ($\alpha_{i,j} > \theta$) from each head that connects pairs of amino acids that are in contact with one another. In TapeBert (a), for example, we can see that 45% of attention in head 12-4 (the 12th layer's 4th head) maps to contacts. The bar plots show the maximum value from each layer. Note that the vertical striping in ProtAlbert (b) is likely due to cross-layer parameter sharing (see Appendix \ref{['sec:other_variants']}).
• Figure 3: Proportion of attention focused on binding sites across five pretrained models. The heatmaps show the proportion of high-confidence attention ($\alpha_{i,j} > \theta$) from each head that is directed to binding sites. In TapeBert (a), for example, we can see that 49% of attention in head 11-6 (the 11th layer's 6th head) is directed to binding sites. The bar plots show the maximum value from each layer.
• Figure 4: Each plot shows the percentage of attention focused on the given property, averaged over all heads within each layer. The plots, sorted by center of gravity (red dashed line), show that heads in deeper layers focus relatively more attention on binding sites and contacts, whereas attention toward specific secondary structures is more even across layers.
• Figure 5: Performance of probing classifiers by layer, sorted by task order in Figure \ref{['fig:layers_combined']}. The embedding probes (orange) quantify the knowledge of the given property that is encoded in each layer's output embeddings. The attention probe (blue), show the amount of information encoded in attention weights for the (pairwise) contact feature. Additional details are provided in Appendix \ref{['sec:probing_methodology']}.