Research map / Architectures and efficiency / RNNs and normalization
Towards Fast, Specialized Machine Learning Force Fields: Distilling Foundation Models via Energy Hessians
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Distills MLFF foundation models (MACE-OFF, MACE-MP-0, JMP) into small specialized student force fields by matching teacher energy Hessian rows, yielding up to 20x faster inference with equal or better accuracy.
Problem
Machine learning force field foundation models are accurate and general but too slow for large-scale MD and are not specialized to the chemical subsets practitioners study; direct-force models also lack energy conservation.
Approach
Train a small student MLFF on a chemical subset with a loss that aligns its force Jacobian (energy Hessian) rows, subsampled per iteration via VJPs, to those of the foundation model teacher, in addition to standard energy and force matching.
Key results
- On SPICE monomers, Hessian-distilled GemNet-dT reaches 6.3 meV/Å force MAE vs 6.6 for the MACE-OFF Large teacher and 11.3 undistilled, at up to 19x max-throughput speedup.
- Distilled PaiNN on solvated amino acids reaches 18.0 meV/Å force MAE vs 50.1 undistilled and runs 20.9x faster than MACE-OFF Large.
- Distilling JMP-L into GemNet-T gives energy conservation over 100 ps NVE simulations of the buckyball catcher, where the JMP-L energy drifts; distillation costs about 4% extra training compute for s=4.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Apply Hessian distillation to specialized students instead of undistilled training | GemNet-dT students distilled from MACE-OFF · SPICE Monomers | Force MAE (meV/Å): 11.3 -> 6.3 | Helped |
| Apply Hessian distillation to specialized students instead of undistilled training | GemNet-dT students distilled from MACE-OFF · SPICE Solvated Amino Acids | Force MAE (meV/Å): 22.4 -> 11.6 | Helped |
| Apply Hessian distillation to specialized students instead of undistilled training | GemNet-dT students distilled from MACE-MP0 · MPtrj Systems with Yttrium | Force MAE (meV/Å): 32.5 -> 21.3 | Helped |
| Use Hessian distillation from JMP-L | GemNet-dT student · MD22 Buckyball Catcher | Force MAE (meV/Å): 8.0 -> 5.1 | Helped |
| Use Hessian distillation from JMP-L | GemNet-dT student · MD22 Double Walled Nanotube | Force MAE (meV/Å): 14.3 -> 10.6 | Helped |
| Replace teacher Hessians with teacher forces for distillation (ablation) | GemNet-dT · SPICE Monomers | Force MAE (meV/A): 6.3 -> 14.2 | Hurt |
| Reduce sampled Hessian rows to s=1 (ablation) | Hessian-distilled students · SPICE Solvated Amino Acids | Force MAE | No clear effect |
| Increase the number of sampled Hessian rows s (ablation) | Hessian-distilled students · SPICE Solvated Amino Acids | Force MAE | No clear effect |
| Add a loss scheduler that halves KD weight after student surpasses teacher accuracy (ablation) | GemNet-dT Hessian distillation · SPICE Solvated Amino Acids | Force MAE: 12.2 -> 11.6 | Helped |
| Add energy gradient supervision to Hessian distillation (ablation) | PaiNN · SPICE Solvated Amino Acids | Energy MAE: 2.9 -> 0.41 | Helped |
| Distill Hessians from JMP-L rather than JMP-S (ablation) | Specialized student models · MD22 Buckyball Catcher | Force MAE: GemNet-dT: 5.1 -> 5.1; GemNet-T: 5.0 -> 4.0 | Mixed |
Surprising findings
- Distilled student models often outperform their teacher foundation model on the specialized subsets, even with the teacher's higher force error.
- Distilling teacher forces instead of Hessians was consistently worse than undistilled training in some cases.
Limitations
- Training with Hessian distillation adds overhead that scales with number of sampled Hessian rows.
- Hessian KD weight is a sensitive hyperparameter that must be tuned per dataset; GemNet-T energy accuracy sometimes worsens.
- Teacher Hessians are precomputed via autodifferentiation, which is costly for very large FMs.
Experimental setup
Datasets: SPICE MPtrj MD22 OC20 OC22 ANI-1x Transition-1x QM9
Benchmarks: SPICE subsets (monomers, solvated amino acids, iodine) MPtrj subsets (Pm-3m spacegroup, Yttrium, band gap >= 5) MD22 buckyball catcher and double-walled nanotube
Metrics: force MAE energy MAE simulation speed (ns/day) NVE energy conservation NVT bond-deviation stability
Baselines: Undistilled n2n a2a MACE-OFF Large MACE-MP-0 JMP-S JMP-L
Builds on: Sobolev training node-to-node feature distillation for MLFFs MACE GemNet JMP
Closest papers by meaning
- EDBench: Large-Scale Electron Density Data for Molecular Modeling (NeurIPS 2025 · similarity 0.88)
- Physics-Informed Weakly Supervised Learning For Interatomic Potentials (ICML 2025 · similarity 0.88)
- Colour me shocked: Exact Molecular Hessians from MLIPs in O(N) time using sparse differentiation! (NeurIPS 2026 · similarity 0.87)
- FlashTP: Fused, Sparsity-Aware Tensor Product for Machine Learning Interatomic Potentials (ICML 2025 · similarity 0.87)
- FlashMD: long-stride, universal prediction of molecular dynamics (NeurIPS 2025 · similarity 0.87)
- The dark side of the forces: assessing non-conservative force models for atomistic machine learning (ICML 2025 · similarity 0.87)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.