Pushing the limits of unconstrained machine-learned interatomic potentials

This paper demonstrates that machine-learned interatomic potentials trained on large datasets without strict physical constraints can outperform traditional constrained models in both accuracy and speed, with simple inference-time corrections ensuring their reliability for practical static simulations.

The Gist

Imagine you are trying to build a digital twin of the physical world, atom by atom. You want to simulate how materials behave, how drugs interact with proteins, or how batteries charge. To do this, you need a "rulebook" that tells every atom how to push and pull on its neighbors. This rulebook is called an interatomic potential.

For decades, scientists had two choices for this rulebook:

1. The Old School Manual: Simple, hand-crafted rules. They are fast to read, but they are often wrong because the real world is too complex for simple math.
2. The Quantum Supercomputer: Extremely accurate, simulating the laws of physics from the ground up. But it's so slow and expensive that you can only simulate a few atoms for a split second.

Enter Machine-Learned Interatomic Potentials (MLIPs).
These are AI models trained on the Supercomputer's data. They promise the speed of the Old School Manual with the accuracy of the Quantum Supercomputer.

The Big Debate: Strict Rules vs. Free Thinking

Traditionally, when building these AI models, scientists forced them to follow strict physical laws, like rotational symmetry (if you spin the whole molecule, the energy shouldn't change) and energy conservation (energy can't just appear or disappear).

Think of it like teaching a child to draw a perfect circle.

• The Constrained Approach: You put a stencil in their hand. They must draw a perfect circle every time. It's safe, but maybe they can't learn to draw anything else, and the stencil makes their hand move slower.
• The Unconstrained Approach: You tell the child, "Just draw what you see." They might draw a wobbly circle at first, but if you show them enough pictures of circles, they eventually learn the concept of a circle on their own. They might draw faster and learn to draw other shapes too.

For a long time, scientists thought the "stencil" (constrained models) was necessary. If you let the AI break the rules, it might make weird, impossible predictions.

What This Paper Discovered

The authors of this paper decided to test the "Free Thinking" approach on a massive scale. They built a new AI model (called PET) that doesn't have a stencil. It doesn't know the rules of rotation or energy conservation; it has to learn them entirely from the data.

Here is the surprising twist: The "Free Thinking" model turned out to be better.

1. It's Faster: Because it doesn't have to stop and check a stencil for every single calculation, it runs significantly faster. It's like a sprinter who doesn't have to check a map at every step.
2. It's Smarter (When Trained Big): When you feed this model a massive amount of data (millions of atomic configurations), it learns the rules of physics so well that it becomes just as accurate as the strict models. In fact, because it has more freedom to learn complex patterns, it sometimes beats the strict models.
3. The "Wobbly Circle" Problem: The only downside is that sometimes, because it learned the rules on its own, it might get slightly "wobbly." For example, if you rotate a molecule, the AI might calculate the energy as slightly different than before.
   • The Fix: The authors found that you can easily fix this "wobble" at the very end. It's like taking a slightly crooked photo and using a simple "straighten" tool in Photoshop. You don't need to rebuild the camera; you just fix the picture after you take it.

Real-World Applications

The team tested this new model on two big challenges:

• Finding New Materials (The "Crystal Hunter"): They used it to scan thousands of crystal structures to find the most stable ones. The unconstrained model was just as good at finding the "gold" as the strict models, but it did it faster.
• Simulating Molecules (The "Molecular Movie"): They used it to simulate how molecules move and vibrate. Even though the model didn't have strict energy conservation built-in, they found that by using a clever trick (mixing the fast, "wobbly" forces with a few slow, "perfect" checks), they could run long, stable simulations that were incredibly accurate.

The Bottom Line

This paper is a game-changer because it suggests we don't need to force AI to follow every single rule of physics from the start. Instead, we can let the AI learn the rules itself, provided we give it enough data.

The Analogy:
Think of the old way as training a robot to walk by strapping it to a rail. It never falls, but it can only walk in a straight line.
The new way is letting the robot walk in an open field. It might stumble a few times, but if you show it enough videos of walking, it learns to balance on its own. And once it learns, it can run faster and jump over obstacles that the robot on the rail never could.

In short: By trusting the data more than the rules, the authors created AI models that are faster, cheaper, and just as accurate as the best models we have today. This opens the door to simulating larger, more complex systems than ever before.

Technical Summary

1. Problem Statement

Machine-learned interatomic potentials (MLIPs) have become essential for bridging the gap between the accuracy of quantum-mechanical methods (like DFT) and the computational efficiency required for large-scale molecular dynamics and materials discovery. However, a prevailing paradigm in MLIP design is the strict enforcement of physical symmetries (rotational invariance, energy conservation, etc.) directly into the model architecture (e.g., using equivariant neural networks).

While these constraints ensure physical consistency, they often limit architectural expressivity and increase computational overhead. The authors investigate a critical, under-explored question: Can fully unconstrained architectures (those that do not hard-code rotational symmetries or energy conservation) scale to massive datasets and outperform state-of-the-art constrained models in both accuracy and speed?

2. Methodology

A. Architecture: Modified PET

