Resumen Técnico: Emyx – Generación de Proteínas de Todo Átomo Rápida y Eficiente

Declaración del Problema

El diseño computacional de enzimas requiere la generación de andamios proteicos que posicionen con precisión los residuos catalíticos y los ligandos. Esta tarea exige tanto precisión geométrica como diversidad estructural. Los modelos generativos de todo átomo actuales, como RFdiffusion3 y Proteína-Complexa, heredan arquitecturas complejas de la predicción de estructuras de proteínas (por ejemplo, representaciones de pares densos, pilas de incrustaciones pesadas). Estas arquitecturas resultan en altos costes de entrenamiento, diversidad de muestras limitada y, a menudo, requieren pasos de plegamiento inverso costosos para generar secuencias de alta calidad. Además, las métricas de evaluación existentes (por ejemplo, el RMSD de autoconsistencia de átomos pesados) pueden sobreestimar el éxito al no verificar la recuperación del plegamiento global, lo que conduce a diseños que satisfacen las restricciones de motivos locales pero fallan en formar proteínas estables y plegables.

Metodología

Arquitectura del Modelo

Emyx es un modelo de ajuste de flujo condicional (conditional flow matching) de 140 millones de parámetros diseñado para la generación de proteínas de todo átomo. A diferencia de sus predecesores que dependen de actualizaciones de pares densos y capas de incrustación específicas de dominio, Emyx concentra su capacidad dentro de bloques transformer estándar.

• Representación: El modelo utiliza una representación unificada de dos niveles. A nivel de átomo, los residuos y ligandos se mapean a una matriz "Rep14" de tamaño fijo (14 ranuras atómicas), donde las posiciones no utilizadas se completan con "átomos fantasma" que heredan las coordenadas del esqueleto (backbone). A nivel de token, los residuos y los HETATM constituyen tokens individuales.
• Conectividad: En lugar de representaciones de pares densos, Emyx utiliza un grafo de conectividad disperso con un presupuesto de aristas fijo. Las aristas se construyen basándose en vecinos de secuencia, enlaces químicos, k-NN de ligandos y conectividad de motivos.
• Diseño del Transformer: La arquitectura emplea un flujo de trabajo "patchify–process–unpatchify":
  1. Codificador de Átomos: Bloques transformer ligeros codifican el detalle atómico local dentro de la representación Rep14.
  2. Downcast: El pooling de atención cruzada con compuerta (gated cross-attention) agrupa las representaciones de átomos en incrustaciones a nivel de token.
  3. Tronco de Tokens: Un tronco transformer profundo (18 capas) procesa el contexto estructural de largo alcance utilizando atención dispersa con sesgos de par aprendidos.
  4. Upcast: Un segundo proceso de atención cruzada transmite el contexto de los tokens de vuelta a los átomos.
  5. Decodificador de Átomos: Bloques transformer finales refinan el detalle atómico local antes de la cabeza de velocidad de salida.
• Modificaciones: Los bloques DiT (Diffusion Transformer) estándar se modifican con proyecciones de cuello de botella para el condicionamiento, compuertas sigmoideas para limitar las actualizaciones residuales y profundidad estocástica (drop-path) para la regularización.

Entrenamiento y Objetivo

Emyx se entrena mediante ajuste de flujo condicional sobre coordenadas cartesianas.

• Objetivo: El modelo aprende un campo de velocidad dependiente del tiempo $v_\theta(x_t, t, c, G)$ para transportar muestras de una distribución de ruido a la distribución de datos. La pérdida de entrenamiento combina una regresión $\ell_2$ sobre las velocidades para los átomos que no son parte del motivo con una pérdida de prueba de diferencia de distancia local (lDDT) diferenciable auxiliar para imponer la calidad estructural.
• Condicionamiento: El modelo se condiciona mediante restricciones geométricas dispersas (geometría del motivo, posiciones de ligandos) y características globales (composición de la estructura secundaria, radio de giro). Estas características se enmascaran estocásticamente durante el entrenamiento para permitir la generación tanto condicional como incondicional.
• Eficiencia: El entrenamiento utiliza un interpolante lineal entre el ruido y los datos, lo que produce gradientes de menor varianza y una convergencia más rápida en comparación con los esquemas de difusión de preservación de varianza. El modelo fue entrenado con datos experimentales de PDB en aproximadamente 3.55 días en 8 GPUs NVIDIA H200 (681.6 horas-GPU).

Muestreo y Reparametrización

Una contribución metodológica clave es la derivación de una reparametrización exacta del interpolante de ajuste de flujo hacia el marco de trabajo de niveles de ruido EDM (Elucidating the Design Space of Diffusion-based Generative Models).

