PaperlayerView exact revision on arXiv

Source-linked AI summary

Equivariant Flow Matching for Electron Density Prediction

Chenxing Liang, Chengdong Wang, Yuchao Lin, Xiaofeng Qian, Shuiwang Ji

arXiv:2610.02651v1physics.chem-phcs.AI

TL;DR

Grid-based density prediction is costly, while basis-set methods’ pointwise coefficient regression misses structural correlations in coefficient space. OrbFlow predicts GTO coefficients with an SE(3)-equivariant flow-matching model trained through a two-phase trajectory curriculum. It achieves state-of-the-art QM9 accuracy and reduces error on every MD molecule against the strongest prior method sharing its basis.

  • Problem

    Grid-based density methods incur high computational cost, while basis-set methods’ pointwise coefficient regression fails to capture structural correlations in coefficient space.

    Source: §1 · p. 4

  • Method

    OrbFlow predicts GTO coefficients through an SE(3)-equivariant flow over a learned probability path, trained with a two-phase trajectory curriculum that fine-tunes the model under its own rollout.

    Source: §1 · p. 4, §3.2 · p. 6

  • Results

    OrbFlow achieves state-of-the-art QM9 accuracy and lowers error on every MD molecule versus the strongest prior method sharing its basis, with SCF acceleration transferring to unseen exchange-correlation functionals.

    Source: §1 · p. 5, §5 · pp. 13–14

  • Takeaways & Limitations

    Predicted densities accelerate SCF across unseen exchange-correlation functionals and recover dipole and quadrupole moments without running SCF.

    Source: §5 · p. 14

    Accuracy holds only at the trained horizon, so OrbFlow is best understood as a trained K-step map.

    Source: §5 · p. 14

Abstract

from arXiv · show

Machine learning surrogates for density functional theory (DFT) have been increasingly used to reduce the cost of first-principles calculations. In this arena, predicting real-space electron densities offers a scalable and transferable initialization for self-consistent field (SCF) procedures. However, current methods face a clear dilemma. That is, grid-based architectures incur a high computational cost, while basis-set methods fail to capture the structural correlations inherent in the coefficient space. Here, we develop OrbFlow, an $\mathrm{SE}(3)$-equivariant generative model that predicts Gaussian-type orbital (GTO) coefficients via flow matching. OrbFlow retains the efficiency of a compact atom-centered basis while replacing pointwise regression with a learned probability path over the full coefficient space. It is trained through a two-phase trajectory curriculum that mitigates discretization drift during numerical integration. OrbFlow achieves state-of-the-art accuracy on QM9, reducing density error by 13.6% relative to the previous best model, and reduces error by 51% to 63% on every molecule of the MD benchmark relative to the strongest prior method sharing its basis. The predicted density also cuts SCF iterations by up to 68% with zero-shot transfer to unseen exchange-correlation functionals and recovers dipole and quadrupole moments to within a few percent of DFT references without any SCF calculation.

1. Introduction

DFT’s cubic SCF cost and limitations of existing density-prediction approaches motivate a compact, correlation-aware alternative. OrbFlow predicts GTO coefficients with equivariant flow matching and a two-phase rollout curriculum, achieving strong density accuracy and downstream utility.

1. Introduction

  • DFT’s iterative SCF calculation scales cubically with electron count, while matrix-centric surrogates can struggle to generalize to larger, unseen geometries.Real-space electron density is described as a smooth, localized, transferable field that can provide a useful initial guess.Source: §1 · p. 4
  • Grid-based density methods incur cubic spatial scaling, while basis-set methods’ pointwise coefficient regression misses structural correlations in coefficient space.OrbFlow addresses this dilemma with flow matching over compact GTO coefficients.Source: §1 · p. 4
  • OrbFlow’s two-phase trajectory curriculum first initializes a stable vector field, then fine-tunes it through differentiable multi-step rollouts to mitigate discretization drift.The contribution summary reports matching single-stage alternatives’ accuracy at 2.5× less training compute.Source: §1 · p. 4
  • OrbFlow achieves state-of-the-art accuracy on QM9 and lowers error on every MD molecule versus the strongest prior model sharing its GTO basis.Its SCF acceleration transfers zero-shot across unseen exchange-correlation functionals.Source: §1 · p. 5

