Retrieval and competition: how a protein foundation model starts a protein

Resumen Técnico: Recuperación y Competencia en Modelos Fundacionales de Proteínas

Planteamiento del Problema
Los modelos de lenguaje de proteínas (PLM) son cada vez más utilizados para guiar decisiones experimentales y clínicas. Sin embargo, persiste una ambigüedad crítica: ¿refleja una predicción confiable el reconocimiento de evidencia biológicamente significativa dentro de la entrada, o simplemente la recuperación de un valor estadístico predeterminado almacenado, aplicado independientemente de la evidencia específica? Aunque ambos mecanismos producen predicciones idénticas, solo el primero justifica la confianza en la generalización. Los métodos existentes de interpretabilidad mecánica han identificado características internas y descompuesto representaciones, pero han tratado en gran medida el mecanismo de atención —el enrutamiento de información entre posiciones— como una caja negra. Específicamente, el papel de la codificación posicional en la dirección del flujo de información ha recibido poca escrutinio. Este artículo aborda esta brecha rastreando la ruta computacional de una regla biológica simple y casi universal: que las proteínas típicamente comienzan con metionina (Met).

Metodología
Los autores analizaron ESM2-8M, un modelo de lenguaje de proteínas de 6 capas y 8 millones de parámetros. Emplearon un conjunto de técnicas de interpretabilidad mecánica para diseccionar la predicción de metionina en la primera posición (posición 0) cuando el residuo inicial está enmascarado:

• Estudios de Ablación: Ablación sistemática por cero de cabezas de atención y módulos MLP en posiciones específicas de tokens (solo enmascarado vs. todas las posiciones) para identificar componentes necesarios.
• Parcheo de Activaciones: Inyección de representaciones intermedias (consultas, claves, valores, salidas de atención) de ejecuciones "limpias" (enmascarado en posición 0) en ejecuciones "corruptas" (enmascarado en posiciones internas) para aislar dependencias causales.
• Descubrimiento de Circuitos: Una búsqueda voraz para identificar conjuntos mínimos de componentes requeridos para la predicción de metionina.
• Descomposición Norma–Dirección: Una descomposición novedosa de las puntuaciones de atención dentro de las bandas de frecuencia de la Incrustación Posicional Rotatoria (RoPE). Los autores separaron la puntuación de atención ($Q \cdot K$) en un término de magnitud (norma de la consulta) y un término de alineación (coseno del ángulo entre la consulta y la clave) para determinar cómo se codifica la información posicional.
• Manipulación de RoPE: Intercambio de rotaciones RoPE entre posiciones y aplicación de desplazamientos de índice uniformes para probar la causalidad geométrica.
• Experimentos de Competencia: Experimentos de alargamiento y destrucción de contexto utilizando secuencias de poli-alanina para observar cómo las señales competidoras afectan el logit de metionina.

Resultados Clave

1. Competencia de Circuitos, No Reconocimiento Local:
   El modelo predice metionina en la posición 0 no mediante la detección local de la identidad de la posición 0, sino a través de una competencia entre circuitos. El circuito de metionina contribuye con una señal aproximadamente constante en la primera posición. En posiciones internas, circuitos dependientes del contexto producen señales competidoras que a menudo superan el prior de metionina. Cuando el contexto se destruye (por ejemplo, enmascarando residuos internos), la señal de metionina vuelve a emerger como dominante. Esto indica que el modelo no "activa" una regla para la posición 0; más bien, mantiene un prior constante que es superado por señales contextuales más fuertes en otros lugares.