• Este mapeo permite que el modelo de ajuste de flujo utilice el esquema de ruido Karras y los mecanismos de agitación estocástica (stochastic churn) sin necesidad de reentrenamiento o cambios arquitectónicos.
• La reparametrización une la eficiencia de entrenamiento del ajuste de flujo con los métodos de muestreo de vanguardia diseñados para modelos de difusión, mejorando significamente las tasas de éxito sobre la integración Euler–Maruyama de SDE estándar.

Contribuciones Clave

1. Generador de Todo Átomo Eficiente: Emyx logra un rendimiento de vanguardia con un recuento de parámetros significativamente menor (140M) y un coste de entrenamiento (4× menor que RFdiffusion3) al prescindir de arquitecturas heredadas de la predicción.
2. Protocolo de Evaluación Estricto: Los autores proponen un protocolo de "sc-RMSD estricto" que requiere tanto la recuperación del plegamiento global (backbone RMSD < 2.0 Å) como la precisión de la geometría catalítica local (motif tip-atom RMSD < 1.5 Å). Demuestran que las métricas anteriores (heavy-atom sc-RMSD) sobreestiman las tasas de éxito en aproximadamente 2× porque no penalizan los diseños que recuperan el sitio activo pero no el plegamiento global.
3. Reparametrización EDM: El artículo introduce un método agnóstico al modelo para convertir modelos de velocidad de ajuste de flujo al marco EDM, permitiendo estrategias de muestreo superiores (esquema Karras, agitación estocástica) sin reentrenamiento.
4. Análisis Espectral: Los autores proporcionan un análisis espectral que revela una brecha fundamental "generador–predictor". Los predictores inicializados con señales ricas de MSA/plantillas exhiben un rango efectivo más alto y un menor sobreentrenamiento, mientras que los generadores inicializados desde coordenadas con ruido operan en un régimen de rango más bajo. Emyx demuestra la fracción de sobreentrenamiento más baja entre los generadores, atribuido a sus cuellos de botella arquitectónicos y profundidad estocástica.

Resultados

Evaluado contra el benchmark AME (Atom-level Motif Enzyme) (41 sitios activos catalíticos):

• Tasa de Éxito: Bajo el protocolo de sc-RMSD estricto, Emyx logra una tasa de éxito del 13.4% (39/41 objetivos resueltos), superando a Proteína-Complexa (8.8%) y RFdiffusion3 (6.7%).
• Novedad Estructural: Emyx produce los andamios más novedosos estructuralmente, con un TM-score mediano de 0.475 respecto al hallazgo más cercano en PDB (menor es más novedoso), comparado con 0.499 para Proteína-Complexa y 0.512 para RFdiffusion3.
• Diversidad: Emyx genera el mayor número de clústeres estructurales únicos (441) entre los diseños exitosos.
• Validez Geométrica: Emyx supera las comprobaciones de validez geométrica (longitudes/ángulos de enlace) en el 82.2% de los diseños, superando significativamente a RFdiffusion3 (0.8%) y Proteína-Complexa (49.9%). Los autores señalan que los bajos puntajes de geometría de RFdiffusion3 se deben probablemente a artefactos de post-procesamiento donde las coordenadas del motivo de entrada sobrescriben las generadas, creando enlaces no físicos.
• Eficiencia: Emyx entrena en ~682 horas-GPU (4× más rápido que RFdiffusion3) y es >2× más rápido en inferencia en una sola GPU A10G en todas las longitudes de cadena.

Significado y Reivindicaciones

El artículo sostiene que la simplicidad arquitectónica es beneficiosa para los generadores de proteínas. A diferencia de los predictores, que requieren señales coevolutivas ricas (MSA, plantillas) para inicializar representaciones de pares complejas, los generadores se condicionan mediante restricciones geométricas dispersas. Por lo tanto, las arquitecturas pesadas heredadas de la predicción de estructuras son innecesarias y perjudiciales para la eficiencia y la diversidad.

Emyx demuestra que concentrar la capacidad en bloques transformer estándar, combinada con un entrenamiento de ajuste de flujo eficiente y un muestreo tipo EDM, puede lograr un rendimiento superior en tareas de diseño de enzimas con recursos computacionales significativamente menores. El trabajo también destaca la importancia de protocolos de evaluación rigurosos (sc-RMSD estricto) para distinguir entre diseños localmente precisos pero globalmente inestables y andamios enzimáticos verdaderamente viables. Finalmente, la reparametrización EDM propuesta ofrece una vía para que otros modelos de ajuste de flujo adopten técnicas de muestreo avanzadas sin necesidad de una reestructuración arquitectónica.