Towards Fast, Specialized Machine Learning Force Fields: Distilling Foundation Models via Energy Hessians
Authors: Ishan Amin
Sanjeev Raja
Aditi S. Krishnapriyan
Venue: ICLR 2025 Poster (2025)
Keywords: machine learning force fields
graph neural networks
knowledge distillation
Primary area: applications to physical sciences (physics, chemistry, biology, etc.)
TL;DR: We distill large machine learning force field foundation models into small, specialized models using knowledge distillation from energy Hessians.

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.

Published as a conference paper at ICLR 2025
TOWARDS FAST, SPECIALIZED MACHINE LEARNING FORCE FIELDS: DISTILLING FOUNDATION MODELS VIA ENERGY HESSIANS
Ishan Amin," Sanjeev Raja", and Aditi S. Krishnapriyan University of Califoriia, Berkeley; LBNL (ishanthewizard) sanjeevr, aditikk (bkkkkeeereeee.edd
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 close the accuracy gap relative to first-principls methods, there is still a strong need for faster inference speed. Additionally, while research 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, MLFFf must also obey physical constraints unique to the downstream for molecular dynamics simulations. This underscores the need for fast, specialized MLFFs relevant specific downstream applications, which preserve test-time physical soundness while maintaining train-time scalability. In this work, we introduuce a method for transferring general-purposse representations from MLFF foundation models smaller, faster MLFFs specialized to specific regions of chemical space. We formulate our approach as an architecture-agnoticc 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, lcage-scale datasets, chemical subsets, and downstream tasks. Our specialized MIFFFs can be up to x 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 student model trained with conservative forces (i.e., computed as derivatives of the potential energy) successfully leverages the representations from the targe-scale teacher for improved accuracy, while maintaining energy conservation during test-time molecular dynamics simulations. More broadly, our work suggests a new paradigm for MIFF development, in which foundation models are released along with smaller, specialized simulation "engines" for common chemical subsets. The implementation of our method is available https://gihhub.com/AAKKK Bttp://bb..oooooBBeeeeeLeLL..... Berkeley/MLFF-dissil
case, such energy conservation
INTRODUCTION
Quantum chemical calculations, such as Density Functional Theory (DFT), underpin a broad range of
applications in computational chemistry, including the discovery of new drugs (Cole & Hine 2016).
materials (Hafner et al. 2006 Jain et al. 20166). and catalysts (Hammer & Norskov 2000). ) Machine
learning force fields (MLFFs) Gasteiger et al. a [22T Batzner al., [222 Musaelian et al. 2022 Batatia
et al. 2022) based on graph neural network (GNN) architectures (Gilmer et al. (017) have recently shown tremendous potential to serve as fast surrogates for these quantum mechanical calculations.
Foundation models (FMs) are general seneral-purpose purpose models trained on large quantities of data, with the ability to generalize to many downstream tasks with little to no fine-tuning. Mirroring advancements in the fields of natural language processing (Achiam et al. (023) and computer vision (Radford et al.) [021] Oquab) ctt al., 2023). the availability of increasingly large and diverse datasets of quantum chemical calculations (Chanussot et al.] 2021 Jain et al. 2020 Eastman et al.) (023)) has enabled the creation of MLFF FMs (Kovács et al.] 2023 [atatii et al. 2023 Shoghi et al. (023). While earlier MLFFss typically trained on relatively narrow datasets (Schütt et al. 2018 Chmiela et all (017), , MLFF FMs trained across a broad
swath chemical space, aiming to perform well across a diverse range of atomic property prediction tasks.
'Equal Contribution.

---

Published as a conference paper at ICLR 2025
Biomolecules Materisla
Catelysts
Specialized MLFFs
Downstream Applications
(Fast)
MD Simulation
Pre-train
Amnu
Adids
Energy Distillation
Foundation Model (Slow)
Perovsxiteg
Geometry Optimization
C Des AKDLEE Kr
Metallolds
Training Data Subset
Teacher
Student Hessian (Rows Computed On-the-Fly)
Free Energy Calculations
Hessian
Subsumple Rows
True Foroes (Procomputed)
Teacher Hessians
Figure 1: Proposed Hessian distillation schematic. In our proposed distillation approach, we start with a machine learning force field (MLFF) foundation model (FM) that has been trained on a large quantity of diverse data. We precompute energy Hessians of the FM over specialized data subset. We then train series of smaller MLFFs on these subsets via our knowledge distillation loss (CKD) which aligns selected rows of the energy Hessian of the smaller (student) models with those of the FM (teacher). We also keep the conventional procedure of training on the ground truth energies and forces (LEF) from the specialized subset. The resulting MLFFs are considerably faster than the FM and can be efficiently used in downstream applications such as MD simulation, geometry optimization, and free energy calculations.
While MLFF FMs trained on large quantities of ab-initio data have begun to approach the accuracy of DFT for some tasks, there are still significant challenges in improving efficiency for modeling large time and length scales. line with the increasing size and diversity of training data, MLFFs have steadily grown in complexity, both in terms of parameter tensor products to enforce higher-order Euclidean symmetries (Batzner et al. 2022 Batatia et al. 2022). Despite efficiency efforts (Luo al. 2024). saate-of-theaart MLFFs remain several orders of magnitude slower than alternatives such as classical force fields (Unke et al. 2021 Wang et al. (2024). As result, MIFF FMS are often prohibitively expensive to use in realistic downstream applications, such molecular dynamics (MD) simulations with > 10 timesteps.
and design choices such as the of expensive
Furthermore, clear tension arises between the need for fast, scalable training of MLFFs, and ensuring physical consistency when the model is deployed in downstream applications such as MD simulations. For example, trying to build constraints into the architecture such as SO(3) equivariance can slow down training. Restricting training to only an equivariant space may also complicate optimization and training dynamics, consequently affecting the quality of the final model. Although there is increasing evidence that SO(3) equivariance constraints, such as energy conservation, are more difficult to enforce without hard constraints (Bigi et al. 2024), but critical in downstream applications such as MD simulations.
be learned accurately directly from data (Qu & Krishnapriyan 2024), other
More broadly, the increasing generality MLFF FMs is at odds with the needs of practitioners, who are often ultimately focused relatively narrow set of systems and downstream applications (perovskites, magnesium-base -based electrolytes, insulators, etc.). Fine-tuning the FM for these downstream applications in principle straightforward, but may be prohibitive for many practitioners and offers no speedup at inference time. These challenges motivate us to ask the following questions:
1. How can we improve the efficiency of MLFFss for specialized tasks while preserving the powerful general- purpose representations learned by FMs?
How ensure that test-time physical soundness, such as energy conservation for molecular dynamics simulations, is maintained while preserving train-time scalability?
To address this, we introduce an approach based on knowledge distillation (KD) which learns fast, specialized MLFFs from large, general-purpose FMs. The core of our approach is training objective that aligns the Hessians of the energy predictions between the foundation MLFF (teacher) and specialized MLFF (student). The method is conceptually simple and efficient: the Hessians of the FM can
2

---

Published as a conference paper at ICLR 2025
be pre-computed once and stored, while the student Hessian computation can be accelerated using approximate sampling techniques. Unlike existing KD methods which align internal features between the teacher and student (Kelvinius et al. 2023). approach is entirely agnostic to model architecture and inductive biases, and can be used out-of-the-box for any student-teacher pairing. We demonstrate our approach on four MLFF FMs: MACE-OFF (Kovács et al. (223)) and EScAIP (Qu & Krishnapriyan 2022) trained on SPICE (Eastman et al. (023)) MACE-MP-0 (Batatia et al. (202)) trained on MPtrj (Deng et al.
2023) from the Materials Project (Jain al. 2020), and JMP (Shoghi et al. (023) finetuned on selected
molecules from MD22 (Chmiela et al. [023). We learn student MLFFs specialized to subsets of the FM's training distribution which mimic realistic downstream applications, such as specifically modeling amino acids or materials containing Yitrium. These specialized student models achieve inference speeds up to 20 times faster than the original FMs. Our approach also achieves substantial improvements in energy and force error, MD simulation stability, energy conservation, and geometry optimization, compared to student models trained without distillation. In most cases, the student models also outperform the original FM. To our knowledge, this is the first approach to create fast, specialized MLFFs from FMs.
2 BACKGROUND AND RELATED WORK
Machine Learning Force Fields. A Machine Learning Force Field (MLFF) is a learnable function approximator Up which maps a molecular configuration potential energy and per-atom forces. Specifically, it takes the positions of atoms, (x(1), ,...,,(())) RTx3, and their atomic numbers, z (2(1), ... z(1)) E R" as inputs, and outputs potential energy Up- I and per-atom forces Fa= (f(") .....f
p!()) ) R MLFFs are typically parameterized as graph neural networks (GNNs), and are
trained via the following regression loss, with supervision from dataset of reference energies and forces;
LEF = No
2 +AF
Σ 11 (z,r) 1, (2.r)
MLFF Inductive Biases. MLFF speed, accuracy, physical soundness, and ease of training are affected by the choice of inductive biases. SO(3) equivariant networks (Geiger & Smidt 2022) guarantee that the forces, and possibly internal features, rotate consistently when the positions of the input rotate, but can be both hard to optimize over and considerably slower than non-equivariant models, often due the reliance on expensive tensor products to handle SO(3) and spherical harmonic representations. Another key inductive bias is conservative force parameterization. MLFFs which obtain conservative forces by differentiating the energy output with respect to the atomic coordinates (Fş = V,U)) guarantee conservation of the model's energy in MD simulations, while MLFFs parameterizing the force separately from the energy lack this guarantee but are faster. There is growing evidence that SO(3) equivariance can be learned without explicit enforcement, given sufficient data and model scaling and/or data augmentation (Neumann et al., 2024 Qu & Krishnapriyan 2024 Pozdnyakov & Ceriotti (024)) while energy conservation is harder to learn without conservative forces (Bigi et al. 20244). In $ we demonstrate that our distillation approach works well with many combinations of teacher and student model inductive biases.
MLFF Foundation Models. While early MLFFs were trained on relatively narrow datasets, foundation models (FMs) trained across a diverse swath of chemical space now becoming increasingly common (Shoghi et al. 2023] Gasteiger et al. [022] Batatia et al. 2023 Kovács et al. 2023). MACE-OFF (Kovács
eet al. 2023) was primarily trained on filtered subset of $51,000 biomolecular lar structures from the SPICE
(Eastman et al. (023)) dataset. MACE-MPO (Batatia et al. (2023) was trained on 1.6 million structures
from the Materials Project (Jain al. 2020). The JMP (Shoghi al. 2023) FM was pre-trained
combined dataset consisting of OC20, OC22, ANI-Ix, and Transition- Ix, and later fine-tuned on several datasets such as QM9, rMD17, and MD22. The promise of MLFFs lies in their ability be used zero-shot or with minimal finetuning acToss many downstream tasks. However, as MLFF FMs have increased in complexity and scale to match the diversity and size of training data, speed has become a limiting factor, particularly in modeling systems with large time and length scales (Unke et al. 2021 Wang et al. 2024).
Knowledge Distillation. Knowledge distillation (KD) (Hinton et al.
2015) aims to transfer knowledge from a larger teacher model to smaller student model, usually by training the student to mimic certain
properties of the teacher (Romero et al. 2015 Sanh et al. 2019 Tang et al.
2020 Gou et al. (021). This is
typically done by minimizing a distillation objective of the form LRD E,, cby (z) ( Pos (z) I where Ør and ds are intermediate features of the teacher and student respectively, and P is linear projection that ac- counts for differences in dimensionality between the two models. Specializing FMs to specific subdomains
3

---

Published as a conference paper at ICLR 2025
has been explored in scage-scale language and vision models QQuu et al. 2024 but, knowledge, is unexplored in the context of MLFFs. Previous work on KD for MLFFs was done in (Kelvinius et al. (023) by aligning node and edge features across models such GemNet-OC, PaiNN, and SchNet the OC-20 and COLL datasets. However, the best-perforning method, referred to as "node-to-node" (n2n), did not specialize the models, and evaluated the student and teacher models on the same data. We demonstrate in that our Hessian distillation approach consistendly outperforms n2n across several datasets and MLFFs.
Learning from Function Derivatives. Sobolev training (Czarnecki et al. 2017) uses function derivatives as supervision train neural networks, including for KD. Their work highlights numerous theoretical ben- efits of training to match function derivatives, including better sample complexity and reduced overfitting. To our knowledge, this form of training has not been used to specialize to a subset of the data distribution.
3 DISTILLING FOUNDATION MODELS WITH ENERGY HESSIANS
After reviewing background energy Hessians in 3.1 introduce our method for producing fast, specialized MLFFs via knowledge distillation (KD) from the energy Hessians of pre-trained foundation models in 3 In our setting, the FM plays the role of the teacher, while the fast, specialized MLFF is the student. We also present our Hessian subsampling strategy to significantly accelerate training ( 333
3.1
BACKGROUND ON ENERGY HESSIANS
The Hessian of the energy is the second derivative of the energy with respect to atomic positions, or equivalently, the negative derivative of the forces with respect to the positions. Given 3.N-dimensional unrolled force and position vectors, the Hessian H H x3NN is given by H OF aoC Accordingly, dr 2r
Hi is the derivative of the tth force with respect to the jth position.
3.2 ENERGY HESSIAN ALIGNMENT OBJECTIVE
Given dataset of N molecular structures paired with quantum- energy and force labels D= [(z,r, F,),F,)) > ) 1, and a MLFF FM T pretrained on D that predicts energies Us and forces Fu. we first precompute the Hessians of the FM energy predictions over the dataset D using automatic differentiation (we demonstrate that finite differences can also be used in (A.1I). This results in an augmented dataset
Dang = [(,,,r,,U U.,FF ).) where H, (2,.). We then train a student MLFF S.g assumed to be small relative to To (i.e., |0] < to a specific downstream application, such as molecules with lodine. Crucially, in addition to matching the energies and forces of DKD, we also train the student Sp to match the FM Hessians over DKD.
Dr" subset of the data, DRD C Daug This subset corresponds
The complete loss function for training the student via knowledge distillation over the subset DKD is:
(()) = R,, LEF(Ø) + H OF (Z,
Or
where Lgr is the standard energy and force- matehinn objective defined in Eq. and AKD D is a hyperparam- eter controlling the strength of KD. We highlight that for direct-force student MLFFs, the energy Hessian is computed as the negative Jacobian of the force, rather than the second derivative of the energy prediction.
In practice, we find that the student can outperform the FM on the original objective (LEF) due to the reduced diversity of 'D D and the regularization effect of teacher supervision. To ensure that the student performance is not bottlenecked by the FM, we reduce the weight of the Hessian distillation loss term, AKD. by a factor of 2 during training once the student's validation loss on the original objective, LgF(@). becomes lower than that of the frozen FM, LEF (1) (see for for more details).
3.3 IMPROVING EFFICIENCY OF HESSIAN COMPUTATIONS WITH SUBSAMPLING
Obtaining the energy Hessian for molecule via autodifferentiation requires 3N backwards passes, one per force value. To mitigate this computational expense, we instead uniformly sample rows from the reference Hessian on which to supervise in each training iteration. We accordingly only compute these rows of the student's Hessian. Each row corresponds to one Euclidean coordinate of single atom. Formally, let J, C [1,....NV] be the set of s randomly sampled indices corresponding to the rows of the Hessian for particular molecular structure: J. - --.ra.-)a) For each sample in the dataset, we supervise the student model on the subset J, of Hessian rows. The modified loss function becomes.
OF
(Z.Tz)
L(6)= EE,, H.DDRD txk
Σ JEJ.
H
(
(3)
Dr
4

---

Published as a conference paper at ICLR 2025
where (1,33) denotes the uniform distribution over subsets of rows from the Hessian, and H.") and F=(") + are the j-th row of the reference Hessian and student forces, respectively. The number of backward passes required to compute the Hessian grows as (ss), so subsampling significantly accelerates training. Reducing the number of sampled rows to s= 1 does not noticeably impact model performance (55)
Computing Hessian Rows via Vector-Jacobian Products. When computing individual rows, we wish to avoid forming the entire Hessian matrix. We achieve this by using vector-Jacobian products (VJPs), the fundamental operation underlying reverse-mode autodifferentiation. Formally, given a function [(x): :i): + Roduut and vector v VoRou reverse-mode autodifferentiation computes the VJP v J in matrix-free manner (e.g., without explicitly forming the Jacobian & ε RRdous Xdin), ). In our setting, f is the MLFF force F; FRS R whose Jacobian is the energy Hessian H R3N33N To extract the jth Hessian row, we construct a one-hot vector e; E R&N to use in the VJP. Let P ER*x3N denote the per- mutation matrix containing one-hot vectors corresponding to the subset of sampled row indices in T: p [e,, er 1T By computing the matrix-Hessial Hessian product P H.H, which is achieved via a vmap over the rows of Par of the force VJP, we can efficiently extract the desired Hessian rows (H ,H()), (I(i,,,.,,HH.)) T
3.4 ENERGY GRADIENT SUPERVISION FOR DIRECT FORCE MODELS
To improve the energy predictions of MLFFs with direct force parameterization, we find it useful to indi- rectly utilize the inductive bias that the force is the negative gradient of the energy. Specifically, we introduce an additional loss term to align the negative gradient of the energy head with the true forces: L = || V.U& with loss weight A Here, F represents the true forces, and V,U, is the gradient of the predicted energy. Importantly, no gradients of the energy head are computed inference time (forces are derived as usual through the separate force head), so there is no impact on inference speed. See $3.33 for more details.
4 EXPERIMENTAL RESULTS
We present the results of our Hessian distillation approach for learning fast, specialized MLFFs from large foundation models. In $4.1 we distill the MACE-OFF FM, which was trained on the SPICE biomolecules dataset. We also distill an EScAIP & Krishnapriyan (2024) FM trained on SPICE in 1 In $4.2 we distill the MACE-MPO model trained on the MPtrj dataset. Finally, in $4.3 we distill the JMP FM, which was pretrained on several large datasets and finetuned on selected MD22 molecules. In each setting, we train fast, specialized MLFF students on data subsets corresponding to realistic downstream applications.
We compare against the following baselines for training student MLFFFs (more details in AA.6
Undistilled: Training on the specialized data subset without Hessian supervision (i.c., ARD 0).
2. n2n:
Training to match the node representations of the FM at the final layer. This is direct comparison to the best pertforrff MLFF KD technique introduced in (Kelvinius et al. 2023
3. aZa: Training on top of a leamed projection of the FM's atom embeddings. We create this baseline in the spirit of self-supervised learning techniques which fine-tune the representations of a large model for specific downstream tasks (Devlin 2018 Radford al. 2021 Caron et al.
2021
We measure MD simulation speed by performing inference with batch size of 64 on a single NVIDIA A6000 GPU and converting to nanoseconds per day by assuming a femtosecond timestep. This mimics real-world use cases of performing parallel MD simulations, vectorized over the batch dimension, to rapidly explore molecular phase space, or for a batch size which maximizes sample throughput for each model. These batch sizes are given in $A.7 We note that production MD simulations, MLFFs would be linked with an optimized platform like LAMMPs (Tompson et al. 2022), potentially leading to significant acceleration beyond the numbers we report.
structure screening. We also measure speed using
4.1 DISTILLING MACE-OFF ON SPICE
Using MACE-OFF (ooáács et al. 2023) as the teacher model, we focus on small subsets of SPICE with limited amounts of data monomess solvated amino acids, and molecules containing iodine- to train small, specialized GemNet-dT, PaiNN, and GemNet-T student MLFFss via our Hessian distillation approach described in (3.2 While the FM and specialized student MLFFs quantities of data, we report force mean absolute error (MAE) on test data not seen by either model during training. Results are shown in Tab. l with additional details and hyperparameter sweeps in A and A.8
trained on different
5

---

Published as a conference paper at ICLR 2025
Table I: Results of distilling the MACE-OFF foundation model trained SPICE into specialized MIFFS, (FM) indicates foundation model. while (S) indicates student model. n2n is the node feature matching baseline from (Kelvinius et al. 2023), while a2a is the atom embedding matching baseline we construct in |A.6 Student models all have identical simulation speeds. Speedups relative to FM are given in parentheses. The "Speed"" column is calculated with a constant batch size of 64 for all models, while the "Maximum Speed" is calculated with the batch size that maximizes throughput for each model.
Force
Energy
Chemical Subgroup
Datasel
Model
Distillation Method
MAE
MAE (me (meVatom))
Speed (nviday) (1)
Maximun Speed (nsddy) (t)
Size
(Parameter Count)
(meV/A)
(U)
(L)
(FM) MACE-OFF Large (4.7M)
666
0.65
38.0
38.1
Monomers
Undistilled
n2n
10.5
1.2
GemNet-dT (0.67M)
a2a
12.9
Hessian (ours)
6.3
16455 (4.3x)
725.2 (19.0k)
Undistilled
25.0
PaiNN (IM)
n2n
20.8
1.5
24.7
Hessian (ours)
8,77
0.48
291.5 (7.7x)
1827 (48.0K)
Undistiled
7.9
3.8
(S) Gemiee-T (0.67M)
n2n
1.7
5.3
Hessinn (ours)
5.1
6.8
66.0 (1.74))
408.9 (10.8s)
(FM) MACE-OFF Luge (4.7M)
19.4
3.8
3.8
Solvated
Amino Acids
805
Undistiled
22.4
2.2
(S) GemNet-dT (0.67M)
n2n
20.7
24.4
Hessian (ours)
1166
0.37
44.4 (11.7))
(11.7)
Undistilled
50.1
3.3
n2n
38.3
(S) PaiNN (1.IM)
aZa
52.4
3.7
Hessian (ours)
18,0
0.41
79.4 (20.9%)
(2.9
Undistiled
18.3
1.7
( GemNet-T (0.67M)
a2a
17.1
1.9
Hessian (curs)
11.2
23.4 (6:2s)
(6.2
(KM) MACE-OHI Large (4.7M)
15.3
1.3
14.8
14.8
Systems
with lodine
11,171
Undistiled
23.4
2,68
() GemNet-TT (0.67M)
23.3
a2a
23.2
2.6
Hessian (ours)
14.7
0.58
148.1 (10.0K)
220.4
Undistilled
51.2
3.3
(S) PAINN (1 (LOIM)
a2a
50.7
3.5
Hessian (ours)
270.2 (18.3x)
(29.8kx
Undistiled
15.9
GemNet-T (0.67M)
n2n
15.9
a2n
15.9
5.6
(ours)
65.4 128.0 (8.6x)
Using the distilled, specialized MLFFs, achieve up to 20x increases in simulation speed relative to the FM, and up to 50x increases for throughpu--maxiiiznng batch sizes. For all splits, our Hessian distillation approach significantly outperforms training without distillation, as well the a2a and n2n baselines, on Energy and Force MAE. In many cases, our distilled models outperform the FM, likely because they can focus all of their expressivity towards learning a narrower slice of chemical space. We report the times required to train the distilled student models relative to that of the original FM in $A.5 These times are generally nominal; when sampling 4 rows of the Hessian (s = 4), training the distilled student model requires an average of 4.0% additional compute beyond FM training. We demonstrate the downstream usefulness of our distilled models in constant-temperaliee student models are more stable than their undistilled counterparts. We also perform geometry optimization in $A.15) and find that our distilled models generally converge to structures with lower energy and force norms. We note that distilling GemNet-T student model occasionally worsens energy accuracy relative to an undistilled model, even though the forces derived from these energies remain unaffected. This would not practically impact MD simulations, for which only forces needed; however, if energies important,
MD simulations $A.14 Our results indicate that our distilled

---

Published as a conference paper at ICLR 2025
Table 2: Results of distilling the MACE-MPO foundation model trained on MPtr into specialized MI FF.. (FM) indicates foundation model, while (S) indicates student model. n2n is the node feature matching baseline from (Kelvinius et al. (023), while a2a is the atom embedding matching baseline we construct in QA.6 Student models all have identical simulation speeds. Speedups relative to FM are given in parentheses. The "Speed" column is calculated with a constant batch size of 64 for all models, while the "Maximum Speed" is calculated with the batch size that maximizes throughput for each model.
Force
Model
Distillation Method
MAE (meV/A) (U)
Speed (uv/day) (
Maximum Speed (ms/day) (t)
Chemical Subgroup
Dataset Size
(Parameter Count)
(FM) MACE-MPO
18.1
101.9
Pm3m Spacegroup
9,725
Undistilled
15.7
GemNettdT (U.67M)
n2n
14.6
a2a
16.9
Hessian (ouns)
11.8
162.7 (1.7x)
260.1
Undistilled
21.9
(8) PaiNN (1.0M)
n2n
19.5
a2n
23.3
Hessian (ours)
15,5
264.4 (2.8x)
451. (4.4x)
(FM) (153MM))
45.2
26.5
27
Undistilled
32.5
Systems Yttiium
30,436
(8) GemNet-dT (0.67M)
n2n
36.5
a2a
36.5
Hessian (ours)
21.3
73 (2.8%)
(2.7g)
Undistilled
55.5
(S) PaiNN (1.0M)
n2n
37.7
aZu
49.8
Hessian (ours)
25.7
215.5 (8.
267.2 (9.9x)
(FM) (A5 M)
31,4
13.4
13.4
Band Gap 2 meV
36,150
Undistilled
17.1
n2n
15.1
(S) GemNet-dT (0.67M)
aZa
16.3
Hession (ours)
12.1
38.7 (2.9x)
(2.9%
Undistilled
32.6
(S) PAINN (LOM)
n2n
27.5
aZa
32.2
Hessiun (ours)
16.3
125.4 (9.4k)
125.4
coefficient rescaling to upweight the energy loss term, or energy fine-tuning as seen in (Kovács et al. 2023). would likely alleviate this issue. In $A.I] we demonstrate distillation with an EScAIP Qu & Krishnapriyan (2024) FM on SPICE, which leads to improved student model performance and eliminates this energy issue.
4.2 DISTILLING MACE-MP-O ON MATERIALS PROJECT
We next consider MACE-MP-0 (Batatia et al. 2023), trained 1.6 million structures from the MPtrj (Deng et al.] (023) dataset, as a teacher model, choosing the following subsets on which to learn specialized Gemnet-dT and PaiNN student MLFFs: materials in the Pm3m spacegroup (which includes cubic perovskites used in photovoltaic devices), materials containing Yttrium (used in lasers and alloys), and materials with band gap of greater than meV (roughly corresponding to insulators). While DFT with the PBE functional is known to underestimate band gap Mori-Sänchez et al. 2008), we assume this delineation sufficient for our purposes of creating broad chemical subgroups. The results are shown in Tab.
We find that the specialized student MLFFs obtained via our Hessian KD approach are up to 10x faster than the original FM, and consistently outperform the undistilled, n2n, and a2a baselines Force MAE across all splits. Interestingly, we find that the GemNet-dT student models outperform the FM even before distillation, and Hessian KD subsequently further improves the student models. We speculate that in this scenario, Hessian KD has a regularizing effect which enables the student to learn better representations despite distilling from a teacher with higher Force MAE, analogous to training with soft or noisy labels (Szegedy et al. 2016; Müller et al. 2019). We also note that the n2n and a2a methods both generally improve over the undistilled baseline, unlike with MACE-OFF on SPICE. This suggests that improvements from KD are not always correlated to the accuracy of the teacher model.
7

---

Published as a conference paper at ICLR 2025
Table 3: Results of distilling JMP-Large (JMP-L) and JMP-Small (JMP-S) foundation models finetuned on selected large MD22 molecules into specialized MLFFs. (FM) indicates foundation model, while (S) indicates student model. Student model speedups relative to the JMP-L FM are given in parentheses. The "Speed" column is calculated with a constant batch size of 1 for all models, while the "Maximum Speed" calculated with the batch size that maximizes throughput for each model.
Energy MAE (meV/ afom) (1)
Force MAE (me//A) (1)
Speed (ns/day) (T)
Maximum Speed (av/day) (T)
Dataset Size
Model (Purumeter Count)
Distillation
Molecule
Method
(FM) JMP S (39.9M) (FM) JMP-1 (220M)
7.8
0.8
1.8 0.6
Buckyball Cascher
4.3
0.4
600
Undistilled
8.0
1.0
(5) GerNNet-TT (0.67M) (Direct-ooress, Invariant)
JMP-S Hessian (puns)
JMP-L Hossian (ours)
5.1
0.15
2,4 (6x)
18.3
Undistilld
8.4
0,08
(5) GemNet-T I (0.57M) (Gradient-Forees Invariant)
JMP-S Hessian (ours)
5.0
0.09
JMP-L Hessian (ours)
4.0
1.4 (3.5%)
9.6 (16x)
Undistiled
8.4
(5) eSCN (0.94M) (Direct-Forees. Equivariant)
JMP-S S Hessian (ouns
9.9
0.79
JMP-L Hossian (ours)
9.9
0.80
1.6 (4x)
135 (2.3x)
(FM) JMP-S S (39.9M) (FM) JMP-L (220M)
23.8
0.5
0.7
0.2
Double Walled
Nanotube
800
Undistileed
(S) GenNet-dT (0.67M) (iirct-Forees, Invarisnt)
JMP-S Hession (ours)
14.3
0.23
JMP-L Hossian (ours)
10.6
0.25
3.2 (16x)
6.4 (32x)
Undistilled
13.6
0.07
() GemNet-T (0.57M) Gradient-Fore Invariant)
JMP-S Hessian (ours)
0.05
JMP- Hessian (ours)
10.8
0.06
188 (9x)
3,3 (16.5x)
Undistilled
19.2
0.50
(S) eSCN (0.94M) (Direct Forces, 1= 2 Equivariant)
JMP'S Hessian (ouns)
16.1
0.40
JMP Hessian (ours)
16.1
0.47
2.6 (13x)
5.4 (27x)
4.3 DISTILLING JMP ON MD22
As a final evaluation of our Hessian distillation approach, we distill JMP (Shoghi et al. 2023 FMs finetuned the largest molecules in MD22-the buckyball catcher and doulle-walled nanotbbe into various student models. This setting presents a number of unique challenges. The selected MD22 molecules, with 148 and 370 atoms respectively. are significantly larger than those considered thus far. Additionally, the JMP models are considerably larger than the previously considered MACE FMs, with approximately 40M and 220M learnable parameters in the small (JMP-S) and large (JMP-L) models respectively. Unlike their MACE counterparts, the JMP FMs must therefore forgo a gradient-based force parameterization to remain within GPU memory limits. The JMP model is also based on a GemNet backbone, which does not utilize built-in higher-order equivariance like the MACE FMs.
Using JMP-S and JMP-L FMS as teachers, we perform Hessian distillation on the selected MD22 molccules to obtain specialized Ge.Net-dT. GemNet-T, and eSCN student MLFFs. GemNet-T uses gradient-base based force parameterizatton, while eSCN uses higher-order 21= equivariance. We find that the our specialized MLFFs are up to 30x times faster than the original FMs, and considerably outperform the undistilled baselines in Force and Energy MAE on both molecules (Table 3). We highlight that our distillation procedure is effective when distilling from a FM with a direct-force parameterization into student model with gradient-based based forces (GemNet-T) T) and higher-order equivariance (eSCN). We capitalize on this property by rumning constant energy (NVE) simulations of the buckyball catcher for 100 ps using the trained GemNet models. We find that distillation produces a GemNet-dT model which conserves energy better than its undistilled counterpart and the original JMP-I. FM (Figure We also find that our GemNet-T student MIFF, employing gradient-basse forces and distilled with JMP-I Hessians, is able to conserve energy and simulate stably for the entire duration of the 100 ps simulation, while the JMP-L energy gradually drifts throughout the simulation (FigureZp), We finally highlight that distilling with Hessians from JMP-L leads to better performance than distilling from JMP-S (see Distillation Method column of Table suggesting that continued scaling of FMs has the potential to further improve student model performance.
8

---

Published as a conference paper at ICLR 2025
(a)
(b)
(c)
Charge Energy (eWitom)
Commies-dt of Errect Foree, Undstilted
Energy (eViatom))
0.08
GamflssT- Grad ur Based Porcas, Undiat bed
Change in (eViatom)
0.08
0.10
Camiie--T Einet Faroe. Distiled
0.06
Nh.. Drrect Forsas (laxcher)
Germent rood Rissed Forres, Rigplled
=Batter)
0.08
4 Grntent Coctendant feed forms sechen
Scmeerof farms
0.06
0.04
0.04
0.04
0.02
0.02
0.00
0.00
0
20
40
60
80
100
80
100
Simulation time (ps)
Simulation time (ps)
Force MAE (meviA) (4 Better))
Figure 2: Energy Conservation in NVE MD Simulations of Buckyball Catcher. We plot the change in the model predicted energy over the trajectory for 5 independent initial conditions. Some simulations become unstable before 100 ps (denoted by x). (a) Hessian distillation improves the energy conservation of GemNet-dT models, which outperforms that of JMP-L. (b) Our student GemNet-T models conserve energy outperrorrrrrnoc
due to using conservative forces, while the JMP-L FM energy steadily drifts, broadly suggesting that large- scale models with few built-in constraints can be effectively distilled into smaller, constrained models. (c) Change in energy plotted against test force MAE. Distillation into a GemNet-T student combines the general- purpose representations and accuracy of JMP-L with the physical inductive biases of conservative forces.
5 ABLATIONS
We conduct ablation studies on various aspects of our approach, namely: the size of the student MLFF model, and the Hessian subsampling frequency used during training. In $A.9 we additionally examine role of teacher Hessians vs. forces in the KD objective, and find that distilling with teacher forces significantly inferior to our approach of distilling with Hessians.
(a)
(b)
30
Undistilled GemNet dl (Student)
Distilled GemNet- (Student)
25
Distilled GerNNet- dT (Student)
048M
25
Distilled PaiNN (Student)
MACE OFF (FM- Teacher)
Force MAE (meV/A)
MACE OFF (FM Teacher)
Force MAE (meW/A)
20
1.6x
6.2x
.11M
20
15
2.3M
.57M .20M
10
15
2.0x
2.7x
6.7x
5
0
200
400
600
800
10
2.5
5.0
7.5
10.0
12.5
15.0
simulation Speed (ns/day)
Rows of Hessian Sampled/ Training Iteration
Figure 3: Parameter count and Hessian subsampling ablations. (a) Force MAE the Monomers split of SPICE as a function of the GemNet-dT student MIFF simulation speed. The size of the dots indicates the relative number of trainable parameters in the each model, and the number of parameters indicated above each dot. Compared to the undistilled model, Hessian distillation improves the speed-accuracy tradeoff. (b) Force MAE on the Solvated Amino Acid split of SPICE a function of the number of rows of the energy Hessian subsampled at each training iteration. The size of the dots and text indicates the time required per step of training. relative to training without distillation. Reducing down not have detrimental effect on model accuracy, and results in more efficient training.
1 does
Student MILF Size. To understand the effect of student MLFF expressivity, we vary the GemNet-dT student model parameter count by reducing the node and edge embedding dimensions from 128 to 8, which yields a total parameter count varying from 2.3 million to 48,000. For each model, we train with Hessian distillation on the Monomers subset of SPICE. We also train the same subset without Hessian distillation for comparison. To measure student MLFF speed, we use a larger inference batch size of 2048, which is the point of memory saturation use a batch size of 128, which maximizes its throughput. We find that Hessian distillation significantly improves the trade-off between speed and accuracy all student MI .F sizes (Fig. a). The differences
NVIDIA A6000 GPU. For the MACE-OF teacher,

---

Published as a conference paper at ICLR 2025
in speed between the FM and student MLFFs are more dramatic at larger batch sizes, suggesting that speed benefits from distillation are magnified when parallelizing across more samples. Interestingly, Hessian KD also appears to unlock better scaling properties: without distillation, the force MAE plateaus after scaling to 0.28M parameters, while using Hessian KD yields continual improvements up to the maximum student size of 2.3M parameters. We speculate that the Hessian supervision term may have regularizing effect on larger models. The improvement becomes more marginal as the student size decreases, indicating that insufficiently expressive students may struggle minimize the multi-term KD objective.
Hessian Subsampling Quantity. We vary the number of rows sampled from the MACE- MACE-OFF FM's reference Hessian during training on the solvated amino acid split of SPICE. We find that increasing does not reliably lead to decrease in Force MAE, and in some cases leads to slight increase. The training cost, measured in GPU-seconds per training step, increases approximately linearly with s (Fig. (b). With s = 1, distillation incurs 1,6x and 2x increase in training cost relative undistilled training for PaiNN and GemNet-dT,, respectively. We speculate that may play similar role as the batch size in conventional training. Small values of s add variance the KD gradient estimates, facilitating escape from local minima in the loss landscape, while large values may lead to generalization gaps, similar to what has been observed for large batch sizes (Keskar et al.) 2016).
CONCLUSION
Key Takeaways. We have presented Hessian distillation, which is to our knowledge the first technique to derive fast, specialized MLFFs by training them to match the energy Hessian of FMs trained on large, diverse datasets. By subsampling rows of the Hessian to supervise the student model at each training iteration, we ensure that Hessian distillation incurs only a nominal cost relative training the original FM. The specialized student MLFFss derived from our distillation approach are up to 20x faster at predicting energies and forces than the original FMs. Despite having far fewer parameters than the teacher FM, our distilled student models are consistently superior to the FM, as well as models trained with other distillation methods or no distillation. All of the demonstrated improvements readily extend to geometry optimizations and MD simulations, where our distilled models are more stable and conserve energy better over time. Our observation that the student MLLFFs often outperform the original FM on the specialized data subset, sometimes even without any distillation, suggests that the field has not yet converged on an effective training recipe for large-scale MIFF FMs, as scaling data and model size should in principle lead to better downstream performance (Kaplan et al. 2020 Hoffmann et al. 2022 We find that our Hessian KD approach still yields improvements over undistilled models in this scenario, suggesting that even when the FM forces are inaccurate, its Hessians still provide effective regularization to make significant improvements in student force accuracy. More generally, the JMP (44.)) and EScAIP ($A.)) results suggest that Hessian distillation will improve in effectiveness as FMs scale and become more performant.
Limitations. The main drawback of our method is that training with Hessian distillation adds training overhead that scales with the number of sampled rows s. However, since choosing s=1 1 has an empirically negligible impact on performance ($5). we can limit the training overhead to around twice that of the undistilled models. Since student models tend to be much smaller than FMs, the training overhead relative to that of the corresponding foundation model is quite small. We also demonstrate that is possible to accelerate Hessian computation using finite difference approximations in $A.11
Future Work and Outlook. In the future, MLFFs could be specialized in ways beyond chemical subgroups, such as performing high-temperature simulations (Stocker et al. 2022) or modeling phase transi- tions (Jinnouchi et al. 2019) using Hessian-based KD. Applying sampling techniques when pre-computing the teacher Hessians would also reduce the upfiront cost of our approach. More broadly, our work sets a precedent for future MLFF development: as training data and model parameter counts continue to grow, new MLFF FM releases should be accompanied by a set of small, specialized "enngine"" for common downstream tasks. We further speculate that in the future, practitioners may rarely, if ever, actually perform inference with MLFF FMs directly. Instead, FMs could serve as a reservoir for general- -urpose representa- tions, which are subsequently distilled into small models specialized for the task at hand. In particular, the JMP and ESc distillation results suggest a recipe in which large FMs are trained with minimal inductive biases to facilitate scalable and general-purpen models with inductive biases tailored to the downstream task (e.g., conservative forces for MD simulations). This paradigm would enable widespread adoption of MLFFs, and move the field closer to the longstanding dream of force fields with the speed of classical methods and the accuracy of quantum mechanical methods.
training. followed by distillation into specialized student
10

---

Published as a conference paper at ICLR 2025
7 ACKNOWLEDGMENTS
The authors thank Eric Qu for assistance with training and distilling the EScAIP model. We also thank Ryan Liu, Toby Kreiman, Rasmus Lindrup, and the FAIR chemistry team for helpful discussions that benefited this paper. The authors acknowledge the support of this work from the U.S. Department of Energy, Office of Science, Energy Earthshot initiatives as part of the Center for lonomer-based Water Electrolysis at Lawrence Berkeley National Laboratory under Award Number DE-ACO2-05CHH 1231. This research used the National Energy Research Scientific Computing Center (NERSC), a U.S. Department of Energy Office of Science User Facility located at Lawrence Berkeley National Laboratory, operated under Contract No. DE-AC02-05CHI122 1231 using NERSC award BES--ERCAPW222133..
7.1 REPRODUCIBILITY STATEMENT
We have built our implementation of Hessian distillation around the Fairchem Github Repository. Our implementation works with any model or dataset compatible with Fairchem. We plan release the code after acceptance of the paper. Details on datasets, models, and hyperparmeters used for training and evaluation are provided in the Appendix.
REFERENCES
Josh Achiam, Steven Adler, Sandhini Agarwal, Lama Ahmad, lige Akkaya, Florencia Leoni Aleman, Diogo Almeida, Janko Altenschmidt, Sam Altman, Shyamal Anadkat, et al. Gpt-4 technical report. arXiv preprint arXiv:2303.088 087744 2023.
Ilyes Batatia, David P Kovacs, Gregor Simm, Christoph Ortner, and Gabor Csanyi. Mace: Higher order equivariant message passing neural networks for fast and accurate force fields. In Koyejo, S. Mohamed, A. Agarwal, D. Belgrave, K. Cho, and A. Oh (eds.), Advances in Neural In- formation Processing Systems, volume 35, pp. 11436. Curran Associates, Inc., 2022.
URL httpps:/prroeeedings neurips cc/paper.pppp...ceeecee wwl.ffills//ee wtllle/paaer/2222//// Pawpeadbbbbboppp-----rrrree Comfere-eo
1136csc51af1eeccf3336
Ilyes Batatia, Philipp Benner, Yuan Chiang, Alin M. Elena, Dávid P. Kovács, Janosh Riebesell, Xavier R. Advincula, Mark Asta, William J. Baldwin, Noam Bernstein, Arghya Bhowmik, Samuel M. Blau, Vlad Cärare, James P. Darby, Sandip De, Flaviano Della Pia, Volker I.. Deringer, Rokas Elijošius, Zakariya El-Machachi, Edvin Fako, Andrea C. Ferrari, Annalena Genreith--Shrieeerr Janine George, Rhys E. A. Goodall, Clare Grey, Shuang Han, Will Handley, Hendrik H. Heenen, Kersti Hermansson, Christian Holm, Jad Jaafar, Stephan Hofmann, Konstantin S. Jakob, Hyunwook Jung, Venkat Kapil, Aaron D. Kaplan, Nima Karimitari, Namu Kroupa, Jolla Kullgren, Matthew C. Kuner, Domantas Kuryla, Guoda Liepuoniute, Johannes T. Margraf, loan-Bogdan Magdău, Angelos Michaelides, J. Harry Moore, Aakash A. Naik, Samuel Niblett, Sam Walton Norwood, Niamh O' Neill, Christoph Ortner, Kristin Persson, Karsten Reuter, Andrew S. Rosen, Lars L. Schaaf, Christoph Schran, Eric Sivonxay. Tamás K. Stenczel, Viktor Svahn, Christopher Sutton, Cas van der Oord, Eszter Martin Vondrák, Yangshuai Wang. William C. Witt, Fabian Zills, and Gábor Csányi A foundation model for atomistic materials chemistry, 2023.
Tejs Vegge,
Simon Batzner, Albert Musaclian, Lixin Sun, Mario Geiger, Jonathan P. Mailoa, Mordechai Kombluth. Nicola Molinari, Tess E. Smidt, and Boris Kozinsky. E(3)-equivariant graph neural networks for data-efficient and accurate interatomic potentials. Nature Communications, 13(1), 5 2022. doi: 10. 10.10388641447--22-20000
Filippo Bigi, Marcel Langer, and Michele Ceriotti. The dark side of the forces: assessing non-conservative force models for atomistic machine learning, 2024.
Erik Bitzek, Pekka Koskinen, Franz Gähler, Michael Moseler, and Peter Gumbsch, Structural relaxation made simple. Physicul review letters, 97(17): (77):11000 170201, 2006.
Mathilde Caron, Hugo Touvron, Ishan Misra, Hervé Jégou, Julien Mairal, Piotr Bojanowski, and Armand Joulin. Emerging properties in self-supervised vision transformers. In Proceedings of the IEEE/CVF international conference computer vision, pp. 9650-9660, 2021.
11

---

Published as a conference paper at ICLR 2025
Lowik Chanussot, Abhishek Das, Siddharth Goyal, Thibaut Lavril, Muhammed Shuaibi, Morgane Riviere, Kevin Tran, Javier and community challenges. Acs Catalysis, L1(10)6059-6013 2021.
Caleb Ho. Weihua Hu, et al. Open catalyst 2020 (oc20) dataset
Stefan Chmiela, Alexandre Tkatchenko, Huziel E. Sauceda, Igor Poltavsky, Kristof T. Schütt, and Klaus-Robert Müller. Machine learning of accurate energy-conserving molecular force fields. Science Advances, 3(5), 5 2017. doi: 10.1 11126//iiav.160000...
Stefan Chmiela, Valentin Vassilev- Galindo, Oliver T Unke, Adil Kabylda, Huziel E Sauceda, Alexandre Tkatchenko, and Klaus-Robert Müller. Accurate global machine learning force fields for molecules with hundreds of atoms. Science Advances, 9())ead(0873, 2023.
Daniel J Cole and Nicholas DM Hine. Applications of large-scale density functional theory in biology. Journal of Physics: Condensed Matter, 28(39):393001, 2016.
Wojeiech M Czamecki, Simon Osindero, Max Jaderberg. Grzegorz Swinszez, and Razvan Pascanu. Sobolev training for neural networks. Advances in neural information processing systems, 30, 2017.
Bowen Deng, Peichen Zhong, KyuJung Jun, Janosh Riebesell, Kevin Han, Christopher J. Bartel, and Gerbrand Ceder. Chgnet as a pretrained universal neural network potential for charge-infornn atomistic modelling. Nature Machine Intelligence, pp. 1-11, 2023. doi: 10. 10.1038//42256-00 01000/60022-0011---0
Jacob Devlin. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint arXiv: arXiv: 1810.0480, 2018.
Peter Eastman, Pavan Kumar Behara, David L Dotson, Raimondas Galvelis, John E Herr, Josh Horton, Yuezhi Mao, John D Chodera, Benjamin P Pritchard, Yuanqing Wang, et al. Spice, a dataset of drug-like molecules and peptides for training machine learning potentials. Scientific Data, 10(1):11, 2023.
Roger Fletcher. Practical methods of optimization. John Wiley & Sons, 2000.
Xiang Fu, Zhenghao Wu, Wujie Wang, Tian Xie, Sinan Keten, Rafael Gomez-Bombarelli, and Tommi Jaakkola. Forces are not enough: Benchmark and critical evaluation for machine learning force fields with molecular simulations. arXiv preprint arXiv:2210.07237. 2022.
Johannes Gasteiger, Florian Becker, and Stephan Günnemann. Gemnet: Universal directional graph neural networks for molecules. Advances in Neural Information Processing Systems, 34:6790 6802, 2021.
Johannes Gasteiger, Muhammed Shuaibi, Anuroop Sriram, Stephan Günnemann, Zachary Ulissi, C. Lawrence Zinick, and Abhishek Das. Gemnet-oc: Developing graph neural networks for large and diverse molecular simulation datasets, 2022.
Mario Geiger and Tess Smidt. e3nn: Euclidean neural networks. arXiv preprint arXiv:2207.09453, 2022.
Justin Gilmer, Samuel S Schoenholz, Patrick F Riley, Oriol Vinyals, and George E. Dahl. Neural message passing for quantum chemistry. In International conference on machine learning, pp. 1263-1272. PMLR, 2017.
Jianping Gou, Baosheng Yu, Stephen J Maybank, and Dacheng Tao. Knowledge distillation: A survey. International Journal of Computer Vision, 12966):189-11899 2021.
impact of density functional theory materials research. MRS bulletin, 31(9):659-663 2006. Jürgen Hafner, Christopher Wolverton, and Gerbrand Ceder. Toward computational materials design: the
Bjørk Hammer and Jens Kehlet Nørskov. Theoretical surface science and catalysis catalysis- calculations and concepts. In Advances in catalysis, volume 45, pp. 71-129. Elsevier, 2000.
Geoffrey Hinton, Oriol Vinyals, and Jeff Dean. Distilling the knowledge in neural network. arXiv preprint arXiv:1503.0253/ 2015.
Jordan Hoffmann, Sebastian Borgeaud, Arthur Mensch, Elena Buchatskaya, Trevor Cai, Eliza Rutherford, Diego de Las Casas, Lisa Anne Hendricks, Johannes Welbl, Aidan Clark, et al. Training compute-optimal large language models. arXiv preprint arXiv:2003. 2022.
12

---

Published as a conference paper at ICLR 2025
Anubhav Jain, Yongwoo Shin, and Kristin A Persson. Computational predictions of energy materials using density functional theory. Nature Reviews Materials, 1(1):1113, 2016.
Anubhav Jain, Joseph Montoya, Shyam Dwaraknath, Nils ER Zimmermann, John Dagdelen, Matthew Horton, Patrick Huck, Donny Winston, Shreyas Cholia, Shyue Ping Ong, et al. The materials project: Accelerating materials design through theory-driven data and tools. Handbook of Materials Modeling: Methods: Theory and Modeling, pp- 1751-1784, 2020.
Frank Jensen. Introduction to computational chemistry. John wiley & sons, 2017.
Ryosuke Jinnouchi, Jonathan Lahnsteiner, Ferenc Karsai, Georg Kresse, and Menno Bokdam. Phase transitions of hybrid perovskites simulated by machine-learning force fields trained on the fly with bayesian inference. Physical review letters, 122(22):22501, 2019.
Jared Kaplan, Sam McCandlish, Tom Henighan, Tom B Brown, Benjamin Chess, Rewon Child, Scott Gray, Alec Radford, Jeffrey Wu, and Dario Amodei. Scaling laws for neural language models. arXiv preprint urXiv:2001.8836/ 2020.
Filip Ekström Kelvinius, Dimitar Georgiev, Artur Petrov Toshev, and Johannes Gasteiger. Accelerating molecular graph neural networks via knowledge distillation. arXiv preprint arXiv:2306.1433 14818, 2023.
Nitish Shirish Keskar, Dheevatsa Mudigere, Jorge Nocedal, Mikhail Smelyanskiy, and Ping Tak Peter Tang. On large-batch training for deep learning: Generalization gap and sharp minima. arXiv preprint arXiv: 1609.0433, 2016.
Dávid Péter Kovács, J. Harry Moore, Nicholas J. Browning, Ilyes Batatia, Joshua T. Horton, Venkat Kapil, William C. Witt, loan-Bogda Magdäu, Daniel J. Cole, and Gábor Csányi Mace-off23; Transferable machine learning force fields for organic molecules, 2023.
Benjamin Lefaudeux, Francisco Massa, Diana Liskovich, Wenhan Xiong, Vittorio Caggiano, Sean Naren, Min Xu, Jieru Hu, Marta Tintore, Susan Zhang, Patrick Labatut, Daniel Haziza, Luca Wehrstedt, Jeremy Reizenstein, and Grigory Sizov. xformers: A modular and hackable transformer modelling library. https::/gittuut com/ wwc.//aaebbooreeeeeee..ooo
xformers] 2022.
Shengjie Luo, Tianlang Chen, and Aditi Krishnapriyan. Enabling efficient equivariant operations in the fourier basis via gaunt tensor products. arXiv preprint arXiv:2401.10211 2024.
Paula Mori-Sánchez, Aron J. Cohen, and Weitao Yang.
Localization and delocalization er- rors in density functional theory and implications for band-gap prediction. Phys. Rev.
Lett., 100:146401, Apr 2008. doi: 10.1103/PhysRewl.ett 160414
URL https
a/iikk.aps.org/doi...
100. 146401
Rafael Müller, Simon Korblithh, and Geoffrey E Hinton. When does label smoothing help? Advances in neural information processing systems, 32, 2019.
Albert Musaelian, Simon Batzner, Anders Johansson, Lixin Sun, Cameron Owen, Mordechai Kombluth, and Boris Kozinsky. Learning local equivariant representations for large-scale atomistic dynamics, 2022.
Mark Neumann, James Gin, Benjamin Rhodes, Steven Bennett, Zhiyi Li, Hitarth Choubisa, Arthur Hussey, and Jonathan Godwin. Orb: A fast, scalable neural network potential, 2024.
Maxime Oquab, Timothée Darcet, Théo Moutakanni, Huy Vo, Mare Szafraniec, Vasil Khalidov, Piere Fernandez, Daniel Haziza, Francisco Massa, Alaaeldin El-Nouby, al. Dinov2: Leaming robust visual features without supervision. arXiv preprint arXi::2004.07193, 2023.
Sergey Pozdnyakov and Michele Ceriotti Smooth, exact rotational symmetrization for deep learning on point clouds. Advances Neural Information Processing Systems, 36, 2024.
Shikai Qiu, Boran Han, Danielle C. Maddix, Shuai Zhang, Yuyang Wang, and Andrew Gordon Wilson. Transferring knowledge from large foundation models to small downstream models, 2024. URL
hittps://arxiv: http:://aaxvvv.oo org/ alss /400 07337
Eric Qu and Aditi S Krishnapriyan. The importance of being scalable: Improving the speed and accuracy of neural network interatomic potentials across chemical domains. arXiv preprint arXiv:2410.44169 2024.
13

---

Published as a conference paper at ICLR 2025
Alec Radford, Jong Wook Kim, Chris Hallacy, Aditya Ramesh, Gabriel Goh, Sandhini Agarwal, Girish Sas- try, Amanda Askell, Pamela Mishkin, Jack Clark, et al. Learning transferable visual models from natural language supervision. In International conference machine learning. pp. 8748 8763. PMLR, 2021.
Sanjeev Raja, Ishan Amin, Fabian Pedregosa, and Aditi S Krishnapriyan. Stability-aware training of neural network interatomic potentials with differentiable boltzmann estimators. arXiv preprint arXiv:2402. 13984, 2024.
Adriana Romero, Nicolas Ballas, Samira Ebrahimi Kahou, Antoine Chassang, Carlo Gatta, and Yoshua Bengio. Fitnets: Hints for thin deep nets. In International Conference on Learning Representations, 2015,
Victor Sanh, Lysandre Debut, Julien Chaumond, and Thomas Wolf. Distilbert, a distilled version of bert: smaller, faster, cheaper and lighter. 5th Workshop on Energy Efficient Machine Learning and Cognitive Computing, 2019.
K. T. Schütt, H. Sauceda, P.J. Kindermans, A. Tkatchenko, and K.-R. Müller. Schnet deep learning architecture for molecules and materials. J. Chem. Phys., 148(24):241222, 6 2018. doi: 10. 10.1069/1.501979. URL https: org/10 11w.1106111110 1tt.://.......1110011010.1... 1019777779
Nima Shoghi, Adeesh Kolluru, John Kitchin, Zachary W. Ulissi, C. Lawrence Zitnick, and Brandon M. Wood. From molecules to materials: Pre-training large generalizable models for atomic property prediction, 2023.
Sina Stocker, Johannes Gasteiger, Florian Becker, Stephan Günnemann, and Johannes T Margraf. How robust are modern graph neural network potentials in long and hot molecular dynamics simulations? Machine Learning: Science and Technology, 3440045010, 11 2022. doi: 10. 00.108882622-2153449955 1088/2632222 21533
URL httthtt:s://ax doi.
Christian Szegedy, Vincent Vanhoucke, Sergey loffe, Jon Shlens, and Zhignnew Wojna. Rethinking the inception architecture for computer vision. In Proceedings of the IEEE conference computer vision and pattemn recognition, pp. 2818-2826 2016.
Yiren Tang, Liangyou Lu, and Graham Neubig. Understanding knowledge distillation in non- non-autoregressive machine translation. In Proceedings of the 58th Annual Meeting of the Association for Compulational Linguistics, pp. 1692-1701, 2020.
Aidan P Thompson, H Metin Aktulga, Richard Berger, Dan S Bolintineanu, W Michael Brown, Paul S Crozier, Pieter In't Veld, Axel Kohlmeyer, Stan G Moore, Trung Dac Nguyen, et Lammps-a flexible simulation tool for particle-based materials modeling at the atomic, meso, and continuum scales. Computer Physics Communications, 271: 100111, 2022.
Justin M Tumey, Andrew Simmonett, Robert M Parrish, Edward G Hohenstein, Francesco Evangelista, Justin T Fermann, Benjamin Mintz, Lori A Burs,, Jeremiah J Wilke, Micah L Abrams, al. Psi4: an open-sounce ab initio electronic structure program. Wiley Interdisciplinary Reviews: Computational Molecular Science, 2(4):556-565 2012.
Oliver T Unke, Stefan Chmiela. Huziel Sauceda, Michael Gastegger, Igor Poltavsky, Kristof T Schutt, Alexandre Tkatchenko, and Klber-Robert Muller. Machine learning force fields. Chemical Reviews, 121(16):11442-0018 2021.
Yuanqing Wang, Kenichiro Takaba, Michael Chen, Marcus Wieder, Yuzhi Xu, John ZH Zhang, Kuang Yu, Xinyan Wang, Linfeng Zhang, Daniel J Cole, et al. On the design space between molecular mechanics and machine learning force fields, 2024.
Edgar Bright Wilson John Courtney Decius, and Paul Cross. Molecular vibrations: the theory of infrared and Raman vibrational spectra. Courier Corporation, 1980.
14

---

Published as a conference paper at ICLR 2025
A APPENDIX
A.1 RESULTS OF DISTILLING ESCAIP ON SPICE
Table 4: Results of distilling the EScAIP model trained on SPICE into specialized GemNet-T student MLFFs. (FM) indicates foundation model, while (S) indicates student model. n2n is the node feature matching baseline from (Kelvinius et al. (202)), while a2a is the atom embedding matching baseline we construct in $A.6) Student models all have identical simulation speeds. Speedups relative to FM are given in parentheses. The "Speed" column is calculated with a constant batch size of 64 for all models, while the "Maximum Speed" calculated with the batch size that maximizes throughput for each model.
Chemical Subgroup
Datasel Size
Model
Distillation Method
MAE (me V/A)
Energy MAE (met/a)om)
Speed (msdday) (t)
Maximum Speed (ns/day) (1)
(Parameter Count)
(PM) EScAIP Large (45M)
3.48
0.41
22.6
22,6
Memorners
14,331
Undistilled
7,9
3.8
n2n
(S) GemNet-T (0.67M)
a2a
8.9
1.9
Hessian (ours)
66.9 (2.9x)
(18.1x)
(PM) ES:AIP Large (45M)
11.52
0.61
19.6
19.6
Solvated Amino Acids
805
Undistilld
18.3
1.7
n2n
(S) GemNei- T (0.67M)
a2a
17.8
1.6
Hessian (ours)
10.0
23.4 (1.28)
23.2
(PM) ESP Large (45M)
11.4
0.73
22,9
22,9
Systems with
11,171
Undistiled
15.9
4.0
lodine
(S) GemNet-T (0.67M))(
a2a
18.6
3.9
Hessian (ours)
10.9
2.2
65.4 (2.9))
128,0 (5.6x)
Complementary to the results in Section 4.1 (distilling the MACE-OFF FM on SPICE), we distill the recent EScAIP Qu & Krishnapriyan (2024) model trained on SPICE into specialized GemNet-T student models. Unlike MACE-OFF, EScAIP does not explicitly build in SO(3) equivariance. We use the version of EScAIP that employs direct force parameterization. As the current tate-of-theeart model on SPICE, EScAIP has lower energy and force errors than MACE-OFF on all three chemical subsets consider. As shown by the results in Tab. 4 we find that this translates improved results with our Hessian distillation procedure: the GemNet-T student models distilled from EScAIP outperform those distilled from MACE-OFF across all chemical subsets in both force and energy MAES. While distillation from MACE-OFF to GemNet-T occasionally led to increases in energy MAE relative to the undistilled model, these issues are eliminated when distilling from EScAIP. This underscores the promise of leveraging large-scale, unconstrained foundation model training procedures and subsequently distilling into specialized student MLFFs with the inductive biases required for the downstream task hand.
A.2 PHYSICAL INTUITION OF ENERGY HESSIAN
The energy Hessian corresponds to the curvature of the energy surface with respect to atomic displacements. The eigenvalues of the energy Hessian are the squares of the normal mode vibrational frequencies, while the eigenvectors represent the amplitudes of motion along each of the 3NV mass-weighted Cartesian coordinates associated with each mode (Jnsenn 2017). These vibrational frequencies are crucial for understanding the thermodynamic properties molecules, such as heat capacity, entropy, and free energy [Jensen] 2017). These frequencies can also be observed experimentally (Wilson et al. 1980). Energy Hessians are also directly used in geometry optimization algorithms to relax molecular structures (Fletcher 2000).
A.3 FOUNDATION MODELS
For the foundation (teacher) models, we use the publicly available, pretrained weights. For MACE-ON and MACE-MP, perform no additional imeetuning/modiification For MACE-OFF, the energy MAE values we obtained using the pretrained checkpoint are slightly higher than what is reported in the original paper (Kováes et al. 2023), while the force MAE values identical. For the JMP models, we finetune
15

---

Published as a conference paper at ICLR 2025
the publicly available, pretrained weights of both JMP-S and JMP-L on the buckyball catcher and double walled nanotube separately, using the exact configurations and hyperparameters provided the JMP repository. While were not able to exactly reproduce the force MAE results reported in the JMP paper (Shoghi et al. 2023) (our obtained losses are slightly higher), we did not perform further tuning experimentation due to limited computational resources. For EScAIP, obtained the pretrained SPICE model checkpoint directly from the authors.
We note that MACE-MP-O was trained on the entirety of MPtrj, and an official held-out test set publicly available. We created our own MPtr train//alidaatonntees splits for the student models. The reported Force MAES are measured on our created test split, which was seen by the FM model during its training, but not by our student models.
Below we provide the links to repositories where the foundation model weights were obtained, as well as the training data:
Mace-OFF repository
Mace-MP repository
JMP repository
Spice Dataset
MPtjj Dataset
We report the foundation model training times in Table A.3 and selected distillation training times in
Table both in GPU hours.
Table Training Times for Foundation Models (trained on A100 GPUs)
Model
Dataset
GPU hours
MACE-OFF Large
SPICE
336
Kovács
2023
MACE-MP Large
MPtrj
1920
Batatia et al
2023
JMP-Small
(Pretraining)
QM9 ANI-Ix + OC20 Translx
5700
Shoghi
2023
IMP-Small
(Finetuning)
Buckyball Catcher
9 (trained ourselves)
JMP-Small (Finetuning)
Double Walled Nanotube
20 (trained ourselves)
JMP-
(Pretraining)
QM9 ANI-Ix + OC20 + Translx
34400 (Shoghi et al.
2023
JMP- Large (Finetuning)
Buckyball Catcher
15 (trained ourselves)
JMP-Large (Finetuning)
Double Walled Nanotube
41 (trained ourselves)
A.4 STUDENT MLFF ARCHITECTURE DETAILS
We provide details on the architectures of the GemNet-dT, GemNet-T, PaiNN, and eSCN MLFFs used student models in this work. A slashed value indicates that different values were used across datasets. In this case, the values correspond to the ordering: SPICE, MD22, MPtrj. Since PaiNN was only used for SPICE and MPtjj, in that case the ordering is SPICE, MPtrj.
Table Hyperparameters for PaiNN student models.
Parameter
Value
Hidden Channels
128
Layers
Radial Basis Functions
128
Cutoff
12.0
Maximum Neighbors
50
Within dataset, hyperparameters for GemNet-dT and GemNet-T T models identical. The only difference that GemNet-T computes forces as the negative gradient of the energy prediction, while GemNet-dT employs direct force parameterization.
16

---

Published as a conference paper at ICLR 2025
Table 7: Hyperparameters for GemNet-dT and GemNet-T student models.
Parameter
Value
Number of Spherical
7
Radial Basis Functions
6
Blocks
4
Atom Emhedding Size
64
Edge Embedding Size
64
Triplet Embedding Size
32
RBF Embedding Size
16
CBF Embedding Size
16
Bilinear Triplet Embedding Size
64
Number Before Skip
1
Number After Skip
Number of Concatenations
1
Number of Atoms
2
Cutoff
5.0/ 6.0
Maximum Neighbors
50
RBF Function
Gaussian
Envelope Function
Polynomial (Exponent: 5)
CBF Function
Spherical Harmonies
Output Initialization
HcOrthogonal
Activation Function
SiLU
A.5 STUDENT MLFF TRAINING DETAILS
We use the FAIR-chem repository https: / coo/FAIRR Het.m//AR-CCe//////ee fairchemm implemented with PyTorch, for model training and evaluation. We include the training details below.
Table 8: Optimization hyperparameters for student models.
Parameter
GemnVe--IT/GeenNNe-TTSSCC
PaiNN
Initial Learning Rate
0.001
0.001
Optimizer
AdamW
AdamW
Weight Decay
0,000002
0.000002
Amsgrad
True
True
Adam epsilon
1.e-7
1.e7
Scheduler
RodlteeL
Ponlaeal
Patience
5
10
Factor
0.8
0.8
Minimum Learning Rate
0.000001
0.000005
EMA Decay
0.999
0.999
Clip Gradient Norm 10
10
We chose batch size and number of Hessian rows primarily based dataset and system sizes and how these would affect training times.
We report training times for Hessian distillation approach a percentage of the original foundation model training time. The caveat is that MACE foundation models were trained on faster, NVIDIA RTX A100 GPUs, while the disilled runs were trained slower NVIDIA RTX A6000 GPUs. Therefore, these percentages are likely overestimates.
A.6 BASELINES
We provide additional details on the n2n and aZa baselines against which we compare our Hessian distillation approach.
17

---

Published as a conference paper at ICLR 2025
Table 9: Loss Weights by Chemical Subset. The same energy and force weights used for undistilled and distilled training, as well as all baselines.
Training Set
Au
Xf
AKD
Avu
Monomers
100
400
Solvated Amino Acids
5
100)
400
5
Structures with lodine
100
400
Pr03m Spacegroup
100
200
Structures with Yttrium
0
100
200
0
Bandgap 2 5meV
100
200
Buckyball Catcher
5
100
400
5
Double Walled Nanotube
100
400
Table 10: Training Batch Size for Student Models by Chemical Subset. The same batch sizes are used for undistilled and distilled training.
Training Set
GemNet-dTT/GeMNee-T
PaiNN eSCN
Monomers Solvated Amino Acids Structures with lodine
4
8
Pm3m Spacegroup
16
16
Structures with Yttrium
16
16
Bandgap 5meV
32
Buckyball Catcher
4
4
Double Walled Nanotube
Table 11: Number of rows sampled from Hessian
Training Set
GemNet-dT/GeeNNe-T
PaiNN
eSCN
Monomers
Solvated Amino Acids
Structures with lodine
4
4
Pm13m Spacegroup
Structures with Yttrium
4
Bandgap 2 5
Buckyball Catcher
Double Walled Nanotube
Table 12:
Training Times in GPU-hours for Selected Distilled Runs. Percentages indicate the fraction of time the training run took compared to the training time of the relevant foundation model.
Training Subset
GemNet-dT
PaiNN
Monomers
72.5 (21.5%)
30.5 (9.1%)
Solvated Amino Acids
15.2 (4.5%)
9.3 (2.8%)
Structures with lodine
68.3 (20.3%)
29.8 (8.8%)
Pm3m Spacegroup
14.8 (0.8%)
7.5 (0.4%)
Structures with Yttrium
57.1 (3.0%)
65.003.4%
Bandgap 5meV
82.2 (4.3 %)
48.7 (2.5%)
Buckyball Catcher
87 (0.25 %)
Double Walled Nanotube
144 (0.42 %)
Node feature supervision (n2n). The n2n approach, introduced in (Kelvinius al. (023), seeks align the node features of the student MLFF with that of the teacher. Specifically, given node features I Rd and (0 E dR from the tth message-passing layer of the teacher and student respectively, the
18

---

Published as a conference paper at ICLR 2025
distillation loss is formulated
LxD =x==Eellhy (z) Ee) (1))1
(=)
where Pe Rexx,, is a learnable linear projection from the student to teacher representation space. The projection weights discarded inference time. (Kelvinius et al., 2023) found that that using the final layer node representation (l= L) yielded the best results, so we chose the same. As in the Hessian distillation setting, we pre-compute and save the teacher's final node features over the dataset prior to training. We found via a sweep over AKD = {10, 100,00000000 1e110000100000 10000, 1000000 that Ax 10000 yields the best results with GemNet-dT on the Solvated Amino Acid split of SPICE (results shown below). Due to limited computed budget, we do not perform sweeps over each individual split, and use this same value for all subsequent splits and models.
Table 13: Validation energy MAE (meV) and force MAE (meV/A) of the n2n baseline the Solvated Amino Acid split of SPICE, using different values of AKD. Energy MAE is total MAE, so there is not one-to-one correponndnnee between the per-atom MAE results reported in the main text.
AKD
Force MAE (me V/A)
Energy MAE (me V)
(Undistilled)
22.4
162
10
22.1
163
100
22.1
169
1000
21.7
142
10000
21.1
117
100000
28.1
144
Atom embedding initialization (a2a).
GNN-based MLFFs, the initial node features (0 are parameterized as learnable dictionary of embeddings for each atomic element. In the a2a approach, we precompute the teacher's atom embeddings h:p and parameterize the student atom embeddings as (0) PhT"' where P TRd, xd, is learnable linear projection from the teacher to student representation space. We train the MLFF using the original energy/force matching objective, with no distillation (i.c
AKD 0). There are no hyperparameters tune for the a2a baseline.
A.7 EVALUATION DETAILS
Reporting of Maximum Simulation Speed. We calculate maximum simulation speeds by performing energy/force inference with the batch size that maximizes throughout for each model. We report these batch sizes below.
Table 14: Throughput-maximizing batch sizes
Training Set
GemNet-dT
PaiNN
MACE
GemNet-T
eSCN
JMP-
JMP-L
Monomers
1024
1024
128
Solvated Amino Acids
64
64
32
Structures with lodine
512
512
64
Pm3m Spacegroup
512
512
512
Structures with Yttrium
128
512
128
Bandgap 2 5meV
64
64
64
Buckyball Catcher
64
24
32
4
Double Walled Nanotube
8
12
We convert inference speed from samples/second to nanoseconds/day by adopting a MD simulation timestep of femtosccond, and assuming that energy/force inference dominates simulation time.
A.8 SENSITIVITY TO KNOWLEDGE DISTILLATION WEIGHT
We assess the sensitivity of our Hessian distillation approach to the knowledge distillation loss weight A used during training. We vary the weight across A&D 400,000,,0,10000 the Solvated Amino
19

---

Published as a conference paper at ICLR 2025
Acids split of SPICE with GemNet-d'T model. We find that value of AKD 400 achieves the best balance of force and energy performance (full results below).
Table 15: Validation energy MAE (meV) and force MAE (me V/A) achieved by Hessian distillation on the Solvated Amino Acid split of SPICE, using different values of ÄKD. Energy MAE is total MAE, so there is not one-to-one correspondence between the per-atom MAE results reported in the main text.
AKD
Force MAE (me V/A)
Energy MAE (meV)
(Undistilled)
22,4
162
10
20.2
145
100
14.0
168
400
12.3
155
1000
13.7
181
We performed a similar sweep for MPTr and found that λKD = 100 was optimal in that setting, Increasing the KD weight generally helps up to a certain point, after which performance saturates and eventually degrades. This is an important hyperparameter that we recommend be tuned for each dataset independently.
A.9 EFFRCT OF TRACHER HESSIANS VERSUS FORCES.
To isolate the benefit of the FM Hessians, we formulate an alternative knowledge distillation objective where the student trained match the FM force predictions instead of its energy Hessians. We simply replace the ground truth forces in Eq. [ with the foundation model (teacher) forces with a weight > 100 (the same weight previously used for the ground truth forces). Concretely, the new loss function becomes,
C=Nr|Uia(Zxx)-Uo(an (o +Ap Σ p (z.) (z.r) (2.))
where fIM denotes the teacher forces and per usual, Use denotes the ground truth energies.
Results on selected splits of the SPICE dataset are presented in Tab. I6 Distilling with the teacher forces consistently inferior distilling with the teacher Hessians, indicating that the richer information contained in the latter helps the model better match the true forces. In fact, force distillation leads to worse performing models than training without distillation, similar to the observation that the n2n and a2a baselines were inferior to undistilled training on SPICE (.4.l)
Table 16; Ablation study looking at distilling with MACE-OFF forces on selected splits of the SPICE
dataset. This approach consistently inferior to our approach of distilling with MACE-OFF Hessians.
Chemical Subgroup
Student Model
Distillation Method
Force MAE (meV/))L)
Undistilled
11.3
GemNet-dT
Forces
14.2
Hessian (ours)
6.3
Monomers
Undistilled
25.0
PaiNN
Forces
25.0
Hessian (ours)
8.77
Undistilled
23.4
GemNet-dT
Forces
26.1
Hessian (ours)
14.7
Systems with lodine
Undistilled
51.2
PaiNN
Forces
51.7
Hessian (ours)
23.7
20

---

Published as a conference paper at ICLR 2025
A.10 COMPARISON OF HESSIAN DISTILLATION TO CONSERVATIVE FORCE TRAINING
We compare the benefits obtained from Hessian distillation of direct force student models (GemNet-dT) with training gyadient-based MI.FFs (GemNet-T) without distillation. Results on SPICE splits are shown in Tab. A0
Table 17:
Force MAE V/A) of conservative force training (GemNet-T) on SPICE test splits, as compared with undistilled and distilled training with GemNet-dT.
Chemical Subset
GemNet-dTT (Undistilled) GemNet-T (Undistilled)
GemNet-dT (Distilled)
Monomers
11.3
8.8
6.3
Solvated Amino Acids
22.4
31.0
11.6
Systems with Todine
22.6
16.0
14.7
Although training with gradient-based forces yields benefits over training with direct forces in the undistilled setting on 2 out of the 3 chemical subsets, we find that the improvements are less than those achieved by Hessian distillation. We hypothesize that while the inductive bias of conservative forces is beneficial, the extra supervision provided by Hessian distillation a stronger learning signal to learn on chemical subsets with potentially limited data. We note that the improvements from adding conservative forces incurs roughly a 3 x increase in inference time due to the extra backpropagation step to compute forces, while the distilled GemNet-dT model has the identical architecture as the undistilled model and incurs no such cost. However, for applications such as constant energy (NVE) MD simulations, gradient- based forces may be required to produce physical and stable simulations. We thus combine both of these elements in $ and $4..3 and distill with a gradient-basse GemNet-T student model to yield greater force MAE improvements while maintaining test-time energy conservation
A.11 COMPUTING HESSIANS WITH FINITE DIFFERENCES
In situations where computing the Hessian via autodifferentiation is unfeasible, such for extremely large foundation models (FMs), models employing attention kernels that are not diice-iifferentialee (Lefaudeux et al.] 2022). or conservative student models where training would require 3rd order gradients, a finite difference approach can be used instead.
In this scheme, molecular structures are perturbed in Fuclidean space and energy/force derivatives are obtained via discretized stencil (in this case, right difference scheme). While this work focuses on cases where autodifferenniation is feasible, we also demonstrate that the finite difference approach is viable alternative. We provide results on the Solvated Amino Acids split in Tab, 18 as proof concept, showing that the method accelerates training for conservative force student models. The difference in final Force MAE achieved by the two methods is negligible.
Table 18: Computing Hessian with Autodifferentiaaion vs Finite Differences with GemNet-T (conservative forces) on Solvated Amino Acids.
Computation Method
Training Speed (Epoch/Min)
Autodifferentiation Finite Differences
1.27
1.67
A.12 EFFECTS OF ADDING A LOSS SCHEDULER TO HESSIAN DISTILLATION
While training via Hessian distillation, we reduce the distillation loss coefficient AKD by 1/2 when the student model exceeds the teacher model's accuracy. This approach dynamically adjusts the training process to focus more matching the ground truth energies and forces rather than the foundation model's I Hessians. Below, we compare Hessian distillation with and without the scheduler on a dataset, demonstrating that the scheduler provides small but consistent improvements in Force MAE. We apply the same scheduling to the n2n and a2a baselines, but in practice the baselines never exceed the teacher accuracy, so the scheduling criterion not met.
21

---

Published as a conference paper at ICLR 2025
Table 19: Comparison of Hessian Distillation with and without Loss Scheduler, training with GemNet-dT on Solvated Amino Acids.
Scheduler
Force MAE
Without Scheduler
12.2
With Scheduler
11.6
A.13 ISOLATING THE EFFECT OF HESSIAN DISTILLATION ON ENERGY ACCURACY
is of interest see how Hessian distillation alone, without energy gradient supervision term described in $3.4) affects energy accuracy. Below in Tab, 20 we compare a training run using only Hessian distillation with one that combines Hessian and energy gradient supervision. While Hessian distillation alone improves Energy MAF, the addition of energy gradient supervision leads to even greater improvements.
Table 20: Ablation investigating the effects of the Hessian and Energy gradient supervision terms when training PaiNN on Solvated Amino Acids.
Distillation Method
Energy MAE
Undistilled
3.3
Hessian Distillation only
2.9
Hessian Distillation Energy Gradient Supervision
0.41
A.14 NVT MD SIMULATIONS WITH STUDENT MODELS
To further evaluate the distilled, specialized MLFFs from $4.1 we run 100 picosecond, constant- temperature (NVT) MD simulations with systems from the Solvated Amino Acid split, with molecules containing 79-96 atoms each. We choose 5 random structures from the held-out test set as initial conditions. We perform Langevin dynamics at a temperature of 300K, a timestep of 1.0 fs, and a friction coefficient of 0.01 fs -, for 100,000 steps, corresponding to 100 ps. Consistent with (Fu et al. 2022), use a maximum bond length deviation metric, which measures unphysical bond stretching or collapse, to measure stability. According this criterion, simulation becomes unstable time 7 if,
(T) bb))) (KJ)
max
where B is the set of all bonds, i. are the two endpoint atoms of the bond, and by.j the equilibrium bond length computed from the training dataset. Following (Fu et al. 2022). we set A = 0.5.4. Since we are simulating non-reactive systems at ambient conditions, bond deviations exceeding this amount are indicative of simulation failure.
Results are shown in Fig. We find that the improvements in Force MAE between the distilled and undis- tilled MLFFs shown in Tab. translate to considerably improved stability over time. We reiterate that both the GemNet-dT and PaiNN student models lack a conservative force parameterization, which generally considered crucial for stable MDD simulations. Our results thus may indicate that distillation from FM teacher employing conservative forces may be sufficient to achieve stable simulation without this inductive bias. We leave more complete analysis of the performance of distilled MLFFs in MD simulations, includ- ing the capturing of observables over long timescales (Fu et al. 2022 Raja et al. 2024), for future work.
A.15 GEOMETRY OPTIMIZATION WITH STUDENT MODELS
As an additional evaluation of the usefulness of our student MLFFs, we perform geometry optimization with our GemNet-dT student models using the FIRE (Bitzek et al. 2006) optimizer the Atomic Simulation Environment (ASE). We select 100 structures from the Monomers split of the SPICE dataset, and run optimization until all the per-atom force norms are below 0.05 eV/A. Finally, we compute the energy and per-atom forces using Density Functional Theory (DFT) at the level of theory (the same level used to generate the dataset in (Eastman et al. (023)). We run DFT calculations with the default settings in the psi4 Python package (Tumey al. 2012).
22

---

Published as a conference paper at ICLR 2025
(a)
(b)
Max bond length deviation (A)
104
Stability Threshold
Stability Threshald
GamNet-dT Undistilled
10a
PaiNN Undistilled
Distilled
PaiNN Distilled
100
10a
10a
10-1
10-1
10"2
20
40
60
80
100
0
20
40
60
80
100
Simulation time (ps)
Simulation time (ps)
Figure 4: Stability Constant Temperature MD Simulations. Results of constant temperature (NVT) MD simulations using the distilled GemNet-dT and PaiNN student MLFFs. We plot the maximum bond length deviation during NVT simulations of 5 selected systems from the SPICE Solvated Amino Acid split. denotes the point which the simulation becomes unstable. Our distilled models considerably more stable than their undistilled counterparts, both for (a) GemNet-dTT and (b) PaiNN.
We find that our distilled GemNet-dT model generally converges to structures with lower energy and mean per-atom force norms than its undistilled counterpart. While the FIRE optimizer does not explicitly use energy Hessians, many quasi-Newton algorithms like BFGS (Bitzek al. (0066) do. It would be interesting to explore whether our Hessian distillation approach leads to even greater gains for these optimizers.
(a)
(b)
Energies After Geometry Optimization
after Geometry Optimization
80
Raame: Fanisctde (meV)
GemWet-dT Unditilld
GemWict-4 4T Ossiled
form (meviA)
40
-20
20
-40
20
40
60
100
40
60
80
Structure
Figure Geometry optimization with GemNet-dTT student MIFFs. (a) Difference in energy the final, relaxed structure obtained via the distilled and undistiled models. On average, the distilled model converges to lower energy structures. (b) Mean per-atom force norm of the final, relaxed structure obtained via the distilled and undistiled models. On average, the distilled model converges to lower structures with lower force norms.
23


Reviews:
[{"review_id": "2WWu97B3rV", "reviewer": "Reviewer_NWzu", "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.", "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?\n- 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?\n- Why don't you compare the results with Kelvinius et al. 2023 on OC20-2M and COLL on which the original work was performed?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The distilled MLFFs are up to 20x faster than their foundation model counterparts, enabling more efficient simulations.\n- The distilled MLFFs achieve comparable or even better force prediction accuracy than the original FMs and demonstrate improved MD stability results.\n- The Hessian distillation method is model architecture agnostic.\n- 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.\n- An anonymous link to the code is not available.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Ja6ffNaELK", "reviewer": "Reviewer_8NSg", "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.", "questions": "See “Weaknesses” for detailed questions and suggestions.\n\nPotential typos:\n* Line 420: “disilling” should be “distilling.”", "limitations": "", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 2, "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.\n* 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.\n* 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.\n\t* 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.\n\t* 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.\n* 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "YFPAYibyRh", "reviewer": "Reviewer_yqHS", "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.", "questions": "* The tables in the paper show only force results, without energy, so I'm curious about the energy results after distillation.\n* 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.\n* 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.\n* 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.\n\n[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.\n\n[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.", "limitations": "", "rating": 6, "confidence": 5, "soundness": 3, "presentation": 3, "contribution": 3, "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.\n* The paper is written in a clear and concise manner, facilitating effortless understanding.", "weaknesses": "Related concerns are discussed in the questions section.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Xz6oAsbJqy", "reviewer": "Reviewer_8TDr", "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.", "questions": "* Why is there a tank in Figure 1?\n* 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.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. Clarity and Readability: The paper is very well-written and accessible, making the proposed method easy to follow.\n2. Practicality: The approach is straightforward to implement and is cost-effective relative to similar knowledge distillation methods.\n3. Experimental Validation: MD simulations performed validate the method, underscoring the practical benefits of the proposed technique.\n4. 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?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)