OmniMol: Transferring Particle Physics Knowledge to Molecular Dynamics with Point-Edge Transformers

Este artículo presenta OmniMol, un potencial interatómico aprendido por máquina de última generación para moléculas pequeñas que aprovecha una arquitectura Transformer de puntos y aristas y la transferencia de conocimiento desde la física de altas energías para lograr un rendimiento excelente con un ajuste fino mínimo y una inferencia excepcionalmente rápida.

Lo esencial

Imagina que tienes dos mundos muy diferentes: uno es el mundo caótico y de alta velocidad de la física de partículas (donde los científicos hacen chocar átomos entre sí para ver qué sale disparado), y el otro es el mundo intrincado y pegajoso de la química molecular (donde los átomos se unen para formar medicamentos, materiales y la vida).

Durante mucho tiempo, los científicos en estos dos campos utilizaron herramientas completamente diferentes para entender sus mundos. Pero en este artículo, los autores presentan OmniMol, una nueva herramienta que intenta enseñar a los expertos en física de partículas a entender la química utilizando un "modelo fundacional" que ya habían construido.

Aquí tienes una explicación sencilla de cómo lo hicieron y qué descubrieron:

1. La analogía del "Chef Maestro"

Piensa en el modelo original, llamado Omnilearned, como un chef maestro que ha pasado años cocinando con chorros de partículas.

• Los ingredientes: En física de partículas, un "chorro" es una lluvia de partículas subatómicas (como protones y neutrones) que salen disparadas de una colisión.
• La habilidad: Este chef aprendió a reconocer patrones en estas lluvias. Sabe cómo interactúan las partículas, cómo se agrupan y cómo predecir qué sucederá a continuación. Fue entrenado con mil millones de diferentes lluvias de partículas.

Ahora, los autores se preguntaron: ¿Puede este mismo chef preparar una comida molecular?

• Los nuevos ingredientes: En lugar de partículas subatómicas, los "ingredientes" son átomos (como Carbono, Oxígeno, Hidrógeno) en una molécula.
• El desafío: Los átomos se comportan de manera diferente a las partículas subatómicas, pero comparten una estructura similar: son simplemente puntos en el espacio con tipos específicos.

2. El "Traductor Universal" (La arquitectura)

Para que esto funcionara, no construyeron un nuevo chef desde cero. Tomaron al "Chef Maestro" existente (Omnilearned) y le dieron un nuevo conjunto de herramientas:

• El Transformador Punto-Borde (PET): Imagina al chef mirando un plato de comida. En lugar de mirar solo un ingrediente a la vez, esta herramienta le permite mirar todos los ingredientes a la vez y ver cómo cada uno se relaciona con todos los demás.
• El "sesgo físico": Este es el secreto. El modelo tiene un "reglamento" incorporado que le dice: "Oye, estas dos partículas/átomos están cerca, así que deberían prestarse más atención entre sí". Esto ayuda al modelo a centrarse en las relaciones más importantes sin confundirse con el ruido.

3. El experimento: Ajuste fino

Los autores tomaron este modelo entrenado en partículas y le dieron un "curso intensivo" en química utilizando un conjunto de datos llamado oMol (una colección de millones de moléculas).

• El objetivo: Querían que el modelo actuara como un Potencial Interatómico Aprendido por Máquina (MLIP). En lenguaje sencillo, esto significa que el modelo necesita predecir dos cosas para cualquier grupo de átomos:
  1. Energía: ¿Cuánta "cola" los mantiene unidos?
  2. Fuerza: Si empujas un átomo, ¿qué tan fuerte empujará de vuelta?

4. Los resultados: Rápido y sorprendentemente bueno

El artículo encontró algunas cosas emocionantes:

• El superpoder de "pocos ejemplos": Por lo general, enseñar química a una computadora requiere cantidades masivas de datos. Pero como OmniMol comenzó con el "conocimiento" de la física de partículas, aprendió química muy rápido. Incluso con una cantidad relativamente pequeña de nuevos datos (como 100,000 moléculas), funcionó casi tan bien como los modelos entrenados con millones. Es como un chef maestro que puede aprender una nueva cocina con solo unas pocas recetas porque ya entiende los fundamentos del sabor y el calor.
• Velocidad: OmniMol es increíblemente rápido. Mientras que otros modelos podrían tardar mucho en calcular cómo se mueve una molécula, OmniMol lo hace en un abrir y cerrar de ojos. Los autores señalan que por cada hora de tiempo de cómputo, OmniMol puede simular tres veces más moléculas que algunos de sus competidores.
• La compensación: Cuando tuvieron enormes cantidades de datos (millones de moléculas), la ventaja de comenzar con el conocimiento de la física de partículas disminuyó un poco. Esto sugiere que el "conocimiento de la física de partículas" actúa como un gran punto de partida, pero si tienes suficiente tiempo y datos para entrenar un modelo desde cero, ese punto de partida importa menos.

5. El panorama general

