Beyond Adam: SOAP and Muon for Faster, Label-Efficient Training of Machine Learning Interatomic Potentials

Título: Más allá de Adam: SOAP y Muon para un entrenamiento más rápido y eficiente en etiquetas de Potenciales Interatómicos de Aprendizaje Automático

Planteamiento del problema
Los Potenciales Interatómicos de Aprendizaje Automático (MLIP, por sus siglas en inglés) se han vuelto esenciales para las simulaciones atomísticas en química y ciencia de materiales, permitiendo estudios que van desde la termodinámica del agua hasta ensamblajes biomoleculares. Si bien el campo ha avanzado significativamente mediante nuevas arquitecturas (por ejemplo, redes neuronales de grafos equivariantes como NequIP y Allegro) y conjuntos de datos a gran escala, el panorama de la optimización del entrenamiento ha permanecido estático. La comunidad recurre predominantemente a Adam y sus variantes (por ejemplo, AdamW) para el entrenamiento. Este artículo postula que la elección del optimizador es un eje de diseño subestimado pero impactante. Específicamente, los autores investigan si los optimizadores de estructura matricial, que explotan la estructura tensorial de los pesos de la red neuronal en lugar de depender del escalado de gradientes diagonal, pueden mejorar la velocidad de convergencia y la precisión final, particularmente en regímenes donde las etiquetas de fuerza son costosas o escasas.

Metodología
Los autores implementan y evalúan sistemáticamente tres optimizadores de estructura matricial recientemente propuestos frente al baseline estándar de AdamW:

1. Muon: Un método de actualización ortogonal utilizando iteraciones de Newton–Schulz para proyectar los gradientes sobre matrices semi-ortogonales.
2. SOAP: Un optimizador que realiza actualizaciones de estilo Adafactor en el espacio de autovectores del precondicionador, seguido de una proyección de vuelta al espacio de parámetros original.
3. SOAP-Muon: Un enfoque híbrido que combina la actualización de SOAP con un paso de ortogonalización inspirado en Muon.

Estos optimizadores fueron integrados en el framework nequip y aplicados a dos sistemas físicos distintos:

• Agua Líquida: Modelada utilizando la arquitectura NequIP.
• Electrolito Ácido Sólido (CsH₂PO₄ o CDP): Modelada utilizando la arquitectura Allegro.

El estudio empleó un protocolo de ajuste sistemático de hiperparámetros para asegurar comparaciones justas. Crucialmente, los autores evaluaron el rendimiento bajo distintos niveles de supervisión de fuerza, que van desde la supervisión completa (100% de etiquetas de fuerza) hasta el entrenamiento solo con energía (0% de etiquetas de fuerza) e intervalos intermedios de supervisión escasa (75%, 50%, 10%, 5%). Esta configuración pone a prueba la capacidad de los optimizadores para aprender superficies de energía potencial (PES) precisas cuando los cálculos de fuerza —a menudo costosos en teorías de alto nivel como el método de acoplamiento de cúmulos o la función de Monte Carlo de difusión— son limitados.

Contribuciones clave

• Integración y Benchmarking: Los autores integraron con éxito Muon, SOAP y SOAP-Muon en el framework nequip, gestionando el remodelado de parámetros específico requerido para las capas e3nn (Linear y FullyConnectedTensorProduct) para aplicar actualizaciones de estructura matricial.
• Análisis de Rendimiento: Demostraron que estos optimizadores pueden superar sustancialmente a AdamW tanto en velocidad de convergencia como en precisión final.
• Eficiencia de Etiquetas: El estudio destaca que SOAP y SOAP-Muon mantienen una alta precisión incluso cuando la supervisión de fuerza se reduce significativamente, lo que sugiere una vía para entrenar MLIP precisos con menos etiquetas de fuerza costosas.
• Validación de Fidelidad Física: Más allá del error absoluto medio (MAE), los autores validaron la corrección física de los potenciales aprendidos ejecutando simulaciones de dinámica molecular (MD) y comparando observables (Funciones de Distribución Radial, Desplazamiento Cuadrático Medio y energías de activación de difusión) contra datos de referencia ab initio y mediciones experimentales.

