Resumen Técnico: NEAT – Transformador de Conjuntos Autoregresivo, Eficiente y Guiado por Vecindades para la Generación Molecular 3D

Enunciado del Problema

Los flujos de trabajo computacionales actuales para el descubrimiento de moléculas bioactivas dependen en gran medida del cribado virtual o de la generación de novo. Si bien el cribado virtual está limitado por el tamaño finito de las bibliotecas químicas y el aumento de los costos computacionales, el modelado generativo ofrece una alternativa prometedora. Sin embargo, la generación de estructuras moleculares 3D presenta desafíos únicos. Las moléculas son inherentemente tridimensionales y sus propiedades a menudo están determinadas por la conformación.

Los modelos de generación 3D existentes más avanzados utilizan predominantemente técnicas de difusión y ajuste de flujo (flow-matching). Estos métodos a menudo sufren de escalado cuadrático en las actualizaciones de paso de mensajes, carecen de adaptabilidad inherente a tamaños moleculares variables (a menudo requiriendo conteos de átomos fijos determinados en la inicialización) y tienen dificultades para gestionar inserciones dinámicas de átomos durante la eliminación de ruido.

Los modelos autoregresivos ofrecen una alternativa flexible que soporta naturalmente conteos variables de átomos y requiere solo una llamada al codificador por token. Sin embargo, los transformadores autoregresivos estándar requieren un ordenamiento secuencial de los tokens. En los grafos moleculares, los átomos no tienen un orden secuencial inherente. Trabajos anteriores han abordado esto utilizando ordenaciones canónicas de átomos, pero estos enfoques no son invariantes a la permutación con respecto a los átomos. Sesgan la predicción del siguiente token hacia convenciones de ordenamiento específicas, limitando la capacidad del modelo para generalizar o realizar tareas como la finalización de prefijos donde el orden de inicio de los átomos puede no alinearse con el orden canónico de la molécula final. Por el contrario, los enfoques de Redes Neuronales de Grafos (GNN) que aseguran la invarianza a la permutación a menudo producen moléculas con baja validez química.

Metodología

Los autores proponen NEAT (Transformador de Conjuntos Autoregresivo, Eficiente y Guiado por Vecindades), el primer modelo transformador autoregresivo para la generación molecular 3D que respeta la invarianza a la permutación de átomos sin depender de ordenaciones canónicas de átomos.

Arquitectura Central

NEAT trata los grafos moleculares como conjuntos de átomos en lugar de secuencias. La arquitectura consta de tres componentes principales:

1. Módulo de Aumento de Datos (Guía de Vecindad): En lugar de codificaciones posicionales secuenciales, el modelo utiliza una estrategia de muestreo para imitar el crecimiento molecular. Dado un grafo molecular, se muestrea un subgrafo conectado (conjunto fuente). El conjunto objetivo se define como los "nodos frontera" (vecinos de 1 salto) del conjunto fuente. Esto proporciona un conjunto de continuaciones de tokens válidas para el siguiente paso, independientemente del orden en que se añadieron los átomos.
2. Tronco del Transformador de Conjuntos: El conjunto fuente se codifica utilizando un transformador de conjuntos (basado en NanoGPT/GPT-2 con atención bidireccional y sin incrustaciones posicionales). Los tipos de átomos se incrustan linealmente y las posiciones se codifican utilizando características de Fourier. La agrupación aditiva agrega estas representaciones de nodos en una representación a nivel de conjunto ($z_n$). Este diseño asegura la invarianza a la permutación con respecto al orden de los átomos en el conjunto fuente.
3. Cabeza de Ajuste de Flujo: La representación a nivel de conjunto $z_n$ se utiliza para predecir el siguiente token, factorizada en tipo de átomo y posición: $p(a_{n+1}, x_{n+1} | z_n) = p(a_{n+1} | z_n) \cdot p(x_{n+1} | a_{n+1}, z_n)$.
   • Tipo de Átomo: Modelado como una distribución categórica mediante una capa lineal y softmax.
   • Posición del Átomo: Modelada mediante Ajuste de Flujo Condicional (CFM). El modelo aprende un campo vectorial para transportar una distribución inicial de posición aleatoria a las posiciones atómicas reales.

Estrategia de Entrenamiento

• Guía de Vecindad: Durante el entrenamiento, el modelo no se supervisa mediante un único siguiente token, sino por el conjunto de todos los nodos frontera válidos. Esto elimina la necesidad de canonicación.
• Forzamiento del Profesor: La predicción del tipo de átomo y la recuperación de la posición se desacoplan. El modelo predice el tipo, y la cabeza de ajuste de flujo aprende la posición condicionada a ese tipo.
• Aumento: Para mejorar la robustez frente a la acumulación de errores, se inyecta ruido gaussiano en una fracción de los átomos del conjunto fuente durante el entrenamiento, y las posiciones se re-centran con respecto al conjunto fuente ruidoso.
• Función de Pérdida: La pérdida total es la suma de la pérdida de entropía cruzada para los tipos de átomos y la pérdida de ajuste de flujo para las posiciones. Se resuelve un problema de asignación lineal para acoplar las distribuciones fuente y objetivo, minimizando los costos de transporte y evitando cruces de trayectorias.

