Equivariant learning of a transferable three-dimensional classical density functional

This paper is available in full, and it's a clean single-author result: a neural network with 1,762 parameters, trained only on time-averaged density fields from canonical MD—no energies, no chemical potentials, no forces—learns the excess free-energy functional of classical density functional theory in full 3D. Once learned, that one scalar functional is minimized and differentiated to reproduce structure factors, the equation of state, liquid–vapor coexistence, interfacial broadening, solvent-mediated forces between colloids, and adsorption isotherms in a gyroid pore—none of which appeared in training. If machine-learned interatomic potentials turned quantum data into reusable force fields, this is the analogous move one level up: turning simulation data into a reusable thermodynamic object.

The missing functional

Classical density functional theory (cDFT) is the statistical-mechanics twin of electronic DFT. Instead of the electron density of a ground state, the central object is the equilibrium number density \rho(\mathbf{r}) of a fluid at temperature T. The theory says: there exists a grand-potential functional

whose minimizer is the true equilibrium density. Here F_{\mathrm{id}} is the ideal-gas term (known exactly, analytic), V_{\mathrm{ext}} is whatever external potential you impose (a wall, a pore, a colloid), and \mu is the chemical potential. Everything interesting about interparticle interactions is buried in the excess free-energy functional F_{\mathrm{exc}}[\rho,T]—and, exactly as with exchange-correlation in electronic DFT, nobody knows what it is for realistic fluids. Hard spheres have an excellent hand-crafted answer (fundamental measure theory); anything with attraction gets mean-field patches of state-dependent quality.

The prize is large. With a good F_{\mathrm{exc}} in hand, every new thermodynamic state or geometry is a cheap variational minimization instead of a fresh multi-million-step MD run. Prior neural attempts (notably "neural functional theory") learned the derivative of this object—the one-body direct correlation c^{(1)}(\mathbf{r};[\rho]) = -\beta\,\delta F_{\mathrm{exc}}/\delta\rho, with \beta = 1/k_BT—but almost entirely in planar (1D-inhomogeneous) or 2D settings, often requiring grand-canonical data where \mu is known and serves as a label. Three dimensions is where the applications live (solvation, nucleation, confinement), and it's also where symmetry becomes unavoidable: a scalar free energy must be invariant, and its derivatives equivariant, under the rotations and reflections of your representation.

The training signal hiding in plain sight

Here's the aha. At equilibrium, define at every grid voxel g the dimensionless local chemical potential

where \Lambda is the thermal wavelength and c^{(1)}_g is the model's automatic derivative of its own free energy. The Euler–Lagrange condition of cDFT says this quantity must be spatially constant—it's the water level in connected vessels. Any spatial variation would be a thermodynamic force pushing density around, contradicting equilibrium.

So the loss is simply the spatial variance:

