Generative Molecular Morphing for Flexible-Size Design via Unbalanced Optimal Transport

Resumen Técnico: Morfismo Molecular Generativo para el Diseño de Tamaño Flexible mediante Transporte Óptimo Desbalanceado

1. Planteamiento del Problema

Los modelos generativos actuales para el diseño molecular en 3D, particularmente los enfoques basados en difusión y flujo, enfrentan una limitación algorítmica crítica: operan sobre un número fijo de átomos ($n$). Esta restricción impide que los modelos aprendan la distribución conjunta de las propiedades moleculares y el tamaño molecular. Dado que muchas propiedades químicas están intrínsecamente ligadas al número de átomos, fijar el tamaño restringe artificialmente la expresividad y la capacidad de control (steerability) del modelo. En consecuencia, estos modelos tienen dificultades para navegar relaciones complejas entre estructura y propiedades, fallan al integrar efectos estructurales previos (como andamios o scaffolds) de manera efectiva y no pueden generar moléculas válidas fuera de la distribución de entrenamiento (OOD) en lo que respecta al tamaño.

2. Metodología

Los autores introducen Morph, un modelo generativo de tamaño flexible para grafos geométricos que ajusta dinámicamente el número de nodos (átomos) durante el proceso de generación. La metodología integra la dinámica continua con saltos discretos, formulada como un proceso de flujo de salto (jump-flow process).

Componentes Principales:

• Espacio de Estados: El modelo opera en un espacio de estados $\mathcal{G}$ definido como la unión disjunta de espacios cocientes para grafos geométricos de diversos tamaños, contabilizando las simetrías $SE(3)$ (traslación/rotación) y $S_n$ (permutación).
• Dinámica Generativa: El proceso combina:
  1. Flujo Continuo: Gobernado por un campo de velocidad $v_t$ que desplaza las posiciones de los átomos.
  2. Saltos Discretos: Gobernados por una Cadena de Markov de Tiempo Continuo (CTMC) con una densidad de tasa $u_t$, que permite la inserción, eliminación y sustitución de tipos de nodos.
• Entrenamiento mediante Coincidencia de Generador Condicional (CGM): Para evitar la intratabilidad del generador marginal, el modelo utiliza CGM. Este aprende a regredir la velocidad condicional y las tasas de salto a lo largo de trayectorias entre pares de extremos $(g_0, g_1)$, donde $g_0$ es un grafo previo y $g_1$ es una molécula objetivo.
• Transporte Óptimo Desbalanceado (UOT) para la Alineación: Una innovación clave es el uso de UOT para alinear grafos geométricos de diferentes tamaños ($n_0 \neq n_1$) antes de definir las trayectorias de interpolación.
  • El objetivo de UOT minimiza una función de costo que involucra costos de transporte, penalizaciones por eliminación de masa ($w_{del}$) y penalizaciones por inserción de masa ($w_{ins}$).
  • Los costos de edición de aristas se establecen en cero, justificado por la capacidad de inferir la topología de los enlaces a partir de la estructura 3D y la valencia.
  • El plan de transporte resultante se discretiza utilizando el algoritmo húngaro en un espacio aumentado para determinar los conjuntos exactos de ediciones (coincidencias, inserciones, eliminaciones).
  • Las posiciones coincidentes se alinean rotacionalmente mediante el algoritmo de Kabsch para asegurar la invarianza $SE(3)$.
• Trayectorias de Interpolación: Basándose en la coincidencia de UOT, el modelo define trayectorias de probabilidad condicional para tres modalidades:
  • Nodos Coincidentes: Las posiciones siguen una trayectoria Gaussiana; los tipos siguen un programador (scheduler).
  • Nodos de Inserción: Los nodos se muestrean en tiempos de evento específicos con posiciones extraídas de una Gaussiana centrada en el objetivo.
  • Nodos de Eliminación: Los nodos persisten hasta un tiempo de evento específico, tras el cual ellos y sus aristas son eliminados.
