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

Authors: Ishan Amin, Sanjeev Raja, Aditi S. Krishnapriyan
Primary area: applications to physical sciences (physics, chemistry, biology, etc.)
Keywords: machine learning force fields, graph neural networks, knowledge distillation

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. At test time, MLFFs must also obey physical constraints unique to the downstream use case, such as energy conservation for molecular dynamics simulations. 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 an architecture-agnostic knowledge distillation procedure, where the smaller "student" MLFF is trained to match the Hessians of the energy predictions of the "teacher" foundation model. We demonstrate our approach across multiple recent foundation models, large-scale datasets, chemical subsets, and downstream tasks. 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. The implementation of our method is available at https://github.com/ASK-Berkeley/MLFF-distill.

Decision: accept (poster)
Reviewer scores: 8, 6, 6, 8

Reviews and discussion:
<thread>
creation date:
1729941507041

thread_type:
official_review

<thread_object>
creator: reviewer_8tdr

<content>
summary:
This paper presents a novel technique for knowledge distillation from foundation model force fields to smaller, faster, and more specialized force fields. The core idea is to align the Hessians of the energies with respect to atomic positions between the teacher (foundation model) and student models, facilitating efficient knowledge transfer. Experiments demonstrate that this approach improves stability and accuracy in the derived force fields compared to simpler distillation methods.

strengths:
1. Clarity and Readability: The paper is very well-written and accessible, making the proposed method easy to follow.
2. Practicality: The approach is straightforward to implement and is cost-effective relative to similar knowledge distillation methods.
3. Experimental Validation: MD simulations performed validate the method, underscoring the practical benefits of the proposed technique.
4. Implementation Insights: The paper also offers practical implementation guidance for practitioners, which is particularly useful for real-world applications.

weaknesses:
The baseline methods used for comparison seem to perform notably poorly, raising questions about fairness. This might be genuine, but the absence of hyperparameter tuning for the baselines undermines this. The authors specifically tuned their method by adjusting the Hessian distillation loss term, "we reduce the weight of the Hessian distillation loss term, λKD, by a factor of 4 during training once the student’s validation loss on the original objective, LEF (φ), becomes lower than that of the frozen FM, LEF (ψ)" (l.207-209). Further clarification on the role of this schedule, as well as the sensitivity to λKD, would be helpful. Would similar tuning or scheduling benefit the alternative approaches as well?

questions:
* Why is there a tank in Figure 1?
* While the Hessian has a nice physical interpretation, as the authors point out, do higher derivates improve transfer further? One way to improve runtime in such cases would be to use Forward on Backward differentiation.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730122240066

thread_type:
official_review

<thread_object>
creator: reviewer_yqhs

<content>
summary:
This paper introduces a method to distillate MLFF foundation models to smaller, faster MLFFs with energy Hessians, achieving significant speed improvements while maintaining performance.

strengths:
* This paper proposes a new method to distill MLFF foundation models into smaller, faster MLFFs with Hessians, which is highly beneficial for simulating realistic systems.
* The paper is written in a clear and concise manner, facilitating effortless understanding.

weaknesses:
Related concerns are discussed in the questions section.

