Resumen Técnico: Modelos Fundacionales Generativos de Moléculas Controlables (CoMole)

Enunciado del Problema

El diseño inverso molecular —la generación de estructuras con propiedades funcionales específicas— sigue siendo un desafío central en el descubrimiento científico. Aunque los métodos recientes de difusión en grafos han avanzado en la generación molecular, operan en gran medida como modelos específicos de tarea con un control limitado sobre las propiedades objetivo. El campo carece de un marco unificado capaz de manejar tareas de diseño heterogéneas (que abarcan desde la ciencia de materiales hasta el descubrimiento de fármacos) con un control fiable.

Los intentos de cerrar esta brecha utilizando Aprendizaje por Refuerzo (RL) para el entrenamiento posterior enfrentan un cuello de botella crítico al aplicarse a la difusión en grafos a nivel de átomo. El vasto espacio de acciones a nivel de átomo está mal alineado con las subestructuras químicamente significativas (motivos) que los químicos utilizan para formular racionalidades estructura-propiedad. Como se demuestra en el artículo, el RL a nivel de átomo a menudo sufre de "colapso de la trayectoria", donde una sola elección atómica inválida desplaza la generación fuera de la variedad química factible, haciendo difícil la recuperación. Esta fragilidad conduce a una validez deficiente y a una asignación de crédito inestable, impidiendo que el modelo asocie de manera fiable los cambios estructurales con las recompensas de propiedades.

Metodología

Los autores proponen CoMole (Modelos Fundacionales Generativos de Moléculas Controlables), un marco unificado construido sobre una tubería de difusión en grafos consciente de motivos. La innovación central es reparametrizar el proceso de RL de difusión en grafos desde un espacio de decisiones a nivel de átomo hacia un espacio de decisiones consciente de motivos.

1. Tokenización de Grafos Consciente de Motivos

CoMole emplea un tokenizador basado en Codificación de Pares de Nodos (NPE) que aprende un vocabulario a partir de datos de preentrenamiento. Este tokenizador mapea grafos moleculares químicamente válidos en grafos de motivos ($z_0$), donde los nodos representan unidades disjuntas de átomos:

• Átomos individuales.
• Unidades de anillo preservadas.
• Subestructuras fusionadas basadas en datos (motivos de nivel superior).

Esta representación preserva la información a nivel de unión y las etiquetas de enlace inter-motivo, permitiendo una reconstrucción sin pérdidas mientras comprime significativamente el estado del grafo.

2. Tubería de Entrenamiento de Tres Etapas

CoMole sigue un paradigma de modelo fundacional que involucra tres etapas:

• Preentrenamiento (PT): El modelo aprende priores estructurales transferibles sobre el espacio de grafos de motivos utilizando difusión incondicional.
• Ajuste Fino Supervisado (SFT): El modelo se condiciona sobre múltiples propiedades (numéricas y categóricas) para introducir una alineación específica de la tarea.
• Aprendizaje por Refuerzo (RL): La política de difusión inversa condicional se optimiza mediante Optimización de Política Proximal (PPO) para alinearse con las recompensas de propiedades objetivo terminales. Críticamente, el RL opera sobre el espacio de acciones consciente de motivos, donde las acciones corresponden a decisiones químicamente significativas (por ejemplo, unir un anillo de benceno) en lugar de ediciones de átomos individuales.

3. Justificación Teórica

El artículo proporciona un análisis teórico que caracteriza el cuello de botella del RL a nivel de átomo. Al definir la complejidad de la decisión como el número de factores categóricos en el kernel inverso, los autores muestran que las representaciones conscientes de motivos reducen el tamaño del espacio de acciones por un factor proporcional al cuadrado de la relación de compresión átomo-motivo ($\chi$). Esta reducción estabiliza la optimización de la política regularizada por KL, permitiendo que las recompensas se asignen más directamente a las acciones relevantes para la propiedad y previniendo el colapso de la trayectoria.

Contribuciones Clave

1. Modelo Fundacional Unificado: CoMole se presenta como la primera familia de modelos fundacionales generativos de moléculas controlables diseñados para el diseño inverso heterogéneo, cerrando la brecha entre los métodos específicos de tarea actuales y las necesidades prácticas.
2. RL Consciente de Motivos: El trabajo introduce un enfoque novedoso para estabilizar el RL para la generación de grafos al desplazar el espacio de decisiones de átomos a motivos químicamente significativos, justificando teóricamente por qué esto reduce la complejidad de la decisión y mejora la asignación de crédito.
3. Controlabilidad Transferible: El marco demuestra la capacidad de transferir la controlabilidad a objetivos de propiedades no vistos optimizando solo los incrustaciones de la tarea mientras se mantiene el generador congelado, logrando un rendimiento competitivo con las líneas base específicas de la tarea.

Resultados Experimentales

CoMole fue evaluado en tres benchmarks heterogéneos: permeabilidad de gases en polímeros, propiedades DFT de polímeros y tareas de descubrimiento de fármacos (FreeSolv, BACE, BBBP).

• Controlabilidad: CoMole ocupó el primer lugar en controlabilidad en los nueve objetivos.
  • En los benchmarks de polímeros, redujo el Error Absoluto Medio (MAE) en más de un 44% en relación con las mejores líneas base.
  • En tareas de fármacos, redujo el MAE de FreeSolv en un 13.1% y logró una precisión de 1.0 en la clasificación de BACE.
  • En comparación con las líneas base más fuertes (por ejemplo, Graph DiT, FREED), CoMole redujo el MAE hasta en un 48.2%.
• Validez: El modelo mantuvo una tasa de validez superior a 0.94 en todos los conjuntos de datos sin corrección basada en reglas o filtrado post-hoc. En contraste, las variantes de RL a nivel de átomo sufrieron caídas en la validez (por ejemplo, de 0.95 a 0.07) y no lograron mantener la controlabilidad.
• Generalización: Cuando se probó en objetivos DFT no vistos ($E_i$, $\epsilon_s$) congelando el generador y aprendiendo solo nuevas incrustaciones de tarea, CoMole logró un rendimiento competitivo con las líneas base entrenadas directamente en esos objetivos.
• Ablación: Reemplazar el tokenizador de motivos con una representación a nivel de átomo causó un fallo catastrófico tanto en la validez como en la controlabilidad, confirmando la necesidad del diseño consciente de motivos.

Significado y Afirmaciones

El artículo afirma que CoMole proporciona una base práctica para el diseño inverso molecular unificado, escalable y controlable. Al transferir priores estructurales preentrenados a una generación controlable mediante optimización de políticas consciente de motivos, el modelo aprende conocimiento estructura-propiedad transferible que soporta diversas tareas de diseño inverso.

Los autores enfatizan que su diseño aborda la fragilidad fundamental del RL a nivel de átomo en la generación de grafos. Al estabilizar la geometría de optimización mediante decisiones químicamente significativas, CoMole logra una alta controlabilidad mientras preserva la validez química, ofreciendo una vía hacia sistemas de diseño basados en modelos fundacionales para moléculas y polímeros que pueden cerrar la brecha entre la generación hipotética y el descubrimiento científico práctico. El trabajo sugiere que tales modelos pueden ayudar a los investigadores a priorizar candidatos prometedores bajo restricciones heterogéneas antes de comprometerse con simulaciones o experimentos costosos.