You never need to know what the water level is—only that it must be flat. The unknown Lagrange multiplier \mu (which canonical MD doesn't give you) is eliminated analytically by subtracting the spatial mean. Every time-averaged density field under a randomized external potential becomes a full-field supervision signal, with the external potential acting as the "question" and the density's shape as the "answer."

This costs one unavoidable gauge freedom: the loss can't see transformations F_{\mathrm{exc}} \to F_{\mathrm{exc}} + b(T)N, i.e., a per-temperature additive shift of the chemical potential. Fixed-N predictions and all second derivatives are unaffected; only absolute reservoir chemical potentials need a one-number calibration per temperature. The paper is refreshingly explicit about this.

The machinery

The representation is where the "equivariant" in the title lives. The density sits on a regular Cartesian grid (spacing 0.5\sigma), so the exact symmetry group is not full O(3) but the cubic point group O_h—the 48 combinations of axis permutations and reflections. For each voxel, the method computes Cartesian moments of the surrounding density within a stencil of 3 grid spacings (123 voxels), A_{\boldsymbol\ell} = \sum_{\mathbf{q}} \rho_{g+\mathbf{q}}\, q_x^{\ell_x} q_y^{\ell_y} q_z^{\ell_z} with total degree \le 3, then symmetrizes products of moments (up to order 2) over all 48 group operations. This is the author's Cartesian atomic cluster expansion (CACE), transplanted from atoms to a density grid, yielding just 15 invariant features. Together with the center density and temperature, these feed a 17{\to}32{\to}16{\to}1 MLP (plus a 2{\to}32{\to}16{\to}1 pointwise baseline), giving the per-particle excess free energy a_{\mathrm{exc}}(\chi_g, T). The functional is then

with \Delta V the voxel volume. Because the same local map is applied everywhere, the functional is extensive by construction and transfers across box sizes.

Two design choices matter beyond the symmetry. First, energy-first: rather than learning c^{(1)} as an output field (the dominant prior approach), the network is the scalar F_{\mathrm{exc}}, and c^{(1)}, c^{(2)} come from first and second autodiff derivatives. The Hessian structure makes c^{(2)} automatically symmetric, and—crucially—you can compare free energies between different states, which is what phase coexistence and adsorption require. This is exactly the conservative-force-field-vs-direct-force-prediction debate from MLIPs, resolved in the same direction. Second, note the receptive fields: a_{\mathrm{exc}} sees 1.5\sigma, but c^{(1)} at a point (a derivative through all environments containing it) reaches 3\sigma—comfortably beyond the 2.5\sigma truncation of the potential.

Training data: 10,957 time-averaged density fields of the truncated-and-shifted Lennard–Jones fluid in 8\sigma periodic cubes on 16^3 grids, under randomized superpositions of Gaussian external potentials, spanning 23 temperatures (T = 0.625–1.8) and mean densities 0.02–0.91\,\sigma^{-3}. All fields at T=0.7, 1.1, 1.5 were held out entirely. Training: 200 epochs, ~1h40m on one L40. Inference is almost comically cheap: a fixed-N equilibrium solve takes 0.3–8 s on a laptop CPU, versus 4.4 million MD steps per reference state.

What you get for free

The benchmarks (Fig. 5 of the paper) confirm the basics: at held-out temperatures the functional inverts densities back to external fields and re-generates densities by minimization; it transfers to boxes 3.375× the training volume; and, trained purely on canonical data, it predicts grand-canonical chemical-potential differences on GCMC test fields.

But the interesting results are the emergent ones—quantities the loss never saw. Differentiating c^{(1)} once more gives the two-body direct correlation c^{(2)}, whose Fourier transform yields the structure factor via the Ornstein–Zernike relation, matching direct MD across three state points. Integrating the zero-wavevector response from the ideal-gas limit reconstructs the full pressure–density equation of state, matching the reference EOS. Most strikingly, at subcritical temperatures the learned homogeneous isotherms spontaneously develop a van der Waals loop—the nonconvexity that encodes liquid–vapor separation—despite the model never being shown pressure, coexistence, or any two-phase supervision. Fixed-N slab solutions then give coexistence densities directly, and a mean-field extrapolation lands the critical point almost on top of the direct-simulation value:

The learned functional also reproduces the progressive broadening of the liquid–vapor interface across four temperatures, when relaxed from MD-initialized profiles at matched geometry and particle number.

Stress tests in genuinely 3D geometry

Two applications, both run without retraining, are where the "transferable generator" claim gets teeth.

Colloid bridging. Two soft repulsive cores with attractive shells sit in the LJ fluid (T=1.2, N=4600, 20\sigma box, 40^3 grid), with zero direct colloid–colloid interaction. The solvent mediates a force anyway: at D=5\sigma the depletion zones merge into a solvent-depleted neck (attraction); near 6.5\sigma the bridge ruptures and solvation layers clash (repulsive maximum); by 9–10\sigma the force dies. The force is extracted from the minimized cDFT density alone via the Feynman–Hellmann identity \langle F\rangle_D = -\Delta V\sum_g \rho_{D,g}\,\partial_D V_{\mathrm{ext},g}—a weighted integral over the whole field, so spatially correlated density errors would show up loudly. The cDFT curve tracks direct MD (including the sign reversal) and does so on a much denser separation grid than MD could affordably sample. A nice control: applying the same quadrature to gridded MD densities reproduces the MD force, isolating discretization effects from model error.

Gyroid adsorption. A triply periodic gyroid pore (16\sigma box, 32^3 grid) has no wall-normal coordinate and is topologically unlike anything in training. The functional reproduces the full 3D density morphology at fixed loading—parity holds from the curved depletion layer into the pore interior—and then, given only particle numbers, predicts the chemical potential of the confined fluid. Matching independently computed bulk and confined branches at equal \mu (after the single per-temperature gauge calibration on the bulk branch) yields the adsorption isotherm against GCMC references, including the strongly nonlinear filling crossover at supercritical T=1.2. This test uses free-energy differences between different geometries—precisely what the energy-first construction buys you and what a learned c^{(1)} field alone would not.

How strong is this, and what to watch

The evidence within its scope is strong and unusually self-consistent: every prediction traces back to one scalar through minimization or differentiation, and the most impressive results (van der Waals loop, coexistence, forces, adsorption) are genuinely out-of-loss. The tiny parameter count is a point in favor, not against—it suggests the symmetry-adapted features carry most of the burden, and it makes overfitting to 8,427 training fields implausible.

Caveats, most of which the paper states plainly:

• One fluid. Everything is the truncated-and-shifted LJ liquid, and only liquid states—crystallizing trajectories were removed. Freezing, which cDFT historically cares about, is untested. The functional is interaction-specific; a new fluid means new training data (though data generation is cheap canonical MD).
• Short-range only. The 3\sigma effective receptive field works because the potential is truncated at 2.5\sigma. Ionic and polar fluids need a long-range branch; the discussion sketches routes (local molecular field theory, reciprocal-space features) but none are demonstrated.
• Fixed resolution. The grid spacing is locked at 0.5\sigma; transfer across volume is tested, transfer across resolution is not.
• Mean-field near criticality. The critical point is a fitted continuation, not a resolution of critical fluctuations, and long-wavelength finite-cell response is flagged as the largest remaining sensitivity. Relatedly, the gyroid test at T=1.2 is supercritical—continuous filling, no capillary condensation with its hysteresis and metastability, which would be a harder test.
• Interface profiles were relaxed from MD-initialized states, not nucleated from uniform density—a fair but gentler test.
• Parity plots in the text come without headline RMSE numbers, so "accurate" is mostly visual for the forward/inverse benchmarks.

On the triage question of variational consistency away from equilibrium: nothing here validates the functional off the equilibrium manifold beyond what the inverse solves probe (which do explore the constrained landscape en route to minima). Its use as a driving force in dynamical DFT is proposed, citing prior work, but not tested.

If it holds and extends—mixtures via multi-channel CACE, electrostatics via a long-range branch, training densities from MLIP-driven MD rather than empirical potentials—this becomes a genuinely new layer in the simulation stack: quantum data → MLIP → equilibrium densities → a mesoscale thermodynamic generator you minimize in seconds. Code and trained models are public (github.com/BingqingCheng/equicdft).

Where to spend your time: Section II.2 (the chemical-potential-balance loss and gauge analysis) is the conceptual core and takes ten minutes. Then go straight to Section III.2 and Figure 3—the colloid force curve is the single most convincing demonstration that a label-free learned free energy encodes real, integrated collective physics rather than pointwise density fitting.