El artículo concluye que OmniMol es la primera vez que un "modelo fundacional" construido para una disciplina científica (física de partículas) se ha transferido con éxito a una completamente diferente (química).

Demostraron que si tienes un modelo inteligente que entiende cómo interactúan los puntos en el espacio en un campo, puede adaptarse para entender cómo interactúan los puntos en el espacio en otro campo, ahorrando tiempo y potencia de cómputo.

En resumen: Los autores tomaron una IA súper inteligente entrenada en colisiones de partículas de alta energía, ajustaron su cerebro para entender átomos en lugar de partículas, y descubrieron que se convirtió en una herramienta extremadamente rápida y altamente precisa para predecir cómo se comportan las moléculas, especialmente cuando los datos son escasos.

Resumen técnico

Resumen Técnico: OmniMol

Enunciado del Problema

El aprendizaje automático (ML) ha transformado la representación y simulación de sistemas físicos complejos, particularmente en física de partículas y química molecular. Aunque estos dominios difieren enormemente en escalas de energía, comparten una estructura de datos fundamental: conjuntos de partículas (o átomos) de tamaño variable en el espacio de fases, que forman efectivamente nubes de puntos estructuradas.

El principal desafío abordado es el desarrollo de Potenciales Interatómicos Aprendidos por Máquina (MLIPs) eficientes. Los métodos tradicionales como la Teoría del Funcional de la Densidad (DFT) son computacionalmente costosos, lo que limita las simulaciones de dinámica molecular (MD) a gran escala y de largo horizonte. Los MLIPs buscan aproximar superficies de energía potencial y fuerzas a una fracción de este costo. Sin embargo, entrenar MLIPs robustos típicamente requiere conjuntos de datos masivos y recursos computacionales significativos. El artículo hipotetiza que un modelo fundamental pre-entrenado en nubes de puntos de física de partículas (específicamente chorros de partículas) podría transferirse a la dinámica molecular, acelerando potencialmente la optimización y mejorando la precisión en regímenes de bajos datos.

Metodología

Arquitectura: Transformer Punto-Borde (PET)

OmniMol se construye adaptando Omnilearned, un modelo fundamental diseñado originalmente para clasificar y generar chorros de partículas en física de altas energías (HEP). La arquitectura central es un Transformer Punto-Borde (PET), que acopla la atención local sobre los $k$-vecinos más cercanos con bloques transformadores globales de todos-a-todos.

Los componentes arquitectónicos clave incluyen:

1. Incrustaciones de Entrada: Los átomos se incrustan en un espacio de tokens que combina información posicional ($\vec{r}$), números atómicos discretos ($Z$) y características adicionales (carga, espín).
2. Bloque de Atención Local: Para cada átomo, se construye un vecindario local utilizando los $K$-vecinos más cercanos ($K=15$ para moléculas, comparado con $K=10$ para chorros). Se calculan características físicas pares, incluyendo términos de distancia, potencias inversas de la distancia y funciones aprendidas de las incrustaciones atómicas. Estas se procesan mediante un pequeño transformador local para crear un vector de incrustación local.
3. Atención Global con Sesgo de Interacción: El mecanismo de auto-atención global incorpora un sesgo explícito derivado de características físicas pares. Los logits de atención se modifican como $A^*_{ij} = A_{ij} + B_{ij}$, donde $B_{ij}$ es un término de sesgo incrustado por un MLP. Este "sesgo de matriz de atención de interacción" inyecta priores físicos pares directamente en el transformador, guiando la red hacia vecindarios físicamente significativos sin sacrificar la expresividad.
4. Cabezas de Salida: La cabeza generativa de Omnilearned se reutiliza para dos tareas:
   • Predicción de Fuerzas: Una cabeza equivariante a permutaciones que predice fuerzas por átomo.
   • Predicción de Energía: Una cabeza que predice correcciones de energía por átomo, las cuales se suman para obtener la energía molecular total, preservando priores extensivos.

Restricciones de Invarianza y Conservación

Para satisfacer las restricciones físicas, los autores abordan dos requisitos:

• Conservación de Energía: Las fuerzas no se predicen directamente, sino que se calculan mediante retropropagación de la salida de energía ($\vec{F}_i = \nabla_{\vec{r}_i} E$). Esto asegura una conservación exacta de la energía, pero aumenta el costo computacional durante el entrenamiento (requiriendo doble retropropagación). En consecuencia, esta restricción se aplica solo a la variante "pequeña" del modelo.
• Equivariancia Rotacional: La arquitectura estándar no es inherentemente equivariante porque las diferencias de coordenadas crudas se alimentan a MLPs. Para remediar esto, los autores introducen una variante "equivariante y conservadora". Esta versión elimina los términos de diferencia de coordenadas directas de las características pares e incorpora en su lugar información angular (cosenos de los ángulos formados por vectores entre átomos vecinos) en el bloque local. Esta modificación mantiene la equivariancia mientras recupera significativamente las pérdidas de rendimiento asociadas con la eliminación de los términos de coordenadas.

