Towards Fast, Specialized Machine Learning Force Fields: Distilling Foundation Models via Energy Hessians

• URL: http://arxiv.org/abs/2501.09009v2
• Date: Fri, 31 Jan 2025 18:12:36 GMT
• Title: Towards Fast, Specialized Machine Learning Force Fields: Distilling Foundation Models via Energy Hessians
• Authors: Ishan Amin, Sanjeev Raja, Aditi Krishnapriyan,
• Abstract summary: We introduce a method for transferring general-purpose representations from MLFF foundation models to smaller, faster MLFFs specialized to specific regions of chemical space. Our specialized MLFFs can be up to 20 $times$ faster than the original foundation model, while retaining, and in some cases, exceeding, its performance and that of undistilled models.
• Score: 0.0
• License:
• Abstract: The foundation model (FM) paradigm is transforming Machine Learning Force Fields (MLFFs), leveraging general-purpose representations and scalable training to perform a variety of computational chemistry tasks. Although MLFF FMs have begun to close the accuracy gap relative to first-principles methods, there is still a strong need for faster inference speed. Additionally, while research is increasingly focused on general-purpose models which transfer across chemical space, practitioners typically only study a small subset of systems at a given time. This underscores the need for fast, specialized MLFFs relevant to specific downstream applications, which preserve test-time physical soundness while maintaining train-time scalability. In this work, we introduce a method for transferring general-purpose representations from MLFF foundation models to smaller, faster MLFFs specialized to specific regions of chemical space. We formulate our approach as a knowledge distillation procedure, where the smaller "student" MLFF is trained to match the Hessians of the energy predictions of the "teacher" foundation model. Our specialized MLFFs can be up to 20 $\times$ faster than the original foundation model, while retaining, and in some cases exceeding, its performance and that of undistilled models. We also show that distilling from a teacher model with a direct force parameterization into a student model trained with conservative forces (i.e., computed as derivatives of the potential energy) successfully leverages the representations from the large-scale teacher for improved accuracy, while maintaining energy conservation during test-time molecular dynamics simulations. More broadly, our work suggests a new paradigm for MLFF development, in which foundation models are released along with smaller, specialized simulation "engines" for common chemical subsets.

Related papers

• Physics-Guided Foundation Model for Scientific Discovery: An Application to Aquatic Science [13.28811382673697]
  We propose a textittextbfPhysics-textbfGuided textbfFoundation textbfModel (textbfPGFM) that combines pre-trained ML models and physics-based models. We demonstrate the effectiveness of this methodology in modeling water temperature and dissolved oxygen dynamics in real-world lakes.
  arXiv  Detail & Related papers  (2025-02-10T00:48:10Z)
• BoostMD: Accelerating molecular sampling by leveraging ML force field features from previous time-steps [3.8214695776749013]
  BoostMD is a surrogate model architecture designed to accelerate molecular dynamics simulations. Our experiments demonstrate that BoostMD achieves an eight-fold speedup compared to the reference model. By combining efficient feature reuse with a streamlined architecture, BoostMD offers a robust solution for conducting large-scale, long-timescale molecular simulations.
  arXiv  Detail & Related papers  (2024-12-21T20:52:36Z)
• Specialized Foundation Models Struggle to Beat Supervised Baselines [60.23386520331143]
  We look at three modalities -- genomics, satellite imaging, and time series -- with multiple recent FMs and compare them to a standard supervised learning workflow. We find that it is consistently possible to train simple supervised models that match or even outperform the latest foundation models.
  arXiv  Detail & Related papers  (2024-11-05T04:10:59Z)
• Recent Advances on Machine Learning for Computational Fluid Dynamics: A Survey [51.87875066383221]
  This paper introduces fundamental concepts, traditional methods, and benchmark datasets, then examine the various roles Machine Learning plays in improving CFD. We highlight real-world applications of ML for CFD in critical scientific and engineering disciplines, including aerodynamics, combustion, atmosphere & ocean science, biology fluid, plasma, symbolic regression, and reduced order modeling. We draw the conclusion that ML is poised to significantly transform CFD research by enhancing simulation accuracy, reducing computational time, and enabling more complex analyses of fluid dynamics.
  arXiv  Detail & Related papers  (2024-08-22T07:33:11Z)
• Masked Particle Modeling on Sets: Towards Self-Supervised High Energy Physics Foundation Models [4.299997052226609]
  Masked particle modeling (MPM) is a self-supervised method for learning generic, transferable, and reusable representations on unordered sets of inputs. We study the efficacy of the method in samples of high energy jets at collider physics experiments.
  arXiv  Detail & Related papers  (2024-01-24T15:46:32Z)
• A Survey on Efficient Federated Learning Methods for Foundation Model Training [62.473245910234304]
  Federated Learning (FL) has become an established technique to facilitate privacy-preserving collaborative training across a multitude of clients. In the wake of Foundation Models (FM), the reality is different for many deep learning applications. We discuss the benefits and drawbacks of parameter-efficient fine-tuning (PEFT) for FL applications.
  arXiv  Detail & Related papers  (2024-01-09T10:22:23Z)
• Improving and generalizing flow-based generative models with minibatch optimal transport [90.01613198337833]
  We introduce the generalized conditional flow matching (CFM) technique for continuous normalizing flows (CNFs) CFM features a stable regression objective like that used to train the flow in diffusion models but enjoys the efficient inference of deterministic flow models. A variant of our objective is optimal transport CFM (OT-CFM), which creates simpler flows that are more stable to train and lead to faster inference.
  arXiv  Detail & Related papers  (2023-02-01T14:47:17Z)
• Forces are not Enough: Benchmark and Critical Evaluation for Machine Learning Force Fields with Molecular Simulations [5.138982355658199]
  Molecular dynamics (MD) simulation techniques are widely used for various natural science applications. We benchmark a collection of state-of-the-art (SOTA) ML FF models and illustrate, in particular, how the commonly benchmarked force accuracy is not well aligned with relevant simulation metrics.
  arXiv  Detail & Related papers  (2022-10-13T17:59:03Z)
• Self-Feature Regularization: Self-Feature Distillation Without Teacher Models [0.0]
  Self-Feature Regularization(SFR) is proposed, which uses features in the deep layers to supervise feature learning in the shallow layers. We firstly use generalization-l2 loss to match local features and a many-to-one approach to distill more intensively in the channel dimension.
  arXiv  Detail & Related papers  (2021-03-12T15:29:00Z)
• Learning Discrete Energy-based Models via Auxiliary-variable Local Exploration [130.89746032163106]
  We propose ALOE, a new algorithm for learning conditional and unconditional EBMs for discrete structured data. We show that the energy function and sampler can be trained efficiently via a new variational form of power iteration. We present an energy model guided fuzzer for software testing that achieves comparable performance to well engineered fuzzing engines like libfuzzer.
  arXiv  Detail & Related papers  (2020-11-10T19:31:29Z)
• Machine Learning Force Fields [54.48599172620472]
  Machine Learning (ML) has enabled numerous advances in computational chemistry. One of the most promising applications is the construction of ML-based force fields (FFs) This review gives an overview of applications of ML-FFs and the chemical insights that can be obtained from them.
  arXiv  Detail & Related papers  (2020-10-14T13:14:14Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.