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

2026-06-25

AI Summary

Existing sparse autoencoders struggle to effectively interpret pairwise representations in Pairformer-like protein co-folding models, often leading to feature explosion and an inability to model the joint distribution of sequence and pairwise features. This work proposes PairSAE, a novel framework that, for the first time, employs N-mode SVD to compress pairwise tensors into token-centric interaction roles and introduces a shared sparse autoencoder to jointly reconstruct both sequence and pairwise representations. By circumventing the quadratic growth inherent in conventional sparse autoencoders along pairwise dimensions, PairSAE extracts highly interpretable features on the PLINDER complex that align closely with UniProt functional annotations and accurately predicts Boltz-2 binding affinities, thereby uncovering structurally meaningful biological concepts learned internally by the model.

Abstract

Foundation models for structural biology have achieved remarkable performance in predicting biomolecular structure and show promise for the design of proteins and small molecules. Yet understanding which internal features drive their outputs remains challenging. Standard sparse autoencoders (SAEs), effective on transformer-style sequence embeddings, do not transfer cleanly to pairformer-like architectures: naively operating on pairwise representations yields a quadratic blow-up of features and obscures concepts distributed jointly across sequence and pair representations. We introduce PairSAE, which summarizes pairwise tensors via an N-mode SVD into token-wise interaction roles, then uses a sparse autoencoder to learn a shared set of token-level features that decode into both sequence and pair representations. Evaluated on Boltz-2 activations for PLINDER protein-ligand complexes, PairSAE yields interpretable features that align with UniProt annotations and predict Boltz-2 affinity values. These results indicate that PairSAE links the latent space of foundation models for structural biology to interpretable structural concepts, clarifying what the model "knows" while avoiding pairformer-induced pitfalls that limit conventional SAEs.

Problem

• mechanistic interpretability
• protein co-folding
• pairwise representations
• sparse autoencoders
• foundation models

Innovation

• PairSAE
• mechanistic interpretability
• pair representations
• N-mode SVD
• sparse autoencoder