Enhancing molecular dynamics with equivariant machine-learned densities

Abstract: Machine-learning interatomic potentials (MLIPs) have enabled molecular dynamics at near ab initio accuracy, yet remain limited to energies and forces by construction, leaving electronic observables such as dipole moments and polarizabilities inaccessible. We introduce DenSNet, a density-first approach to machine-learned electronic structure that learns the Hohenberg--Kohn map from nuclear configurations to the ground-state electron density. Our approach employs an SE(3)-equivariant neural network to predict density coefficients of a flexible atom-centered Gaussian basis, combined with a $Δ$-learning strategy that uses superposed atomic densities as a prior to accelerate training. A second equivariant network then maps the predicted density to the total energy, providing a unified framework for molecular dynamics and electronic structure. We validate DenSNet on ethanol, ethanethiol, and resorcinol, where infrared spectra from machine-learned trajectories show excellent agreement with experimental gas-phase measurements. To test scalability, we train on polythiophene oligomers with 1--6 monomers and extrapolate to chains of up to 12 monomers, generating stable long-time trajectories whose infrared spectra agree with reference density functional theory calculations. Here, we show that reinstating the electron density as the central learned quantity opens a practical route to transferable prediction of spectroscopic and electronic observables in large-scale molecular simulations.

• The paper demonstrates that DenSNet achieves a tenfold improvement in density accuracy with absolute fractional errors of 5.1–6.6×10⁻⁴ compared to conventional methods.
• It employs a modular SE(3)-equivariant architecture with delta-learning, decoupling the learning of electron density from energy prediction for enhanced simulation fidelity.
• Transferability tests on small molecules and polythiophene oligomers confirm stable MD simulations and reliable spectroscopic properties, validating its practical utility.

Equivariant Machine-Learned Electron Densities for Molecular Dynamics: The DenSNet Framework

Motivation and Theoretical Foundations

The DenSNet framework introduces a density-centric approach to machine-learned electronic structure for molecular dynamics (MD). Traditional MLIPs focus exclusively on energies and forces, omitting access to electronic observables such as dipole moments, polarizability tensors, and IR/Raman spectral properties, all critical for accurate spectroscopy and for probing electronic phenomena. The central motivation draws from the Hohenberg–Kohn theorem in density functional theory (DFT): the ground-state electron density $\rho(\mathbf{r})$ is a unique functional of the external potential defined by nuclear charges and positions. DenSNet leverages this correspondence by directly learning the map from nuclear configuration $(Z, \mathbf{R})$ to $\rho(\mathbf{r})$, using an SE(3)-equivariant neural architecture to guarantee physical symmetry invariance as molecules are rotated or translated.

Architecture and Methodological Innovations

DenSNet consists of two modular SE(3)-equivariant neural networks:

1. Density Prediction Network: The architecture builds upon PhiSNet, using rotationally equivariant message passing to predict expansion coefficients for an atom-centered Gaussian–spherical harmonics basis. This flexible basis adapts both radial and angular parameters per atom, enhancing expressivity compared to the standard fixed density fitting basis.
2. Energy Network: The second network maps the learned density coefficients to the molecular energy, generating forces by automatic differentiation. This modularity decouples the learning of electronic structure and energetics.

An important methodological advance is delta ($\Delta$)-learning: the model learns corrections to superpositions of free-atom densities (SAD). Rather than learning the full density ab initio, DenSNet focuses capacity on the physically relevant residual bonding density, accelerating convergence and improving transferability.

Physical constraints are imposed post-training: electron count conservation is ensured via analytic projection of $l=0$ spherical harmonics coefficients, and density positivity enforced via the softplus transformation, with subsequent fine-tuning to optimize accuracy for observables sensitive to these constraints.

Figure 1: Schematic of DenSNet’s SE(3)-equivariant representation module, detailing atomic spherical harmonics features, residual block structure, and message-passing interaction.

Benchmarking and Accuracy: Small Molecules

DenSNet’s performance was benchmarked on ethanol, ethanethiol, and resorcinol, targeting both density-derived and spectroscopic observables.

Density absolute fractional error (AFE) for DenSNet ($\Delta$-learning) achieves $5.1\times10^{-4}$–$6.6\times10^{-4}$, noticeably lower than LDA and Hartree–Fock densities and a tenfold improvement over conventional density fitting. Dipole moment magnitude and angular errors outperform direct ML approaches, approaching DF-basis accuracy. Energy and force mean absolute errors are consistently below the quantum chemical threshold (1 kcal/mol/Å), matching the best MLIPs.

A crucial outcome is that DenSNet correctly reproduces spatial gradients of the density, evidenced by low exchange–correlation energy and von Weizsäcker kinetic energy errors, validating fidelity to GGA-level DFT functionals.

