PFT: Phonon Fine-tuning for Machine Learned Interatomic Potentials

This paper introduces Phonon Fine-tuning (PFT), a scalable method that directly supervises machine learned interatomic potentials with DFT-derived force constants to significantly improve the accuracy of vibrational and thermal properties by correcting curvature errors in the potential energy surface.

The Gist

Imagine you are trying to teach a robot chef to cook a perfect meal. You show the robot thousands of recipes (data) and tell it, "Make sure the final dish tastes right (energy), the ingredients are chopped to the right size (force), and the pot isn't too heavy (stress)."

The robot gets very good at this. It can predict the taste and weight of the dish almost perfectly. However, there's a problem: the robot doesn't quite understand the texture or the snap of the food. If you ask it to predict how the food vibrates when you tap it, or how much heat it holds, the robot fails. It's because the robot learned the result of the cooking, but not the curvature of the recipe itself—how the flavors change if you tweak an ingredient just a tiny bit.

This paper introduces a new training method called Phonon Fine-tuning (PFT) to fix exactly that problem for materials science.

The Problem: The "Flat" Map

In the world of materials, scientists use a "Potential Energy Surface" (PES). Think of this as a giant, 3D topographic map of a mountain range.

• The Valley: The bottom of the valley is where a material is stable (like a ball sitting at the bottom of a bowl).
• The Slope: How steep the sides are tells you how hard it is to push the material (Force).
• The Curvature: How "bowl-shaped" the bottom is tells you how the material vibrates.

Standard AI models for materials are great at finding the bottom of the valley and measuring the slope. But they often get the curvature wrong. They might think the bowl is flat when it's actually deep and round, or vice versa. Because of this, they can't accurately predict how the material vibrates (phonons), how much heat it holds, or how well it conducts electricity.

The Solution: PFT (The "Vibration Coach")

The authors created a new training technique called Phonon Fine-tuning (PFT). Instead of just showing the robot the final dish, they now show it the vibrations of the ingredients.

1. Direct Supervision: They take the AI model and force it to match the "curvature" of the map directly. They compare the AI's math against a super-accurate reference (called DFT) that calculates exactly how the atoms push and pull on each other when shaken.
2. The "Stochastic" Shortcut: Calculating the curvature for a giant crystal (a supercell with thousands of atoms) is usually like trying to measure every single grain of sand on a beach. It's too slow and expensive.
   • The Analogy: PFT is like hiring a scout to walk the beach and randomly pick a few handfuls of sand to measure, rather than measuring the whole beach. By doing this randomly but smartly, the AI learns the shape of the whole beach without needing to count every grain. This makes the training fast enough to run on standard computers.
3. The "Co-training" Safety Net: There's a risk that if you teach the robot too much about vibrations, it might forget how to cook the basic meal (this is called "catastrophic forgetting").
   • The Fix: The authors use a "co-training" strategy. They alternate between teaching the robot about vibrations (PFT) and teaching it the original basic recipes (standard data). This keeps the robot sharp at both tasks, ensuring it doesn't lose its original skills.

The Results: Sharper Predictions

When they tested this new method on a model called Nequix MP:

• Vibrations: The model's ability to predict how materials vibrate improved by 55% on average.
• Heat: It got much better at predicting heat capacity and thermal conductivity (how well heat moves through the material).
• The "Third-Degree" Bonus: Even though they only trained the model on second-order vibrations (the "bowl shape"), the model accidentally got better at predicting third-order effects (how the bowl shape changes if you push it really hard). This is like learning to balance a ball in a bowl and suddenly getting better at juggling three balls.

Why It Matters

This isn't just about making a better math model; it's about making materials discovery faster and more accurate. By fixing the "curvature" of the AI's understanding, scientists can now trust these models to predict real-world properties like:

• How much a material expands when heated.
• How well a battery material conducts heat.
• Whether a new material will be stable or fall apart.

In short, PFT takes a smart AI that knows where things are, and teaches it to understand how things move and vibrate, all without forgetting what it already knew.

Technical Summary

Technical Summary: Phonon Fine-tuning for Machine Learned Interatomic Potentials (PFT)

Problem Statement

Machine Learned Interatomic Potentials (MLIPs) have become essential surrogates for Density Functional Theory (DFT) in high-throughput materials screening. However, standard MLIPs are typically trained on energy, forces, and stress (EFS) from equilibrium or near-equilibrium structures. While effective for predicting ground-state energies, this training regime often fails to accurately capture the curvature of the Potential Energy Surface (PES). Since vibrational properties (phonons) and thermodynamic quantities depend on the second-order derivatives of the PES (force constants), errors in curvature lead to significant inaccuracies in phonon dispersions, vibrational entropy, and heat capacity.

Furthermore, training directly on phonon data presents computational challenges. Phonon calculations require large supercells to avoid self-interaction of displaced atoms and to capture long-range interactions. The resulting force constant matrix (Hessian) scales quadratically ($O(N^2)$) with the number of atoms, making full Hessian computation and training infeasible for the large supercells necessary for accurate phonon modeling. Additionally, fine-tuning exclusively on phonon data (which consists largely of equilibrium geometries) risks "catastrophic forgetting," degrading the model's performance on the diverse, non-equilibrium configurations present in its original pre-training datasets.

Methodology

The authors introduce Phonon Fine-tuning (PFT), a procedure designed to directly align the curvature of an MLIP's PES with DFT-derived force constants while maintaining scalability and preserving upstream performance.

