Resumen Técnico: Cofolding de pocos pasos con Mapas de Flujo de Átomos Completos (DECAF)

Planteamiento del Problema

El modelado computacional preciso de complejos biomoleculares 3D, particularmente para el cofolding de proteína-ligando, se ha convertido en un paradigma dominante en el diseño de fármacos basado en la estructura (SBDD). Si bien los modelos generativos de átomos completos modernos (por ejemplo, AlphaFold 3, Boltz-1) pueden predecir estructuras con resolución atómica, dependen de ejecuciones de difusión iterativas costosas. Generar estructuras de alta fidelidad suele requerir cientos de evaluaciones de funciones neuronales (NFE), lo que crea un cuello de botella computacional significativo tanto para el despliegue como para las técnicas de búsqueda en tiempo de inferencia.

Además, generar estructuras físicamente válidas a nivel atómico es crítico; un modelado local inexacto conduce a fallos catastróficos como choques estéricos, longitudes/ángulos de enlace incorrectos y cadenas laterales tensionadas. Las restricciones físicas estrictas suelen desplazar la carga computacional hacia métodos en tiempo de inferencia, donde se utiliza la búsqueda guiada por recompensas (por ejemplo, Monte Carlo Secuencial, MCTS) para refinar las muestras. Sin embargo, estos métodos de búsqueda aumentan el costo de inferencia, ya que requieren múltiples consultas de recompensa sobre múltiples partículas durante la trayectoria. La pregunta de investigación central es si es posible entrenar un modelo de cofolding de átomos completos que genere muestras optimizadas por recompensa de manera eficiente utilizando solo un pequeño número de NFEs en el tiempo de inferencia.

Metodología: El Marco DECAF

Los autores presentan DENOISER COFOLDING ALL-ATOM FLOWMAP (DECAF), un marco diseñado para destilar modelos de difusión de cofolding de átomos completos de vanguardia en mapas de flujo de pocos pasos. La metodología aborda desafíos técnicos específicos al adaptar los modelos de difusión estándar a los mapas de flujo para el modelado de proteínas de átomos completos.

1. Reparametrización del Mapa de Flujo en el Espacio-$\sigma$

Los modelos de difusión estándar suelen utilizar la parametrización EDM (Elucidating the Design Space), que condiciona en un único nivel de ruido $\sigma$. Esta formulación no se extiende limpiamente a los mapas de flujo de tiempo dual, y la reparametrización de tiempo no lineal $\sigma(t)$ utilizada en EDM conduce a inestabilidad numérica al calcular las derivadas parciales temporales ($\partial \sigma / \partial t$) requeridas para los objetivos del mapa de flujo.

Para evitar esto, DECAF define el mapa de flujo directamente en el espacio de ruido $\sigma$. En lugar de mapear el tiempo $t$, mapea entre niveles de ruido $\rho \to \sigma$ (donde $\rho > \sigma$).

• $\sigma$-Velocidad: El marco define una $\sigma$-velocidad instantánea $v_\sigma(x_\sigma) = dx_\sigma/d\sigma$.
• Velocidad Promedio: El mapa de flujo se parametriza como un desplazamiento utilizando la velocidad $\sigma$ promedio entre dos niveles de ruido, permitiendo que el modelo realice "saltos" a lo largo de la trayectoria de la ODE de flujo de probabilidad directamente en el espacio de ruido.
• Estabilidad: Esta reparametrización elimina el problemático factor $\partial \sigma / \partial t$, permitiendo un entrenamiento estable y una destilación directa desde modelos de cofolding preentrenados de estilo EDM.

2. Parametrización del Denoiser con Alineación SE(3)

Una innovación crítica es el uso de una formulación basada en denoiser en lugar de un predictor de velocidad directa.

• Predicción de Punto Final: El modelo predice el punto final objetivo $\hat{x}_{tgt}$ (la estructura limpia) en lugar de la velocidad instantánea.
• Alineación Rígida SE(3): La función de pérdida incorpora una minimización sobre el grupo SE(3) (usando el algoritmo de Kabsch) para alinear la estructura predicha $\hat{x}_{tgt}$ con la verdad de campo (o la predicción del profesor). Esto impone suavemente la simetría SE(3), lo cual es crítico para entrenar modelos estables y reducir la varianza.
• Función de Pérdida: El objetivo de entrenamiento iguala el punto final predicho contra la predicción del profesor tras la alineación rígida, soportando tanto el entrenamiento diagonal (ajuste de puntuación) como el fuera de la diagonal (salto).

3. DECAF-SEARCH: Marco de Tiempo de Inferencia

Para aprovechar el mapa de flujo para la optimización de la recompensa, los autores introducen DECAF-SEARCH.

• Mecanismo de Lookahead: A diferencia de los modelos de difusión estándar donde la estimación de la recompensa requiere un denoising completo, el mapa de flujo de DECAF proporciona una predicación de "lookahead" del punto final limpio $\hat{x}_0$ desde cualquier estado ruidoso intermedio $x_\sigma$. Esto permite una estimación de recompensa de mayor fidelidad ($R(\hat{x}_0)$) sin completar la trayectoria completa.
• Algoritmo de Búsqueda: El marco adapta los métodos de dirección en tiempo de inferencia (Feynman-Kac, MCTS) al entorno del mapa de flujo. Utiliza un muestreador estocástico $\gamma$ para generar partículas, evalúa sus recompensas de lookahead y realiza el ascenso de gradiente en el espacio limpio o el remuestreo para dirigir la generación hacia configuraciones de alta recompensa y físicamente válidas.