Figure 2: Comparative density error metrics and IR spectra for ethanol, ethanethiol, and resorcinol; ML-predicted spectra align with experimental data, demonstrating accuracy in both density and force fields.

Spectroscopic Validation and Stability in MD

DenSNet's densities enable direct computation of IR spectra from MD trajectories by evaluating dipole moments as spatial integrals over $\rho(\mathbf{r})$. The machine-learned spectra for all three benchmark molecules display accurate peak positions and intensities relative to experimental gas-phase measurements (NIST), including O–H/S–H stretching, C–H stretching clusters, and fingerprint region modes characteristic of functional group identity and aromatic substitution patterns.

MD trajectories of hundreds of picoseconds reveal no instability, confirming DenSNet’s robust force fields and reinforcing the practical utility in spectroscopic property prediction from large-scale simulations.

Figure 3: Density prediction accuracy and IR spectra overlay for small organics; ML-derived spectra overlay experimental benchmarks.

Transferability and Scaling to Extended Systems

DenSNet was trained on polythiophene oligomers (1–6 monomers) and extrapolated to chains up to 12 monomers, testing size transferability and robustness in complex, extended conjugated systems. The density AFE remains within $1.1\times10^{-3}$ across the extrapolation regime, with no degradation as chain length increases. Force errors are stable and sub-chemical accuracy, enabling 100 ps stable MD trajectories for 8–12 mer systems.

Dipole moment MAE increases predictably with system size, following $(Z, \mathbf{R})$0 scaling, but per-monomer error is invariant—consistent with random-walk addition of dipole vectors in head-to-tail stacked thiophenes.

IR spectra computed via ML-MD closely match DFT reference calculations for 8–12 mers, reproducing diagnostic C=C ring stretching and C–H out-of-plane bending modes. Peak position errors are within experimental IR resolution limits, and integrated peak area errors remain in the 4–18% regime for most features, validating DenSNet transferability to spectroscopic observables in extended systems.

Figure 4: Transferability to polythiophene oligomers: density isosurfaces and IR spectra for chain lengths up to 12, with ML-MD results matching DFT references.

Figure 5: Quantitative IR peak errors for polythiophene 8–12-mers: peak heights and FWHMs anti-correlate; integrated areas remain within 4–18% error for dominant modes.

Practical and Theoretical Implications

DenSNet unifies molecular dynamics with access to electronic observables at every simulation step, bypassing the computational bottleneck of on-the-fly DFT. The framework’s equivariant architecture minimizes the need for data augmentation and realignment, targeting physical symmetry directly and aiding interpretability and transfer learning. Practical implications include:

• Direct vibrational spectroscopy (IR intensity, dipole derivatives, frequency-dependent dielectric response) from large-scale MD, bridging a gap left by MLIPs.
• Applicability to extended π-conjugated systems, hydrogen-bond networks, and reaction coordinates involving electronic redistribution.
• Potential as a Kohn–Sham initial density guess to accelerate SCF convergence.
• Compatibility with large quantum chemical datasets (Open Catalyst, Open Molecules, QCML, QuantumChem-200K), facilitating pre-training or few-shot fine-tuning.

On the theoretical side, DenSNet shifts the modeling paradigm from learning isolated observables to reinstating electron density as the central, universal molecular descriptor. The accuracy and transferability demonstrated substantiate the assertion that $(Z, \mathbf{R})$1 reliably encodes information sufficient for both MD and spectroscopic predictions.

Future Directions

Limitations remain: the density AFE loss does not directly optimize dipole moment accuracy, particularly challenging for large systems where errors accumulate geometrically. Explicit inclusion of dipole-regularized loss functions or multi-task learning (dipole, density, energy, force) may further improve electrostatic property prediction.

Future extensions may include learning excited-state densities, response densities for polarizabilities, and full one-particle density matrices, enabling non-adiabatic and photochemical MD. Utilizing curated datasets and physics-informed architectures will catalyze broader fine-tuning and transfer to unseen chemistry.

Conclusion

DenSNet leverages SE(3)-equivariant machine learning, delta-learning, and atom-centered flexible bases to produce accurate electron densities, energies, forces, and spectroscopic properties throughout MD simulations. The unified density-first approach repositions $(Z, \mathbf{R})$2 at the center of atomistic modeling, advancing both practical simulation and theoretical understanding of electronic structure in large-scale molecular environments. The demonstrated transferability, interpretability, and spectroscopic reliability position DenSNet as a key framework in data-driven quantum chemistry and molecular simulation.

Figure 6: Density-gradient-dependent errors for $(Z, \mathbf{R})$3-learning versus density fitting, validating DenSNet’s accuracy for spatial density gradients critical to GGA-level functional evaluation.