questions:
* The tables in the paper show only force results, without energy, so I'm curious about the energy results after distillation.
* A major concern is that the primary use of the distilled MLFF model is for molecular dynamics simulations, where conservation properties are crucial for scientists in physics, chemistry, biology, and materials science. I understand the authors avoided second-order derivatives to calculate the Hessian by directly predicting forces, using JVP calculations. However, a pretrained model might predict forces directly to save computation due to its large size, but the student model should compute forces using autograd, similar in the SFT in JMP[1], which makes more sense. Although Fig. 2 shows stable molecular dynamics in the NVT ensemble, following [2], energy will not be conserved in the NVE ensemble.
* The paper claims that the student model outperforms the teacher model, which is confusing. I suspect this is because the energy and force labels used in training come from the dataset itself. While the inclusion of Hessian loss is shown to be better than using only energy and force loss, this highlights the importance of derivatives. Since the Hessian matrix introduces force derivatives, could training a traditional MLFF from scratch, with forces computed via autograd, achieve similar or better results? Additionally, the statement in Fig. 3b about "speculating that s may play a similar role as the batch size" is akin to conclusions from traditional MLFF training, suggesting that direct autograd training without Hessian distillation might yield similar outcomes. The authors could compare such models to illustrate the Hessian's impact.
* Regarding the appendix experiment using force for distillation, what is the specific loss function? If I'm correct in understanding that the Hessian term in Eq. 3 is replaced by the force term, could the poor results from force distillation be due to the inherent force label loss in the data, where the teacher model's force predictions contradict the data labels? This implies that the force distillation setup might be flawed.

[1] Shoghi N, Kolluru A, Kitchin J R, et al. From Molecules to Materials: Pre-training Large Generalizable Models for Atomic Property Prediction[C]//The Twelfth International Conference on Learning Representations.

[2] Fu X, Wu Z, Wang W, et al. Forces are not Enough: Benchmark and Critical Evaluation for Machine Learning Force Fields with Molecular Simulations[J]. Transactions on Machine Learning Research.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730570801501

thread_type:
official_review

<thread_object>
creator: reviewer_8nsg

<content>
summary:
This paper introduces a method for transferring general-purpose representations from large ML force field (MLFF) foundation models to smaller, faster MLFFs specialized for specific regions of chemical space, with the aim of improving inference speed. The approach is formulated as a knowledge distillation (KD) process, where the smaller “student” MLFF learns to match the Hessians of energy predictions made by the “teacher” foundation model. By selectively subsampling rows of the Hessian corresponding to individual atomic coordinates, the “student” MLFF achieves a training process that is computationally efficient relative to foundation models and demonstrates improved force prediction on downstream datasets.

strengths:
* The approach is well-motivated, leveraging knowledge from large MLFF foundation models and adapting it to specific chemical space regions using knowledge distillation. The method also achieves promising results across organic and material molecular systems.
* The paper is well-organized and easy to follow.

weaknesses:
* Although the method shows promising results, the rationale for using Hessian information as a distillation signal is unclear, which may impact the perceived technical contribution. Additional theoretical or intuitive insights on this choice would clarify the method’s grounding.
* Foundation models are often trained with energy and force supervision, possibly derived from various electronic structure methods, making the physical reliability of Hessians from pre-trained foundation models questionable.
	* Notably, the authors mention that some student models outperform foundation models in specialized chemical spaces even before distillation, suggesting that foundation models may not fully converge in certain cases. This raises questions about the significance and reliability of using Hessians from foundation models as distillation targets.
	* Accurate Hessians are crucial for tasks like geometry optimization (as referenced in Figure 1). It remains unclear how potentially inaccurate Hessians from foundation models could affect the student model's performance in such applications.
* It is uncertain whether the proposed method can be extended to MLFF architectures designed with energy-conserving forces or high-order equivariance, which are often crucial factors for stable and transferable ML force fields. Discussing the impact of these inductive biases on the Hessian-driven KD approach would strengthen the work.

questions:
See “Weaknesses” for detailed questions and suggestions.

Potential typos:
* Line 420: “disilling” should be “distilling.”
</content>
</thread_object>

</thread>
<thread>
creation date:
1730716291022

thread_type:
official_review

<thread_object>
creator: reviewer_nwzu

<content>
summary:
This work proposes a new method for improving the efficiency of Machine Learning Force Fields (MLFFs) by utilizing a knowledge distillation technique. The method distills knowledge from large, general-purpose foundation MLFF models (FMs) into smaller, faster MLFFs specialized for specific regions of chemical space. This is accomplished by aligning the energy Hessians, which are the second derivatives of the energy with respect to atomic positions, between the teacher FM and the student MLFF. By strategically subsampling rows of the Hessian, the authors significantly reduce the computational cost of the distillation process. The authors demonstrate that their approach can achieve speedups of up to 20 times compared to the original FMs while retaining, and in some cases exceeding, the accuracy of the foundation models.

