Понимание и смягчение сдвигов распределения в машинном обучении Силовые поля

Understanding and Mitigating Distribution Shifts For Machine Learning Force Fields

March 11, 2025

Авторы: Tobias Kreiman, Aditi S. Krishnapriyan

cs.AI

Аннотация

Машинное обучение силовых полей (MLFF) представляет собой перспективную альтернативу дорогостоящим квантово-механическим молекулярным симуляциям ab initio. Учитывая разнообразие химических пространств, представляющих интерес, и стоимость генерации новых данных, важно понимать, как MLFF обобщаются за пределы своих обучающих распределений. Чтобы охарактеризовать и лучше понять сдвиги распределений в MLFF, мы проводим диагностические эксперименты на химических наборах данных, выявляя общие сдвиги, которые представляют значительные трудности даже для крупных базовых моделей, обученных на обширных данных. На основе этих наблюдений мы выдвигаем гипотезу, что текущие методы обучения с учителем недостаточно регуляризуют MLFF, что приводит к переобучению и плохому представлению систем, выходящих за пределы распределения. Затем мы предлагаем два новых метода в качестве начальных шагов для смягчения сдвигов распределений в MLFF. Наши методы сосредоточены на стратегиях уточнения во время тестирования, которые требуют минимальных вычислительных затрат и не используют дорогостоящие эталонные метки ab initio. Первая стратегия, основанная на спектральной теории графов, изменяет ребра тестовых графов, чтобы они соответствовали структурам графов, наблюдаемым во время обучения. Наша вторая стратегия улучшает представления для систем, выходящих за пределы распределения, во время тестирования, делая шаги градиента с использованием вспомогательной цели, такой как дешевый физический априор. Наши стратегии уточнения во время тестирования значительно снижают ошибки на системах, выходящих за пределы распределения, что позволяет предположить, что MLFF способны и могут двигаться в сторону моделирования разнообразных химических пространств, но не обучаются эффективно для этого. Наши эксперименты устанавливают четкие ориентиры для оценки обобщающих способностей следующего поколения MLFF. Наш код доступен по адресу 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/.