Force Free Molecular Dynamics With Equivariant Networks

The tyranny of the tiny time step

Molecular dynamics is one of the workhorses of modern science. Give a computer the rules for how atoms attract and repel each other, and it can play the movie of a material forward in time, revealing how a protein folds, how a liquid freezes, or how a crystal responds to pressure. The method has produced insights that no experiment could reach, including simulations of viruses with tens of millions of atoms.

The trouble is the clock. A standard simulation advances by computing the force pulling on each atom, then moving every atom a small step in response, then recomputing all the forces. That step has to stay around half a femtosecond, which is five ten millionths of a nanosecond, because a larger step lets errors snowball until the simulation blows up. To watch anything interesting you need billions of these steps, and each one demands a full force calculation. That is the wall the field keeps hitting.

Most of the recent progress has attacked one half of the problem, the cost of the forces. Machine learning interatomic potentials learn to predict forces at the accuracy of expensive quantum calculations but far faster, and they have been a real success. Yet they leave the other half untouched. However cheap the force becomes, you are still stuck taking those minuscule steps. TrajCast goes after the step itself.

What force free actually means

The idea sounds almost too bold. Rather than compute a force and integrate it, TrajCast looks at the current state of the system, the position and velocity of every atom along with its chemical identity, and predicts the next state a much larger interval later. No force is ever calculated. No equation of motion is integrated. The network has learned the net effect of all that machinery and shortcuts straight to the answer.

This matters for a subtle reason the authors are careful about. Some earlier learning methods generate snapshots of a system without tracking how it actually got from one to the next, which throws away the dynamics and makes it impossible to measure how fast things move or diffuse. TrajCast keeps the thread intact. It works autoregressively, feeding each prediction back in as the starting point for the next, so the full trajectory it builds is a real sequence in time rather than a bag of disconnected frames. That is what lets it recover properties like diffusion that depend on the ordering of events.

Why the network has to respect symmetry

You cannot just bolt any neural network onto this task. Physics has symmetries, and the most important one here is rotation. If you spin the entire system, the forces and motions should spin with it, no more and no less. A network that ignores this would have to learn the same physics separately for every possible orientation, wasting data and inviting error.

TrajCast is built from equivariant message passing networks, which are a flavour of graph neural network engineered to obey rotation by construction. Each atom is a node, its neighbours within a cutoff distance are connected by edges, and information passes along those edges in a form that transforms correctly when the whole system rotates. Spin the molecule and the prediction spins with it, guaranteed, not merely learned. For readers who want the nuts and bolts of the graph neural network machinery underneath all this, our tour of graph neural network foundations is a good companion, and a practical entry point sits in our look at GNN tooling for the Julia language.

The architecture draws on the same lineage as the leading machine learning force fields, families known as NequIP and MACE, but it adds a crucial ingredient. It feeds the atomic velocities into the network as well as the positions, encoding each velocity through a basis that depends on the chemical element. Handing the network the velocities is what lets it forecast momentum rather than just geometry, and the authors show this extra information is a big part of why the model learns so much from so little data.

The equations behind the machine

At the heart of any message passing network are two steps repeated in layers. Each node gathers messages from its neighbours, then updates its own state using what it gathered.

The equivariance requirement is a precise mathematical statement. Rotate the atomic positions and the network features must transform by the matching rotation, captured by a representation of the rotation group.

Training is refreshingly direct. The network is asked to match both the true displacement and the true velocity of every atom over the forecast interval, with the two errors weighted equally.

Two safeguards keep the long rollouts honest. A constraint step removes any spurious drift of the whole system by conserving total momentum, and a gentle thermostat rescales the velocities after each prediction to hold the temperature steady, which also stops small errors from accumulating over millions of steps.

Three systems, three kinds of matter

The authors did not settle for one easy demonstration. They tested TrajCast on a small molecule, a crystalline solid, and a liquid, three regimes that stress a simulator in different ways.

The small molecule test used paracetamol purely as a convenient object with four different chemical elements, a standard physics benchmark rather than anything medical. TrajCast reproduced its vibrational fingerprint without the artefacts that larger steps usually introduce, matched the spread of its energy, and even captured the tall energy barrier that separates two twisted forms of the molecule. That last point is telling, because it means the model learned the rare, difficult transitions and not just the easy jiggling.

The quartz crystal showed off scale. Because the atoms in a lattice move in a more constrained, correlated way, the team could push the forecast interval all the way to thirty times the base step and still get stable, accurate trajectories. The throughput is the headline here. A single model generated close to twenty nanoseconds of trajectory per day for a crystal of several thousand atoms, where a leading machine learning force field would manage roughly one nanosecond per day for a thousand atoms. Our analysis of how language models are entering materials science covers a very different angle on the same broad goal of computing over matter faster.

Liquid water is the trickiest of the three, because a fluid has no fixed structure and its molecules tumble and drift. TrajCast still matched the reference on the vibrational spectrum almost perfectly and reproduced the water structure well, with only a slight softening of one feature and a diffusion coefficient modestly higher than the reference. For a liquid at room temperature generated without any forces, that is a strong showing.

