Compreendendo e Mitigando Deslocamentos de Distribuição para Aprendizado de Máquina Campos de Força

Understanding and Mitigating Distribution Shifts For Machine Learning Force Fields

March 11, 2025

Autores: Tobias Kreiman, Aditi S. Krishnapriyan

cs.AI

Resumo

Campos de Força de Aprendizado de Máquina (MLFFs) são uma alternativa promissora às caras simulações moleculares quânticas ab initio. Dada a diversidade de espaços químicos de interesse e o custo de gerar novos dados, é importante entender como os MLFFs generalizam além de suas distribuições de treinamento. Para caracterizar e compreender melhor as mudanças de distribuição em MLFFs, realizamos experimentos diagnósticos em conjuntos de dados químicos, revelando mudanças comuns que representam desafios significativos, mesmo para modelos de base grandes treinados com dados extensos. Com base nessas observações, hipotetizamos que os métodos atuais de treinamento supervisionado regularizam inadequadamente os MLFFs, resultando em overfitting e aprendizado de representações pobres para sistemas fora da distribuição. Em seguida, propomos dois novos métodos como passos iniciais para mitigar as mudanças de distribuição em MLFFs. Nossos métodos focam em estratégias de refinamento em tempo de teste que incorrem em custo computacional mínimo e não utilizam rótulos de referência ab initio caros. A primeira estratégia, baseada na teoria espectral de grafos, modifica as arestas dos grafos de teste para alinhá-las com estruturas de grafos vistas durante o treinamento. Nossa segunda estratégia melhora as representações para sistemas fora da distribuição em tempo de teste ao dar passos de gradiente usando um objetivo auxiliar, como um prior físico de baixo custo. Nossas estratégias de refinamento em tempo de teste reduzem significativamente os erros em sistemas fora da distribuição, sugerindo que os MLFFs são capazes de e podem avançar na modelagem de espaços químicos diversos, mas não estão sendo efetivamente treinados para isso. Nossos experimentos estabelecem benchmarks claros para avaliar as capacidades de generalização da próxima geração de MLFFs. Nosso código está disponível em 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/.