Resumen Técnico: KRONOS – Difusión Latente Autorregresiva para la Generación de Moléculas 3D

Planteamiento del Problema

La generación de moléculas tridimensionales (3D) es una tarea crítica en el descubrimiento de fármacos y el diseño molecular. Los enfoques actuales de vanguardia dependen principalmente de modelos de difusión y flow matching. Si bien estos métodos logran una alta calidad de generación, típicamente requieren que el tamaño molecular (número de átomos) se especifique a priori, lo que limita su flexibilidad para la generación de longitud variable.

Por el contrario, los modelos autorregresivos (AR) soportan naturalmente la generación de longitud variable y el condicionamiento sobre estructuras moleculares parciales (por ejemplo, la elaboración de andamios o scaffolds conocidos), lo cual es altamente valioso en la química medicinal. Sin embargo, los modelos AR existentes enfrentan dos desafíos principales:

1. Brecha de Rendimiento: Históricamente, los modelos AR se han quedado rezagados frente a los modelos de difusión en cuanto a calidad de generación, particularmente en términos de fidelidad geométrica y validez química.
2. Equilibrio de Objetivos: Lograr un alto rendimiento tanto en la generación incondicional (diseño de novo) como en la generación condicionada por fragmentos (completado molecular) dentro de una misma arquitectura sigue siendo difícil. Muchos modelos AR tienen dificultades para generalizar entre estos paradigmas sin modificaciones arquitectónicas o compromisos en el rendimiento.

Además, los generadores moleculares latentes existentes han dependido casi exclusivamente de la difusión, dejando el potencial de la generación autorregresiva en el espacio latente en gran medida inexplorado.

Metodología: KRONOS

Los autores presentan KRONOS, un marco de difusión autorregresiva latente que genera moléculas 3D directamente en el espacio latente de un Autoencoder Unificado (UAE) preentrenado. El método combina tres componentes principales:

1. Generación Autorregresiva Latente

En lugar de operar en el espacio de datos brutos (prediciendo tipos de átomos discretos/enlaces y coordenadas continuas simultáneamente), KRONOS opera en el espacio latente continuo de un UAE preentrenado (Luo et al., 2025).

• Representación: Una molécula 3D se codifica en una secuencia de tokens latentes $Z = \{z_1, \dots, z_n\}$, donde cada token codifica conjuntamente la identidad atómica, el entorno de enlace y la geometría 3D.
• Factorización: El modelo aprende una distribución autorregresiva de izquierda a derecha:
  $$p_\theta(z_{1:n}) = \prod_{i=1}^n p_\theta(z_i | z_{<i})$$
• Longitud Variable: Dado que el espacio latente carece de un token natural de fin de secuencia, KRONOS emplea un clasificador de parada dedicado (un MLP ligero) entrenado junto con el generador para predecir la terminación de la secuencia dinámicamente, eliminando la necesidad de un recuento de átomos predefinido.

2. Predicción de Tokens Basada en Difusión

Para modelar la distribución condicional de los tokens latentes continuos $p_\theta(z_i | z_{<i})$, KRONOS utiliza una formulación de Pérdida de Difusión (Li et al., 2024).

• Un transformador causal codifica los tokens precedentes en una representación de contexto $c_i$.
• Una red de eliminación de ruido de difusión $\epsilon_\theta$ predice el ruido añadido al token latente actual $z_i$, condicionado en $c_i$ y el paso de tiempo de la difusión.
• Este enfoque reemplaza la predicción de tokens discretos con un objetivo de eliminación de ruido, permitiendo el modelado flexible de variables continuas sin discretización.

3. Estrategia de Entrenamiento Mixto (Generación Condicionada por Fragmentos)

Para soportar tanto la generación incondicional como la condicionada por fragmentos dentro de un solo modelo, los autores introducen una estrategia de entrenamiento inspirada en el paradigma Fill-in-the-Middle (FIM).

• Descomposición de Fragmentos: Las moléculas se descomponen en fragmentos químicamente significativos utilizando un esquema basado en BRICS.
• Construcción de Secuencias: Durante el entrenamiento, cada molécula se transforma en:
  • Una secuencia de orden por defecto (para la generación incondicional).
  • Una secuencia condicionada por fragmentos donde un subconjunto aleatorio de fragmentos actúa como "semilla" (contexto), y los fragmentos restantes son el objetivo de generación.