strengths:
- The distilled MLFFs are up to 20x faster than their foundation model counterparts, enabling more efficient simulations.
- The distilled MLFFs achieve comparable or even better force prediction accuracy than the original FMs and demonstrate improved MD stability results.
- The Hessian distillation method is model architecture agnostic.
- Subsampling hessian rows significantly reduces computational costs without sacrificing performance. Its also interesting that subsampling quality doesn't impact the performance much.

weaknesses:
- Training with hessian distillation increases the computational cost compared to undistilled training.
- An anonymous link to the code is not available.

questions:
- Some models use conservative forces and some don't. Do you have a sense of how much that impacts when you distill Hessians from a non-conservative model instead or if the student model is conservative? Or do you expect that not to impact performance as it seems like hessian quality doesn't?
- Why are the rows sampled so different for GemNet across training set and the same for PaiNN? What would be the general suggestion to set this hyperparameter? Or should everyone be iterating on this for every dataset, model architecture, etc?
- Why don't you compare the results with Kelvinius et al. 2023 on OC20-2M and COLL on which the original work was performed?
</content>
</thread_object>

</thread>
<thread>
creation date:
1732134179679

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
General Response

comment:
We thank the reviewers for their helpful comments on our work, which introduces a new Hessian distillation approach to produce fast, specialized machine learning force fields (MLFFs) from large, general-purpose foundation models (FMs). We appreciate that the reviewers found the paper to be clear, well-written, and showed strong results. Below, we highlight the major improvements and changes we have made to the paper. We have also uploaded a new version of the paper with changes indicated in blue text.

1. **Distilling from a non-conservative teacher model to a conservative or higher-order equivariant student model.** Several reviewers were interested in how our Hessian distillation approach performs when the student model uses gradient-based/conservative forces, while the teacher model uses direct/non-conservative forces. We have addressed this by adding a Section (4.3) to our paper focusing on JMP [1], a large, non-conservative foundation model, as the teacher model, with conservative GemNet-T models as the students. We achieve strong improvements in force MAE compared to undistilled models on the buckyball catcher and double walled nanotube molecules from MD22. These are molecules which are sufficiently large that running the original JMP foundation model with conservative forces is memory-prohibitive on a standard GPU. Crucially, when we run the distilled models in constant energy (NVE) MD simulations, the JMP model’s energy gradually drifts, while our GemNet-T student model conserves energy by design. Although we found the expense of distilling conservative student models to be manageable in practice, it is possible to accelerate training using finite differences (see point 4 below). We also demonstrate strong results when using eSCN [2], which employs higher-order (l=2) equivariance, as the student model. This clearly suggests that our approach is useful in the relevant setting where the student models have more inductive biases than the teacher FM. 

2. **Energy results.** In addition to showing results on Force MAE, we have now added results on Energy MAE for the SPICE dataset (Table 1). We find that our Hessian distillation approach, along with a new loss term involving the gradient of the energy prediction head of non-conservative models (defined in Section 3.4), yields strong improvements on Energy MAE relative to undistilled student MLFFs and our n2n and a2a baselines. We also provide an ablation study in the Appendix (Section A.11) to show that pure Hessian distillation without this additional loss term also leads to improvements in energy MAE.

3. **Geometry optimization**. As another demonstration of the downstream usefulness of our method in realistic use cases, we perform geometry optimization with our undistilled and distilled student models on selected systems from the SPICE dataset, and find that our distilled models converge to structures with lower energies and force norms (evaluated with DFT) than their undistilled counterparts (Section A.13).