2. Related Work

Prior density-prediction work spans grid-based probes and several basis-set approaches, each with different representational and computational characteristics. Flow matching has been used for quantum-mechanical operators, while OrbFlow extends it to GTO coefficient spaces and trains under its own rollout.

Electron Density Prediction

  • Probe-based methods evaluate density on dense 3D grids [23] [30] [16], while basis-set methods predict expansion coefficients using hybrid plane-wave/GTO [25], Gaussian-mixture [11], or auxiliary GTO bases [8] [14] [44].Grid methods face cubic scaling with voxel resolution; auxiliary GTO regression uses a fixed atom-centered basis.Source: §2 · p. 5

Flow Matching and Exposure Bias

  • Flow matching [36] [54] has been applied to molecular tasks and Hamiltonian prediction in QHFlow [26]; OrbFlow extends it to GTO coefficient spaces with trajectory-regularized multi-step Euler training.The cited work frames the mismatch between single-step training and compounding multi-step inference error as exposure bias [4] [46].Source: §2 · p. 5
  • Prior approaches to exposure bias include teacher-based distillation [49] and shortcut models [13]; OrbFlow’s curriculum uses neither a separate teacher nor extra step-size conditioning.It continues training the same network under its own K-step rollout, with identical weights across both phases.Source: §2 · p. 5

3. Methods

OrbFlow represents continuous molecular density with atom-centered GTOs and predicts their coefficients using a geometry-conditioned equivariant flow. Its training curriculum aligns the learned field with inference rollouts, while the architecture preserves symmetry in the coefficient space.

3.1. Charge Density Representation via Gaussian-Type Orbitals

  • The density is represented as a linear combination of localized, atom-centered GTOs, with basis enrichment using bond-midpoint virtual nodes and even-tempered shell augmentations.The basis starts from uncontracted def2-QZVPPD elemental configurations.Source: §3.1 · p. 6

3.2. Equivariant Flow Matching

  • OrbFlow learns a transport from a prior over coefficients to target coefficients, using a geometry-conditioned vector field in a continuous normalizing flow.Targets are obtained by projecting DFT densities onto the basis via regularized least squares; the predicted endpoint is supervised through synthesized density error on subsampled spatial probes.Source: §3.2 · pp. 6–7
  • The two-phase curriculum initializes a field at isolated timesteps, then fine-tunes it from pure noise under the inference-time K-step Euler rollout.Phase 2 penalizes accumulated drift directly, and continues training the same network without a separate teacher.Source: §3.2 · p. 7
  • An equivariant network paired with an irrep-isotropic prior preserves rotational symmetry through the flow and yields a symmetry-preserving density.The prior draws each angular-momentum subspace independently from a zero-mean isotropic Gaussian.Source: §3.2 · p. 8

3.3. Model Implementation

  • The model uses stacked EquiformerV3 blocks with equivariant attention and local SO(2) convolutions, followed by a bilinear e3nn tensor-product readout for per-atom GTO coefficients.The backbone processes both molecular geometry and the current flow state.Source: §3.3 · p. 8

4. Experiments

Across QM9 and MD, OrbFlow lowers density-prediction error; the experiments also assess SCF initialization, density-derived observables, inference cost, and training ablations.

4.1. QM9 Dataset

  • 0.152% NMAE on the primary QM9 test set is 13.6% below ELECTRA (0.176%); on Split 2, OrbFlow reaches 0.165%, 13.6% below SCDP + VO.OrbFlow achieves state-of-the-art accuracy on both partitions.Source: §4.1 · pp. 8–9, Table 1 · p. 9

