Mechanistic Interpretability of Antibody Language Models Using SAEs

Resumen Técnico: Interpretabilidad Mecanicista de Modelos de Lenguaje de Anticuerpos mediante SAEs

1. El Problema

El desarrollo de fármacos basados en anticuerpos depende de modelos de lenguaje de proteínas (PLMs) para optimizar la generación de secuencias, la humanización y la "desarrollabilidad" (solubilidad, estabilidad, etc.). Sin embargo, estos modelos operan como "cajas negras". La falta de interpretabilidad plantea tres riesgos críticos en el descubrimiento de fármacos:

• Falta de confianza: Dificultad para validar si las predicciones se basan en características biológicas reales.
• Detección de sobreajuste (overfitting): Incapacidad para saber si el modelo está memorizando datos en lugar de aprender reglas biológicas.
• Control limitado: La dificultad para dirigir (steer) la generación de secuencias hacia propiedades específicas de manera precisa y racional.

2. Metodología

El estudio utiliza Autoencoders Dispersos (Sparse Autoencoders - SAEs) para descomponer las activaciones densas de los modelos de lenguaje en un espacio latente de características interpretables. Se evalúan dos arquitecturas de SAE:

• TopK SAEs: Limitan el número de latentes activos a un valor $k$, lo que facilita la identificación de conceptos, pero puede sufrir de problemas como la "división de características" (feature splitting).
• Ordered SAEs (O-SAEs): Implementan una arquitectura de agrupamiento anidado (nested) que impone una jerarquía estricta. Esto obliga a los primeros latentes a capturar estructuras globales y abstractas, mientras que los posteriores capturan detalles finos.

Modelos evaluados:

• p-IgGen: Modelo principal (decoder-only, 17M parámetros).
• IgLM y ProGen2-OAS: Modelos adicionales para validar la generalización de los resultados a través de diferentes escalas y arquitecturas.

Tareas de validación:
Se utilizaron sondas lineales (linear probes) para identificar latentes correlacionados con conceptos biológicos: la identidad de las regiones determinantes de la complementariedad (CDR) y la identidad de los genes germinales (V/J). Además, se aplicó la técnica de direccionamiento (steering), que consiste en intervenir el estado oculto del modelo sumando el vector del decodificador del latente para alterar la generación.

3. Contribuciones Clave

• Desacoplamiento entre Interpretación y Direccionabilidad: El estudio demuestra que una alta correlación entre un latente y un concepto biológico (interpretabilidad) no garantiza que ese latente pueda controlar la generación de ese concepto (direccionabilidad/steerability).
• Validación de O-SAEs para Control Generativo: Identifica que los SAEs ordenados son superiores para el control preciso de la generación, a pesar de tener patrones de activación menos intuitivos visualmente.
• Generalización de la Técnica: Demuestra que la interpretabilidad mediante SAEs es aplicable y robusta en múltiples arquitecturas de modelos de anticuerpos.

4. Resultados Principales

• TopK SAEs: Lograron identificar con éxito características biológicas (como la identidad de los genes IGHJ y las regiones CDR) con una precisión muy alta (F1-score elevado). Sin embargo, fallaron al intentar dirigir la generación; aumentar la activación de estos latentes no alteraba de forma fiable la proporción de genes en las secuencias generadas.
• Ordered SAEs (O-SAEs): Aunque sus patrones de activación son más difusos y menos localizados espacialmente que los de TopK, permitieron un direccionamiento exitoso y predecible. Por ejemplo, al aumentar el latente asociado al gen IGHJ4, la proporción de este gen en la biblioteca generada aumentó significativamente (Pearson $R = 0.939$).
• Calidad de la Generación: El direccionamiento mediante O-SAEs produjo secuencias con mayor novedad (distancia de Hamming) y diversidad, pero manteniendo una "biologicidad" intacta (probabilidad de ESM-2 similar a las secuencias naturales).
• Transferibilidad: Los experimentos en IgLM y ProGen2-OAS confirmaron que los SAEs pueden recuperar la estructura de identidad de genes germinales independientemente del modelo utilizado.

5. Significado e Impacto

Este trabajo marca un avance en el diseño racional de anticuerpos mediante IA. Al demostrar que los Ordered SAEs pueden actuar como "controles" para la generación de secuencias, se abre la puerta a la creación de bibliotecas de anticuerpos con propiedades predefinidas (como la afinidad o la estabilidad) de manera controlada.

El estudio concluye que, para la ciencia de proteínas, no basta con "entender" qué está viendo el modelo (interpretabilidad); es crucial poder "manipular" esos conceptos para el diseño de fármacos (direccionabilidad), y los SAEs ordenados son la herramienta más prometedora para lograr este objetivo.