PairSAE: Mechanistic Interpretability from Pair Representations in Protein Co-Folding

Technical Summary: PairSAE – Mechanistic Interpretability from Pair Representations in Protein Co-Folding

Problem Statement
Foundation models for structural biology, such as AlphaFold and the Boltz models, have achieved remarkable accuracy in predicting biomolecular structures and complexes. However, understanding the internal features driving these outputs remains a significant challenge. While Sparse Autoencoders (SAEs) have successfully disentangled "superposed" features in transformer-based language models (LLMs) and protein language models (pLMs) by mapping them to monosemantic concepts, they do not transfer cleanly to structure prediction models.

Unlike standard LLMs, structure prediction models utilize pairformer architectures that maintain both sequence embeddings and pairwise interaction tensors. Naïvely applying SAEs to these models presents two critical issues:

1. Quadratic Blow-up: Learning separate dictionaries for pairwise representations scales quadratically with sequence length, making analysis computationally prohibitive.
2. Obscured Concepts: Features may be superposed jointly across sequence and pair spaces. Treating them separately fails to capture concepts distributed across both representations, while treating them independently obscures the joint nature of the learned features.

Methodology: PairSAE
The authors introduce PairSAE, a framework designed to learn a shared set of token-level features that jointly reconstruct both sequence and pair representations. The method consists of three primary stages:

1. Summarizing Pairwise Interactions via N-mode SVD:
   To avoid the quadratic scaling of pairwise tensors ($Z \in \mathbb{R}^{N_{tok} \times N_{tok} \times n_z}$), the authors perform an N-mode Singular Value Decomposition (SVD). They extract the mode-1 and mode-2 matrices of left singular vectors ($U^{(1)}$ and $U^{(2)}$), which encode the row and column interaction roles of each token. These are truncated to the top $r$ singular vectors and concatenated to form a token-wise summary embedding $m_i$. This step compresses the pairwise information into a token-level representation that preserves interaction roles without quadratic complexity.

2. Shared Autoencoder Architecture:
   The PairSAE encoder takes the concatenation of the original sequence embedding ($s_i$) and the SVD-derived interaction summary ($m_i$) as input. It produces a sparse latent feature vector $h_i$ using a sparsity-inducing nonlinearity (a composition of BatchTopK and ReLU).
   Crucially, the decoder uses this shared feature set $h$ to reconstruct both:

   • The sequence embedding: $\hat{s}_i = D_s h_i + b_{dec}^s$
   • The pairwise embedding: $\hat{z}_{i,j} = D_z^{row} h_i + D_z^{col} h_j + b_{dec}^z$
     This forces the model to learn a unified feature basis that explains both local sequence context and pairwise interactions.

3. Loss Function and Training:
   The model is trained using a Matryoshka SAE loss, which slices the decoder and feature vectors across nested widths ($c_1 < c_2 < c_3 = D$) to ensure robustness across dictionary sizes. The total loss includes reconstruction terms for both sequence and pair representations, weighted equally, plus an auxiliary loss to prevent "dead" features. Training is accelerated by computing a Monte Carlo approximation for the pairwise loss term, sampling a single $j$ for each $i$.

Key Contributions

• PairSAE Architecture: A novel method to apply sparse dictionary learning to pairformer-based models by summarizing pairwise tensors via N-mode SVD and decoding into both sequence and pair spaces from a shared feature set.
• Interpretable Feature Discovery: The method successfully recovers features that align with biophysical concepts, avoiding the pitfalls of conventional SAEs when applied to structural biology models.
• Application to Binding Affinity: The framework is applied to Boltz-2 activations on protein-ligand complexes (PLINDER dataset) to link latent representations to binding affinity predictions.

Results
The authors evaluated PairSAE on Boltz-2 representations for protein-ligand complexes:

• Linear Probing: PairSAE features demonstrated high predictive power for UniProt residue annotations. Features from layers 33 and 64 (third recycling step) achieved F1 scores of 29.1% and 24.0% at the token level, and 53.2% and 49.6% at the complex level, significantly outperforming classifiers based on ESM2-650M neurons (which achieved 0.8% and 19.7% respectively).
• Concept Alignment: Visualizations showed specific features activating on distinct structural motifs, such as transmembrane regions, disulfide bonds, and protease active sites, with high F1 scores against UniProt ground truth.
• Affinity Prediction: Using LASSO regression on max-pooled PairSAE features, the authors predicted Boltz-2 affinity values with an $R^2$ of 0.528 on a test set (PoseBusters). Specific features (e.g., feature 2299) were identified that strongly correlate with high-affinity complexes, activating specifically on ligands and showing significant group differences compared to inactive complexes.

Significance and Claims
The paper claims that PairSAE successfully links the latent space of foundation models for structural biology to interpretable structural concepts. By addressing the specific architectural constraints of pairformer models, PairSAE clarifies "what the model knows" regarding folding and binding without succumbing to the quadratic complexity or feature superposition issues that limit conventional SAEs.

The authors position this work as a step toward principled evaluation and discovery in structural biology, enabling researchers to assess prediction reliability and biophysical plausibility. They note that while the method effectively maps sequence and pair features to known annotations, a limitation remains in mapping ligand-activating features to specific concepts due to a lack of detailed ligand annotations. Future work is proposed to scale the study to more layers, integrate these features into automated pipelines, and develop steering methods for interpretable protein design.