The authors utilize and modify the PET (Point-Edge-Transformer) architecture, a Graph Neural Network (GNN) that processes atomic neighborhoods using Transformer layers. Key modifications to the original PET include:

• Modern Transformer Components: Implementation of Root Mean Square Layer Normalization (RMSNorm), SwiGLU activation, and pre-normalization.
• Asymmetric Feature Scaling: Increasing the dimensionality of node features by a factor of four relative to edge features. Since chemical structures typically have 30–50 times more edges than nodes, this allows for a massive increase in parameter count with negligible computational overhead.
• Direct Force Heads: The model predicts forces directly as an additional output head rather than deriving them via automatic differentiation of the energy. This avoids a backward differentiation step, offering a 2–3× speedup.
• Adaptive Cutoff: An adaptive cutoff strategy is employed to maintain a roughly constant number of neighbors across materials with varying densities, ensuring uniform computational cost.

B. Training Strategy

• Datasets: The models were trained on three massive, diverse datasets: MPtrj, Alexandria (subsampled), and OMat24, resulting in a 730M parameter model (PET-OAM). A smaller 190M parameter model was trained on the SPICE molecular dataset.
• Two-Stage Training:
  1. Pre-training: Models are first trained using non-conservative forces and stresses (direct prediction). This stage is faster and allows the model to learn the underlying data distribution and approximate symmetries.
  2. Fine-tuning: The models are then fine-tuned using conservative counterparts (forces derived from energy gradients) to ensure physical consistency for specific applications.
• Symmetry Handling: Instead of hard-coding symmetries, the models rely on data augmentation (rotations and inversions) during training to learn invariance. At inference time, rotational symmetrization (averaging predictions over a grid of rotations) is applied to recover exact physical symmetries for static calculations.

C. Inference-Time Corrections

The paper proposes practical workflows to mitigate symmetry breaking in unconstrained models:

• Molecular Dynamics: Using random rotations at every time step or multiple time-stepping algorithms to average out errors.
• Geometry Optimization: Using permissive symmetry detection thresholds or projecting incompatible force components to restore symmetry.
• Phonon Calculations: Symmetrizing the Hessian matrix and averaging over rotations to eliminate spurious soft modes.

3. Key Contributions

1. Scalability of Unconstrained Models: Demonstrated that unconstrained architectures can be scaled to hundreds of millions of parameters and trained on massive datasets (billions of configurations) without sacrificing accuracy.
2. Superior Efficiency: Showed that unconstrained models, particularly those using direct force heads, achieve significantly faster inference times (2–3× speedup) compared to equivariant counterparts while maintaining competitive accuracy.
3. Empirical Validation of "Learned" Symmetries: Provided evidence that with sufficient data, models can effectively learn rotational invariance and energy conservation from the data distribution, rather than requiring them to be hard-coded.
4. Practical Workflows: Established robust protocols for using unconstrained models in static simulations (geometry optimization, lattice dynamics) by applying simple inference-time corrections (averaging, symmetrization).

4. Results

A. Benchmark Performance

• Matbench-Discovery: The PET-OAM model achieved state-of-the-art performance on materials discovery metrics (DAF, Accuracy, F1 score), matching or exceeding top equivariant models like eSEN, NequIP, and SevenNet.
• SPICE (Molecular Domain): On the SPICE dataset, the unconstrained PET model outperformed leading models (MACE, eSEN) in both energy and force accuracy (e.g., 0.09 meV/atom energy MAE on PubChem vs. 0.88 for MACE-S).
• Accuracy-Speed Trade-off: Pareto front analysis showed that unconstrained models occupy the optimal region, offering the best balance between high accuracy and low inference latency.

B. Static Simulations (Geometry & Phonons)

• Geometry Optimization: Unconstrained models successfully relaxed unstable high-symmetry structures (e.g., BCC Titanium) into stable low-symmetry phases (e.g., FCC) that equivariant models often get "trapped" in due to zero-force constraints.
• Phonon Calculations: While unconstrained models initially produced slightly distorted geometries, applying rotational averaging and Hessian symmetrization yielded phonon density of states and dispersion curves indistinguishable from symmetric cases.

C. Transferability and Fine-tuning

• Models pre-trained on large datasets (OMat24) could be fine-tuned to smaller, specialized datasets (e.g., MAD, SPICE) in <1/20th of the epochs required for training from scratch, achieving significantly lower test errors.

5. Significance and Conclusion

This work challenges the dogma that physical symmetries must be hard-coded into MLIP architectures. The authors conclude that:

• Unconstrained models are viable and superior for large-scale applications where inference speed is critical.
• Data is the key: With sufficiently large and diverse datasets, models can learn symmetries effectively, and the "cost" of learning these symmetries is offset by the efficiency gains of simpler architectures.
• Practicality: Simple inference-time modifications (averaging, symmetrization) are sufficient to recover physical observables consistent with symmetry laws, making unconstrained models safe for production use in geometry optimization and lattice dynamics.

The paper establishes rotationally unconstrained architectures as a promising alternative for next-generation universal MLIPs, offering better trade-offs between parameter budget, accuracy, and computational cost. The authors have made their pre-trained models and code publicly available to facilitate further research and application.