Graph neural networks for molecular dynamics: learning force fields from quantum data
Classical molecular dynamics runs on hand-fitted force fields — fast but crude approximations of quantum reality. Ab initio ("from first principles") methods like DFT (Density Functional Theory, a way of computing quantum energies without solving the full many-electron equation) are accurate but too slow to simulate more than a few hundred atoms for picoseconds. Graph neural networks split the difference: trained on quantum-accurate data, they predict energies and forces almost as fast as classical force fields while approaching quantum-chemical accuracy.
Molecules as graphs
A molecule is naturally a graph: atoms are nodes carrying features (element type, charge), and bonds — or simply nearby atoms within a cutoff radius — are edges carrying geometric features (distance, sometimes bond order). Unlike PageRank's web graph, this graph lives embedded in 3D space, so its edges encode continuous geometry, not just discrete connectivity:
Edge (i, j): r_ij = ‖x_i − x_j‖, e_ij = RBF(r_ij)
RBF = radial basis function expansion of the interatomic distance
This graph representation is exactly what makes GNN potentials size-transferable: a model trained on 100-atom molecules can, in principle, run inference on a 10,000-atom system, because it only ever operates on local neighbourhoods, not the whole system at once.
Message passing: the core operation
Every message-passing GNN layer repeats the same two steps: each node gathers "messages" from its neighbours, then updates its own hidden state from the aggregated message:
h_i' = φ_update(h_i, m_i) update node i's hidden state
Stacking L such layers lets information propagate L bond-hops away from each atom — enough layers, and every atom's final representation implicitly encodes its whole local chemical environment, from immediate bonds to longer-range steric and electrostatic effects.
Why equivariance matters
A molecule's physics doesn't change if you rotate, translate, or reflect the whole system — energy is invariant to these symmetries, and forces (vectors) must rotate along with the molecule (equivariant). A GNN that ignores this has to learn rotational invariance from data alone, wasting capacity and generalizing poorly to unseen orientations.
| Architecture family | Symmetry handling | Example |
|---|---|---|
| Invariant GNNs | Use only rotation-invariant features (distances, angles) | SchNet |
| Equivariant GNNs | Maintain vector/tensor features that transform correctly under rotation | NequIP, MACE |
| Message-passing + attention | Combines equivariant geometry with transformer-style attention | Equiformer |
Predicting energy-conserving forces
For a molecular dynamics trajectory to be physically valid, forces must be the negative gradient of a single scalar potential energy — otherwise energy leaks or accumulates unphysically over long simulations. The standard trick: predict a scalar total energy E, then get forces by automatic differentiation, never by predicting force vectors directly:
F_i = −∂E / ∂x_i forces via autodiff — automatically energy-conserving
Because E is guaranteed to be a real scalar function of positions, differentiating it can never produce a force field with "curl" — the model is architecturally incapable of violating energy conservation, which naive direct force prediction cannot guarantee.
A message-passing layer in JavaScript
function messagePassingLayer(nodeFeatures, edges, mlpMsg, mlpUpdate) {
const n = nodeFeatures.length;
const messages = new Array(n).fill(null).map(() => zeros(mlpMsg.outDim));
// Step 1: each edge produces a message, summed at the receiving node
for (const { i, j, rij } of edges) {
const eij = radialBasisExpand(rij);
const msg = mlpMsg.forward([...nodeFeatures[i], ...nodeFeatures[j], ...eij]);
messages[i] = addVec(messages[i], msg);
}
// Step 2: combine aggregated message with old state to get the new one
return nodeFeatures.map((h, i) => mlpUpdate.forward([...h, ...messages[i]]));
}
function predictEnergyAndForces(positions, atomicNumbers, layers) {
let h = atomicNumbers.map(embed);
const edges = buildNeighborList(positions, 5.0); // 5 Å cutoff
for (const layer of layers) h = messagePassingLayer(h, edges, layer.msg, layer.update);
const perAtomEnergy = h.map(hi => layers.readout.forward(hi)[0]);
const totalEnergy = perAtomEnergy.reduce((a, b) => a + b, 0);
const forces = autodiffGradient(totalEnergy, positions).map(g => scale(g, -1));
return { energy: totalEnergy, forces };
}Training against quantum reference data
These models are trained supervised on datasets of (geometry, energy, forces) triples computed with DFT or coupled-cluster methods — datasets like MD17, ANI-1, or OC20. The loss combines energy and force error, with forces usually weighted much more heavily since they directly drive the dynamics:
λ_F ≫ λ_E typically, since force errors compound over a trajectory
Applications beyond drug and materials design
- Drug discovery: fast, accurate binding-energy estimates for virtual screening of candidate molecules.
- Materials science: predicting properties of battery electrolytes, catalysts, and novel alloys at near-DFT accuracy.
- Protein folding dynamics: long-timescale simulations that would be infeasible with ab initio methods.
- Catalysis: mapping reaction pathways and transition states on learned potential energy surfaces.
- Climate and atmospheric chemistry: simulating reactive processes across large ensembles of molecules.
🧪 Explore molecular dynamics live
Watch atoms interact under a force field, from Lennard-Jones potentials to the learned potentials GNNs predict