Comprensión y Mitigación de Cambios en la Distribución para el Aprendizaje Automático Campos de Fuerza

Understanding and Mitigating Distribution Shifts For Machine Learning Force Fields

March 11, 2025

Autores: Tobias Kreiman, Aditi S. Krishnapriyan

cs.AI

Resumen

Los Campos de Fuerza de Aprendizaje Automático (MLFFs, por sus siglas en inglés) representan una alternativa prometedora frente a las costosas simulaciones moleculares cuánticas ab initio. Dada la diversidad de espacios químicos de interés y el costo asociado a la generación de nuevos datos, es crucial comprender cómo los MLFFs generalizan más allá de sus distribuciones de entrenamiento. Para caracterizar y entender mejor los cambios de distribución en los MLFFs, realizamos experimentos de diagnóstico en conjuntos de datos químicos, revelando desplazamientos comunes que plantean desafíos significativos, incluso para modelos base entrenados con grandes volúmenes de datos. Basándonos en estas observaciones, planteamos la hipótesis de que los métodos de entrenamiento supervisado actuales regularizan de manera insuficiente los MLFFs, lo que resulta en sobreajuste y en la obtención de representaciones deficientes para sistemas fuera de distribución. A continuación, proponemos dos nuevos métodos como pasos iniciales para mitigar estos cambios de distribución en los MLFFs. Nuestros métodos se centran en estrategias de refinamiento en tiempo de prueba que implican un costo computacional mínimo y no utilizan etiquetas de referencia ab initio costosas. La primera estrategia, basada en la teoría espectral de grafos, modifica las aristas de los grafos de prueba para alinearlos con las estructuras de grafos observadas durante el entrenamiento. Nuestra segunda estrategia mejora las representaciones para sistemas fuera de distribución en tiempo de prueba mediante la toma de pasos de gradiente utilizando un objetivo auxiliar, como un prior físico de bajo costo. Nuestras estrategias de refinamiento en tiempo de prueba reducen significativamente los errores en sistemas fuera de distribución, sugiriendo que los MLFFs son capaces de modelar espacios químicos diversos y pueden avanzar en esa dirección, pero no están siendo entrenados de manera efectiva para lograrlo. Nuestros experimentos establecen puntos de referencia claros para evaluar las capacidades de generalización de la próxima generación de MLFFs. Nuestro código está disponible en https://tkreiman.github.io/projects/mlff_distribution_shifts/.

English

Machine Learning Force Fields (MLFFs) are a promising alternative to expensive ab initio quantum mechanical molecular simulations. Given the diversity of chemical spaces that are of interest and the cost of generating new data, it is important to understand how MLFFs generalize beyond their training distributions. In order to characterize and better understand distribution shifts in MLFFs, we conduct diagnostic experiments on chemical datasets, revealing common shifts that pose significant challenges, even for large foundation models trained on extensive data. Based on these observations, we hypothesize that current supervised training methods inadequately regularize MLFFs, resulting in overfitting and learning poor representations of out-of-distribution systems. We then propose two new methods as initial steps for mitigating distribution shifts for MLFFs. Our methods focus on test-time refinement strategies that incur minimal computational cost and do not use expensive ab initio reference labels. The first strategy, based on spectral graph theory, modifies the edges of test graphs to align with graph structures seen during training. Our second strategy improves representations for out-of-distribution systems at test-time by taking gradient steps using an auxiliary objective, such as a cheap physical prior. Our test-time refinement strategies significantly reduce errors on out-of-distribution systems, suggesting that MLFFs are capable of and can move towards modeling diverse chemical spaces, but are not being effectively trained to do so. Our experiments establish clear benchmarks for evaluating the generalization capabilities of the next generation of MLFFs. Our code is available at https://tkreiman.github.io/projects/mlff_distribution_shifts/.