Resultados

• Supervisión Completa: Bajo supervisión completa de energía y fuerza, SOAP y SOAP-Muon superaron consistentemente a AdamW.
  • En el sistema CDP (Allegro), SOAP logró la mejor precisión de energía (reducción del 9% en MAE), mientras que SOAP-Muon logró la mejor precisión de fuerza (reducción del 14%).
  • En el sistema de agua (NequIP), SOAP-Muon proporcionó la mejor precisión de energía (reducción del 24%), y SOAP la mejor precisión de fuerza (reducción del 19%).
  • Muon por sí solo mostró resultados mixtos, rindiendo por debajo de AdamW en el sistema de agua, pero mejorando en CDP.
• Velocidad de Convergencia: Los optimizadores de estructura matricial alcanzaron el mínimo de MAE de fuerza de validación de AdamW significativamente más rápido. En CDP, SOAP fue 4.9× más rápido; en agua, fue 5.8× más rápido. Esta aceleración ocurrió a pesar del mayor costo computacional por paso de optimización, lo que implica una reducción sustancial en el número de épocas requeridas.
• Supervisión de Fuerza Escasa: Las ventajas de los optimizadores de estructura matricial se hicieron más pronunciadas a medida que se reducían las etiquetas de fuerza.
  • Entrenamiento Solo con Energía: En ausencia de etiquetas de fuerza, SOAP y SOAP-Muon superaron sustancialmente a AdamW. Por ejemplo, en CDP, SOAP-Muon redujo el MAE de energía y fuerza en aproximadamente un 47% y 60% respectivamente en comparación con AdamW.
  • Etiquetas Escasas: En CDP, el modelo entrenado con SOAP-Muon usando solo el 50% de las etiquetas de fuerza alcanzó una precisión comparable a la de AdamW entrenado con el 100%.
  • Estabilidad: Se observó un hallazgo crítico en la supervisión del 5% de fuerza en CDP. El modelo AdamW se volvió catastróficamente inestable durante las simulaciones de MD, produciendo resultados no físicos. En contraste, el modelo SOAP-Muon se mantuvo estable y reprodujo con precisión las RDFs ab initio y los MSDs.
• Observables Físicos: Los modelos entrenados con SOAP-Muon (incluso con solo un 5% de etiquetas de fuerza) reprodujeron las energías de activación de la difusión de protones (0.42 eV) en buen acuerdo con el rango experimental (0.39–0.43 eV), mientras que el modelo inestable de AdamW no logró alcanzar el régimen de difusión lineal.

Significancia y Reivindicaciones
El artículo concluye que la elección del optimizador es un componente crítico, aunque a menudo descuidado, del diseño de MLIP. Los autores argumentan que las técnicas de entrenamiento deben tratarse como un "eje de diseño de primera clase" junto con la selección de arquitectura y conjunto de datos.

• Robustez: SOAP emerge como el optimizador por defecto más robusto a través de diferentes sistemas y niveles de supervisión. SOAP-Muon muestra los resultados más fuertes en configuraciones específicas, pero requiere un ajuste de hiperparámetros más cuidadoso (particularmente respecto a los coeficientes de momentum y la potencia de los valores singulares) para asegurar la estabilidad.
• Impacto Práctico: La capacidad de estos optimizadores para acelerar el entrenamiento y mantener la precisión con etiquetas de fuerza escasas ofrece una solución práctica para entrenar MLIP de alta fidelidad con datos de referencia donde los cálculos de fuerza son prohibitivamente costosos (por ejemplo, más allá de DFT).
• Limitaciones: Los autores reconocen que su trabajo está restringido a dos sistemas y dos arquitecturas. Señalan que, si bien SOAP parece efectivo a la escala actual, se necesita un estudio sistemático adicional para confirmar su eficacia en el entrenamiento de potenciales base universales. También destacan que la inestabilidad de Muon en el sistema de agua requiere mayor investigación, sugiriendo que el paso de ortogonalización puede ser una fuente primaria de degradación en ciertos contextos.

Los autores declaran que su código se liberará en futuras versiones del framework de código abierto nequip, facilitando la adopción más amplia de estas estrategias de optimización.