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

Technische Samenvatting: PairSAE – Mechanistische Interpreteerbaarheid vanuit Paarrepresentaties in Eiwit Co-folding

Probleemstelling
Fundamentele modellen voor structurele biologie, zoals AlphaFold en de Boltz-modellen, hebben een opmerkelijke nauwkeurigheid bereikt in het voorspellen van biomoleculaire structuren en complexen. Het begrijpen van de interne kenmerken die deze outputs aansturen, blijft echter een grote uitdaging. Hoewel Sparse Autoencoders (SAEs) succesvol "superponeerde" kenmerken hebben ontward in transformer-gebaseerde taalmodellen (LLMs) en eiwittaalmodellen (pLMs) door deze te mappen naar monosemantische concepten, laten ze zich niet zonder problemen vertalen naar structurele voorspellingsmodellen.

In tegenstelling tot standaard LLMs maken structurele voorspellingsmodellen gebruik van pairformer-architecturen die zowel sequentie-embeddings als paarinteractietensoren behouden. Het naïef toepassen van SAEs op deze modellen brengt twee kritieke problemen met zich mee:

1. Kwadratische Explosie: Het leren van afzonderlijke woordenboeken voor paarrepresentaties schaalt kwadratisch met de sequentielengte, wat analyse computationeel onhaalbaar maakt.
2. Obscure Concepten: Kenmerken kunnen gezamenlijk superponeerd zijn over zowel sequentie- als paarruimtes. Het apart behandelen van deze ruimtes faalt in het vastleggen van concepten die over beide representaties verdeeld zijn, terwijl het onafhankelijk behandelen de gezamenlijke aard van de geleerde kenmerken verhult.

Methodologie: PairSAE
De auteurs introduceren PairSAE, een raamwerk dat is ontworpen om een gedeelde set token-niveau kenmerken te leren die zowel de sequentie- als de paarrepresentaties gezamenlijk reconstrueren. De methode bestaat uit drie primaire stadia:

1. Samenvatten van Paarinteracties via N-mode SVD:
   Om de kwadratische schaling van paar-tensoren ($Z \in \mathbb{R}^{N_{tok} \times N_{tok} \times n_z}$) te vermijden, voeren de auteurs een N-mode Singular Value Decomposition (SVD) uit. Ze extraheren de mode-1 en mode-2 matrices van de linker singuliere vectoren ($U^{(1)}$ en $U^{(2)}$), die de rol van de rij- en kolominteractie voor elk token coderen. Deze worden afgekort naar de top $r$ singuliere vectoren en geconcateneerd om een token-wise samenvattende embedding $m_i$ te vormen. Deze stap comprimeert de paar-informatie naar een token-niveau representatie die interactierollen behoudt zonder kwadratische complexiteit.

2. Gedeelde Autoencoder Architectuur:
   De PairSAE-encoder neemt de concatenatie van de oorspronkelijke sequentie-embedding ($s_i$) en de via SVD verkregen interactie-samenvatting ($m_i$) als input. Het produceert een ijle (sparse) latente feature vector $h_i$ met behulp van een sparsity-inducerende niet-lineariteit (een compositie van BatchTopK en ReLU).
   Cruciaal is dat de decoder deze gedeelde feature set $h$ gebruikt om zowel de sequentie-embedding als de paar-embedding te reconstrueren:

   • De sequentie-embedding: $\hat{s}_i = D_s h_i + b_{dec}^s$
   • De paar-embedding: $\hat{z}_{i,j} = D_z^{row} h_i + D_z^{col} h_j + b_{dec}^z$
     Dit dwingt het model om een verenigde feature-basis te leren die zowel de lokale sequentiecontext als de paarinteracties verklaart.

3. Verliesfunctie en Training:
   Het model wordt getraind met een Matryoshka SAE loss, die de decoder en feature-vectoren over geneste breedtes ($c_1 < c_2 < c_3 = D$) snijdt om robuustheid over verschillende dictionary-groottes te garanderen. De totale loss bevat reconstructietermen voor zowel de sequentie- als de paarrepresentaties, die even zwaar wegen, plus een hulpverlies (auxiliary loss) om "dode" features te voorkomen. De training wordt versneld door een Monte Carlo-benadering voor de paar-loss term te berekenen, waarbij voor elke $i$ een enkele $j$ wordt gesampled.

Belangrijkste Bijdragen

• PairSAE Architectuur: Een nieuwe methode om sparse dictionary learning toe te passen op pairformer-gebaseerde modellen door paar-tensoren samen te vatten via N-mode SVD en te decoderen naar zowel sequentie- als paarruimtes vanuit een gedeelde feature set.
• Interpreteerbare Ontdekking van Kenmerken: De methode slaagt er succesvol in om kenmerken te herstellen die overeenkomen met biofysische concepten, waarbij de valkuilen van conventionele SAEs bij toepassing op structurele biologie modellen worden vermeden.
• Toepassing op Bindingsaffiniteit: Het raamwerk wordt toegepast op Boltz-2 activaties voor eiwit-ligand complexen (PLINDER dataset) om de link tussen latente representaties en bindingsaffiniteit-voorspellingen te leggen.

Resultaten
De auteurs evalueerden PairSAE op Boltz-2 representaties voor eiwit-ligand complexen:

• Lineaire Probing: PairSAE-features vertoonden een hoge voorspellende kracht voor UniProt residu-annotaties. Features uit lagen 33 en 64 (derde recyclingstap) behaalden F1-scores van 29,1% en 24,0% op token-niveau, en 53,2% en 49,6% op complex-niveau, wat de classifiers gebaseerd op ESM2-650M neuronen (die respectievelijk 0,8% en 19,7% behaalden) aanzienlijk overtrof.
• Conceptuele Alignement: Visualisaties toonden specifieke features die activeren op onderscheidende structurele motieven, zoals transmembraanregio's, disulfidebruggen en protease actieve sites, met hoge F1-scores tegenover de UniProt grondwaarheid.
• Affiniteitsvoorspelling: Door LASSO-regressie te gebruiken op max-gepoolde PairSAE-features, voorspelden de auteurs Boltz-2 affiniteitswaarden met een $R^2$ van 0,528 op een testset (PoseBusters). Specifieke features (bijv. feature 2299) werden geïdentificeerd die sterk correleren met hoog-affiene complexen, waarbij ze specifiek activeren op liganden en significante groepsverschillen vertonen vergeleken met inactieve complexen.

Significantie en Claims
Het artikel claimt dat PairSAE er succesvol in slaagt de latente ruimte van fundamentele modellen voor structurele biologie te koppelen aan interpreteerbare structurele concepten. Door de specifieke architecturale beperkingen van pairformer-modellen aan te pakken, verheldert PairSAE "wat het model weet" met betrekking tot vouwen en binding, zonder ten onder te gaan aan de kwadratische complexiteit of de problemen met feature-superpositie die conventionele SAEs beperken.

De auteurs positioneren dit werk als een stap naar een principiële evaluatie en ontdekking in de structurele biologie, wat onderzoekers in staat stelt om de betrouwbaarheid van voorspellingen en biofysische plausibiliteit te beoordelen. Ze merken op dat hoewel de methode effectief sequentie- en paar-features naar bekende annotaties mapt, er een beperking blijft in het mappen van ligand-activerende features naar specifieke concepten vanwege een gebrek aan gedetailleerde ligand-annotaties. Toekomstig werk wordt voorgesteld om de studie op te schalen naar meer lagen, de features te integreren in geautomatiseerde pipelines en methoden voor sturing (steering) te ontwikkelen voor interpreteerbaar eiwitontwerp.