DCFold: Efficient Protein Structure Generation with Single Forward Pass

Resumen Técnico: DCFold – Generación Eficiente de Estructuras de Proteínas con una Única Paso Adelante

Enunciado del Problema
Aunque AlphaFold3 (AF3) ha establecido un nuevo estado del arte para la predicción de estructuras de proteínas y la generación de complejos biomoleculares al utilizar una arquitectura basada en difusión de átomos completos, su despliegue práctico se ve obstaculizado por altos costos computacionales. AF3 depende del reciclaje iterativo del Pairformer y del muestreo de difusión en múltiples pasos, lo que resulta en tiempos de inferencia medidos en minutos para secuencias largas. Esta latencia limita su utilidad en aplicaciones de alto rendimiento aguas abajo, como el cribado virtual y el diseño de proteínas (por ejemplo, la alucinación de ligantes). Los intentos existentes para mitigar esto, como la reducción manual de las iteraciones de reciclaje, sacrifican precisión por eficiencia. Además, los modelos de consistencia estándar (CM), que buscan acelerar la difusión reduciendo los pasos de muestreo, enfrentan desafíos significativos cuando se aplican a estructuras de proteínas: (1) asumen datos de dimensión fija y utilizan la distancia euclidiana para el emparejamiento de pasos de tiempo, lo que no logra acomodar secuencias de proteínas de longitud variable y conduce a un entrenamiento inestable; y (2) no abordan la sobrecarga iterativa del módulo Pairformer.

Metodología
Los autores proponen DCFold, un modelo generativo de un solo paso que logra una precisión a nivel de AlphaFold3 mediante un marco de entrenamiento de Consistencia Dual. Este marco aborda los dos principales cuellos de botella de AF3: el proceso de difusión y el reciclaje del Pairformer.

1. Marco de Consistencia Dual:

   • Consistencia de Difusión: El modelo impone consistencia entre la salida de un solo paso y la trayectoria de múltiples pasos del módulo de difusión. Un muestreador de un paso se estabiliza desactivando la inyección de ruido estocástico y normalizando los tamaños de paso. El objetivo de entrenamiento minimiza el Error Cuadrático Medio (MSE) entre la predicción del modelo en el paso de tiempo $t$ y un paso de tiempo de referencia $r$.
   • Consistencia del Pairformer: Para eliminar la sobrecarga iterativa del Pairformer, los autores introducen una función de pérdida de consistencia cíclica. Dado que el Pairformer actualiza las representaciones internas de forma iterativa, una única pasada hacia adelante contiene inherentemente representaciones correspondientes a diferentes profundidades de ciclo. El modelo se entrena para minimizar el MSE entre las representaciones en ciclos adyacentes ($z_i$ frente a $z_{i+1}$), extrayendo efectivamente el refinamiento de múltiples ciclos en una única pasada hacia adelante.

2. Emparejamiento Geodésico Temporal (TGM):
   Para resolver la inestabilidad de aplicar modelos de consistencia a secuencias de proteínas de longitud variable, los autores introducen TGM. A diferencia de los programadores convencionales que emparejan pasos de tiempo basándose en intervalos euclidianos fijos, TGM empareja pasos de tiempo basándose en su distancia geodésica en la variedad de información temporal.

   • El método define una métrica temporal utilizando Información de Fisher ($I(t)$), que mide la tasa de cambio en la distribución a lo largo de la trayectoria de difusión.
   • La distancia geodésica $d_g(t, r)$ se calcula como la integral de $\sqrt{I(\tau)}$.
   • TGM asegura que los pares de entrenamiento $(t, r) mantengan una distancia geodésica constante en relación con el progreso del entrenamiento. Esto adapta el currículo a la dimensionalidad de los datos (longitud de la secuencia), previniendo la "disparidad de información" que causa el colapso del entrenamiento en secuencias largas.

Contribuciones Clave

• Arquitectura DCFold: Un predictor de un solo paso que elimina la sobrecarga iterativa tanto del muestreador de difusión como del reciclaje del Pairformer, permitiendo velocidades de inferencia comparables a AlphaFold2 mientras mantiene la precisión de AlphaFold3.
• Entrenamiento de Consistencia Dual: Un marco novedoso que impone conjuntamente la consistencia a través del módulo de difusión y el Pairformer, permitiendo una inferencia unificada de una sola pasada hacia adelante.
• Emparejamiento Geodésico Temporal (TGM): Un mecanismo de programación principista para la destilación de consistencia en datos de longitud variable. Al utilizar la Información de Fisher para definir una métrica riemanniana en la variedad temporal, TGM estabiliza la dinámica de entrenamiento y previene el colapso de pesos en tareas de modelado de proteínas de alta dimensión y longitud variable.
• Aplicabilidad Aguas Abajo: El método permite una inferencia totalmente diferenciable de un solo paso, haciéndolo adecuado para tareas de optimización basadas en gradientes como la alucinación de ligantes, que anteriormente eran difíciles con modelos iterativos de múltiples pasos.

Resultados Experimentales
Los autores evaluaron DCFold en benchmarks de predicción de estructuras y diseño de ligantes:

• Predicción de Estructuras (Posebusters V2 & PDB Reciente): DCFold logra una precisión comparable al modelo completo de AlphaFold3. En el benchmark Posebusters V2, alcanza una tasa de éxito de RMSD < 2 Å del 78,6% (mejor caso) y del 71,43% (peor caso), coincidiendo estrechamente con el rendimiento de AlphaFold3 mientras es 15 veces más rápido (tiempo de inferencia promedio reducido de 133,3 s a 8,9 s).
• Calidad de la Distribución: DCFold remodela la distribución de salida, mejorando significativamente el RMSD del peor caso y las tasas de éxito en comparación con la ODE de AF3 de línea base (un solo paso, sin destilación), lo que indica predicciones más fiables y estables.
• Diseño de Ligantes: En tareas de alucinación de ligantes (utilizando objetivos como IL2-Rα, TrkA, H3), DCFold supera la línea base BindCraft basada en AF2 en las tasas de éxito in silico. Por ejemplo, en el objetivo H3, DCFold logró una tasa de éxito del 23% bajo restricciones basadas en la física en comparación con el 16% para BindCraft.
• Estudios de Ablación: Los experimentos confirman que tanto la Consistencia de Difusión como la Consistencia del Pairformer contribuyen al rendimiento. Se demuestra que TGM es superior a otras líneas base de consistencia (CD, sCM, ECM) en tareas de proteínas, con ECM fallando en entrenar efectivamente en secuencias largas debido a la varianza del gradiente.

Significado y Afirmaciones
El artículo afirma que DCFold cierra exitosamente la brecha entre la eficiencia de AlphaFold2 y la precisión de AlphaFold3. Al comprimir la arquitectura iterativa en una única pasada hacia adelante sin sacrificar la fidelidad predictiva, DCFold permite el diseño y cribado de proteínas escalables y de alto rendimiento. La introducción de TGM proporciona una solución generalizable para aplicar modelos de consistencia a datos secuenciales de longitud variable, abordando una limitación fundamental de los métodos anteriores de aceleración basados en consistencia. Los autores posicionan a DCFold como un modelo fundamental capaz de apoyar diversas aplicaciones aguas abajo, incluido el cribado virtual y el diseño de proteínas diferenciable, que anteriormente estaban limitados por la sobrecarga computacional de los modelos de difusión iterativos.