Source-linked AI summary
BERTology Meets Biology: Interpreting Attention in Protein Language Models
Jesse Vig, Ali Madani, Lav R. Varshney, Caiming Xiong, Richard Socher, Nazneen Fatema Rajani
TL;DR
Protein Transformers learn useful representations, but their attention mechanisms are difficult to interpret. This paper analyzes attention across five pretrained models using protein-specific structural and functional annotations, finding alignment with contacts, binding sites, substitution relationships, and increasingly complex properties across layers. The authors also provide three-dimensional visualization while emphasizing that the analyses are associative rather than causal.
Problem
Transformer representations perform well for protein tasks but present interpretability challenges that can obscure model bias, spurious correlations, and biological information.
Method
The study analyzes attention and hidden-state representations across five pretrained protein Transformers using contact maps, binding-site annotations, probing tasks, and three-dimensional visualization.
Results
Attention captures sequence-distant structural contacts, targets binding sites, aligns with substitution relationships, and represents progressively higher-level protein properties with increasing layer depth.
Takeaways & Limitations
Attention can expose structural and functional protein knowledge learned from sequence-only pretraining and may support scientific discovery when presented in appropriate biological context.
Takeaways & Limitations
The analyses are associative and do not establish causal links between attention and model behavior or explain model predictions.
Abstract
from arXiv · showhide
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.
1 INTRODUCTION
The paper adapts Transformer interpretability methods to protein sequences, using attention to investigate structural and functional properties. It addresses interpretability challenges while exploring how learned representations may support biological understanding.
- Transformer models offer strong protein representations but remain difficult to interpret, leaving potential bias and spurious correlations hidden.
- The study analyzes protein Transformer models through attention-based interpretability methods tailored to protein structure and function.
- The analysis also compares knowledge in attention weights with knowledge encoded in hidden-state representations.
- Attention connects sequence-distant amino acids that are spatially close, targets binding sites, and captures progressively higher-level properties with increasing layer depth.
- Three-dimensional visualization contextualizes attention within protein structure, supporting inspection of learned structural and functional relationships.
2 BACKGROUND: PROTEINS
Proteins are amino-acid sequences that form physical three-dimensional structures and carry out functions through features such as contacts and binding sites. Biological properties also include substitution relationships and post-translational modifications.
- Amino acids: Proteins consist of sequences drawn from a vocabulary of 20 commonly observed amino acids.
- Substitution matrix: Substitution matrices score how readily amino-acid pairs can replace one another while maintaining protein viability.
- Protein structure: Protein structure includes local secondary structure and large-scale tertiary structure that determines overall shape and function.
- Protein structure: Contact maps characterize tertiary structure by describing amino-acid pairs that lie within 8 angstroms in the folded protein.
- Binding sites: Binding sites are regions that bind other molecules to carry out specific functions and can serve as drug-development targets.
- Post-translational modifications: Post-translational modifications, including phosphorylation, occur after translation and contribute to protein structure and function.
3 METHODOLOGY
The study evaluates five pretrained protein Transformers with attention-alignment metrics, probing tasks, and checks against background frequencies and shuffled-attention null models. Analyses cover pairwise and token-level protein properties while excluding selected special-token attention and low-data heads.
- Model: Five pretrained models include TapeBert, ProtBert, ProtBert-BFD, ProtAlbert, and ProtXLNet across BERT, ALBERT, and XLNet architectures.TapeBert used Pfam with 31M sequences; ProtBert-BFD used BFD with 2.1B sequences, while other ProtTrans models used UniRef100 with 216M sequences.
- Transformer attention: Attention weights α_i,j represent influence from token i to token j, and heads are identified by layer and head index.
- Attention analysis: For token pairs, the analysis measures the proportion of high-attention pairs exhibiting a target property such as contact.High-confidence attention is selected using threshold θ.
- Attention analysis: For individual-token properties such as binding sites, the metric measures the proportion of attention directed toward tokens with that property.
- Validation checks: Alignment results are checked against property background frequencies and models with randomly shuffled attention weights to assess whether findings are due to chance.
- Probing tasks: The study uses embedding and attention probes to measure property knowledge encoded in model representations.
- Datasets: The datasets provide amino-acid sequences annotated with spatial coordinates, secondary structure, and binding sites, without using those annotations for model training.
- Experimental details: The analyses exclude attention to [SEP] and [CLS] tokens and retain heads with at least 100 high-confidence attention arcs.The attention threshold is set to θ = 0.3.
4 WHAT DOES ATTENTION UNDERSTAND ABOUT PROTEINS?
Attention in protein Transformers captures structural, functional, and biochemical properties across layers and models. It aligns with contacts and binding sites, concentrates on selected amino acids and PTM positions, and reflects progressively higher-level information in deeper layers.
- 4.1 PROTEIN STRUCTURE: Up to 63.2% of attention in the most aligned heads focused on contacts, versus a 1.3% background contact frequency.The strongest alignment occurred in deepest layers across the five evaluated models.
- 4.2 BINDING SITES AND POST-TRANSLATIONAL MODIFICATIONS: Binding-site attention was strongest in ProtAlbert, where 22 heads focused over 50% of attention on binding sites versus a 4.8% background frequency.The other BERT models also showed heads focusing up to 48.2%, 50.7%, and 45.6% of attention on binding sites.
- 4.2 BINDING SITES AND POST-TRANSLATIONAL MODIFICATIONS: ProtXLNet showed weaker binding-site targeting, with its top head at 15.1% and average head at 6.2%, though the cause of this disparity remained unclear.The authors note differences in architecture and pretraining objective as possible explanations.
- 4.2 BINDING SITES AND POST-TRANSLATIONAL MODIFICATIONS: A small number of heads strongly targeted PTM-associated positions, including 64% of TapeBert head 11-6 despite PTMs occurring at only 0.8% of sequence positions.The same head also targeted binding sites at 49%.
- 4.3 CROSS-LAYER ANALYSIS: Deeper layers focused relatively more on binding sites and contacts, while lower layers represented local secondary structure and attention encoded contact knowledge mainly in the final 1–2 layers.Embedding probes accumulated contact knowledge gradually, unlike attention weights in this analysis.
- 4.4 AMINO ACIDS AND THE SUBSTITUTION MATRIX: At least one head focused over 25% of attention on 16 amino acids in TapeBert, and attention distributions correlated with substitution relationships across models.ProtTrans correlations ranged from 0.60 to 0.75, compared with randomized-model correlations ranging from -0.05 to 0.21.
5 RELATED WORK
Prior work established Transformer interpretability in NLP and protein modeling, while this paper extends attention-based analysis to biophysical properties and compares attention with embeddings across layers.
- Protein-model interpretability has examined output embeddings, but this work adopts an interpretability-first focus on internal representations.
- NLP Transformer interpretability commonly probes embeddings, syntax, and semantics using classifiers and other analyses.
- Attention studies in language find correlations with syntactic and semantic relationships, with explanatory power varying by task and architecture.
- This work extends attention interpretability to protein-specific biophysical properties and relationships.
- A unified metric enables direct comparison of information in attention weights and layer embeddings across layers.
- The paper introduces a tool for visualizing attention in three-dimensional protein structure.
6 CONCLUSIONS AND FUTURE WORK
The paper frames protein language models as interpretable tools for biological investigation and visualizes attention in structural context, using models spanning multiple architectures and datasets.
- Attention visualized within three-dimensional protein structure is presented as a way to make learned representations accessible to domain experts.
- The authors suggest attention could help uncover novel relationships or nuanced forms of existing measures such as contact maps.
- Table 1 summarizes five pretrained models by architecture, depth, heads, parameters, pretraining source, and dataset size.
A.2 BERT TRANSFORMER ARCHITECTURE
BERT applies stacked encoder layers with multi-head self-attention to amino-acid sequences, while ALBERT and XLNet modify parameterization, embeddings, or pretraining and positional mechanisms.
- BERT: BERT embeds an amino-acid sequence and repeatedly applies encoder layers combining multi-head self-attention, feedforward networks, normalization, and residual connections.
- BERT: Each BERT layer outputs contextualized embeddings for the sequence positions.
- Self-Attention: Self-attention assigns normalized positive weights to token pairs, with each layer and head learning a distinct attention mechanism.
- Self-Attention: Attention weights arise from scaled query-key dot products followed by softmax and produce weighted sums of value vectors.
- Self-Attention: Multi-head attention linearly projects queries, keys, and values in parallel before concatenating the resulting representations.
- ALBERT and XLNet: ALBERT shares parameters across layers and uses factorized embeddings, whereas XLNet uses bidirectional autoregressive pretraining, segment recurrence, and relative positional encoding.
B ADDITIONAL EXPERIMENTAL DETAILS
The analysis measures how high-attention token pairs or destinations align with biological properties and uses significance tests and null models to assess whether alignments exceed chance.
- An indicator function marks whether a token pair has a property, and the metric computes the dataset-level proportion of high-attention pairs possessing it.
- Statistical tests compare high-confidence attention alignment with the property’s background frequency using a two-proportion z-test.
- Null-model checks compare findings against randomly initialized or shuffled-sequence models, although neither produced attention above the analysis threshold.
B.3 PROBING METHODOLOGY
The study probes protein-model embeddings and attention weights to quantify how much information they encode about structural and functional properties. It uses task-specific classifiers and metrics for token-level and token-pair predictions.
- Embedding probe: Embedding probes apply linear classifiers to layerwise token representations for binding-site and secondary-structure prediction.For contact prediction, pairwise features concatenate elementwise differences and products of two token vectors.
- Evaluation metrics: Secondary-structure probing uses F1, while contact and binding-site probing use precision@L/5 and precision@L/20, respectively.The binding-site metric reflects that approximately 4.8% of sequence positions are binding sites in the dataset.
- Attention probe: Attention probes treat pairwise attention weights as features for predicting whether amino acids are in contact.This directly quantifies the attention mechanism’s knowledge of the contact property.
- Attention probe: In the multi-head setting, attention weights across all heads in a layer form a feature vector for probing property knowledge.The classifier assesses information available across the entire layer.
B.4 DATASETS
The analysis uses two TAPE protein-sequence datasets for different structural tasks and adds a third dataset for binding-site and post-translational-modification analysis.
- Datasets: ProteinNet supports the analysis of amino acids and contact maps.ProteinNet is one of two sequence datasets drawn from the TAPE repository.
- Datasets: The Secondary Structure dataset supports the analysis of secondary structure.It is the second TAPE repository dataset used in the study.
- Datasets: A third dataset was created for binding-site and post-translational-modification analysis.The supplied passage identifies its purpose but does not provide further dataset details.
C ADDITIONAL RESULTS OF ATTENTION ANALYSIS
Additional attention analyses examine secondary-structure classes, contact-map alignment, binding sites, post-translational modifications, and amino-acid preferences. Several captions compare model heads with background frequencies and randomized null models using corrected significance tests.
- Secondary structure: Figures 8–10 report the percentage of each head’s attention focused on helix, strand, and turn/bend secondary structures.These figures organize attention concentration by secondary-structure category.
- Contact maps: Figure 11 identifies the top 10 heads per model by attention aligned with contact maps, with significant differences from background contact frequency.The caption reports p < 0.00001 and Bonferroni correction for confidence intervals and tests.
- Contact maps: Figure 12 provides the top-10 contact-aligned heads for null models.The caption directs readers to Appendix B.2 for methodological details.
- Binding sites: Figure 13 identifies the top 10 heads per model by attention focused on binding sites, with all differences from background frequency statistically significant.The caption reports p < 0.00001 and Bonferroni correction.
- Post-translational modifications: Figures 15–16 examine attention focused on post-translational modifications, including top heads whose enrichment differs significantly from background frequency.Figure 16 reports p < 0.00001 with Bonferroni correction; Figure 17 gives corresponding null-model heads.
- Amino-acid preferences: Figures 18–19 report each head’s attention percentage for individual amino acids, while Table 3 compares maximally attentive heads in standard and randomized TapeBert.Table 3 states that attention percentages differ significantly from amino-acid background frequencies after Bonferroni correction.