Estrategias de Entrenamiento y Ajuste Fino

El modelo se ajusta fino en el conjunto de datos oMol (específicamente los subconjuntos oMol-25, oMol-4M, oMol-100M y oMol-140M). Se exploran dos estrategias:

1. LoRA (Adaptación de Bajo Rango): Los pesos de la columna vertebral PET pre-entrenada se congelan. Se introducen adaptadores de bajo rango solo para las matrices del cuerpo del transformador ($W_Q, W_K, W_V, W_O, W_{MLP}$), junto con el entrenamiento de los codificadores de entrada molecular, el MLP de sesgo y las cabezas de tarea. También se añade una capa de "adaptación de incrustación" para modificar las incrustaciones aprendidas.
2. Ajuste Fino Completo: Todos los pesos en el cuerpo y los codificadores de entrada se desbloquean y entrenan, mientras que las cabezas de tarea se entrenan desde cero.

El objetivo de entrenamiento minimiza la suma de los Errores Absolutos Medios (MAE) para energías y fuerzas, con las fuerzas ponderadas más pesadamente ($\lambda_F = 10$).

Resultados Clave

Rendimiento en oMol

• Entrenamiento Completo: Cuando se entrena en conjuntos de datos grandes (oMol-4M y oMol-100M/140M), OmniMol logra un rendimiento competitivo con los MLIPs más avanzados. Por ejemplo, en oMol-140M, el modelo OmniMol-large logra un MAE de energía de 1.04 meV/átomo y un MAE de fuerza de 13.59 meV/Å.
• Régimen de Bajos Datos: Las ganancias más significativas se observan cuando los datos de entrenamiento son limitados. Cuando se ajusta fino en solo 100k moléculas o con muy pocas épocas (2 pasadas) sobre oMol-4M, las variantes pre-entrenadas de OmniMol superan significativamente a los modelos entrenados desde cero.
  • En un subconjunto de 100k, el pre-entrenamiento mejoró el MAE de energía hasta un 29.4% y el MAE de fuerza en un 26.9% para el modelo mediano.
  • Con solo dos épocas en oMol-4M, el modelo mediano mostró una mejora del 54.6% en el MAE de energía y del 56.9% en el MAE de fuerza en comparación con su contraparte no pre-entrenada.
• Variante Equivariante/Conservadora: La variante del modelo equivariante y conservadora muestra un rendimiento significativamente mejorado (especialmente para fuerzas) en regímenes de bajos datos, aunque esta ventaja disminuye a medida que aumenta el tamaño del conjunto de datos.

Escalado y Velocidad de Inferencia

• Escalado: OmniMol sigue una ley de potencias limpia con el tamaño del modelo, sin mostrar signos de saturación hasta 1 mil millones de parámetros, consistente con hallazgos recientes sobre MLIPs basados en transformadores.
• Velocidad de Inferencia: A pesar de grandes conteos de parámetros, OmniMol demuestra velocidades de inferencia únicas debido a optimizaciones de hardware para transformadores. En una GPU A100 para sistemas de ~100 átomos, OmniMol-mediano es aproximadamente 3 veces más rápido que las líneas base comparables de Redes Neuronales de Grafos (GNN) (eSEN-md-d y AllScAIP-md) mientras mantiene una precisión competitiva (solo ~0.7 meV/átomo de error de energía más alto que AllScAIP-md).

Significado y Afirmaciones

El artículo afirma presentar la primera demostración de transferencia interdisciplinaria para modelos fundamentales de nubes de puntos científicos. Al adaptar un modelo pre-entrenado en chorros de partículas de física de altas energías a la dinámica molecular, los autores demuestran que:

1. La Transferencia Transdominio es Viable: Un modelo fundamental construido para física de partículas puede transferirse efectivamente a la química molecular, sugiriendo que las estructuras subyacentes de nubes de puntos comparten características aprendibles a través de escalas físicas enormemente diferentes.
2. El Sesgo Inductivo Acelera el Aprendizaje: El pre-entrenamiento actúa como un fuerte sesgo inductivo. Similar a cómo la equivariancia ayuda cuando los datos son escasos, la "lección amarga" del pre-entrenamiento permite una optimización rápida y una precisión mejorada cuando los datos de entrenamiento son limitados.
3. Eficiencia: La transferencia arquitectónica permite velocidades de inferencia únicas, lo cual es crítico para aplicaciones que requieren una exploración rápida de espacios de diseño, como el descubrimiento de fármacos de moléculas pequeñas.

Los autores concluyen que, aunque el estudio se centra en MLIPs, las lecciones sobre modelos fundamentales de nubes de puntos pueden tener utilidad generalizada en dominios científicos donde los sistemas se describen como conjuntos desordenados de cuerpos interactuantes. No afirman una superioridad universal sobre todos los métodos existentes en todos los regímenes, sino que destacan las ventajas específicas en escenarios de bajos datos y velocidad de inferencia.