4.2. MD Dataset

  • 51% to 63% lower test error on every MD molecule than SCDP; averaged NMAE falls from 1.81% to 0.798% (2.3×).OrbFlow also outperforms GPWNO on all six molecules and approaches ELECTRA, which has roughly 55× more parameters.Source: §4.2 · pp. 9–10, Table 2 · p. 10

4.3. DFT Acceleration and Functional Transferability

  • PBE0 reduces SCF iterations by 68.13% and wall-clock time by 60.98%; B3LYP5's RIC is 77.04% against PBE's 54.48%.OrbFlow was trained at the GGA level and evaluated on B3LYP5 and PBE0 without retraining; the PBE0 SAD baseline converges unusually slowly.Source: §4.3 · p. 11, Table 3 · p. 12

4.4. Direct Physical Properties from Predicted Densities

  • Relative MAE is 1.0% for dipole magnitude and 1.1% for quadrupole norm, with no SCF calculation; accuracy is uniform across molecules.The evaluation uses 500 held-out QM9 molecules against PBE references.Source: §4.4 · p. 12

4.5. Inference Efficiency

  • 0.066s end-to-end inference cost is on par with ELECTRA (0.061s) and 1.6× SCDP, while remaining negligible against a typical SCF iteration.Network cost scales close to linearly with rollout budget K, making latency predictable and controllable.Source: §4.5 · p. 12

4.6. Ablation Studies

  • 0.240 ± 0.005% NMAE with noise plus rollout beats direct regression (0.280 ± 0.003%), noise alone (0.295 ± 0.016%), and rollout alone (0.268 ± 0.003%).The configurations share backbone, parameter count, and step budget; rollout supervision trains on the model's own intermediate states.Source: §4.6 · pp. 12–13, Table 4 · p. 13
  • 74 GPU-h for the curriculum is comparable in accuracy to 81 without a warm start and 182 for joint optimization.The curriculum reuses one Phase-1 checkpoint across rollout horizons; optimization-step budgets are matched, not wall-clock budgets.Source: Table 5 · p. 13, §4.6 · p. 13
  • 0.240% NMAE at two inference steps worsens to 0.539% at four, because Phase 2 sharpens the network at states visited by the two-step rollout.The train-matched budget is both the most accurate and the cheapest.Source: §4.6 · p. 13

5. Conclusion and Limitations

OrbFlow reaches state-of-the-art QM9 accuracy and lowers error on every MD molecule relative to the strongest prior method sharing its basis. The appendices detail its GTO basis and coefficient projection, while noting that accuracy holds only at the trained horizon.

5. Conclusion and Limitations

  • OrbFlow reaches state-of-the-art QM9 accuracy, lowers error on every MD molecule versus the strongest same-basis prior method, and matches a Gaussian-mixture model 55× its size in accuracy.The rollout adds a small inference overhead, and accuracy holds only at the trained horizon; the authors describe OrbFlow as a trained K-step map.Source: §5 · pp. 13–14

A. GTO Basis Set Details and Coefficient Projection

  • Ground-truth DFT density is projected onto coefficients of a localized auxiliary GTO basis to define the model targets.The section covers the basis formulation, even-tempered radial exponents, and regularized least-squares coefficient projection.Source: A. GTO Basis Set Details and… · p. 19

A.1. Even-Tempered Gaussian-Type Orbital (GTO) Basis Set

  • The atom-centered GTO expansion separates radial and angular structure, with even-tempered exponents spanning core-like and diffuse valence densities.The exponents use an element-specific base and scaling factor; hyperparameters are optimized per element.Source: §A.1 · p. 19

A.2. Element-Specific Hyperparameter Configurations

  • Element-specific basis configurations use the full element table for QM9 and the H/C/O subset for MD, with oxygen’s basis assigned to virtual nodes.The configurations specify maximum angular momentum, radial shells, progression boundaries, and scaling factors.Source: §A.2 · p. 19, Table 6 · p. 20

A.3. L2 Least-Squares Coefficient Projection Mechanics

  • Ground-truth DFT densities are projected onto GTO coefficients by minimizing a grid-weighted least-squares objective with Tikhonov regularization.The grid formulation uses integration weights and a regularizer to prevent numerical instability from near-linear dependencies in the basis.Source: §A.3 · p. 21

