Comprendre et atténuer les décalages de distribution pour l'apprentissage automatique Champs de force

Understanding and Mitigating Distribution Shifts For Machine Learning Force Fields

March 11, 2025

Auteurs: Tobias Kreiman, Aditi S. Krishnapriyan

cs.AI

Résumé

Les champs de force par apprentissage automatique (MLFFs) constituent une alternative prometteuse aux simulations moléculaires quantiques ab initio coûteuses. Compte tenu de la diversité des espaces chimiques d'intérêt et du coût lié à la génération de nouvelles données, il est important de comprendre comment les MLFFs généralisent au-delà de leurs distributions d'entraînement. Afin de caractériser et de mieux comprendre les décalages de distribution dans les MLFFs, nous menons des expériences diagnostiques sur des ensembles de données chimiques, révélant des décalages courants qui posent des défis significatifs, même pour les grands modèles de base entraînés sur des données étendues. Sur la base de ces observations, nous émettons l'hypothèse que les méthodes actuelles d'entraînement supervisé régularisent insuffisamment les MLFFs, entraînant un surajustement et l'apprentissage de représentations médiocres pour les systèmes hors distribution. Nous proposons ensuite deux nouvelles méthodes comme premières étapes pour atténuer les décalages de distribution dans les MLFFs. Nos méthodes se concentrent sur des stratégies de raffinement au moment du test qui engendrent un coût computationnel minimal et n'utilisent pas d'étiquettes de référence ab initio coûteuses. La première stratégie, basée sur la théorie des graphes spectraux, modifie les arêtes des graphes de test pour les aligner avec les structures de graphes observées pendant l'entraînement. Notre deuxième stratégie améliore les représentations pour les systèmes hors distribution au moment du test en effectuant des pas de gradient à l'aide d'un objectif auxiliaire, tel qu'un a priori physique peu coûteux. Nos stratégies de raffinement au moment du test réduisent significativement les erreurs sur les systèmes hors distribution, suggérant que les MLFFs sont capables de modéliser des espaces chimiques divers et peuvent progresser dans cette direction, mais ne sont pas efficacement entraînés pour le faire. Nos expériences établissent des références claires pour évaluer les capacités de généralisation de la prochaine génération de MLFFs. Notre code est disponible à l'adresse suivante : 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/.