Scaling Beyond the Atomic Scale with Coarse-Grained Machine Learning

Summary

Most coarse-grained molecular simulations treat each molecule as a single point, defined by a position and nothing else. For a liquid like water, this discards exactly what makes it behave like water, namely the molecule's shape and orientation.

We instead represent each molecule as an ellipsoidal bead carrying shape, volume, and orientation, and train an equivariant neural network to learn its energy, forces, and torques directly from atomistic data. The model reproduces the structure of liquid water down to how neighboring molecules align with one another, which a point-based representation cannot capture at all. When we remove the orientation and train an otherwise identical model, this structure degrades, isolating the contribution of the anisotropic features. Notably, the added geometry incurs minimal overhead, and we still observe 7-27x speedups over all-atom simulation.

Why geometry matters

Much of the behavior that decides how a material performs is emergent. It does not live in any single molecule, but in how millions of them organize collectively. Polymers entangle and phase-separate, electrolytes structure themselves around an electrode, surfactants and lipids self-assemble, and these collective arrangements are what determine whether a battery lasts, a coating holds up, or a plastic has the mechanical properties we want. They play out over microseconds and across length scales far larger than a single molecule, and atomistic simulation cannot reach them. Even with machine-learning interatomic potentials (MLIPs), which deliver near-quantum accuracy at a fraction of the usual cost, simulations remain practically capped at tens of thousands of atoms and short timescales. This leaves a wide gap between the resolution at which we can model chemistry accurately and the scale at which the behavior we care about emerges.

Coarse-graining (CG) is the established route that bridges this gap. By grouping atoms into beads and reducing the number of particles to model, simulations can reach sizes and timescales that all-atom methods cannot. A CG simulation, however, is only as good as the potential that drives the beads’ behavior. Traditionally, these models have been hand-tuned for one system and transfer poorly. Data-driven CG potentials have begun to address this the same way MLIPs did for atomistic modeling – by learning the potential from data rather than fitting it by hand.

Most coarse-grained potentials inherit a point-particle representation from atomistic models, where each bead is a dimensionless site defined only by its position. For systems in which the underlying molecule is roughly isotropic, this may be a reasonable approximation, but for polar or asymmetric molecules, it is a poor assumption. A water molecule has a well-defined axis, and its neighbors arrange around that axis rather than around a featureless point, so reducing it to a position discards exactly the information that gives the liquid its structure. The limit of a structured liquid is its solid state, where water molecules align into different arrangements of ice crystals, which point particle representations cannot model.

Anisotropic sites address this – ellipsoidal beads with orientation-dependent interactions go back decades, and more recent work has even learned anisotropic CG potentials with neural networks. Importantly, these approaches encode the local environment with hand-crafted descriptors such as symmetry functions or SOAP-style expansions, the same kind of fixed representation that learned, equivariant message-passing networks have largely replaced in atomistic modeling. We solve this longstanding problem by pioneering equivariant representations for anisotropic CG sites that the network learns directly from data.

What we built

We map each molecule to a single ellipsoid through its inertia tensor, which fixes the bead's shape, volume, and orientation at every frame. The mapping is set by the molecule's own mass distribution, so there are no parameters to fit. We then aggregate the atomistic energy, forces, and torques onto that bead and train on them directly. Orientation enters the network through spherical-harmonic projections of the bead's principal axes, and the E(3)-equivariant architecture handles the rotational equivariance that otherwise makes orientation difficult to learn. Forces and torques are both obtained by differentiating the total energy, the latter through the quaternion gradient, so the dynamics conserve energy by construction. Training runs inside an active-learning loop that iteratively builds the training dataset, with a high-fidelity atomistic potential serving as the ground-truth oracle. For the water results below, that oracle is the MACE universal potential, and the reference is a 64-molecule system run at 300 K.

×

Isolating what the geometry buys

The first question is whether the model is accurate at all. Measured against the all-atom reference, the anisotropic model reproduces the structure of liquid water across the board, with a radial distribution function that matches almost exactly in both peak positions and heights.

Then, to quantify what shape and orientation contribute separate from everything else in the model, we trained a second model identical in every regard except the shape and orientation inputs were removed, collapsing each bead back to an isotropic, point-site on the molecular center of mass. If the geometry were unnecessary, the two models would agree. We find that the isotropic baseline shifts the first two coordination peaks and lowers the first, missing the correct nearest-neighbor distances. The same pattern of discrepancies appears in the angular correlations – the anisotropic model tracks the reference while the baseline misses peak positions and amplitudes precisely where directional packing dominates, and recovers only at long range where those correlations have already washed out.

×

Given that the beads carry orientation, the model also resolves observables a point particle cannot represent at all, namely how neighboring molecules align relative to one another. We probe this by projecting each bead's principal axes onto the intermolecular vector, and onto each other, across many neighbor pairs. These distributions are strongly non-uniform, reflecting the preferred relative orientations of nearby molecules, and the anisotropic model reproduces them. We take this as evidence that the model has learned the rotational structure of the liquid alongside its positional structure.

×

Speed

With accuracy addressed, the efficiency of the approach takes center stage. Reducing computational cost is the point of coarse-graining, and we find that the added geometry does not compromise it. Both models run far faster than all-atom, roughly 7x for 64 molecules and 27x for 512, and the anisotropic features add no meaningful overhead relative to the isotropic baseline. Given that both models carry the same number of sites and the orientation features are inexpensive to evaluate, confirming that anisotropy is not bought at the expense of speed.

Moreover, we note that water is close to the worst case here, since with only three atoms in each bead, the reduction in degrees of freedom is modest. For a polymer or protein, where tens of atoms collapse onto a single site, the speedup scales with the coarse-graining ratio.

What this unlocks

The main benefit of anisotropy is accuracy. Shape and orientation increase the expressibility of the model to be able to accurately capture structural disorder and directional information. With learned equivariant anisotropy in hand, the path to general-purpose CG models is an exercise in scaling to more molecular classes through learned atom-to-bead mapping and tightening uncertainty-driven active learning to drive down data cost. Each of these is tractable, and together they point toward a CG modeling paradigm that is as general and as data-driven as MLIPs have become for atomistic simulation but for systems 100x larger in size.

What this work settles is concrete: given a learned equivariant representation, anisotropic coarse-graining is not just practical but measurably better than the isotropic alternative trained the same way, and is ready to scale.

Full methods and results are in our paper: https://arxiv.org/abs/2607.10002