A.4. Empirical Fitting Floor

OrbFlow’s QM9 error is above the GTO fitting floor, while its remaining gap varies across MD molecules. Its predicted electron count is already close to the valence count before renormalization for SCF initialization.

A.4. Empirical Fitting Floor

  • 0.152% OrbFlow QM9 NMAE exceeds the 0.061% fitting residual; versus SCDP + VO at 0.178%, it removes 22% of the error above the floor.Table 7 reports the GTO density fitting floor versus OrbFlow NMAE; the floor is reconstructed from stored ridge ground-truth coefficients without a network.Source: §A.4 · pp. 21–22, Table 7 · p. 22
  • On MD, benzene and ethane lie near the basis fitting floor, whereas ethanol and malonaldehyde remain well above it.The passage characterizes compact rigid molecules as well represented by the basis and flexible oxygenates as leaving a larger residual for the model to close.Source: §A.4 · p. 22

A.5. Electron Count Conservation

  • 0.014 e is OrbFlow’s average QM9 electron-count difference from Nval, or 0.030% of Nval, before renormalization for SCF initialization.The mean valence count is 48.28 e; charge conservation is not imposed during training, and the predicted density is renormalized to the PAW valence count for an SCF initial guess.Source: §A.5 · p. 22

B. Proof of Prior Rotational Invariance

The proof establishes rotational invariance of the irrep-isotropic prior by using the unitarity of Wigner D-matrices. It concludes that the density is unchanged under arbitrary three-dimensional rotations.

B. Proof of Prior Rotational Invariance

  • Unitary Wigner D-matrices preserve each angular-momentum subspace’s Gaussian quadratic form, yielding rotational invariance of the density.The proposition concerns invariance of the irrep-isotropic prior under rotations in SO(3).Source: B. Proof of Prior Rotational Invariance · p. 22

C. Implementation Details

The appendix consolidates shared settings for the model architecture, target construction, and optimization procedure, with dataset-specific settings given separately.

C. Implementation Details

  • Architecture, target construction, and optimization settings are shared across QM9 and MD except where noted; dataset-specific configurations are listed separately.The appendix identifies Appendix D.1 as the location of per-dataset hyperparameters and training configuration.Source: C. Implementation Details · p. 23

C.1. Network Architecture and Coefficient Parameterization

OrbFlow uses an equivariant graph architecture to predict GTO coefficients and trains a conditional flow from structured noise to target coefficients. Its appendix records the optimization and experimental protocols.

C.1. Network Architecture and Coefficient Parameterization

  • OrbFlow represents molecules with nuclear-coordinate and bond-midpoint orbital nodes, then uses an equivariant EquiformerV3 backbone and maps latent tensors to per-atom GTO coefficients.The graph uses a 6 Å cutoff; the QM9 and MD backbones use 8 and 3 blocks, respectively.Source: §C.1 · p. 23

C.2. Flow-Matching Objective

  • Training uses conditional flow matching with an endpoint parameterization and an irrep-structured Gaussian prior in normalized coefficient space.Only the scalar channel has a nonzero mean; the MD prior’s scalar mean and standard deviation are estimated from training coefficients and held fixed.Source: §C.2 · p. 23, Table 8 · p. 24

C.3. Optimization Protocol

  • Production runs use Adam, a 1 × 10^-3 base learning rate, warm-up followed by power decay, EMA, gradient clipping, and FP32 training.Checkpoints are selected by validation density error after two-step Euler integration; jobs ran on NVIDIA H200 GPUs.Source: §C.3 · pp. 23–24

D. Experiment Details

  • The appendix collects supporting procedures, dataset-specific training settings, SCF solver settings, latency measurement protocols, and multipole definitions.These protocols support the paper’s SCF-acceleration, efficiency, and direct-property evaluations.Source: D. Experiment Details · p. 24

D.1. Training Configuration