Inferencia

La generación comienza con un átomo inicial aleatorio. En cada paso, el conjunto actual de átomos se codifica, se predice el tipo del siguiente átomo y se muestrea una posición inicial aleatoria. Luego, el campo vectorial de ajuste de flujo se integra (utilizando métodos de Euler o Euler-Maruyama) para determinar las coordenadas 3D finales. Este proceso se repite hasta que se predice un token de parada.

Contribuciones Clave

1. Generación Autoregresiva Invariante a la Permutación: NEAT introduce una estrategia de entrenamiento novedosa que utiliza guía de vecindad y transformadores de conjuntos para lograr una generación molecular 3D que es invariante al ordenamiento de átomos, superando las limitaciones del ordenamiento canónico en los modelos autoregresivos.
2. Eficiencia: Al tratar la entrada como un conjunto y utilizar una única llamada al codificador por token, NEAT evita el escalado cuadrático del paso de mensajes basado en difusión y la necesidad de múltiples pasos de integración por paso de generación de átomos encontrados en algunas líneas base de ajuste de flujo.
3. Capacidad de Finalización de Prefijos: La naturaleza agnóstica al orden de NEAT permite una finalización de prefijos de alta calidad (generación de moléculas alrededor de un andamio dado) sin requerir que el prefijo se alinee con un ordenamiento canónico específico, una tarea donde los modelos basados en secuencias como QUETZAL luchan.
4. Rendimiento de Vanguardia: El modelo logra resultados competitivos o superiores en puntos de referencia estándar (QM9 y GEOM-Drugs) en comparación tanto con líneas base autoregresivas como con líneas base de difusión/ajuste de flujo.

Resultados Experimentales

El modelo se evaluó en los conjuntos de datos QM9 (moléculas orgánicas pequeñas) y GEOM-Drugs (moléculas similares a fármacos).

• Calidad de Generación: En QM9, NEAT superó a los modelos autoregresivos anteriores (G-SchNet, G-SphereNet, Symphony, Geo2Seq) en términos de validez y unicidad de xyz2mol. Logró una validez de xyz2mol del 95,2 % y una unicidad del 91,6 %, ocupando el segundo lugar en validez y unicidad de búsqueda. En GEOM, NEAT logró la mayor validez de búsqueda (99,7 %) y la segunda mejor validez de xyz2mol (93,9 %) y unicidad (91,6 %) entre los modelos autoregresivos, igualando el rendimiento de modelos de ajuste de flujo de vanguardia como FlowMol3.
• Finalización de Prefijos: En una tarea que generaba 1.000 moléculas a partir de 100 prefijos diversos, NEAT superó significativamente a QUETZAL. NEAT logró una validez de búsqueda del 97,8 % y una validez de xyz2mol del 92,7 %, mientras que QUETZAL logró solo el 27,8 % y el 28,1 % respectivamente. Esto demuestra la capacidad de NEAT para generalizar características estructurales independientemente de las convenciones de ordenamiento de átomos.
• Eficiencia de Tiempo de Ejecución: NEAT demuestra ventajas significativas de velocidad. Funciona aproximadamente 5 veces más rápido que SemlaFlow y 16 veces más rápido que FlowMol3. En escenarios del peor caso, la complejidad de NEAT es $O(BN^3)$ (donde $N$ es el conteo de átomos), mientras que los modelos de difusión/flujo escalan como $O(BTN^2)$ (donde $T$ son los pasos de integración). Dado que $N$ es típicamente menor que $T$ para moléculas similares a fármacos, NEAT ofrece una ventaja práctica de velocidad.

Significado y Afirmaciones

El artículo afirma que NEAT proporciona una base sólida para aprovechar las arquitecturas de transformadores en tareas complejas de diseño de fármacos, específicamente para la generación condicional de moléculas 3D. Al eliminar la dependencia de ordenaciones canónicas de átomos, el modelo permite tareas flexibles como la generación basada en andamios (finalización de prefijos) que anteriormente eran difíciles para los enfoques autoregresivos.

Los autores señalan que, si bien la arquitectura actual no es explícitamente equivariante a SE(3), el uso de aumento de rotación aleatoria durante el entrenamiento compensa efectivamente esto, resultando en solo una variación menor en las predicciones. El trabajo sugiere que integrar NEAT con un transformador equivariante a SE(3) es una dirección prometedora para futuras investigaciones.

Los autores enfatizan que, si bien el modelo genera estructuras 3D de alta calidad, la generación de moléculas es solo un paso en una compleja tubería que requiere validación extensa, síntesis y pruebas. También reconocen limitaciones, como la dependencia de información de bordes para la guía de vecindad (lo cual puede limitar la aplicabilidad a cristales inorgánicos) y la falta de extrapolación más allá del conteo máximo de átomos visto durante el entrenamiento.