Contribuciones Clave

1. Primer Mapa de Flujo de Proteínas de Átomos Completos: El diseño de DECAF-Boltz, el primer mapa de flujo de proteínas de átomos completos que destila exitosamente un profesor Boltz-1.
2. Innovaciones Metodológicas:
   • Una novedosa reparametrización en el espacio-$\sigma$ que permite el entrenamiento estable de mapas de flujo para el modelado de átomos completos.
   • Una parametrización basada en denoiser que permite la alineación rígida ponderada por SE(3), crucial para la estabilidad y la precisión.
3. Búsqueda en Tiempo de Inferencia: La introducción de DECAF-SEARCH, un marco que utiliza el lookahead del mapa de flujo para mejorar la estimación de la recompensa y la alineación, permitiendo una búsqueda eficiente con menos NFEs.
4. Validación Empírica: Demostración de que DECAF supera a los baselines basados en difusión en regímenes de bajo NFE y iguala el rendimiento de presupuesto completo con significativamente menos computación.

Resultados Experimentales

Los autores evaluaron DECAF en los benchmarks Runs N' Poses (RnP) y PoseBusters, comparándolo contra Boltz-1x (una variante de potencial físico de Boltz-1) y otros modelos de vanguardia como AlphaFold 3 y Pearl.

Rendimiento en Presupuestos de Bajo NFE

• Benchmark RnP: Con presupuestos estrictos de NFE (por ejemplo, 20–50 NFEs), DECAF-SEARCH supera significativamente a Boltz-1x (ajustado) en todas las métricas (PB-Valid, Tasa de Éxito, lDDT-PLI). Por ejemplo, a 20 NFEs, DECAF logra un puntaje PB-Valid del 83.5% frente al 75.9% de Boltz-1x.
• Generalización: La ganancia de rendimiento se mantiene en todos los cuartiles de similitud de bolsillo, incluyendo objetivos fuera de la distribución, confirmando la robustez del modelo.
• Eficiencia: DECAF iguala el rendimiento de Boltz-1x a 600 NFEs utilizando solo 30 NFEs (una reducción de computación de 20x) en PoseBusters.

Destilación de Modelos de Vanguardia

• DECAF-Pearl: El marco se aplicó para destilar el modelo de cofolding de vanguardia Pearl. DECAF-Pearl iguala la tasa de éxito del profesor utilizando 5 veces menos NFEs (40 NFEs frente a la simulación completa).
• Comparación: DECAF-Pearl supera a otros baselines potentes (incluyendo la simulación completa de Pearl) en métricas de validez física y tasa de éxito en presupuestos de bajo NFE.

Frontera de Pareto

El análisis de la frontera de cómputo óptimo muestra que DECAF-SEARCH domina a Boltz-1x en todos los presupuestos de computación de inferencia. La estrategia de búsqueda óptima varía según el presupuesto:

• Bajo NFE ($\le$ 30): El SMC basado en partículas (dirección FK) es óptimo.
• NFE Moderado (100–250): La estimación de gradiente de Monte Carlo (MC-GRAD) es más efectiva.
• Alto NFE ($\ge$ 142): La búsqueda basada en MCTS produce los mejores resultados.

Significado y Reivindicaciones

El artículo afirma que DECAF representa un paso significativo para hacer que la generación de biomoléculas de átomos completos sea computacionalmente factible para aplicaciones de alto rendimiento.

• Aceleración: Al comprimir las trayectorias de difusión en mapas de flujo de pocos pasos, DECAF reduce la latencia de inferencia en órdenes de magnitud (por ejemplo, reducciones de 5x a 20x en NFEs) sin sacrificar la precisión estructural o la validez física.
• Optimización de Recompensa: El lookahead del mapa de flujo permite una búsqueda guiada por recompensa más eficiente y precisa, abordando el "impuesto de inferencia" asociado con múltiples consultas de recompensa en los modelos de difusión estándar.
• Impacto Práctico: La aceleración permite nuevos flujos de trabajo, como el cofolding de librerías enteras de ligandos contra un objetivo para el cribado virtual y la generación de conjuntos de datos sintéticos a gran escala para entrenar modelos posteriores, lo que anteriormente era impracticable debido al costo de las simulaciones de difusión completas.
• Modestia: Los autores reconocen que los modos de fallo (por ejemplo, enlaces sp2 no planos) son heredados de las señales de recompensa de los modelos profesores (Boltz-1x) más que del procedimiento de búsqueda en sí. También señalan que la relación entre el presupuesto de NFE y la tasa de éxito no es monótona, lo que sugiere trabajos futuros en búsqueda adaptativa y entrenamiento conjunto con recompensas específicas de la tarea.

El código se ha publicado en https://github.com/genesistherapeutics/decaf.