The appendix specifies training and SCF evaluation protocols, explains why iteration savings translate only partly to runtime savings, and documents multipole and inference-latency measurements. It also describes the deterministic two-step ablation baseline.

D.1. Training Configuration

  • The QM9 training configuration uses the standard 123,835/50/10,000 train/validation/test split, while MD uses separate models for six trajectories.The appendix gives probe counts, batch sizes, and phase lengths for both datasets.Source: §D.1 · p. 24, Table 10 · p. 25

D.2. DFT Solver Settings

  • SCF acceleration is evaluated by matched SAD and OrbFlow-density initializations, recording iteration counts and a relative time ratio that includes OrbFlow inference latency.The evaluation uses 500 PBE test molecules and 200 molecules for unseen functionals; other solver settings are held identical between the matched runs.Source: §D.2 · p. 26

D.3. Why Iteration Savings Only Partially Translate to Wall-Clock Savings

  • PBE’s fixed-cost share is 77.0%, versus 26.5% for B3LYP5 and 10.5% for PBE0; fixed setup limits PBE’s wall-clock savings despite iteration reductions.The appendix notes that a 46% iteration reduction on PBE molecules of this size buys only about 10% in wall-clock time; hybrid functionals’ exact exchange makes iteration costs more dominant.Source: §D.3 · pp. 26–27

D.4. Multipole Moments

  • Held-out QM9 dipole and quadrupole prediction accuracy is uniform across molecules rather than driven by a favorable average over poorly predicted outliers.Figure 3 compares predicted and DFT-reference multipoles per molecule; its panels report the corresponding Table 11 MAEs.Source: §D.4 · p. 27, Table 11 · p. 27, Figure 3 · p. 28

D.5. Efficiency Measurement Protocol

  • OrbFlow’s network cost scales close to linearly with integration steps, while density-synthesis latency differs with the density representation rather than flow matching.At two steps, OrbFlow’s network cost is 0.054s; synthesis costs 0.013s for OrbFlow and 0.015s for SCDP, versus 0.003s for ELECTRA.Source: §D.5 · pp. 27, 29

E.1. Deterministic Two-Step Baseline

  • The deterministic two-step baseline uses the same EquiformerV3 configuration and training budget, replacing sampled noise and continuous time with a zero-state initialization and learned step-index embedding.This design isolates the flow formulation’s stochastic prior and bridge interpolation.Source: §E.1 · p. 29, E. Extended Ablation Results · p. 29

E.2. Curriculum Variants and Training Compute

The curriculum’s shared Phase-1 warm start reduces the incremental compute needed to fine-tune a rollout horizon.

E.2. Curriculum Variants and Training Compute

  • 39 GPU-h of Phase-2 fine-tuning brings the K=2 curriculum to 74 GPU-h total, with its 35 GPU-h Phase-1 warm start shared with K=3.The incremental fine-tuning cost is roughly half the 81 GPU-h required to train the K=2 objective from scratch.Source: §E.2 · p. 29

E.3. Training-Time Rollout Horizon

Training and inference perform best when rollout supervision and integration budgets match: the appendix reports no measurable gain from training beyond two steps and degraded accuracy when inference exceeds the trained horizon.

E.3. Training-Time Rollout Horizon

  • K=1 reaches 0.296 ± 0.005%, statistically indistinguishable from Phase-1-only training at 0.295 ± 0.016%; K=2 and K=3 reach 0.240 ± 0.005% and 0.240 ± 0.002%.Because training compute grows with K, K=2 is adopted as the smallest informative rollout budget.Source: §E.3 · pp. 29–30, Table 13 · p. 30

E.4. Inference-Time Integration Budget

  • At four inference steps, QM9 NMAE rises to 0.539% from 0.240% at two steps; relative inference cost reaches 3.98× at sixteen steps.For the K=2-trained model, two steps are both the accuracy-optimal and most compute-efficient inference choice.Source: Table 14 · p. 30, §E.4 · p. 30

Related papers

Loading 2610.02651v1…