TrajCast can generalize in a zero shot manner to unseen regions of phase space, producing physically meaningful ensembles in metastable equilibrium and out of equilibrium regimes beyond the training data, without compromising accuracy. Thiemann and colleagues, Nature Machine Intelligence, 2026

The glass that was never in the training set

The most striking result is one the model was never taught. The team took the water network, trained only at ordinary room temperature, and asked it to simulate what happens when water is cooled fast enough to freeze into a glass rather than ice. This is a hard, subtle process that lives far outside anything the model had seen.

It worked. Across four independent runs, each cooling a few hundred molecules from three hundred down to one hundred and eighty kelvin over twelve nanoseconds and two point four million prediction steps, the model stayed stable and reproduced the glass transition. The diffusion slowed and vanished around two hundred and ten kelvin, exactly where the liquid seizes up into a glassy state, and the structural signatures sharpened in the way real experiments report. Predicting behaviour in a temperature range the model never trained on, and getting the physics right, is the kind of generalization that separates a genuine model of the dynamics from a lookup table.

Why it learns so much from so little

Underneath the demonstrations sits a quieter finding about data. TrajCast needs remarkably little training data, often just a few hundred picoseconds of trajectory where competing methods need hundreds of nanoseconds. The authors trace this to the velocities. By learning from the motion of atoms and not only their arrangement, each training snapshot carries far more information, and they measure the effect precisely. The rate at which accuracy improves with more data is up to four times steeper than for comparable force field methods on the same systems.

They are also honest about the limits of the free lunch. How large a forecast interval you can get away with depends on the system. The rigid quartz tolerated a thirty fold step, while water became unstable at more modest settings, because a jostling liquid decorrelates faster and gives the model less to hold onto. Bigger steps always trade accuracy for speed, and the sweet spot has to be found per system rather than assumed.

Where it falls short

The paper is candid about what TrajCast cannot yet do, and a fair reading has to sit with the caveats.

The current version works at a single, fixed forecast interval chosen before training. A more flexible model that mixes several step sizes, as recent weather forecasting networks do, would likely generalize further, but that is future work rather than a solved feature. The method also lives in specific thermodynamic settings. It was built for simulations at constant energy or constant temperature, and extending it to constant pressure, which many real problems need, would require coupling it to a separate mechanism for controlling pressure.

Accuracy is less forgiving than in force field work. A machine learning force field can tolerate fairly large errors in its predicted forces and still produce a stable simulation, because the integration smooths things out. TrajCast has no such cushion. Its velocity errors must stay small for the rollout to remain stable, and the tolerance is tight enough that a couple of percent can tip a run into instability. The thermostat helps, but it is doing real work to keep long trajectories from drifting.

Finally, every system here was trained on trajectories from classical force fields rather than from the most expensive quantum calculations, and each model handles one chemistry at a time. Whether the approach holds up when trained on many chemically diverse materials at once, the direction that would make it a general foundation model, remains an open and interesting question.

Why it matters

Step back and the significance is clear. For decades the small time step was treated as an immovable fact of molecular simulation, a tax you simply paid. TrajCast reframes it as a modelling choice that a network can learn around, and in doing so it opens a path to longer simulations of larger systems at a fraction of the cost. The gains are largest exactly where they are most needed, in slow processes like glass formation that unfold over timescales traditional methods struggle to reach.

The deeper appeal is architectural. By marrying the symmetry aware design of modern force fields with the autoregressive forecasting style of sequence models, the authors built something that keeps the physics honest while shedding the integration bottleneck. It is the kind of recombination of existing ideas that tends to travel, and the authors point toward materials discovery and the study of phenomena beyond the reach of current simulation as the natural next destinations. If the approach scales to many chemistries at once, it could become a genuinely new tool for computing over matter.

A reference implementation

The code below is a compact, runnable sketch of the core idea. It builds a small message passing network that reads atomic positions, velocities, and types, predicts a displacement and a velocity update over a forecast interval, removes any net momentum to keep the system honest, rescales the velocities toward a target temperature, and rolls the prediction forward autoregressively. A production model uses fully equivariant features through libraries such as e3nn, which this sketch leaves out for clarity. The authors full code and models are linked under the block.

# TrajCast, compact reference sketch of force free molecular dynamics
# A message passing network predicts the next positions and velocities
# directly, with momentum removal and a simple thermostat for stability.

import torch
import torch.nn as nn

torch.manual_seed(0)


class MLP(nn.Module):
    def __init__(self, sizes):
        super().__init__()
        layers = []
        for i in range(len(sizes) - 1):
            layers.append(nn.Linear(sizes[i], sizes[i + 1]))
            if i < len(sizes) - 2:
                layers.append(nn.SiLU())
        self.net = nn.Sequential(*layers)

    def forward(self, x):
        return self.net(x)


