Resumen Técnico: ENSEMBITS

Enunciado del Problema

Los modelos de lenguaje de proteínas (PLMs) y las herramientas de predicción de estructura han dependido en gran medida de representaciones estructurales estáticas o de secuencias simples de aminoácidos. Aunque los tokenizadores de estructura de proteínas (PST) existentes capturan eficazmente la geometría local de instantáneas estáticas, no logran codificar los movimientos correlacionados y los estados conformacionales alternativos inherentes a los conjuntos de proteínas. La función biológica a menudo no se codifica en una única estructura, sino en la distribución de conformaciones que una proteína muestrea (por ejemplo, bucles catalíticos que regulan el acceso al sustrato). El campo se está desplazando hacia la generación de conjuntos, sin embargo, existe una falta de un vocabulario discreto capaz de tokenizar estos conjuntos dinámicos para su uso en arquitecturas basadas en transformadores. Además, los datos de dinámica suelen ser escasos, y los métodos existentes o bien no logran modelar conjuntamente el movimiento correlacionado entre residuos, o requieren trayectorias completas que pueden no estar disponibles en el momento de la inferencia.

Metodología: ENSEMBITS

ENSEMBITS se presenta como el primer tokenizador diseñado específicamente para conjuntos conformacionales de proteínas. Mapea un multiconjunto desordenado de marcos estructurales (conformeros) a un token discreto mediante una arquitectura de Autoencoder Variacional con Cuantización Vectorial Residual (RVQ-VAE).

1. Formulación del Problema

El modelo trata la tokenización de conjuntos como un problema de compresión de multiconjunto a token.

• Entrada: Un multiconjunto desordenado $P = \{x_1, \dots, x_P\}$ de $P$ marcos para una sola proteína, donde cada marco contiene coordenadas por átomo.
• Salida: Un token discreto de un código $\mathcal{C}$ que representa el motivo dinámico local de un residuo a través del conjunto.
• Restricción: El tokenizador debe ser agnóstico a $P$, aceptando conjuntos de cualquier cardinalidad positiva y siendo invariante a la permutación respecto al orden de los marcos.

2. Diseño del Descriptor

Para capturar la dinámica, ENSEMBITS calcula descriptores invariantes a SE(3) a través de los marcos. Los autores utilizan un descriptor de marco relativo estilo ESM3 (el entorno de producción) que codifica el entorno local de la columna vertebral como transformaciones SE(3) entre el marco de un residuo y sus $K$ vecinos espaciales más cercanos.

• Modo Dinámico: Crucialmente, las identidades de los vecinos ($N^p_r$) se vuelven a calcular de forma independiente para cada marco $p$. Esto permite que el descriptor codifique naturalmente la formación y ruptura de contactos a lo largo de la trayectoria, una señal invisible para los enfoques de vecinos fijos.

3. Arquitectura

• Codificador de Conjuntos: Una arquitectura estilo PerceiverIO que mapea el multiconjunto de descriptores por residuo a un único vector latente $z$. Utiliza una MLP compartida por elemento seguida de atención cruzada con tokens de consulta aprendibles, garantizando la invarianza a la permutación por construcción.
• Cuantizador Residual: Un RVQ-VAE multi-etapa con $K=3$ niveles discretiza el latente $z$ en una tupla de índices de código $(c_1, \dots, c_K)$. El token final es la suma de las incrustaciones de estos niveles.
• Decodificador y Pérdida: El decodificador reconstruye el multiconjunto de descriptores. La pérdida de reconstrucción emplea emparejamiento húngaro para alinear los descriptores predichos y objetivo, garantizando la invarianza a la permutación sin requerir un orden fijo.

4. Destilación de Marco Único a Token (SFTD)

Para abordar la escasez de datos de dinámica y permitir la inferencia en estructuras estáticas individuales, los autores proponen SFTD.

• Mecanismo: Durante el entrenamiento, el modelo procesa dos conjuntos de la misma proteína: un conjunto completo ($P_{max}$) y un subconjunto muestreado aleatoriamente (hasta $P=1$).
• Objetivo: La representación latente del subconjunto se acerca a la del conjunto completo (con un gradiente detenido en el conjunto completo).
• Resultado: Esto permite que el tokenizador prediga un "token de dinámica" a partir de una única estructura predicha en el momento de la prueba, destilando eficazmente la información a nivel de conjunto en una representación de marco único.

Contribuciones Clave