4. **Accelerating Hessian computation with finite differences**. We show in Appendix Section A.9 that we can replace auto-differentiation with finite differences to accelerate Hessian computation by approximately 40% for gradient-based student models.

5. **Hyperparameter sweeps/baseline tuning.** We have performed sweeps over selected hyperparameters, including the knowledge distillation weight and loss scheduling, for our Hessian distillation method and the baselines we compare to. This has improved the performance of our baselines, making them in line with results reported in [3] but still considerably inferior to our approach. 

[1] Shoghi, Nima, et al. "From molecules to materials: Pre-training large generalizable models for atomic property prediction." International Conference on Learning Representations, 2024.

[2] Passaro, Saro, and C. Lawrence Zitnick. "Reducing SO (3) convolutions to SO (2) for efficient equivariant GNNs." International Conference on Machine Learning. PMLR, 2023.

[3] Ekström Kelvinius, Filip, et al. "Accelerating molecular graph neural networks via knowledge distillation." Advances in Neural Information Processing Systems 36 (2024).
</content>
</thread_object>

</thread>


Metareview:
metareview:
The submission presents a distillation method for machine learning force field (MLFF). The notable feature of the method is that the guidance from the teacher MLFF is provided in the form of the energy Hessian. A few dedicated techniques are introduced for better efficiency, e.g., row subsampling and the Jacobian-vector product formulation. I'm glad to see a work that considers a practical and promising situation of distilling a specialized, light-weighted MLFF from an MLFF foundation model, and noting the value of the energy Hessian. The reviewers also agree on this contribution.

Reviewers raised a few concerns and insufficiencies of the paper:
* The common ask is the effectiveness of the method when the inductive bias does not match: using a non-conservative MLFF teacher models and student model that is conservative. The authors have provided additional results in such cases, which appropriately demonstrate the effectiveness.
* The authors provided more details on hyperparameter choices (in response to Reviewers 8TDr and NWzu) for the consideration to support that the better results do not come from unfair tuning on baselines. This seems reasonably convincing.
* Reviewer 8NSg challenged the reliability of the Hessian of the teacher model. In response, the authors mentioned a related empirical evidence in one case, and refer to scaling law to relief the problem, which seem to have satisfied the reviewer.
* The authors provided energy evaluation results in one case in response to Reviewer yqHS, which are supportive.
* Reviewer yqHS challenged the effectiveness of the method under the spirit of derivative supervision by asking for a comparison with autograd force MLFF model which already has its derivatives supervised. The authors provided supportive results for the proposed Hessian distillation method in one case.
* The authors provided details on additional computational cost in response to Reviewer NWzu, and the results seem acceptable.
* I came up with the question that, while the energy Hessian does hold physical relevance, it is also relevant to provide interpretations on why it is also a good choice for distilling from a teacher MLFF. In the rebuttal to Reviewer 8NSg, they provided additional interpretations in terms of Sobolev training, which sounds a reasonable supplement.

After the rebuttal period, all reviewers increased their scores, and all the scores are positive. This suggests that all the major concerns and insufficiencies are addressed. I hence recommend accepting the paper. To make the paper stronger, I would suggest the followings:
* Regarding the reliability of the Hessian of the teacher model, the authors may consider more discussions since the Hessian of a neural network could still largely depend on a specific architecture: the Hessian of the teacher model is neither trained by data, and certain architectures may restrict the expressiveness of Hessian. The authors may consider discussing what kind of architectures can be considered.
* The authors may be interested in discussing the relation to distillation methods that align the latent representation of the two models.
* Reviewer yqHS asked about the setting for the comparison with teacher force distillation. In the rebuttal, although it is clear that there is no conflict in the force labels, but this setting seems to lose the ground-truth force supervision. The authors may consider further clarifying the setting in the section or table caption, and how it is a fair comparison with Hessian distillation.
* Some issues asked by reviewers are only demonstrated in one case. Hope the authors could provide general explanations on these issues, or provide results in more cases, if possible.