class TrajCast(nn.Module):
    # Two message passing layers over a radius graph, then a readout
    # that predicts a displacement and a velocity change per atom.
    def __init__(self, n_types, hidden=64, cutoff=2.5):
        super().__init__()
        self.embed = nn.Embedding(n_types, hidden)
        # message uses neighbour features plus the relative position and velocity
        self.msg = MLP([hidden * 2 + 6, hidden, hidden])
        self.upd = MLP([hidden * 2, hidden, hidden])
        self.readout = MLP([hidden, hidden, 6])   # 3 for displacement, 3 for velocity
        self.cutoff = cutoff

    def forward(self, pos, vel, types):
        h = self.embed(types)
        # build a radius graph from pairwise distances
        diff = pos.unsqueeze(0) - pos.unsqueeze(1)          # [n, n, 3]
        dist = diff.norm(dim=-1)
        mask = (dist < self.cutoff) & (dist > 0)
        vdiff = vel.unsqueeze(0) - vel.unsqueeze(1)

        for _ in range(2):
            n = pos.shape[0]
            hi = h.unsqueeze(1).expand(n, n, -1)
            hj = h.unsqueeze(0).expand(n, n, -1)
            edge = torch.cat([hi, hj, diff, vdiff], dim=-1)
            m = self.msg(edge) * mask.unsqueeze(-1)         # zero out non neighbours
            agg = m.sum(dim=1)                              # pool over neighbours
            h = h + self.upd(torch.cat([h, agg], dim=-1))   # residual update

        out = self.readout(h)
        return out[:, :3], out[:, 3:]                       # delta_r, delta_v


def remove_momentum(vel, mass):
    # subtract the centre of mass velocity so total momentum stays zero
    p = (mass.unsqueeze(-1) * vel).sum(dim=0)
    return vel - p / mass.sum()


def thermostat(vel, mass, target_temp, kb=1.0):
    # rescale velocities to match a target kinetic energy
    ke = 0.5 * (mass.unsqueeze(-1) * vel ** 2).sum()
    n_dof = 3 * vel.shape[0] - 3
    cur_temp = 2.0 * ke / (n_dof * kb)
    scale = (target_temp / (cur_temp + 1e-8)) ** 0.5
    return vel * scale


def rollout(model, pos, vel, types, mass, steps=50, target_temp=1.0):
    traj = [pos.clone()]
    for _ in range(steps):
        dr, dv = model(pos, vel, types)
        pos = pos + dr
        vel = remove_momentum(vel + dv, mass)
        vel = thermostat(vel, mass, target_temp)
        traj.append(pos.clone())
    return torch.stack(traj)


if __name__ == "__main__":
    n_atoms, n_types = 40, 3
    pos = torch.rand(n_atoms, 3) * 5.0
    vel = torch.randn(n_atoms, 3) * 0.1
    types = torch.randint(0, n_types, (n_atoms,))
    mass = torch.ones(n_atoms)

    model = TrajCast(n_types)
    traj = rollout(model, pos, vel, types, mass, steps=50)
    drift = (traj[-1].mean(0) - traj[0].mean(0)).norm()
    print(f"rolled out {traj.shape[0]} frames, centre of mass drift {drift.item():.4f}")

Conclusion

The core achievement here is a molecular dynamics engine with no forces in it. TrajCast learns to jump straight from the current arrangement and motion of a set of atoms to their next arrangement and motion, over an interval many times larger than the step a traditional simulation could survive. It did this convincingly across a small molecule, a crystal, and a liquid, matching reference simulations on the structural, dynamic, and energetic properties that actually matter.

The conceptual shift worth remembering is that the tiny time step was never a law of nature, only a limitation of integrating forces one nudge at a time. By learning the net effect of that machinery and preserving the sequence of states, the model keeps the true dynamics while discarding the bottleneck. The choice to feed velocities into an equivariant network turns out to be the quiet engine behind both its accuracy and its striking hunger for so little training data.

The approach travels because the design is principled rather than tuned to one case. The same architecture handled three very different states of matter, and, most impressively, a water model reached into the physics of glass formation it was never trained on and got it right. That kind of zero shot reach is the strongest evidence that the network learned something real about how atoms move rather than memorizing a narrow slice of behaviour.

The honest limits keep the enthusiasm in check. The method commits to one forecast interval at a time, it lives in fixed thermodynamic settings that would need extending for pressure controlled problems, and its accuracy tolerances are tighter than a force field can afford, so a small velocity error can destabilize a long run. Each model also handles a single chemistry, trained on classical rather than quantum reference data, which leaves the leap to a general purpose foundation model as unfinished business.

Future directions follow naturally. Let the model mix several step sizes, teach it pressure control, and, most ambitiously, train it across many materials at once to see whether one network can forecast the motion of matter in general. If those steps hold up, the slow, patient grind of molecular simulation could speed up enough to reach the long, slow processes that have always sat just out of computational reach, and that would change what questions scientists can afford to ask.

Frequently asked questions