1. Formulación y Pipeline: La primera formulación de extremo a extremo de la tokenización de dinámica de proteínas como un problema de compresión de multiconjunto a token, utilizando descriptores invariantes a SE(3) y codificadores de conjuntos invariantes a la permutación.
2. Objetivo SFTD: Una estrategia de entrenamiento novedosa que alinea las incrustaciones de subconjuntos con sus contrapartes de conjuntos completos, permitiendo que el modelo funcione como un tokenizador consciente de la dinámica incluso cuando solo está disponible una única estructura estática.
3. Validación Empírica: Validación exhaustiva que muestra que ENSEMBITS captura la dinámica a nivel de conjunto y se transfiere eficazmente a tareas funcionales posteriores, superando a menudo a los tokenizadores estáticos a pesar de utilizar significativamente menos datos de preentrenamiento.

Resultados

1. Representación de la Dinámica (RMSF y ANOVA)

• Predicción de RMSF: ENSEMBITS domina los benchmarks de predicción de Fluctuación Cuadrática Media (RMSF).
  • Multi-marco: ENSEMBITS ($P=full$) supera a la línea base más fuerte (ProtProfileMD) en aproximadamente 2.5 puntos de Spearman en mdCATH-div y entre 11 y 18 puntos en los cortes MISATO.
  • Marco único: Incluso con $P=1$ (destilado), ENSEMBITS supera a todos los demás tokenizadores de marco único (incluyendo AminoAseed y ESM3struct) en la predicción de RMSF.
• Prueba ANOVA: Un ANOVA de un factor sobre la amplitud del movimiento ($s_1$) revela que los tokens de ENSEMBITS explican el 37.1% de la varianza en la amplitud del movimiento local ($\eta^2 = 0.371$). Esto es aproximadamente 3 veces mayor que la siguiente mejor línea base de múltiples marcos (Vote_3Di, $\eta^2 \approx 0.128$) y 7 veces mayor que los tokenizadores estáticos de marco único. Esto confirma que los tokens codifican dinámicas locales distinguibles y no solo la clase de pliegue estática o la posición.

2. Tareas Funcionales Posteriores

A pesar de utilizar muchos menos datos de preentrenamiento (solo conjuntos de MD, sin preentrenamiento de plegamiento inverso a escala PDB), ENSEMBITS iguala o supera a los tokenizadores estáticos en varias tareas:

• Predicción del Efecto de Mutación: En el benchmark PROTEINGYM, ENSEMBITS ($P=1$) combinado con ESM2 supera a todos los demás tokenizadores estructurales combinados con ESM2, logrando una ganancia del 6.9% sobre ESM2 solo.
• Sitio de Unión y Afinidad: ENSEMBITS lidera en la predicción de sitios de unión (AUROC 0.750) y la regresión de afinidad de unión en todos los cortes. Cabe destacar que supera a ESM3struct y AminoAseed en la detección de sitios de unión, lo que sugiere que la flexibilidad conformacional local es una señal crítica para los residuos de interfaz que los tokenizadores estáticos pasan por alto.
• Predicción de EC y GO: ENSEMBITS es competitivo con ESM3struct en las tareas de Comisión de Enzimas (EC) y Ontología de Genes (GO), aunque ESM3struct mantiene una ventaja en tareas fuertemente correlacionadas con la clase de pliegue (probablemente debido a su preentrenamiento de plegamiento inverso).

Significado y Afirmaciones

El artículo posiciona a ENSEMBITS como un paso necesario en la evolución del modelado de proteínas, desde la predicción de estructura estática hacia la generación de conjuntos.

• Cerrando la Brecha: Proporciona el vocabulario discreto necesario para integrar la dinámica de proteínas en el modelado de lenguaje y el diseño.
• Eficiencia de Datos: Demuestra que representaciones dinámicas de alta calidad pueden aprenderse a partir de corpus de MD relativamente pequeños (en comparación con los masivos corpus de secuencias/estructuras utilizados por modelos como ESM3), siempre que la representación esté diseñada explícitamente para la dinámica.
• Utilidad Práctica: El mecanismo SFTD alivia el cuello de botella de la disponibilidad de datos de dinámica, permitiendo consultar el modelo en estructuras estáticas individuales (por ejemplo, predicciones de AlphaFold) mientras se aprovechan los priores de dinámica aprendidos.

Los autores notan modestamente que la calidad de los tokens está limitada por la cobertura y la fidelidad de las trayectorias de entrenamiento subyacentes y que el objetivo es introducir una metodología para inculcar información dinámica en tokens discretos, en lugar de afirmar haber resuelto toda la dinámica de proteínas. El trabajo futuro implica escalar a corpus de conjuntos más grandes e integrar estos tokens directamente en modelos de lenguaje de proteínas para la generación de estructuras de múltiples marcos.