• Función de Pérdida: El objetivo de entrenamiento minimiza una divergencia de Bregman entre los generadores condicionales predichos y el objetivo, comprendiendo términos para movimiento (MSE), sustitución (entropía cruzada), inserción (NLL de Poisson y GMM), eliminación (entropía cruzada binaria) y predicción de carga.

3. Contribuciones Clave

1. Extensión a Grafos Geométricos 3D: Este trabajo extiende las técnicas de generación de secuencias de tamaño flexible a grafos geométricos 3D, permitiendo la integración fluida de estructuras previas (como scaffolds) en el proceso generativo.
2. Algoritmo de Coincidencia basado en UOT: Los autores proponen un algoritmo de coincidencia basado en el Transporte Óptimo Desbalanceado que alinea grafos geométricos de diferentes tamaños para definir las trayectorias de probabilidad objetivo para el entrenamiento, manejando la naturaleza transdimensional del problema.
3. Mejora en el Control y Generación OOD: El artículo demuestra que Morph iguala a los modelos de tamaño fijo de vanguardia (SOTA) en la generación incondicional, mientras ofrece un control de propiedades superior. Crucialmente, permite la generación exitosa Fuera de la Distribución (OOD) en regímenes (por ejemplo, tamaños de moléculas más grandes) donde los modelos previos fallan.

4. Resultados

• Diseño De-novo: En el conjunto de datos QM9, Morph logra métricas competitivas (Estabilidad de Átomos, Estabilidad Molecular, Validez, Unicidad) comparadas con modelos de tamaño fijo SOTA (FlowMol, MiDi, EQGAT-diff, SemlaFlow), utilizando a menudo un número menor o igual de Evaluaciones de Número de Funciones (NFE). También se desempeña bien en el más complejo conjunto de datos GEOM-Drugs.
• Control de Propiedades (Steering):
  • Condicionamiento de Tamaño: Cuando se condiciona a conteos de átomos significativamente superiores a la distribución de entrenamiento (OOD), Morph mantiene una alta validez, generando moléculas con hasta un 40% más de átomos que el ejemplo más grande de entrenamiento. Los modelos de tamaño fijo fallan al extrapolar en este régimen.
  • Condicionamiento de LogP: El modelo logra dirigir con éxito el tamaño de las moléculas generadas basándose en el coeficiente de partición logarítmico (logP), recuperando el margen de tamaño de la distribución de entrenamiento y generando moléculas válidas en regiones de bajo soporte o de logP OOD.
• Decoración de Scaffolds: En una tarea de optimización de líderes (lead optimization), Morph extiende con éxito pequeños andamios químicos hacia moléculas completas preservando la subestructura central, logrando una alta validad y unicidad.

5. Significación y Reivindicaciones

El artículo afirma que Morph aborda una limitación fundamental en el diseño molecular generativo actual: la incapacidad de adaptar el tamaño molecular durante la generación. Al aprender la distribución conjunta de tamaño y propiedades, Morph desbloquea:

• Flexibilidad de Muestreo Sin Precedentes: La capacidad de generar moléculas de diversos tamaños a partir de un solo modelo.
• Mayor Potencial de Descubrimiento: La capacidad de explorar espacios químicos fuera de la distribución de entrenamiento (OOD), lo cual es esencial para descubrir nuevas moléculas con propiedades específicas y no vistas.
• Integración de Estructuras Previas: La capacidad de utilizar andamios (scaffolds) como puntos de partida para la expansión flexible, un requisito común en el descubrimiento de fármacos.

Los autores reconocen limitaciones, incluyendo el aumento de los costos de entrenamiento debido a la complejidad combinatoria, la necesidad de más épocas para alcanzar el rendimiento SOTA, y la dependencia actual de altos pasos de integración debido a la naturaleza de la CTMC de su modelo. Posicionan a Morph como un bloque de construcción prometedor para acelerar el descubrimiento en las ciencias químicas, específicamente para tareas que requieren la generación de tamaño flexible.