2. Arquitectura del Circuito de Metionina:
   La predicción está mediada por un circuito distribuido y multicapa en lugar de una regla local simple:

   • Lectura (Capa 6): La cabeza L6H8 actúa como la lectura final. Atiende casi exclusivamente al token de Inicio de Secuencia (BOS) y escribe una señal específica de metionina en la salida. Su vector de salida-valor (OV) está altamente alineado con el gradiente del logit de metionina.
   • Composición de Consultas (Capas Previas): L6H8 no actúa de forma independiente. Depende de una consulta construida aguas arriba (principalmente en las Capas 1, 4 y 5) por cabezas de atención que atienden al token BOS. Estas cabezas aguas arriba utilizan señales posicionales mediadas por RoPE para ensamblar una consulta que, al ser rotada por el RoPE de L6H8, se alinea con la clave BOS.
   • Construcción de Representaciones (MLP): Los MLP de capas tempranas (L1, L2, L4) son esenciales para construir las representaciones clave/valor en BOS y la consulta en la posición enmascarada. Sin ellos, el circuito falla.

3. El Papel Activo del Token BOS:
   Contrario a la visión del BOS como un delimitador pasivo, los autores encuentran que funciona como un nodo computacional activo. Almacena una representación de referencia estable e invariante a la posición (clave y valor) que el circuito recupera. La especificidad de la predicción (por qué ocurre en la posición 0 y no en otros lugares) está determinada enteramente por la consulta formada en el token enmascarado, no por cambios en la representación BOS.

4. Mecanismo de Selectividad Posicional (RoPE):
   Los autores introducen una descomposición norma–dirección de las puntuaciones de atención a través de bandas de frecuencia RoPE. Descubren que la selectividad posicional surge de cambios acoplados en la norma de la consulta y la alineación angular:

   • Banda Dominante ($f_0$): Opera en ambos ejes. En posiciones internas, la norma de la consulta para esta banda colapsa aproximadamente siete veces, y la alineación de coseno con la clave BOS cae de ~0.8 a cerca de cero.
   • Bandas Secundarias ($f_1, f_2, f_6$): Actúan principalmente como selectores de alineación, donde los valores de coseno invierten su signo o colapsan, mientras que las normas permanecen relativamente estables.
   • RoPE como Puerta: En un circuito mínimo aislado (sin circuitos contextuales competidores), eliminar RoPE hace que el modelo prediga metionina en todas las posiciones. Por lo tanto, RoPE no solo habilita la predicción en la posición 0; activa la puerta del circuito para suprimir el prior en posiciones internas.

Significado y Afirmaciones
El artículo afirma que distinguir entre el reconocimiento biológico y la recuperación estadística requiere resolución a nivel de circuitos individuales, bandas de frecuencia y composición de consultas.

• Verificación Mecánica: Incluso para la regla biológica más simple (las proteínas comienzan con Met), la predicción del modelo está mediada por un circuito computacional distribuido en lugar de un reconocimiento directo. Esto sugiere que para tareas más complejas, la relación entre la confianza del modelo y la evidencia biológica subyacente será cada vez más oscurecida.
• Interpretabilidad de la Confianza: Una alta confianza no garantiza el uso de evidencia biológica correcta; puede indicar simplemente que un circuito (el prior estadístico) está dominando a otros.
• Motivos Computacionales Generales: Los hallazgos destacan el papel activo de los tokens BOS como marcos de referencia y el uso de la composición de consultas para el enrutamiento posicional, motivos que pueden ser generales a las arquitecturas de transformadores más allá de los modelos de proteínas.
• Limitaciones: Los autores señalan que, aunque la arquitectura del circuito se identifica en ESM2-8M, sigue siendo una pregunta abierta si estos motivos específicos persisten en modelos más grandes o a través de diferentes familias de PLM. También identifican que la composición de consultas aguas arriba involucra tanto cabezas impulsadas puramente por RoPE como al menos una cabeza dependiente del contexto (L4H3), cuyo contenido específico permanece por identificar.

El estudio concluye que la verificación mecánica es necesaria y desafiante para predicciones donde las apuestas biológicas son altas, ya que el "conocimiento" del modelo sobre la biología puede ser un subproducto de regularidades estadísticas recuperadas a través de circuitos complejos y distribuidos.