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

Sintesi Tecnica: PairSAE – Interpretabilità Meccanicistica dalle Rappresentazioni di Coppia nella Co-Piegatura delle Proteine

Problema
I modelli di fondazione per la biologia strutturale, come AlphaFold e i modelli Boltz, hanno raggiunto un'accuratezza straordinaria nella previsione di strutture e complessi biomolecolari. Tuttavia, comprendere le caratteristiche interne che guidano questi output rimane una sfida significativa. Sebbene gli Autoencoder Sparsi (SAE) abbiano con successo disaccoppiato le caratteristiche "sovrapposte" nei modelli transformer-based per i modelli linguistici (LLM) e i modelli linguistici proteici (pLM), mappandoli a concetti monosemantici, essi non si trasferiscono agevolmente ai modelli di previsione strutturale.

A differenza degli standard LLM, i modelli di previsione strutturale utilizzano architetture pairformer che mantengono sia gli embedding di sequenza che i tensori di interazione a coppie. L'applicazione ingenua degli SAE a questi modelli presenta due criticità fondamentali:

1. Esplosione Quadratica: L'apprendimento di dizionari separati per le rappresentazioni di coppia scala quadraticamente con la lunghezza della sequenza, rendendo l'analisi computazionalmente proibitiva.
2. Concetti Offuscati: Le caratteristiche possono essere sovrapposte congiuntamente tra gli spazi di sequenza e di coppia. Trattarle separatamente non riesce a catturare i concetti distribuiti tra entrambe le rappresentazioni, mentre trattarle indipendentemente oscura la natura congiunta delle caratteristiche apprese.

Metodologia: PairSAE
Gli autori introducono PairSAE, un framework progettato per apprendere un insieme condiviso di caratteristiche a livello di token che ricostruisce congiuntamente sia le rappresentazioni di sequenza che quelle di coppia. Il metodo consiste in tre fasi primarie:

1. Riassunto delle Interazioni di Coppia tramite SVD N-mode:
   Per evitare la scalabilità quadratica dei tensori di coppia ($Z \in \mathbb{R}^{N_{tok} \times N_{tok} \times n_z}$), gli autori eseguono una Decomposizione dei Valori Singolari (SVD) N-mode. Estraggono le matrici del modo-1 e del modo-2 dei vettori singolari sinistri ($U^{(1)}$ e $U^{(2)}$), che codificano i ruoli di interazione riga e colonna di ciascun token. Questi vengono troncati ai primi $r$ vettori singolari e concatenati per formare un embedding di riassunto per token $m_i$. Questo passaggio comprime l'informazione di coppia in una rappresentazione a livello di token che preserva i ruoli di interazione senza la complessità quadratica.

2. Architettura dell'Autoencoder Condiviso:
   L'encoder PairSAE prende come input la concatenazione dell'originale embedding di sequenza ($s_i$) e del riassunto di interazione derivato da SVD ($m_i$). Produce un vettore di caratteristiche sparse $h_i$ utilizzando una non-linearità che induce la sparsità (una composizione di BatchTopK e ReLU).
   Fondamentalmente, il decoder utilizza questo set di caratteristiche condivise $h$ per ricostruire sia:

   • L'embedding di sequenza: $\hat{s}_i = D_s h_i + b_{dec}^s$
   • L'embedding di coppia: $\hat{z}_{i,j} = D_z^{row} h_i + D_z^{col} h_j + b_{dec}^z$
     Ciò costringe il modello a imparare una base di caratteristiche unificata che spieghi sia il contesto locale della sequenza che le interazioni di coppia.

3. Funzione di Perdita e Addestramento:
   Il modello è addestrato utilizzando una perdita SAE Matryoshka, che sezionizza il decoder e i vettori di caratteristiche attraverso larghezze annidate ($c_1 < c_2 < c_3 = D$) per garantire la robustezza rispetto alle dimensioni del dizionario. La perdita totale include termini di ricostruzione sia per le rappresentazioni di sequenza che di coppia, pesati equamente, più una perdita ausiliaria per prevenire le caratteristiche "morte". L'addestramento è accelerato calcolando un'approssimazione Monte Carlo per il termine di perdita di coppia, campionando un singolo $j$ per ogni $i$.

Contributi Chiave

• Architettura PairSAE: Un metodo innovativo per applicare l'apprendimento di dizionari sparsi ai modelli basati su pairformer, riassumendo i tensori di coppia tramite SVD N-mode e decodificando sia nello spazio di sequenza che in quello di coppia da un set di caratteristiche condiviso.
• Scoperta di Caratteristiche Interpretabili: Il metodo recupera con successo caratteristiche che si allineano con concetti biofisici, evitando le insidie dei convenzionali SAE quando applicati ai modelli di biologia strutturale.
• Applicazione all'Affinità di Legame: Il framework viene applicato alle attivazioni di Boltz-2 su complessi proteina-ligando (dataset PLINDER) per collegare le rappresentazioni latenti alle previsioni di affinità di legame.

Risultati
Gli autori hanno valutato le rappresentazioni di Boltz-2 su complessi proteina-ligando:

• Linear Probing: Le caratteristiche di PairSAE hanno dimostrato un alto potere predittivo per le annotazioni di residui UniProt. Le caratteristiche dei layer 33 e 64 (terzo step di recycling) hanno raggiunto F1 score del 29,1% e 24,0% a livello di token, e 53,2% e 49,6% a livello di complesso, superando significativamente i classificatori basati sui neuroni di ESM2-650M (che hanno ottenuto rispettivamente lo 0,8% e il 19,7%).
• Allineamento dei Concetti: Le visualizzazioni hanno mostrato caratteristiche specifiche che si attivano su distinti motivi strutturali, come regioni transmembrana, ponti disolfuro e siti attivi delle proteasi, con alti punteggi F1 rispetto alla verità di base (ground truth) di UniProt.
• Previsione dell'Affinità: Utilizzando la regressione LASSO sulle caratteristiche max-pooled di PairSAE, gli autori hanno predetto i valori di affinità di Boltz-2 con un $R^2$ di 0,528 su un set di test (PoseBusters). Sono state identificate caratteristiche specifiche (ad esempio, la feature 2299) che correlano fortemente con i complessi ad alta affinità, attivandosi specificamente sui ligandi e mostrando differenze di gruppo significative rispetto ai complessi inattivi.

Significatività e Rivendicazioni
Il documento afferma che PairSAE collega con successo lo spazio latente dei modelli di fondazione per la biologia strutturale ai concetti strutturali interpretabili. Affrontando i vincoli architettonici specifici dei modelli pairformer, PairSAE chiarisce "cosa sa il modello" riguardo al ripiegamento e al legame senza soccombere alla complessità quadratica o ai problemi di sovrapposizione delle caratteristiche che limitano i convenzionali SAE.

Gli autori pongono questo lavoro come un passo verso la valutazione e la scoperta principiata nella biologia strutturale, consentendo ai ricercatori di valutare l'affidabilità della previsione e la plausibilità biofisica. Notano che, sebbene il metodo mappi efficacemente le caratteristiche di sequenza e di coppia alle annotazioni note, rimane una limitazione nel mappare le caratteristiche attivate dai ligandi a concetti specifici a causa della mancanza di annotazioni dettagliate dei ligandi. Il lavoro futuro proposto mira a scalare lo studio su più layer, integrare queste caratteristiche in pipeline automatizzate e sviluppare metodi di steering per il design proteico interpretabile.