VFUSE: Virulent Feature Understanding with Sparse autoEncoders

Resumen Técnico: VFUSE – Comprensión de Características Virulentas mediante Autoencoders Dispersos

Planteamiento del Problema

Los modelos generativos para el diseño de proteínas, particularmente los modelos de difusión de todos los átomos como RoseTTAFold3 (RF3) y RFDiffusion3 (RFD3), poseen la capacidad de doble uso para diseñar proteínas peligrosas (p. ej., neurotoxinas, proteínas de evasión inmunitaria) mediante el andamiaje de nuevas estructuras alrededor de plantillas peligrosas. Los métodos actuales de detección de bioseguridad dependen de la homología de secuencias o de clasificadores a nivel de secuencia. Estos enfoques sufren dos limitaciones críticas: operan post-hoc (marcando secuencias después de la generación) y carecen de interpretabilidad mecánica, lo que significa que no pueden explicar qué características estructurales específicas activaron una señal de peligro ni intervenir durante el proceso de muestreo. Además, aunque los Autoencoders Dispersos (SAEs) se han aplicado con éxito a modelos de lenguaje de proteínas (PLMs) y modelos de visión, aún no se han utilizado para sondear las representaciones internas de los modelos de difusión de todos los átomos.

Metodología

Los autores presentan VFUSE, un marco de interpretabilidad mecánica que entrena Autoencoders Dispersos en las activaciones de modelos de difusión-transformador para auditar la presencia de características asociadas a peligros.

1. Modelo y Recolección de Activaciones:

   • El estudio se centra en RFDiffusion3 (un modelo de difusión de todos los átomos de 18 bloques) y RoseTTAFold3 (su contraparte de plegamiento).
   • Las activaciones ($d=768$) se capturan mediante hooks de PyTorch en bloques específicos del flujo residual: bloques 6, 8 y 12 para RFD3, y bloques 12 y 16 para RF3.
   • Para RFD3, las entradas se generan mediante difusión parcial (en $\partial t = 5$, aprox. 5Å de ruido) sobre estructuras PDB de plantillas peligrosas para simular el proceso de diseño del modelo. Las entradas de RF3 se pliegan directamente.

2. Construcción del Conjunto de Datos:

   • Secuencias Virulentas: Obtenidas de SafeProtein (275 pares) y ToxinPred3 (1,200 secuencias solo para RF3).
   • Secuencias Benignas: Secuencias de longitud emparejada (100–300 aa) de UniProt, filtradas para excluir palabras clave de toxina/virulencia.
   • Control de Homología: Para evitar que las sondas memoricen la identidad de la familia de plegamiento en lugar de aprender señales de virulencia, las secuencias se agrupan por identidad del 30% usando mmseqs2. Las divisiones de entrenamiento/prueba aseg로ran que no aparezcan homólogos cercanos en ambos conjuntos.

3. Arquitectura y Entrenamiento de SAE:

   • Arquitectura: Se emplea un Matryoshka BatchTopK SAE. Este combina:
     • Matryoshka: Tamaños de prefijo anidados que permiten la reconstrucción a múltiples granularidades.
     • BatchTopK: Fuerza la dispersión manteniendo solo los $K$ pre-activaciones superiores a través de un lote, reemplazando las penalizaciones L1.
   • Hiperparámetros: Tamaño del diccionario $m=12,288$ (expansión de 16x), $K=80$ latentes activos, entrenado durante 20,000 pasos.
   • Rendimiento: Los SAEs entrenados explican el 96.9% de la varianza de la activación de la muestra retenida con una dispersión $L_0$ promedio de 79.9.

4. Protocolo de Evaluación:

   • Sondeo (Probing): Se utilizan vectores latentes de SAE promediados por media (mean-pooled) ($z$) y activaciones puras ($x$) para entrenar sondas de regresión logística con regularización L2.
   • Métricas: El AUROC se evalúa bajo dos regímenes: divisiones estratificadas aleatorias y divisiones agrupadas por homología.
   • Puntuación de Características: Se calcula el AUROC univariante para cada característica frente a las etiquetas de clase. La significancia se determina mediante pruebas de Mann–Whitney U con corrección de Benjamini–Hochberg FDR ($q < 0.05$).

Resultados Clave

1. Los Latentes de SAE Superan a las Activaciones Puras:

   • En el Bloque 12 de RFD3, las características codificadas por SAE superan significativamente a las activaciones puras en el entorno de homología agrupada, logrando un AUROC de 0.817 frente a 0.763 para las características puras (una mejora de +0.054).
   • Esta mejora es específica del Bloque 12; los bloques anteriores (6, 8) muestran ganancias neutras o negativas con los SAEs.
   • Para RF3 (conjunto de datos ToxinPred3), las características del SAE del Bloque 12 también muestran una mejora de +0.041 sobre las características puras en las divisiones agrupadas.

2. Reducción de la Memorización de la Homología:

   • La brecha entre el AUROC de la división aleatoria y el de la división agrupada por homología (un indicador de la memorización de familias) es la más pequeña para las características de SAE del Bloque 12 de RFD3 (brecha de 0.13) en comparación con el Bloque 6 (brecha de 0.22). Esto sugiere que los SAEs del Bloque 12 capturan señales estructurales generalizables a la clase y discriminativas de la clase, en lugar de memorizar familias de plegamiento.

3. Identificación de Características de Peligro Monosemánticas:

   • El estudio identifica características individuales de SAE que se activan selectivamente en diseños peligrosos.
   • Característica Superior: La Característica 639 en el Bloque la 12 de RFD3 logra un AUROC de 0.819 ($q \approx 10^{-13}$) para detectar neurotoxinas de víboras.
   • Localización: Estas características exhiben localización por residuo. Para la Característica 639, los 3 residuos con mayor activación (99, 104, 109) se agrupan dentro de una única hélice alfa de la estructura de Ammodytoxin A. Otras características se localizan de manera similar en regiones específicas de proteínas de evasión inmunitaria de la viruela del mono, proteínas inactivadoras de ribosomas y toxinas salivales de mosquitos.

4. Rendimiento Competitivo:

   • La mejor sonda de VFUSE (RFD3 Bloque 12 SAE, división aleatoria, AUROC 0.877) es competitiva con los clasificadores de solo secuencia de última generación (p. ej., DTVF con 0.92), a pesar de utilizar una sonda logística simple y un conjunto de datos más pequeño. Crucialmente, VFUSE proporciona justificaciones por residuo, algo que los modelos de solo secuencia no poseen.

Significado y Reivindicaciones

El artículo afirma ser el primer SAE entrenado en un modelo de difusión de todos los átomos y la primera auditoría de peligros a nivel de característica de un modelo de diseño de proteínas.

• Perspectiva Mecanicista: El trabajo demuestra que el aprendizaje de diccionarios en profundidades específicas (Bloque 12) desenreda conceptos estructurales polisemánticos en direcciones monosemánticas que son más aptas para la detección de peligros de baja capacidad que las representaciones puras del modelo.
• Seguridad e Interpretabilidad: VFUSE establece una vía para el monitoreo en tiempo de ejecución y la auditoría interpretable de los modelos generativos de proteínas. Al identificar residuos y motivos estructurales específicos asociados con la virulencia, el método va más allá de la detección de "caja negra" hacia controles de seguridad explicables.
• Dirección Futura: Aunque los autores señalan limitaciones (conjunto de datos pequeño, falta de éxito en el control mediante dirección en pruebas preliminares), el trabajo allana el camino para un diseño de proteínas seguro e interpretable al permitir la detección de características peligrosas a nivel de característica en lugar de solo a nivel de secuencia.