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

Resumen Técnico: PairSAE – Interpretabilidad Mecanística a partir de Representaciones de Pares en el Co-plegamiento de Proteínas

Declaración del Problema
Los modelos fundacionales para la biología estructural, como AlphaFold y los modelos Boltz, han logrado una precisión notable en la predicción de estructuras y complejos biomoleculares. Sin embargo, comprender las características internas que impulsan estos resultados sigue siendo un desafío significativo. Si bien los Autoencoders Dispersos (SAEs) han logrado desenredar características "superpuestas" en modelos de transformadores basados en lenguaje (LLMs) y modelos de lenguaje de proteínas (pLMs) al mapearlas hacia conceptos monosemánticos, estos no se transfieren limpiamente a los modelos de predicción de estructuras.

A diferencia de los LLMs estándar, los modelos de predicción de estructuras utilizan arquitecturas pairformer que mantienen tanto las incrustaciones (embeddings) de secuencia como los tensores de interacción por pares. La aplicación ingenua de SAEs a estos modelos presenta dos problemas críticos:

1. Explosión Cuadrática: El aprendizaje de diccionarios separados para las representaciones de pares escala cuadráticamente con la longitud de la secuencia, lo que hace que el análisis sea computacionalmente prohibitivo.
2. Conceptos Obscurecidos: Las características pueden estar superpuestas conjuntamente a través de los espacios de secuencia y de pares. Tratar los espacios por separado no logra capturar conceptos distribuidos en ambos, mientras que tratarlos de forma independiente oscurece la naturaleza conjunta de las características aprendidas.

Metodología: PairSAE
Los autores presentan PairSAE, un marco diseñado para aprender un conjunto compartido de características a nivel de token que reconstruya conjuntamente las representaciones de secuencia y de pares. El método consta de tres etapas principales:

1. Resumen de Interacciones de Pares mediante SVD de N-modos:
   Para evitar la escala cuadrática de los tensores de pares ($Z \in \mathbb{R}^{N_{tok} \times N_{tok} \times n_z}$), los autores realizan una Descomposición en Valores Singulares (SVD) de N-modos. Extraen las matrices de modo-1 y modo-2 de los vectores singulares de la izquierda ($U^{(1)}$ y $U^{(2)}$), que codifican los roles de interacción de fila y columna de cada token. Estos se truncan a los $r$ vectores singulares superiores y se concatenan para formar una incrustación de resumen por token $m_i$. Este paso comprime la información de pares en una representación a nivel de token que preserva los roles de interacción sin la complejidad cuadrática.

2. Arquitectura de Autoencoder Compartido:
   El codificador PairSAE toma la concatenación de la incrustación de secuencia original ($s_i$) y el resumen de interacción derivado de SVD ($m_i$) como entrada. Produce un vector de características dispersas $h_i$ utilizando una no linealidad que induce la dispersión (una composición de BatchTopK y ReLU).
   Crucialmente, el decodificador utiliza este conjunto de características compartido $h$ para reconstruir tanto:

   • La incrustación de secuencia: $\hat{s}_i = D_s h_i + b_{dec}^s$
   • La incrustación de pares: $\hat{z}_{i,j} = D_z^{row} h_i + D_z^{col} h_j + b_{dec}^z$
     Esto obliga al modelo a aprender una base de características unificada que explique tanto el contexto local de la secuencia como las interacciones de pares.

3. Función de Pérdida y Entrenamiento:
   El modelo se entrena utilizando una pérdida SAE de Matrioshka, que segmenta el decodificador y los vectores de características a través de anchos anidados ($c_1 < c_2 < c_3 = D$) para asegurar la robustez a través de diferentes tamaños de diccionario. La pérdida total incluye términos de reconstrucción para las representaciones de secuencia y de pares, ponderados equitativamente, más una pérdida auxiliar para prevenir características "muertas". El entrenamiento se acelera mediante el cálculo de una aproximación de Monte Carlo para el término de pérdida de pares, muestreando un único $j$ para cada $i$.

Contribuciones Clave

• Arquitectura PairSAE: Un método novedoso para aplicar el aprendizaje de diccionarios dispersos a modelos basados en pairformer mediante el resumen de tensores de pares vía SVD de N-modos y la decodificación en espacios de secuencia y de pares desde un conjunto de características compartido.
• Descubrimiento de Características Interpretables: El método recupera con éxito características que se alinean con conceptos biofísicos, evitando los inconvenientes de los SAEs convencionales cuando se aplican a modelos de biología estructural.
• Aplicación a la Afinidad de Unión: El marco se aplica a las activaciones de Boltz-2 en complejos proteína-ligando (conjunto de datos PLINDER) para vincular las representaciones latentes con la predicción de la afinidad de unión.

Resultos
Los autores evaluaron las representaciones de Boltz-2 en complejos proteína-ligando:

• Probing Lineal: Las características de PairSAE demostraron un alto poder predictivo para las anotaciones de residuos de UniProt. Las características de las capas 33 y 64 (tercer paso de reciclamiento) alcanzaron puntuaciones F1 de 29.1% y 24.0% a nivel de token, y 53.2% y 49.6% a nivel de complejo, superando significativamente a los clasificadores basados en neuronas de ESM2-650M (que alcanzaron 0.8% y 19.7% respectivamente).
• Alineación de Conceptos: Las visualizaciones mostraron características específicas que se activan en motivos estructurales distintos, como regiones transmembrana, puentes disulfuro y sitios activos de proteasas, con puntuaciones F1 elevadas frente a la verdad de terreno de UniProt.
• Predicción de Afinidad: Utilizando regresión LASSO sobre las características de max-pooling de PairSAE, los autores predijeron los valores de afinidad de Boltz-2 con un $R^2$ de 0.528 en un conjunto de prueba (PoseBusters). Se identificaron características específicas (por ejemplo, la característica 2299) que se correlacionan fuertemente con complejos de alta afinidad, activándose específicamente en ligandos y mostrando diferencias de grupo significativas en comparación con complejos inactivos.

Significado y Reivindicaciones
El artículo afirma que PairSAE vincula con éxito el espacio latente de los modelos fundacionales para la biología estructural con conceptos estructurales interpretables. Al abordar las restricciones arquitectónicas específicas de los modelos pairformer, PairSAE esclarece "qué sabe el modelo" respecto al plegamiento y la unión sin sucumbir a la complejidad cuadrática o a los problemas de superposición de características que limitan a los SAEs convencionales.

Los autores posicionan este trabajo como un paso hacia la evaluación y el descubrimiento con principios en la biología estructural, permitiendo a los investigadores evaluar la fiabilidad de la predicción y la plausibilidad biofísica. Señalan que, si bien el método mapea eficazmente las características de secuencia y pares a anotaciones conocidas, persiste una limitación en el mapeo de las características que activan ligandos hacia conceptos específicos debido a la falta de anotaciones detalladas de ligandos. Se propone como trabajo futuro escalar el estudio a más capas, integrar estas características en procesos automatizados y desarrollar métodos de control (steering) para el diseño de proteínas interpretable.