Verständnis und Minderung von Verteilungsverschiebungen im maschinellen Lernen Kraftfelder

Understanding and Mitigating Distribution Shifts For Machine Learning Force Fields

March 11, 2025

Autoren: Tobias Kreiman, Aditi S. Krishnapriyan

cs.AI

Zusammenfassung

Machine Learning Force Fields (MLFFs) stellen eine vielversprechende Alternative zu teuren ab-initio-Quantenmechanik-Molekülsimulationen dar. Angesichts der Vielfalt der chemischen Räume, die von Interesse sind, und der Kosten für die Generierung neuer Daten, ist es wichtig zu verstehen, wie MLFFs über ihre Trainingsverteilungen hinaus generalisieren. Um Verschiebungen in den Verteilungen von MLFFs zu charakterisieren und besser zu verstehen, führen wir diagnostische Experimente mit chemischen Datensätzen durch, die häufige Verschiebungen aufdecken, die erhebliche Herausforderungen darstellen, selbst für große Foundation-Modelle, die mit umfangreichen Daten trainiert wurden. Basierend auf diesen Beobachtungen stellen wir die Hypothese auf, dass aktuelle überwachte Trainingsmethoden MLFFs unzureichend regularisieren, was zu Overfitting und schlechten Repräsentationen von Out-of-Distribution-Systemen führt. Wir schlagen dann zwei neue Methoden als erste Schritte zur Minderung von Verteilungsverschiebungen für MLFFs vor. Unsere Methoden konzentrieren sich auf Testzeit-Verfeinerungsstrategien, die mit minimalem Rechenaufwand verbunden sind und keine teuren ab-initio-Referenzlabels verwenden. Die erste Strategie, basierend auf der spektralen Graphentheorie, modifiziert die Kanten von Testgraphen, um sie mit den während des Trainings gesehenen Graphenstrukturen in Einklang zu bringen. Unsere zweite Strategie verbessert die Repräsentationen für Out-of-Distribution-Systeme zur Testzeit, indem Gradientenschritte unter Verwendung eines Hilfsziels, wie eines kostengünstigen physikalischen Priors, durchgeführt werden. Unsere Testzeit-Verfeinerungsstrategien reduzieren die Fehler bei Out-of-Distribution-Systemen erheblich, was darauf hindeutet, dass MLFFs in der Lage sind und sich darauf zubewegen können, diverse chemische Räume zu modellieren, aber nicht effektiv dafür trainiert werden. Unsere Experimente etablieren klare Benchmarks für die Bewertung der Generalisierungsfähigkeiten der nächsten Generation von MLFFs. Unser Code ist verfügbar unter 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/.