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

Résumé Technique : PairSAE – Interprétabilité Mécaniste à partir des Représentations de Paires dans le Co-repliement des Protéines

Énoncé du Problème
Les modèles de fondation pour la biologie structurale, tels qu'AlphaFold et les modèles Boltz, ont atteint une précision remarquable dans la prédiction des structures et des complexes biomoléculaires. Cependant, comprendre les caractéristiques internes qui pilotent ces résultats reste un défi important. Bien que les Autoencodeurs Creux (Sparse Autoencoders - SAE) aient réussi à désentrelacer les caractéristiques « superposées » dans les modèles de transformeurs basés sur le langage (LLM) et les modèles de langage protéiques (pLM) en les cartographiant vers des concepts monosémantiques, ils ne se transfèrent pas proprement aux modèles de prédiction de structure.

Contrairement aux LLM standards, les modèles de prédiction de structure utilisent des architectures pairformer qui maintiennent à la fois des plongements de séquence et des tenseurs d'interaction par paires. L'application naïve des SAE à ces modèles présente deux problèmes critiques :

1. Explosion Quadratique : L'apprentissage de dictionnaires séparés pour les représentations par paires augmente de manière quadratique avec la longueur de la séquence, rendant l'analyse trop coûteuse en termes de calcul.
2. Concepts Obscurcis : Les caractéristiques peuvent être superposées conjointement à travers les espaces de séquence et de paires. Les traiter séparément échoue à capturer les concepts distribués à travers les deux représentations, tandis que les traiter indépendamment occulte la nature conjointe des caractéristiques apprises.

Méthodologie : PairSAE
Les auteurs introduisent PairSAE, un cadre conçu pour apprendre un ensemble partagé de caractéristiques au niveau des tokens qui reconstruisent conjointement les représentations de séquence et de paires. La méthode se compose de trois étapes principales :

1. Résumé des Interactions par Paires via la SVD N-mode :
   Pour éviter l'échelle quadratique des tenseurs de paires ($Z \in \mathbb{R}^{N_{tok} \times N_{tok} \times n_z}$), les auteurs effectuent une Décomposition en Valeurs Singulières (SVD) N-mode. Ils extraient les matrices de mode-1 et de mode-2 des vecteurs singuliers de gauche ($U^{(1)}$ et $U^{(2)}$), qui encodent les rôles d'interaction de ligne et de colonne de chaque token. Ces matrices sont tronquées selon les $r$ vecteurs singuliers supérieurs et concaténées pour former un plongement de résumé par token $m_i$. Cette étape compresse l'information de paire en une représentation par token qui préserve les rôles d'interaction sans la complexité quadratique.

2. Architecture d'Autoencodeur Partagé :
   L'encodeur PairSAE prend en entrée la concaténation du plongement de séquence original ($s_i$) et du résumé d'interaction dérivé de la SVD ($m_i$). Il produit un vecteur de caractéristique creux $h_i$ utilisant une non-linéarité induisant la parcimonie (une composition de BatchTopK et ReLU).
   Crucialement, le décodeur utilise ce jeu de caractéristiques partagé $h$ pour reconstruire à la fois :

   • Le plongement de séquence : $\hat{s}_i = D_s h_i + b_{dec}^s$
   • Le plongement de paire : $\hat{z}_{i,j} = D_z^{row} h_i + D_z^{col} h_j + b_{dec}^z$
     Cela force le modèle à apprendre une base de caractéristiques unifiée qui explique à la fois le contexte local de la séquence et les interactions par paires.

3. Fonction de Perte et Entraînement :
   Le modèle est entraîné à l'aide d'une perte SAE Matryoshka, qui découpe le décodeur et les vecteurs de caractéristiques à travers des largeurs imbriquées ($c_1 < c_2 < c_3 = D$) pour assurer la robustesse à travers les tailles de dictionnaires. La perte totale inclut des termes de reconstruction pour les représentations de séquence et de paires, pondérés également, plus une perte auxiliaire pour empêcher les caractéristiques « mortes ». L'entraînement est accéléré en calculant une approximation de Monte Carlo pour le terme de perte de paire, en échantillonnant un seul $j$ pour chaque $i$.

Contributions Clés

• Architecture PairSAE : Une nouvelle méthode pour appliquer l'apprentissage de dictionnaire creux aux modèles basés sur les pairformers en résumant les tenseurs de paires via la SVD N-mode et en décodant dans les espaces de séquence et de paires à partir d'un ensemble de caractéristiques partagées.
• Découverte de Caractéristiques Interprétables : La méthode réussit à récupérer des caractéristiques qui s'alignent sur des concepts biophysiques, évitant les pièges des SAE conventionnels lorsqu'ils sont appliqués aux modèles de biologie structurale.
• Application à l'Affinité de Liaison : Le cadre est appliqué aux activations de Boltz-2 sur des complexes protéine-ligand (jeu de données PLINDER) pour lier les représentations latentes aux prédictions d'affinité de liaison.

Résultats
Les auteurs ont évalué PairSAE sur les représentations de Boltz-2 pour des complexes protéine-ligand :

• Probing Linéaire : Les caractéristiques de PairSAE ont démontré un fort pouvoir prédictif pour les annotations de résidus UniProt. Les caractéristiques des couches 33 et 64 (troisième étape de recyclage) ont atteint des scores F1 de 29,1 % et 24,0 % au niveau du token, et 53,2 % et 49,6 % au niveau du complexe, surpassant de manière significative les classificateurs basés sur les neurones d'ESM2-650M (qui ont atteint respectivement 0,8 % et 19,7 %).
• Alignement de Concept : Les visualisations ont montré des caractéristiques spécifiques s'activant sur des motifs structurels distincts, tels que les régions transmembranaires, les ponts disulfures et les sites actifs des protéases, avec des scores F1 élevés par rapport à la vérité de terrain UniProt.
• Prédiction d'Affinité : En utilisant la régression LASSO sur les caractéristiques max-poolées de PairSAE, les auteurs ont prédit les valeurs d'affinité de Boltz-2 avec un $R^2$ de 0,528 sur un ensemble de test (PoseBusters). Des caractéristiques spécifiques (par exemple, la caractéristique 2299) ont été identifiées comme étant fortement corrélées avec les complexes à haute affinité, s'activant spécifiquement sur les ligands et montrant des différences de groupe significatives par rapport aux complexes inactifs.

Signification et Revendications
L'article affirme que PairSAE parvient à lier l'espace latent des modèles de fondation pour la biologie structurale à des concepts structurels interprétables. En répondant aux contraintes architecturales spécifiques des modèles pairformer, PairSAE clarifie « ce que le modèle sait » concernant le repliement et la liaison sans succomber à la complexité quadratique ou aux problèmes de superposition de caractéristiques qui limitent les SAE conventionnels.

Les auteurs positionnent ce travail comme une étape vers une évaluation et une découverte principled de la biologie structurale, permettant aux chercheurs d'évaluer la fiabilité des prédictions et la plausibilité biophysique. Ils notent que bien que la méthode cartographie efficacement les caractéristiques de séquence et de paire vers des annotations connues, une limite demeure dans la cartographie des caractéristiques activées par les ligands vers des concepts spécifiques en raison d'un manque d'annotations détaillées des ligands. Des travaux futurs sont proposés pour étendre l'étude à davantage de couches, intégrer ces caractéristiques dans des pipelines automatisés et développer des méthodes de pilotage pour le design protéique interprétable.