• Mecanismo: El modelo es entrenado para generar autorregresivamente los fragmentos objetivo condicionados en los fragmentos semilla. Esto se logra asignando un único orden determinista a los fragmentos y átomos dentro de ellos, en lugar de marginalizar sobre todas las permutaciones (como en los modelos AR Agnósticos al Orden). Esto permite que el modelo aprenda capacidades de completado molecular sin cambios arquitectónicos.

Contribuciones Clave

1. Marco KRONOS: Un novedoso marco de difusión autorregresiva latente que modela conjuntamente la topología del grafo molecular y la geometría en un espacio latente unificado.
2. Estrategia de Entrenamiento Unificada: Un enfoque de entrenamiento mixto que permite la generación tanto incondicional como condicionada por fragmentos dentro de un único modelo autorregresivo de izquierda a derecha, inspirado en el paradigma FIM.
3. Generación de Longitud Variable: Un marco unificado que soporta naturalmente la generación de longitud variable mediante un clasificador de parada, evitando la necesidad de especificar el tamaño a priori.
4. Rendimiento: Demostración de que la generación autorregresiva latente puede alcanzar un rendimiento líder entre los métodos autorregresivos mientras se mantiene competitiva con los modelos de difusión, sin sacrificar la calidad incondicional al habilitar el condicionamiento por fragmentos.

Resultados Experimentales

Los autores evaluaron KRONOS en los conjuntos de datos QM9 y GEOM-Drugs.

Generación Incondicional

• QM9: KRONOS logró la mayor Validez 3D (97.3%) y el rendimiento general más fuerte entre los métodos autorregresivos. Se mantuvo competitivo con los modelos de difusión líderes (ej. JODO, EDM) en métricas que incluyen Estabilidad Atómica 3D, Longitud de Enlace y MMD de Ángulo de Enlace.
• GEOM-Drugs: KRONKOS logró el rendimiento general más fuerte entre todos los baselines evaluados (tanto autorregresivos como de difusión).
• Plausibilidad Física: Utilizando optimización de geometría semiempírica (GFN2-xTB), las conformaciones generadas por KRONOS mostraron una baja Distancia de Raíz Cuadrática Media (RMSD) y diferencias de energía ($\Delta E$) respecto a las estructuras optimizadas, indicando que se encuentran cerca de regiones de espacio conformacional físicamente plausibles.

Generación Condicionada por Fragmentos

• En un conjunto de datos de aproximadamente 15k pares molécula-fragmento, KRONOS superó sustancialmente al baseline autorregresivo QUETZAL tanto en las métricas Mean@3 como Best@3 para validez y estabilidad.
• Crucialmente, habilitar la generación condicionada por fragmentos tuvo un impacto insignificante en la generación incondicional, demostrando que ambos paradigmas pueden ser soportados dentro de una misma arquitectura.

Significado y Reivindicaciones

El artículo afirma que KRONOS proporciona una alternativa prometedora a los enfoques actuales de generación autorregresiva y de difusión en el espacio de datos. Su importancia radica en:

• Cerrar la Brecha: Reduce sustancialmente la brecha de rendimiento entre los modelos autorregresivos y los de difusión, logrando resultados líderes entre los métodos AR mientras iguala a los baselines de difusión.
• Flexibilidad: Unifica las fortalezas de los modelos autorregresivos (longitud variable, condicionamiento parcial) con la alta calidad de generación de los modelos de difusión, todo dentro de una formulación de espacio latente.
• Utilidad Práctica: Al soportar la generación condicionada por fragmentos sin modificaciones arquitectónicas, ofrece una herramienta práctica para tareas de química medicinal donde las moléculas suelen diseñarse mediante la elaboración de andamios conocidos.

Los autores señalan limitaciones, incluyendo la dependencia de la expresividad del espacio latente preentrenado y el uso de un esquema de descomposición BRICS fijo, sugiriendo que el trabajo futuro podría explorar el aprendizaje de abstracciones de fragmentos de forma conjunta con el modelo generativo o extender el marco hacia el diseño guiado por propiedades.