1. Loss Function Formulation:
   PFT modifies the standard training objective by adding a term that minimizes the error between the MLIP's analytical energy Hessian and the DFT-computed force constants ($\Phi$). The total loss is a weighted sum:
   $$L_{PFT} = \lambda_E L_E + \lambda_F L_F + \lambda_\sigma L_\sigma + \lambda_\Phi L_\Phi$$
   where $L_\Phi$ measures the Mean Absolute Error (MAE) between the predicted second derivatives and the target force constants.

2. Scalable Training via Hessian-Vector Products (HVP):
   To address the $O(N^2)$ scaling of the Hessian in large supercells, PFT employs stochastic sampling. Instead of computing the full Hessian, the method samples a single atom and Cartesian direction per structure in a batch, effectively selecting one column of the Hessian.

   • This column is computed using a Hessian-Vector Product (HVP): $\nabla^2_r \hat{E}(r)v = \nabla_r (\nabla_r \hat{E}(r)^\top v)$.
   • This approach reduces the computational complexity of a training step from quadratic to linear ($O(N)$) with respect to the number of atoms.
   • The implementation utilizes forward-mode Jacobian-vector products through reverse-mode gradients (e.g., in JAX), allowing the calculation of loss across a batch with a single HVP call, enabling training on GPUs with limited memory.

3. Co-training Strategy:
   To mitigate catastrophic forgetting, the authors propose a co-training scheme. During the fine-tuning process, training steps on the phonon dataset ($D_{phonon}$) are interleaved with steps on the original upstream dataset ($D_{up}$, e.g., Materials Project relaxation trajectories). The ratio of upstream to phonon steps ($K$) is tuned to balance the preservation of general EFS performance with the improvement of curvature accuracy.

4. Analytical vs. Finite Displacement:
   The framework supports computing force constants analytically via automatic differentiation (AD) rather than relying solely on finite displacement calculations, ensuring consistency with the model's energy conservation properties.

Key Contributions

• Direct Curvature Supervision: A fine-tuning objective that explicitly aligns MLIP energy Hessians with DFT force constants, addressing the "softening" of PES curvature often observed in standard EFS-trained models.
• Scalable Hessian Training: A strategy enabling training on large supercells by sampling Hessian columns and utilizing HVPs, reducing computational cost from quadratic to linear.
• Catastrophic Forgetting Mitigation: A simple co-training recipe that interleaves upstream data to maintain performance on general tasks while improving vibrational properties.
• Empirical Correlation: Demonstration that Hessian error strongly correlates with errors in multiple phonon thermodynamic metrics, validating the approach of direct curvature supervision.

Results

The authors evaluated PFT by fine-tuning the Nequix MP foundation model (trained on MPtrj) on the MDR Phonon benchmark.

• Phonon Properties: PFT improved the average Mean Absolute Error (MAE) across phonon thermodynamic properties (maximum frequency, vibrational entropy, Helmholtz free energy, heat capacity) by 55% compared to the base Nequix MP model. The PFT-tuned model achieved state-of-the-art accuracy among models trained on Materials Project trajectories.
• Third-Order Derivatives: Despite only supervising second-order constants, PFT generalized to improve third-order force constants, reducing their MAE by 20–30% across different base models.
• Thermal Conductivity: Leveraging the improved third-order derivatives, PFT improved thermal conductivity predictions on the Matbench Discovery benchmark. The Nequix MP PFT model achieved a symmetric relative mean error ($\kappa_{SRME}$) of 0.307, surpassing other MPtrj-trained models and approaching the performance of models trained on significantly larger datasets.
• Generalization and Stability:
  • Co-training Effect: Without co-training, fine-tuning on phonon data degraded the model's performance on the upstream MPtrj validation set (energy/force/stress errors increased significantly). Co-training largely eliminated this degradation with only a marginal trade-off in Hessian accuracy.
  • Matbench Discovery: On the geometry optimization and stability classification tasks (which rely heavily on energy accuracy rather than curvature), PFT models maintained performance comparable to the base model. Notably, the non-co-trained PFT model suffered a 60% drop in F1 score for stability classification, whereas the co-trained model reduced this drop to less than 1%.
• Efficiency: The fine-tuning process required approximately 35 A100 GPU hours (with co-training), a fraction of the cost required for the initial pre-training of the base models.

Significance and Claims

The paper claims that PFT provides a model-agnostic, scalable, and efficient method to significantly enhance the accuracy of MLIPs for vibrational and thermal properties without sacrificing their general utility. By directly supervising the curvature of the PES, the authors demonstrate that it is possible to correct the specific deficiencies of standard EFS-trained models regarding phonon predictions.

The work highlights that:

1. Curvature matters: Directly training on force constants yields superior results compared to indirect training via perturbed structures.
2. Scalability is achievable: Stochastic sampling of Hessians makes high-fidelity phonon fine-tuning feasible for large supercells.
3. Transferability exists: Improvements in second-order derivatives naturally generalize to third-order properties (thermal conductivity), suggesting a robust representation of the PES anharmonicity.
4. Practicality: The method preserves the model's original capabilities on non-equilibrium tasks through co-training, making it a viable strategy for refining foundation models in materials science workflows.

The authors conclude that while PFT significantly improves properties dependent on higher-order derivatives, the evaluation is currently limited to materials near equilibrium, and further research is needed to assess performance in highly off-equilibrium regimes.