Comprensione e Mitigazione degli Spostamenti di Distribuzione per il Machine Learning Campi di Forza

Understanding and Mitigating Distribution Shifts For Machine Learning Force Fields

March 11, 2025

Autori: Tobias Kreiman, Aditi S. Krishnapriyan

cs.AI

Abstract

I Campi di Forza basati su Machine Learning (MLFFs) rappresentano un'alternativa promettente alle costose simulazioni molecolari quantomeccaniche ab initio. Data la diversità degli spazi chimici di interesse e il costo associato alla generazione di nuovi dati, è fondamentale comprendere come gli MLFFs generalizzino al di là delle loro distribuzioni di addestramento. Per caratterizzare e comprendere meglio gli spostamenti di distribuzione negli MLFFs, conduciamo esperimenti diagnostici su dataset chimici, rivelando spostamenti comuni che pongono sfide significative, anche per modelli di base di grandi dimensioni addestrati su dati estesi. Sulla base di queste osservazioni, ipotizziamo che i metodi di addestramento supervisionati attuali regolarizzino in modo inadeguato gli MLFFs, portando a overfitting e all'apprendimento di rappresentazioni scarse per sistemi fuori distribuzione. Proponiamo quindi due nuovi metodi come passi iniziali per mitigare gli spostamenti di distribuzione negli MLFFs. I nostri metodi si concentrano su strategie di raffinamento al momento del test che comportano un costo computazionale minimo e non utilizzano costose etichette di riferimento ab initio. La prima strategia, basata sulla teoria spettrale dei grafi, modifica gli archi dei grafi di test per allinearli alle strutture dei grafi osservate durante l'addestramento. La nostra seconda strategia migliora le rappresentazioni per sistemi fuori distribuzione al momento del test eseguendo passi di gradiente utilizzando un obiettivo ausiliario, come un precedente fisico a basso costo. Le nostre strategie di raffinamento al momento del test riducono significativamente gli errori sui sistemi fuori distribuzione, suggerendo che gli MLFFs sono in grado di e possono avvicinarsi alla modellazione di spazi chimici diversi, ma non vengono addestrati efficacemente per farlo. I nostri esperimenti stabiliscono benchmark chiari per valutare le capacità di generalizzazione della prossima generazione di MLFFs. Il nostro codice è disponibile all'indirizzo 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/.