Paper deep dive
Equivariant Flow Matching for Electron Density Prediction
Chenxing Liang, Chengdong Wang, Yuchao Lin, Xiaofeng Qian, Shuiwang Ji
Intelligence
Status: succeeded | Model: Gemma-4-26B-A4B | Prompt: intel-v1 | Confidence: 91%
Last extracted: 10/5/2026, 6:39:00 AM
Summary
The paper introduces OrbFlow, an SE(3)-equivariant generative model that predicts Gaussian-type orbital (GTO) coefficients for electron density prediction using flow matching. It addresses the trade-off between computational cost and accuracy in existing methods by learning a probability path over the coefficient space, trained via a two-phase trajectory curriculum. OrbFlow achieves state-of-the-art accuracy on the QM9 dataset and significantly reduces error on the MD benchmark, while also accelerating Self-Consistent Field (SCF) iterations and recovering physical properties like dipole moments.
Entities (10)
Relation Signals (7)
OrbFlow → achievesstateofthearton → QM9
confidence 97% · OrbFlow achieves state-of-the-art accuracy on QM9
OrbFlow → usesmethod → Flow Matching
confidence 95% · OrbFlow, an SE(3)-equivariant generative model that predicts Gaussian-type orbital (GTO) coefficients via flow matching.
OrbFlow → hasproperty → SE(3)-equivariant
confidence 94% · OrbFlow, an SE(3)-equivariant generative model
OrbFlow → predicts → Gaussian-type orbital
confidence 92% · predicts Gaussian-type orbital (GTO) coefficients via flow matching.
OrbFlow → accelerates → Self-Consistent Field
confidence 88% · The predicted density also cuts SCF iterations by up to 68%
OrbFlow → outperforms → ELECTRA
confidence 85% · reducing density error by 13.6% relative to the previous best model... ELECTRA
OrbFlow → outperforms → SCDP
confidence 85% · Compared with SCDP, the strongest prior GTO-based method, OrbFlow reduces test error on every molecule
Cypher Suggestions (0)
No Cypher suggestions yet.
Abstract
Abstract: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.
Tags
Links
- Source: https://arxiv.org/abs/2610.02651v1
- Canonical: https://arxiv.org/abs/2610.02651v1
Trouble viewing inline? Open PDF directly →
Full Text
82,326 characters extracted from source content.
Expand or collapse full text
October 5, 2026 Equivariant Flow Matching for Electron Density Prediction Chenxing Liang 1 , Chengdong Wang 2 , Yuchao Lin 1 , Xiaofeng Qian 2,3,4∗ , Shuiwang Ji 1,2,5∗ 1 Department of Computer Science and Engineering, Texas A&M University 2 Department of Materials Science and Engineering, Texas A&M University 3 Department of Electrical and Computer Engineering, Texas A&M University 4 Department of Physics and Astronomy, Texas A&M University 5 Department of Mechanical Engineering, Texas A&M University Abstract: 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 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. Keywords: Equivariant Flow Matching, Electron Density Prediction, Generative Models, Quantum Chemistry * Equal senior authorship Github Repository: https://github.com/divelab/AIRS/tree/main/OpenDFT/OrbFlow Model Checkpoints: https://huggingface.co/divelab/OrbFlow Dataset: https://huggingface.co/datasets/divelab/OrbFlow-data arXiv:2610.02651v1 [physics.chem-ph] 2 Oct 2026 Equivariant Flow Matching for Electron Density Prediction Contents 1 Introduction4 2 Related Work5 3 Methods6 3.1 Charge Density Representation via Gaussian-Type Orbitals . . . . . . . . . . . . . . . . . . . 6 3.2 Equivariant Flow Matching . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 6 3.3 Model Implementation . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 8 4 Experiments8 4.1 QM9 Dataset . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 8 4.2 MD Dataset . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 9 4.3 DFT Acceleration and Functional Transferability . . . . . . . . . . . . . . . . . . . . . . . . . 10 4.4 Direct Physical Properties from Predicted Densities . . . . . . . . . . . . . . . . . . . . . . . 12 4.5 Inference Efficiency . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 12 4.6 Ablation Studies . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 12 5 Conclusion and Limitations13 A GTO Basis Set Details and Coefficient Projection19 A.1 Even-Tempered Gaussian-Type Orbital (GTO) Basis Set . . . . . . . . . . . . . . . . . . . . . 19 A.2 Element-Specific Hyperparameter Configurations . . . . . . . . . . . . . . . . . . . . . . . . 19 A.3 L2 Least-Squares Coefficient Projection Mechanics . . . . . . . . . . . . . . . . . . . . . . . . 19 A.4 Empirical Fitting Floor . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 21 A.5 Electron Count Conservation . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 22 B Proof of Prior Rotational Invariance22 C Implementation Details23 C.1 Network Architecture and Coefficient Parameterization . . . . . . . . . . . . . . . . . . . . . 23 C.2 Flow-Matching Objective . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 23 C.3 Optimization Protocol . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 23 2 Equivariant Flow Matching for Electron Density Prediction D Experiment Details24 D.1 Training Configuration . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 24 D.2 DFT Solver Settings . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 26 D.3 Why Iteration Savings Only Partially Translate to Wall-Clock Savings . . . . . . . . . . . . . 26 D.4 Multipole Moments . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 27 D.5 Efficiency Measurement Protocol . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 27 E Extended Ablation Results29 E.1 Deterministic Two-Step Baseline . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 29 E.2 Curriculum Variants and Training Compute . . . . . . . . . . . . . . . . . . . . . . . . . . . 29 E.3 Training-Time Rollout Horizon . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 29 E.4 Inference-Time Integration Budget . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 30 3 Equivariant Flow Matching for Electron Density Prediction 1. Introduction Density functional theory (DFT) is a cornerstone of computational quantum chemistry, materials science, and solid-state physics [20,29,21,28,39]. It maps the many-body Schrödinger equation onto many one-electron Kohn-Sham equations, in which the total energy is a functional of electron density rather than many-electron wavefunctions, enabling ab initio evaluation of molecular energies and electronic structure. The iterative self-consistent field (SCF) calculation that solves the Kohn-Sham equations scales cubically (O(N 3 )) with the number of electronsN, making it a bottleneck for large molecular structures, amorphous systems, and long-timescale molecular dynamics. Recently, rapid advances have been made in developing deep learning surrogates to accelerate or bypass DFT, most commonly by predicting electronic structure operators directly from atomic coordinates [57], such as Kohn-Sham Hamiltonian matrixH[53,33,17,55,56,58,18,59,34,38,26] or density matrix P[50,19,12] in localized basis sets. However, these matrix-centric predictions face systemic scalability and transferability bottlenecks.His numerically difficult to predict owing to its large, dense entrywise structure [37]. Models trained on small clusters therefore often fail to generalize to larger, unseen geometries, sometimes causing numerical divergence in the downstream DFT SCF calculations. In contrast, the real-space electron density is a physically smooth, localized, and universally transferable scalar field, and a high-quality initial guess for it can consistently reduce the number of SCF iterations under same convergence threshold. Existing density prediction methods generally fall into two paradigms [14]. Probe-based methods evaluate density on dense 3D grids and offer strong expressiveness but incur heavy overhead from cubic spatial scaling. In contrast, basis-set methods are more efficient but rely on pointwise coefficient regression that fails to capture the structural correlations of the coefficient space. We propose to address this dilemma by flow matching (FM) [36, 54], which parameterizes a time-dependent vector field that defines efficient ordinary differential equation (ODE) trajectories and yields rapid convergence while capturing intricate structural dependencies. FM has already achieved state-of-the-art accuracy in predicting quantum mechanical operators such as Kohn-Sham Hamiltonians [26], but its capabilities have not yet been applied to other continuous physical fields such as Gaussian-type orbital (GTO) basis spaces. Motivated by the above considerations, we introduce OrbFlow, a geometry-conditioned,SE(3)-equivariant flow matching framework operating within a compact GTO coefficient space. Rather than a rigid, one-shot regression mapping, OrbFlow models the distribution of optimal basis spaces by learning a probability path that deforms a structured noise prior into the exact target distribution. To eliminate discretization drift during ODE integration rollouts, we propose a two-phase trajectory curriculum training strategy. The model is first initialized with an independent flow matching objective to build a stable global vector field. It is then dynamically fine-tuned via an end-to-end differentiable multi-step Euler rollout loss that directly suppresses numerical truncation errors. Figure 1 presents an overview of OrbFlow. Our primary contributions are summarized below. • Few-Step Flow Matching for Basis-Set Density Prediction. We formulate GTO coefficient prediction as conditional flow matching, and show that its benefit is realized when the flow is trained under its own few-step rollout. •Trajectory Curriculum Training. We introduce a two-phase curriculum that fine-tunes the flow under its own differentiable few-step rollout, requiring no teacher network and matching the accuracy of single-stage alternatives at 2.5× less training compute. •Symmetry-Preserving Equivariant Integration. We pair a high-order equivariant graph transformer with an irrep-isotropic prior, so that symmetry holds at every state along the rollout and the predicted 4 Equivariant Flow Matching for Electron Density Prediction 푐 ! 푐 " 푐 # SCF / DFT acceleration Downstream Properties 푣 $ 푡,푀푑푡 푐 ! ~ 푝 ! 푐 # ~ 푝 # 휌(푐 ! )휌(푐 # ) 푀 Figure 1: Overview of OrbFlow. Given a molecular graphM, OrbFlow initializes Gaussian-type orbital (GTO) coefficients from a structured noise priorc 0 ∼ p 0 and integrates a learned, geometry-conditioned vector fieldv θ (t, M)overt∈ [0, 1]to obtain the target coefficientsc 1 ∼ p 1 . The recovered coefficients are converted into a continuous real-space electron densityρ(r)through the auxiliary GTO basis expansion, providing a fast surrogate for DFT initialization and enabling direct prediction of physical properties. density transforms equivariantly under rigid motions. •Accuracy and Downstream Utility. OrbFlow achieves state-of-the-art accuracy on QM9 and reduces error on every MD molecule relative to the strongest prior model sharing its GTO basis, with SCF acceleration transferring zero-shot across unseen exchange-correlation functionals. 2. Related Work Electron Density Prediction. Continuous real-space electron density prediction generally falls into two paradigms [14]. Probe-based models evaluate density as spatial queries on dense 3D grids [23,30,16], but cubic scaling with voxel resolution imposes severe memory and throughput costs at scale. Basis-set methods instead predict expansion coefficients of localized analytical functions, and split further by basis choice. Hybrid models combine a plane-wave basis with a localized GTO basis [25]. Gaussian mixture models such as ELECTRA [11] instead use a learnable set of Gaussian primitives rather than a fixed atom-centered basis. Auxiliary GTO regression predicts coefficients directly against a fixed atom-centered basis [8, 14, 44, 9]. Flow Matching and Exposure Bias. Flow Matching (FM) [36,54] has emerged as a powerful alternative to diffusion models for molecular generative tasks [22,40,24,52], most relevantly in QHFlow [26], which frames Hamiltonian matrix regression as an equivariant distribution-matching problem. We extend this paradigm to GTO coefficient spaces, pairing a flow over compact GTO coefficients with a trajectory-regularized multi-step Euler objective. The mismatch between single-step training and the compounding error of multi- step inference is known as exposure bias [4,46]. Diffusion and flow-based models remove it in two ways. Distillation trains a student to reproduce the multi-step behavior of a separately trained, frozen teacher, as in progressive distillation [49] and consistency models [51]. Shortcut models [13] drop the teacher instead, conditioning the network on the step size so that one model stays accurate at many step budgets. Our curriculum takes a third route and needs neither. It continues training the same network under its own K-step rollout, with identical weights across both phases, so it adds no model capacity, no teacher, and no extra conditioning input. 5 Equivariant Flow Matching for Electron Density Prediction 3. Methods 3.1. Charge Density Representation via Gaussian-Type Orbitals To represent the continuous molecular electron densityρ(r)compactly, we adopt an atom-centered Gaussian- type orbital (GTO) basis expansion [10], following Fu et al.[14]. Each basis functionφ α, ℓ, m, R i (r) is a normalized spherical Gaussian centered at node coordinateR i , characterized by radial exponentα, angular momentumℓ, and magnetic quantum numberm. The full functional form and the even-tempered scheme used to construct the radial exponents are given in Appendix A.1. Given a molecular geometryR i N i=1 containing both nuclear coordinates and bond-midpoint virtual nodes, we parameterize the spatial density field as a linear combination of these localized GTOs, ρ(r) = N ∑︁ i=1 N b i ∑︁ j=1 ℓ i,j ∑︁ m=−ℓ i,j c i,j,m φ α i,j , ℓ i,j , m, R i (r),(1) whereN b i denotes the number of unique radial shells assigned to nodei. Following the basis enrichment recipe of Fu et al.[14], we construct our basis by starting from uncontracteddef2-QZVPPDelemental configurations [47], embeddingSE(3)-equivariant virtual nodes at bond midpoints to capture interatomic shared densities, and applying even-tempered geometric shell augmentations [2] for radial flexibility. Element- specific hyperparameters are detailed in Appendix A.2, and the representation floor of this basis, the error of the best fit it can express, is reported in Appendix A.4. 3.2. Equivariant Flow Matching Problem Formulation. Our goal is to predict the target GTO coefficient vectorc ⋆ ∈ R d from a molecular graphG = (R, H), whereR∈ R N×3 andH∈ R N×D denote the node coordinates and features. The density is reconstructed linearly from these coefficients on a discrete grid as ρ(c) =O(c) = sΦ c,(2) whereΦ pj = φ j (r p )evaluates thej-th GTO function at grid pointp,s > 0is a scale factor, andO(·)is our differentiable density synthesis operator. We train onD =(G, c ⋆ ), where target coefficientsc ⋆ are obtained by projecting ground-truth DFT densities onto the basis via a regularized least-squares fit (Appendix A.3). Instead of mappingG 7→ c ⋆ in a single step, our model learns a transport from a tractable priorc 0 ∼ p 0 (·|G) to the target c ⋆ . Conditional Flow Matching in GTO Space. We formulate a continuous normalizing flow (CNF) within this compact, auxiliary GTO coefficient space: a time-dependent, geometry-conditioned vector fieldv θ (x t , t) defining an ODEdx t /dt = v θ (x t , t)overt∈ [0, 1]that deforms the priorp 0 (c 0 )att = 0into a distribution on c ⋆ at t = 1. Following Kim et al. [26], we use a linear probability bridge, x t = (1− t) c 0 + t c ⋆ ,(3) with conditional field u t (x t | c ⋆ , c 0 ) = c ⋆ − c 0 . We parameterize v θ via an endpoint scheme, v θ (x t , t) = f θ (G, x t , t)− x t 1− t ,(4) 6 Equivariant Flow Matching for Electron Density Prediction wheref θ (G, x t , t)predicts the clean target endpoint ˆ c ⋆ ≈ c ⋆ . Since our physical objective is real-space density accuracy rather than raw coefficient error, we supervise the predicted endpoint through the synthesis operator O(·), evaluating the MAE over P ′ subsampled spatial probes: L stage1 (θ) = E t,G, c 0 [︁ 1 P ′ P ′ ∑︁ p=1 ⃒ ⃒ [︀ O( f θ (G, x t , t)) ]︀ p −ρ lab p ⃒ ⃒ ]︁ ,(5) Two-Phase Trajectory Curriculum Learning. Training solely on isolated, random timesteps via(5)does not align the vector field with the multi-step trajectories used at inference. We address this with a two-phase curriculum. Phase 1 first establishes a good vector field at arbitrary, isolated timesteps. Phase 2 then fine- tunes that field under the same multi-step rollout used at inference, so training conditions match evaluation conditions directly. Phase 1 (Field Initialization). We minimizeL stage1 forT 1 steps, teaching the network to map arbitrary perturbed states(G, x t , t)to a target endpoint whose synthesized density matches the true DFT field. This establishes a stable continuous normalizing flow without the overhead of unrolling the ODE, giving a stable warm-start vector field v θ . Phase 2 (Trajectory Regularization). Starting from the Phase-1 weights, we continue training forT 2 steps under a fullK-step Euler rollout that mirrors how the model is actually used at inference. Each rollout begins at the pure noise priorx t 0 = c 0 att 0 = 0, rather than at an interpolated point along the bridge, and is unrolled end to end through the network in a differentiable manner. We discretize the interval[0, 1]intoK equal steps of size∆t = 1/K, with t k = k·∆t, and advance the state at each step with the Euler update x t k+1 = x t k +∆t· v θ (x t k , t k ), for k = 0, . . . , K− 1.(6) The final state ˆ c 1 = x t K is the model’s simulated prediction of the target coefficients, usingK = 2in our primary experiments. We train the model by comparing the density reconstructed from this simulated endpoint to the true DFT density, giving the Phase-2 objective L stage2 (θ) = E G, c 0 [︁ 1 P ′ P ′ ∑︁ p=1 ⃒ ⃒ [︀ O( ˆ c 1 ) ]︀ p −ρ lab p ⃒ ⃒ ]︁ .(7) In Phase 1, the input statex t is built by interpolating with the ground-truth target itself via(3), so the network is partially shown the answer during training. In Phase 2, the rollout starts from pure noise with no such guidance, so the model must navigate the full trajectory on its own, and any small per-step error can accumulate into discretization drift. Phase 2 directly penalizes this accumulated drift. Penalizing this drift acts as a trajectory regularizer that straightens the flow field, which is why Phase 2 improves performance at the trained rollout horizon even though it starts from Phase-1 weights that already minimizeL stage1 . Phase 2 achieves this by continuing to train the same network under its own rollout, rather than by introducing a separate teacher network or a distillation step. Equivariant Flow and Isotropic Irrep Priors. For the generated density to be invariant under rigid translations and rotations across both phases, the flow must be strictlySE(3)-equivariant. Atomic coordinates are centered at the origin (r i ← r i − 1 N ∑︀ j r j ), removing translational dependence. Under a rotation R∈ SO(3), GTO basis functions (built from spherical harmonicsY m ℓ ) of angular momentumℓtransform via the(2ℓ + 1)-dimensional Wigner D-matrixD (ℓ) (R), so the coefficient vectorctransforms covariantly as a 7 Equivariant Flow Matching for Electron Density Prediction direct sum of irreducible representations, c R −→ D(R)c = ⨁︁ ℓ D (ℓ) (R)c (ℓ) .(8) We therefore require both the networkf θ and the priorp 0 (c 0 )to respect this symmetry. We parameterizef θ with a high-order equivariant graph transformer (Sec. 3.3) satisfying f θ (︀ R·G, D(R)x t , t )︀ = D(R) f θ (G, x t , t),(9) and requirep 0 (c 0 ) = p 0 (D(R)c 0 )for allR ∈ SO(3), satisfied by an irrep-isotropic Gaussian prior in which each angular momentum subspaceℓis drawn independently from a zero-mean isotropic Gaussian, c (ℓ) 0 ∼N (0, σ 2 ℓ I 2ℓ+1 ), with σ 2 ℓ matched to the empirical variance of the target coefficients in subspace ℓ. Proposition 1 Letc 0 ∼ p 0 (c 0 )be sampled from the irrep-isotropic prior. Under any rotationR∈ SO(3), the prior density is invariant, p 0 (c 0 ) = p 0 (D(R)c 0 ). The proof is in Appendix B. Pairing this rotationally invariant prior with our equivariant networkf θ guarantees that integrating the velocity field yields a physically consistent, symmetry-preserving density. 3.3. Model Implementation We parameterizef θ using a stack ofLEquiformerV3 blocks [35], leveragingSE(3)-equivariant multi-head attention and localSO(2)convolutions [42] to jointly process the molecular geometryGand the current flow statex t . The backbone readout is mapped to per-atom GTO coefficients via a bilineare3nntensor-product readout [15], serving directly as the flow head f θ in (4). Full architectural details are in Appendix C. 4. Experiments We evaluate OrbFlow on the QM9 and MD benchmarks, building on the GTO auxiliary basis formulation of SCDP [14] and a state-of-the-art EquiformerV3 [35] backbone. We evaluate predicted real-space densities ρ pred =O( ˆ c 1 )against ground-truth DFT referencesρ lab on the full spatial grid. All QM9 and MD results are averaged over three independently trained random seeds. We report global reconstruction accuracy via the normalized Mean Absolute Error, NMAE = 100× ∥ρ pred −ρ lab ∥ 1 ∥ρ lab ∥ 1 . 4.1. QM9 Dataset Dataset. QM9 [48,45] contains∼134,000 stable organic molecules with up to nine heavy atoms (C, O, N, F) plusH. We use the real-space density files curated by Jørgensen and Bhowmik[23](VASP, PBE functional, PAW method), adopting the standard 123,835/50/10,000 train/val/test partition. To compare directly with ChargE3Net [30], InfGCN [8], and GPWNO [25], we also evaluate on “Split 2,” the last 1,600 structures of the standard test set, matching these baselines’ protocol. Results. OrbFlow achieves state-of-the-art accuracy on both partitions (Table 1). It reaches 0.152% NMAE on the primary test set, a 13.6% reduction over the strongest baseline ELECTRA (0.176%) and 14.6% over 8 Equivariant Flow Matching for Electron Density Prediction Table 1: QM9 charge density prediction error, grouped by model representation. The first column reports NMAE [%] on the standard 10,000-molecule test set, and the second on “Split 2,” the last 1,600 structures of that test set. Results for OrbFlow and ELECTRA were obtained in this work as means over three training seeds, all other entries are quoted from the cited papers. Lower NMAE is better (↓). Bold marks the best result per column, underlinethe second-best. ModelNMAE [%]↓ NMAE, Split 2 [%]↓ Probe-based i-DeepDFT [23]0.357± 0.001– e-DeepDFT [23]0.284± 0.001– ChargE3Net [30]0.196± 0.0010.203± 0.003 Hybrid PW & GTO GPWNO [25]–0.73 Gaussian Mixture ELECTRA [11]0.176± 0.001– GTO-based InfGCN [8]0.869± 0.0020.93 SCDP, L = 3 [14]0.432± 0.0010.438± 0.003 SCDP + VO, L = 6 [14]0.178± 0.0010.191± 0.002 OrbFlow (ours)0.152± 0.0010.165± 0.001 its closest GTO-based predecessor SCDP + VO (0.178%), and 0.165% on Split 2, 13.6% below SCDP + VO and 18.7% below ChargE3Net (0.203%). These margins reflect both the flow matching formulation and the EquiformerV3 backbone. Sec. 4.6 separates the two at fixed architecture and compute, and finds that the flow matching formulation contributes on its own. They are also best read against the representation floor of GTO basis, since a direct least-squares fit of the ground-truth densities reaches 0.061% NMAE (Appendix A.4) and OrbFlow removes 22% of the error SCDP + VO leaves above it. Figure 2 visualizes predicted isosurfaces and spatial error. Training Details. We train with Adam under a linear warm-up followed by a power-decayed schedule, following the two-phase curriculum of Sec. 3.2 with rollout budgetK = 2. Full hyperparameters are in Appendix D.1. 4.2. MD Dataset Dataset. We benchmark on molecular dynamics (MD) trajectories spanning six molecules, ethanol, benzene, phenol, resorcinol, ethane, and malonaldehyde (MDA) [7,6], with reference densities computed at the PBE level via plane-wave Quantum ESPRESSO calculations, sampled across diverse geometries on uniform 3D grids. Results. Table 2 reports NMAE across the six MD trajectories. Compared with SCDP, the strongest prior GTO-based method, OrbFlow reduces test error on every molecule by 51% to 63% in relative terms, with the largest gains on benzene (1.13% to 0.421%) and ethane (2.05% to 0.871%). Averaged over the six 9 Equivariant Flow Matching for Electron Density Prediction Table 2: NMAE [%] on MD test sets for OrbFlow, SCDP [14], InfGCN [8], ELECTRA [11], and GPWNO [25], grouped by model representation. Bold marks the best result per row,underlinemarks the second-best. Results for OrbFlow and ELECTRA were obtained in this work and are reported as mean±std over three training seeds. SCDP, InfGCN and GPWNO numbers are taken from Fu et al.[14]. ELECTRA’s ethanol standard deviation is below 5× 10 −4 and rounds to 0.000. Molecule GTO-basedGaussian Mixture Hybrid PW & GTO OrbFlow (ours) SCDP InfGCNELECTRAGPWNO # Params13M16M–723M– MD-ethanol1.078± 0.0232.34± 0.25 8.431.030± 0.0004.00 MD-benzene0.421± 0.0021.13± 0.06 5.110.472± 0.0132.45 MD-phenol0.575± 0.0081.29± 0.07 5.510.573± 0.0122.68 MD-resorcinol0.663± 0.0141.35± 0.08 5.950.619± 0.0052.73 MD-ethane0.871± 0.0082.05± 0.12 7.010.970± 0.0553.67 MD-malonaldehyde1.177± 0.0122.71± 0.60 10.340.826± 0.0255.32 molecules, NMAE falls from 1.81% to 0.798%, a2.3×reduction at comparable model size. On benzene and ethane OrbFlow comes within 0.02 and 0.06 percentage points of the basis fitting floor (Appendix A.4), leaving little room for further gains in this basis. Against ELECTRA, a 723M-parameter Gaussian-mixture model with roughly55×more parameters, OrbFlow remains competitive. OrbFlow is more accurate on benzene and ethane and statistically indistinguishable on phenol (0.575% against 0.573%, well within seed variance), while ELECTRA leads on ethanol, resorcinol, and malonaldehyde, by a small margin on the first two and a clear one on malonaldehyde. Therefore, OrbFlow approaches the accuracy of a much larger Gaussian-mixture model while improving substantially over every prior method operating in the same GTO basis, and it outperforms the hybrid plane-wave method GPWNO on all six molecules by wide margins. Figure 2 visualizes predicted density isosurfaces and spatial error for both benchmarks. Training Details. Full details are in Appendix D.1. 4.3. DFT Acceleration and Functional Transferability A key downstream application of machine-learned charge density is to use it as improved initial guess for SCF calculations, thereby reducing the number of SCF iterations and accelerating DFT calculations. Standard Kohn-Sham DFT calculations typically initialize the SCF loop using a superposition of atomic densities (SAD), which may provide a relatively poor approximation to the converged electronic density and therefore require a larger number of computationally expensive SCF iterations. To improve the initialization, we instead use the real-space charge density ρ pred predicted by OrbFlow as the initial density for the SCF calculation. The resulting SCF acceleration was evaluated using two complementary metrics. The Relative Iteration Count (RIC), defined asRIC = (N pred /N SAD )× 100%, compares the number of SCF iterations required to reach the specified energy convergence criterion when initialized with the OrbFlow-predicted density and the conventional SAD density, respectively. Here,N pred andN SAD denote the corresponding numbers of SCF iterations. Similarly, the Relative Time Ratio (RTR), defined asRTR = (T pred /T SAD )× 100%, compares the corresponding total wall-clock times, whereT pred andT SAD denote the wall-clock times for calculations initialized with ρ pred and SAD, respectively. 10 Equivariant Flow Matching for Electron Density Prediction (a) (b) Prediction Prediction Error Error Figure 2: Visualization of OrbFlow-predicted electron densities and prediction error. (a) Six representative QM9 test molecules. (b) Six representative MD test molecules. Prediction rows show the predicted density isosurface, rendered atρ = 0.69e/Å 3 for QM9 andρ = 1.12e/Å 3 for MD. Error rows show an isosurface of the absolute error|ρ pred − ρ DFT |, at|∆ρ| = 1.2× 10 −3 e/Å 3 for QM9 and|∆ρ| = 1.1× 10 −2 e/Å 3 for MD. Per-example NMAE is reported below for each molecule. Functional Transferability. To assess whether OrbFlow learns transferable physical representations rather than memorizing artifacts specific to a single level of theory, we evaluate its zero-shot generalization across unseen exchange-correlation (XC) functionals using held-out QM9 molecules. All reference calculations are performed using VASP [31] with a fixed plane-wave basis, with the complete computational settings provided in Appendix D.2. Although trained exclusively at the GGA level, OrbFlow is evaluated across higher rungs of Jacob’s ladder to hybrid functionals (B3LYP5 [3, 32], PBE0 [1, 41]), without any retraining. As shown in Table 3, at the training-distribution functional PBE, OrbFlow reduces the required SCF iterations to 54.48% of the SAD baseline and the wall-clock time to 89.51%. The wall-clock saving is smaller because SCF iterations are only one component of a DFT run (Appendix D.3). This acceleration carries over to the two functionals never seen during training, both of which evaluate exact exchange rather than the semi-local exchange-correlation OrbFlow was trained on. PBE0 exhibits the strongest speedup on both metrics, reducing SCF iterations by 68.13% and wall-clock time by 60.98%, though its SAD baseline converges unusually slowly (43.02 iterations against 21.33 for PBE). B3LYP5 also benefits, with a larger wall-clock reduction than at the training-distribution functional (RTR 83.12% against 89.51%) though a higher relative iteration count (RIC 77.04% against 54.48%). Taken together, these results indicate that OrbFlow’s predicted density is a 11 Equivariant Flow Matching for Electron Density Prediction Table 3: SCF acceleration from OrbFlow initializations, evaluated on QM9 molecules.N SAD andN pred are SCF iterations to convergence from the SAD and OrbFlow initializations, respectively. LowerRIC↓and lower RTR↓both indicate greater acceleration. Iteration counts are means over the evaluated molecules and the ± values are standard deviations across those molecules. RIC and RTR are computed from the means. Functional (Family: Name)N SAD N pred RIC [%]↓ RTR [%]↓ Training-distribution functional GGA: PBE21.33± 1.03 11.62± 1.6554.4889.51 Zero-shot transfer to unseen XC functionals Hybrid GGA: B3LYP521.60± 1.73 16.64± 1.4277.0483.12 Hybrid GGA: PBE043.02± 4.77 13.71± 2.4731.8739.02 genuinely informative initial guess across levels of theory rather than a narrow fit to PBE-specific dynamics. 4.4. Direct Physical Properties from Predicted Densities Beyond accelerating SCF, a predicted density also serves as a DFT surrogate. Density-derived observables follow directly fromρ pred , with no self-consistent calculation. We evaluate the molecular dipole momentμ and the traceless quadrupole tensorQ ab on 500 held-out QM9 molecules, against PBE references computed with VASP. Relative MAE is 1.0% on the dipole magnitude and 1.1% on the quadrupole norm, and 1.7% and 2.5% on their individual components (Table 11, Appendix D.4). The accuracy is uniform across molecules rather than an average over well-predicted and poorly-predicted ones (Figure 3, Appendix D.4). Accuracy on both indicates that OrbFlow captures the density’s shape at multiple spatial scales. One prediction therefore supports two downstream uses, an optimized SCF initial guess (Sec. 4.3) and a stand-alone property surrogate. 4.5. Inference Efficiency Table 12 reports per-molecule inference latency. OrbFlow’s end-to-end cost of0.066s is on par with ELECTRA (0.061s), and1.6×that of the single-shot GTO baseline SCDP, reflecting the two Euler steps of the rollout. The absolute overhead is a few hundredths of a second per molecule, negligible against a single SCF iteration, which typically costs on the order of seconds. Because the network cost scales close to linearly withK, OrbFlow’s latency is predictable and directly controllable through the rollout budget. The full protocol is in Appendix D.5. 4.6. Ablation Studies All ablations use a random 10% subset of QM9, subsampled in both the training and test splits, held fixed across configurations, with three seeds and the protocol of Sec. 4.1. Absolute NMAE is therefore higher than in Table 1, where models train on the full set. We ask three questions, whether the flow formulation contributes beyond the backbone, whether the two-phase curriculum is necessary, and how many Euler steps are needed. What the Flow Formulation Contributes. Our flow formulation differs from direct regression in two ways. It starts from noise drawn from the prior rather than from a fixed input, and it is supervised on the output of its own multi-step rollout ((7)) rather than one timestep at a time ((5)). The four cells in Table 4 run from 12 Equivariant Flow Matching for Electron Density Prediction direct regression, which does neither, to OrbFlow, which does both. Table 4: Ablations as a2× 2over where the input state comes from and what the loss supervises. NMAE (%), mean±std over three seeds. All cells share the backbone, parameter count, basis, real-space loss, and step budget, and all use two network evaluations per prediction except direct regression, which uses one. Input stateSingle-step loss Rollout loss Deterministic (from 0)0.280± 0.0030.268± 0.003 Stochastic (noise prior)0.295± 0.016 0.240± 0.005 Neither ingredient works alone. Direct re- gression reaches0.280 ± 0.003%. Adding noise alone (Phase-1-only training) reaches 0.295± 0.016%, slightly worse. Adding roll- out alone, a deterministic two-step chain ˆ c = g θ (︀ G, g θ (G, 0) )︀ , reaches0.268± 0.003%. Both together reach0.240± 0.005%. Starting from noise exposes the network to a wide range of input states, and rollout supervision makes that exposure count by training on the states inference actually visits, the model’s own in- termediates rather than points interpolated toward the target. All four cells share the backbone, parameter count, and9× 10 4 -step budget, both rollout cells use two network evaluations, and the deterministic chain keeps a learned step-index embedding (Appendix E.1), so neither architecture, compute, nor step conditioning explains the gap. That gap is what few-step methods such as consistency [51] and shortcut [13] models exist to close. Table 5: Curriculum variants on the same subset, with total training compute. The rollout-budget sweep and full compute accounting are in Ap- pendix E.2. ConfigurationNMAE (%)↓ GPU-h Phase 1 only0.295± 0.016 35 Phase 2 only0.250± 0.027 81 Joint single-stage0.241± 0.002 182 Full curriculum (Ours) 0.240± 0.005 74 Two-Phase Training Curriculum. Both ingredients are needed, so the question becomes how to train them. Table 5 compares three routes to the sameK=2 objective. Training directly atK=2with no warm start reaches0.250± 0.027%, and optimizing the single- step and rollout losses jointly from random initializa- tion reaches0.241± 0.002%, within seed variance of our0.240± 0.005%. Accuracy is therefore not what separates the rollout-trained variants. Cost is, 74 GPU- h for the curriculum against 81 without a warm start and 182 for joint optimization, and the curriculum reuses one Phase-1 checkpoint across rollout horizons while the alternatives retrain for each. Budgets here are matched in optimization steps rather than wall-clock, since a Phase-2 step evaluates the network twice. How Many Euler Steps. We sweep the step count twice, over the Phase-2 rollout horizonKused during training (Appendix E.3) and over the integration budget used at inference by the resultingK=2model (Appendix E.4). In training,K=1matches Phase-1-only (0.296% against 0.295%), since under the endpoint parameterization a single step fromt=0lands where the single-step loss already supervises, whileK=2 andK=3tie at a cost that grows withK. We therefore train atK=2. At inference, integrating more finely degrades accuracy, from0.240%at two steps to0.539%at four, because Phase 2 sharpens the network only at the states the two-step rollout visits. The train-matched budget is thus both the most accurate and the cheapest. 5. Conclusion and Limitations We introduced OrbFlow, anSE(3)-equivariant flow matching framework for electron density prediction. OrbFlow reaches state-of-the-art accuracy on QM9 and lowers error on every MD molecule against the strongest prior method sharing its basis, while matching a Gaussian-mixture model55×its size in accuracy 13 Equivariant Flow Matching for Electron Density Prediction and latency. Its predicted densities accelerate SCF across unseen exchange-correlation functionals and recover dipole and quadrupole moments without running SCF. The rollout adds a small inference overhead, and accuracy holds only at the trained horizon, so OrbFlow is best understood as a trainedK-step map. Supervision across a range of horizons is a natural route to discretization-free behavior. Acknowledgments This work was supported in part by the National Science Foundation under Grants IIS-2551110, MOMS- 2331036, and CMMI-2226908; the Advanced Research Projects Agency for Health (ARPA-H) under grant 1AY1AX000053; the Texas A&M University Division of Research Targeted Proposal Teams Funding Program; and the Texas A&M Institute of Data Science Thematic Labs Program. First-principles calculations by C.W. were supported by the U.S. Department of Energy Office of Basic Energy Sciences under grant DE-SC0023866. We thank Jonas Elsborg and Arghya Bhowmik for their assistance in reproducing the results of the ELECTRA model. We are grateful to Shubham Parashar, Lianhao Zhou, and Wendi Yu for insightful discussions. References [1] Carlo Adamo and Vincenzo Barone. Toward reliable density functional methods without adjustable parameters: The PBE0 model. The Journal of Chemical Physics, 110(13):6158–6170, April 1999. ISSN 1089-7690. doi: 10.1063/1.478522. URL http://dx.doi.org/10.1063/1.478522. [2] Richard D Bardo and Klaus Ruedenberg. Even-tempered atomic orbitals. VI. Optimal orbital exponents and optimal contractions of Gaussian primitives for hydrogen, carbon, and oxygen in molecules. The Journal of Chemical Physics, 60(3):918–931, 1974. [3]Axel D. Becke. Density-functional thermochemistry. i. the role of exact exchange. The Journal of Chemical Physics, 98(7):5648–5652, April 1993. ISSN 1089-7690. doi: 10.1063/1.464913. URL http://dx.doi.org/10.1063/1.464913. [4] Samy Bengio, Oriol Vinyals, Navdeep Jaitly, and Noam Shazeer. Scheduled sampling for sequence prediction with recurrent neural networks. Advances in Neural Information Processing Systems, 28, 2015. [5]P. E. Blöchl. Projector augmented-wave method. Phys. Rev. B, 50:17953–17979, Dec 1994. doi: 10.1103/PhysRevB.50.17953. URL https://link.aps.org/doi/10.1103/PhysRevB.50.17953. [6]Mihail Bogojeski, Leslie Vogt-Maranto, Mark E Tuckerman, Klaus-Robert Müller, and Kieron Burke. Quantum chemical accuracy from density functional approximations via machine learning. Nature Communications, 11:5223, 2020. [7]Felix Brockherde, Leslie Vogt, Li Li, Mark E Tuckerman, Kieron Burke, and Klaus-Robert Müller. Bypassing the kohn-sham equations with machine learning. Nature Communications, 8:872, 2017. [8]Chaoran Cheng and Jian Peng. Equivariant neural operator learning with graphon convolution. Advances in Neural Information Processing Systems, 36:61960–61984, 2023. [9]Beatriz G del Rio, Brandon Phan, and Rampi Ramprasad. A deep learning framework to emulate density functional theory. npj Computational Materials, 9:158, 2023. 14 Equivariant Flow Matching for Electron Density Prediction [10]Karin Eichkorn, Oliver Treutler, Holger Öhm, Marco Häser, and Reinhart Ahlrichs. Auxiliary basis sets to approximate coulomb potentials. Chemical Physics Letters, 240(4):283–290, 1995. [11]Jonas Elsborg, Luca Thiede, Alán Aspuru-Guzik, Tejs Vegge, and Arghya Bhowmik. ELECTRA: A Cartesian network for 3D charge density prediction with floating orbitals. Advances in Neural Information Processing Systems, 38:28092–28121, 2025. [12]Pol Febrer, Peter Bjørn Jørgensen, Miguel Pruneda, Alberto Garcia, Pablo Ordejón, and Arghya Bhowmik. Graph2mat: universal graph to matrix conversion for electron density prediction. Machine Learning: Science and Technology, 6(2):025013, 2025. [13]Kevin Frans, Danijar Hafner, Sergey Levine, and Pieter Abbeel. One step diffusion via shortcut models. In International Conference on Learning Representations, volume 2025, pages 34668–34684, 2025. [14]Xiang Fu, Andrew Rosen, Kyle Bystrom, Rui Wang, Albert Musaelian, Boris Kozinsky, Tess Smidt, and Tommi Jaakkola. A recipe for charge density prediction. Advances in Neural Information Processing Systems, 37:9727–9752, 2024. [15]Mario Geiger and Tess Smidt. e3n: Euclidean neural networks. arXiv preprint arXiv:2207.09453, 2022. [16] Sheng Gong, Tian Xie, Taishan Zhu, Shuo Wang, Eric R Fadel, Yawei Li, and Jeffrey C Grossman. Predicting charge density distribution of materials using a local-environment-based graph convolutional network. Physical Review B, 100(18):184103, 2019. [17] Xiaoxun Gong, He Li, Nianlong Zou, Runzhang Xu, Wenhui Duan, and Yong Xu. General framework for E(3)-equivariant neural network representation of density functional theory Hamiltonian. Nature Communications, 14:2848, 2023. [18]Qiangqiang Gu, Zhanghao Zhouyin, Shishir Kumar Pandey, Peng Zhang, Linfeng Zhang, and Weinan E. Deep learning tight-binding approach for large-scale electronic simulations at finite temperatures with ab initio accuracy. Nature Communications, 15:6772, 2024. [19]Suman Hazra, Urvesh Patil, and Stefano Sanvito. Predicting the one-particle density matrix with machine learning. Journal of Chemical Theory and Computation, 20(11):4569–4578, 2024. [20]Pierre Hohenberg and Walter Kohn. Inhomogeneous electron gas. Physical Review, 136(3B):B864, 1964. [21]Anubhav Jain, Yongwoo Shin, and Kristin A Persson. Computational predictions of energy materials using density functional theory. Nature Reviews Materials, 1:15004, 2016. [22]Bowen Jing, Bonnie Berger, and Tommi Jaakkola. Alphafold meets flow matching for generating protein ensembles. In Forty-first International Conference on Machine Learning, 2024. URL https: //openreview.net/forum?id=rs8Sh2UASt. [23]Peter Bjørn Jørgensen and Arghya Bhowmik. Equivariant graph neural networks for fast electron density estimation of molecules, liquids, and solids. npj Computational Materials, 8:183, 2022. [24]Nayoung Kim, Seongsu Kim, Minsu Kim, Jinkyoo Park, and Sungsoo Ahn. MOFFlow: Flow matching for structure prediction of metal-organic frameworks. In The Thirteenth International Conference on Learning Representations, 2025. URL https://openreview.net/forum?id=dNT3abOsLo. 15 Equivariant Flow Matching for Electron Density Prediction [25]Seongsu Kim and Sungsoo Ahn. Gaussian plane-wave neural operator for electron density estimation. In Forty-first International Conference on Machine Learning, 2024. URL https://openreview.net/forum? id=TvoG41N1Y3. [26]Seongsu Kim, Nayoung Kim, Dongwoo Kim, and Sungsoo Ahn. High-order Equivariant Flow Matching for Density Functional Theory Hamiltonian Prediction. In The Thirty-ninth Annual Conference on Neural Information Processing Systems, 2025. URL https://openreview.net/forum?id=iFIjNXb0Y5. [27]Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.6980, 2014. [28] James Kirkpatrick, Brendan McMorrow, David HP Turban, Alexander L Gaunt, James S Spencer, Alexander GDG Matthews, Annette Obika, Louis Thiry, Meire Fortunato, David Pfau, et al. Pushing the frontiers of density functionals by solving the fractional electron problem. Science, 374(6573): 1385–1389, 2021. [29] Walter Kohn and Lu Jeu Sham. Self-consistent equations including exchange and correlation effects. Physical Review, 140(4A):A1133, 1965. [30]Teddy Koker, Keegan Quigley, Eric Taw, Kevin Tibbetts, and Lin Li. Higher-order equivariant neural networks for charge density prediction in materials. npj Computational Materials, 10:161, 2024. [31]G. Kresse and J. Furthmüller. Efficient iterative schemes forab initiototal-energy calculations using a plane-wave basis set. Physical Review B, 54:11169–11186, 1996. [32] Chengteh Lee, Weitao Yang, and Robert G. Parr. Development of the colle-salvetti correlation-energy formula into a functional of the electron density. Phys. Rev. B, 37:785–789, Jan 1988. doi: 10.1103/ PhysRevB.37.785. URL https://link.aps.org/doi/10.1103/PhysRevB.37.785. [33]He Li, Zun Wang, Nianlong Zou, Meng Ye, Runzhang Xu, Xiaoxun Gong, Wenhui Duan, and Yong Xu. Deep-learning density functional theory Hamiltonian for efficient ab initio electronic-structure calculation. Nature Computational Science, 2(6):367–377, 2022. [34]Yunyang Li, Zaishuo Xia, Lin Huang, Xinran Wei, Samuel Harshe, Han Yang, Erpai Luo, Zun Wang, Jia Zhang, Chang Liu, Bin Shao, and Mark Gerstein. Enhancing the Scalability and Applicability of Kohn-Sham Hamiltonians for Molecular Systems. In The Thirteenth International Conference on Learning Representations, 2025. URL https://openreview.net/forum?id=twEvvkQqPS. [35]Yi-Lun Liao, Alexander J Hoffman, Sabrina C Shen, Alexandre Duval, Sam Walton Norwood, and Tess Smidt. EquiformerV3: Scaling efficient, expressive, and general SE(3)-equivariant graph attention transformers. arXiv preprint arXiv:2604.09130, 2026. [36]Yaron Lipman, Ricky T. Q. Chen, Heli Ben-Hamu, Maximilian Nickel, and Matthew Le. Flow matching for generative modeling. In The Eleventh International Conference on Learning Representations, 2023. URL https://openreview.net/forum?id=PqvMRDCJT9t. [37] Zhe Liu, Yuyan Ni, Zhichen Pu, Qiming Sun, Siyuan Liu, and Wen Yan. Towards a transferable acceleration method for density functional theory. In The Fourteenth International Conference on Learning Representations, 2026. URL https://openreview.net/forum?id=JNuk3yGDKE. 16 Equivariant Flow Matching for Electron Density Prediction [38]Erpai Luo, Xinran Wei, Lin Huang, Yunyang Li, Han Yang, Zaishuo Xia, Zun Wang, Chang Liu, Bin Shao, and Jia Zhang. Efficient and Scalable Density Functional Theory Hamiltonian Prediction through Adaptive Sparsity. In Forty-second International Conference on Machine Learning, 2025. URL https: //openreview.net/forum?id=K3lykWhXON. [39]Nicola Marzari, Andrea Ferretti, and Chris Wolverton. Electronic-structure methods for materials design. Nature Materials, 20(6):736–749, 2021. [40]Benjamin Kurt Miller, Ricky T. Q. Chen, Anuroop Sriram, and Brandon M Wood. FlowMM: Generating Materials with Riemannian Flow Matching. In Forty-first International Conference on Machine Learning, 2024. URL https://openreview.net/forum?id=W4pB7VbzZI. [41]Joachim Paier, Robin Hirschl, Martijn Marsman, and Georg Kresse. The Perdew–Burke–Ernzerhof exchange-correlation functional applied to the G2-1 test set using a plane-wave basis set. The Journal of Chemical Physics, 122(23), 2005. ISSN 1089-7690. doi: 10.1063/1.1926272. URL http://dx.doi. org/10.1063/1.1926272. [42]Saro Passaro and C Lawrence Zitnick. Reducing SO(3) convolutions to SO(2) for efficient equivariant GNNs. In International Conference on Machine Learning, pages 27420–27438. PMLR, 2023. [43]John P. Perdew, Kieron Burke, and Matthias Ernzerhof. Generalized gradient approximation made simple. Phys. Rev. Lett., 77:3865–3868, Oct 1996. doi: 10.1103/PhysRevLett.77.3865. URL https: //link.aps.org/doi/10.1103/PhysRevLett.77.3865. [44] Zhuoran Qiao, Anders S Christensen, Matthew Welborn, Frederick R Manby, Anima Anandkumar, and Thomas F Miller I. Informing geometric deep learning with electronic interactions to accelerate quantum chemistry. Proceedings of the National Academy of Sciences, 119(31):e2205221119, 2022. [45]Raghunathan Ramakrishnan, Pavlo O Dral, Matthias Rupp, and O Anatole Von Lilienfeld. Quantum chemistry structures and properties of 134 kilo molecules. Scientific Data, 1:140022, 2014. [46] Marc’Aurelio Ranzato, Sumit Chopra, Michael Auli, and Wojciech Zaremba. Sequence level training with recurrent neural networks. arXiv preprint arXiv:1511.06732, 2015. [47] Dmitrij Rappoport and Filipp Furche. Property-optimized Gaussian basis sets for molecular response calculations. The Journal of Chemical Physics, 133(13), 2010. [48]Lars Ruddigkeit, Ruud Van Deursen, Lorenz C Blum, and Jean-Louis Reymond. Enumeration of 166 billion organic small molecules in the chemical universe database GDB-17. Journal of Chemical Information and Modeling, 52(11):2864–2875, 2012. [49] Tim Salimans and Jonathan Ho. Progressive distillation for fast sampling of diffusion models. In International Conference on Learning Representations, 2022. URL https://openreview.net/forum?id= TIdIXIpzhoI. [50] Xuecheng Shao, Lukas Paetow, Mark E Tuckerman, and Michele Pavanello. Machine learning electronic structure methods based on the one-electron reduced density matrix. Nature Communications, 14: 6281, 2023. [51]Yang Song, Prafulla Dhariwal, Mark Chen, and Ilya Sutskever. Consistency models. In International Conference on Machine Learning, pages 32211–32252. PMLR, 2023. 17 Equivariant Flow Matching for Electron Density Prediction [52]Hannes Stark, Bowen Jing, Regina Barzilay, and Tommi Jaakkola. Harmonic self-conditioned flow matching for joint multi-ligand docking and binding site design. In Forty-first International Conference on Machine Learning, 2024. [53]Zechen Tang, He Li, Peize Lin, Xiaoxun Gong, Gan Jin, Lixin He, Hong Jiang, Xinguo Ren, Wenhui Duan, and Yong Xu. A deep equivariant neural network approach for efficient hybrid density functional calculations. Nature Communications, 15:8815, 2024. [54]Alexander Tong, Kilian FATRAS, Nikolay Malkin, Guillaume Huguet, Yanlei Zhang, Jarrid Rector- Brooks, Guy Wolf, and Yoshua Bengio. Improving and generalizing flow-based generative models with minibatch optimal transport. Transactions on Machine Learning Research, 2024. ISSN 2835-8856. URL https://openreview.net/forum?id=CD9Snc73AW. [55]Haiyang Yu, Zhao Xu, Xiaofeng Qian, Xiaoning Qian, and Shuiwang Ji. Efficient and Equivariant Graph Networks for Predicting Quantum Hamiltonian. In International Conference on Machine Learning, pages 40412–40424. PMLR, 2023. [56]Haiyang Yu, Yuchao Lin, Xuan Zhang, Xiaofeng Qian, and Shuiwang Ji. Efficient prediction of SO(3)- equivariant hamiltonian matrices via SO(2) local frames. In Forty-third International Conference on Machine Learning, 2026. URL https://openreview.net/forum?id=f9i6VLHT5v. [57] Xuan Zhang, Limei Wang, Jacob Helwig, Youzhi Luo, Cong Fu, Yaochen Xie, ..., and Shuiwang Ji. Artificial intelligence for science in quantum, atomistic, and continuum systems. Foundations and Trends® in Machine Learning, 18(4):385–912, 2025. [58]Yang Zhong, Hongyu Yu, Mao Su, Xingao Gong, and Hongjun Xiang. Transferable equivariant graph neural networks for the Hamiltonians of molecules and solids. npj Computational Materials, 9:182, 2023. [59] Zhanghao Zhouyin, Zixi Gan, Shishir Kumar Pandey, Linfeng Zhang, and Qiangqiang Gu. Learning local equivariant representations for quantum operators. In The Thirteenth International Conference on Learning Representations, 2025. URL https://openreview.net/forum?id=kpq3IIjUD3. 18 Equivariant Flow Matching for Electron Density Prediction A. GTO Basis Set Details and Coefficient Projection To train our machine learning model to predict the continuous real-space electron densityρ(r), we project the continuous ground-truth density obtained from density functional theory (DFT) onto a localized, auxiliary basis set of Gaussian-Type Orbitals (GTOs). This section details the mathematical formulation of this basis set, the even-tempered scheme used to generate the radial exponents, and the regularized least-squares projection framework used to extract the target coefficient vectors. A.1. Even-Tempered Gaussian-Type Orbital (GTO) Basis Set For an atomilocated at coordinateR i , the real-space densityρ(r)is expanded using a linear combination of atom-centered GTO basis functionsφ μ (r− R i ). The basis functions are decomposed into radial and angular components, φ n,ℓ,m (r− R i ) =N n,ℓ |r− R i | ℓ exp (︀ −α n,ℓ |r− R i | 2 )︀ Y m ℓ (︂ r− R i |r− R i | )︂ (10) whereY m ℓ represents the real spherical harmonics for angular momentumℓand magnetic quantum number m∈ [−ℓ, ℓ], andN n,ℓ is the normalization constant given by N n,ℓ = √︃ 2 2ℓ+7/2 α ℓ+3/2 n,ℓ (2ℓ + 1)!! √ π (11) where (2ℓ + 1)!! denotes the double factorial of 2ℓ + 1. To systematically cover the radial coordinate space, the exponentsα n,ℓ are parameterized using an even- tempered geometric progression. For each angular momentum channelℓup toℓ max , we constructN ℓ radial shells. The exponents are determined by a base exponent α 0,ℓ and a scaling factor β ℓ , α n,ℓ = α 0,ℓ · β n−1 ℓ , for n = 1, 2, . . . , N ℓ (12) The hyperparametersα 0,ℓ , β ℓ , N ℓ are optimized per element to span both the sharp, core-like densities near the nuclei (large α) and the diffuse, bonding valence densities (small α). A.2. Element-Specific Hyperparameter Configurations For each chemical element present in our dataset, Table 6 reports the maximum angular momentum ℓ max , the number of radial shellsN ℓ , the boundaries of the even-tempered progression (α min = α 0,ℓ and α max = α 0,ℓ β N ℓ −1 ℓ ), and the geometric progression scaling factorβ ℓ , all instantiated from an uncontracted def2-QZVPPDparent basis with a shared even-tempered scaleβ = 1.3. QM9 uses the full table across H, C, N, O, F. MD, whose six trajectories span onlyH, C, O(Sec. 4.2), uses the corresponding subset. Bond- midpoint virtual nodes (Sec. 3.1) are assigned the oxygen basis. A.3. L2 Least-Squares Coefficient Projection Mechanics Letφ p (r)represent a shorthand for the set of all active basis functions centered on all atoms (and virtual bond-midpoint nodes) in the system, indexed by a single flattened indexp = (i, n, ℓ, m). Incorporating the scale factors > 0defined in the forward continuous density synthesis operatorO(c) = sΦ c, the continuous electron density ρ(r) is approximated by the expansion ρ(r)≈ ̃ ρ(r) = s ∑︁ p c p φ p (r)(13) 19 Equivariant Flow Matching for Electron Density Prediction Table 6: Even-tempered GTO parameters used for density projection (def2-QZVPPDparent basis, global β = 1.3). α max = α min β N ℓ −1 . MD uses the H/C/O rows. QM9 uses the full table. Virtual nodes use O. Element Channel N ℓ α min α max β Hs (ℓ = 0) 23 0.1606 51.59 1.3 p (ℓ = 1) 18 0.1643 14.22 1.3 d (ℓ = 2) 13 0.1681 3.917 1.3 Cs (ℓ = 0) 42 0.08996 4224 1.3 p (ℓ = 1) 32 0.1169 398.2 1.3 d (ℓ = 2) 27 0.1176 107.9 1.3 f (ℓ = 3) 18 0.1529 13.22 1.3 g (ℓ = 4) 13 0.1538 3.583 1.3 Ns (ℓ = 0) 42 0.1243 5837 1.3 p (ℓ = 1) 32 0.1630 555.2 1.3 d (ℓ = 2) 27 0.1641 150.5 1.3 f (ℓ = 3) 18 0.2151 18.61 1.3 g (ℓ = 4) 13 0.2165 5.044 1.3 Os (ℓ = 0) 42 0.1393 6540 1.3 p (ℓ = 1) 37 0.1102 1394 1.3 d (ℓ = 2) 32 0.08720 297.0 1.3 f (ℓ = 3) 23 0.1479 47.51 1.3 g (ℓ = 4) 14 0.2509 7.600 1.3 Fs (ℓ = 0) 42 0.1753 8229 1.3 p (ℓ = 1) 37 0.1280 1619 1.3 d (ℓ = 2) 33 0.09354 414.2 1.3 f (ℓ = 3) 23 0.1728 55.51 1.3 g (ℓ = 4) 14 0.3194 9.673 1.3 20 Equivariant Flow Matching for Electron Density Prediction To obtain the optimal target coefficientsc p corresponding to the ground-truth DFT densityρ DFT (r), we minimize the L 2 error over the continuous space, E(c) = ∫︁ R 3 ⎡ ⎣ ρ DFT (r)− s ∑︁ p c p φ p (r) ⎤ ⎦ 2 dr(14) Setting the partial derivatives with respect to each coefficient to zero, ∂E ∂c p = 0 , yields a system of linear equations, the normal equations, ∑︁ q S pq c q = y p (15) where S pq is the analytic overlap matrix between GTO basis functions, S pq = ∫︁ R 3 φ p (r)φ q (r)dr(16) andy p represents the normalized projection of the true DFT density onto the basis functions, factoring out the scale factor s, y p = 1 s ∫︁ R 3 ρ DFT (r)φ p (r)dr(17) In practice, the ground-truth DFT reference densityρ lab is represented on a discrete numerical grid(r k , w k ), wherer k is the coordinate of grid pointkandw k is the corresponding integration weight. Incorporating the diagonal spatial mask matrixW = diag(w 1 , w 2 , . . . , w P ), the discrete least-squares objective formulated on the grid is expressed as c ⋆ = arg min c ⃦ ⃦ W 1/2 (︀ Φc− y )︀ ⃦ ⃦ 2 2 + λ∥c∥ 2 2 + γ c ⊤ diag(Φ ⊤ Φ) c(18) wherey =ρ lab /s,Φ kj = φ j (r k )is the forward GTO evaluation design matrix, andλ = 10 −6 is a Tikhonov regularizer to prevent numerical instabilities caused by basis near-linear dependencies. The auxiliary parameter γ regularizes the coefficients scaling with respect to the continuous basis functions’ self-overlap. Solving this regularized objective directly yields the target coefficients c ⋆ , c ⋆ = (︁ Φ ⊤ WΦ + λI + γ diag(Φ ⊤ Φ) )︁ −1 Φ ⊤ Wy(19) During preprocessing we setλ = 10 −6 andγ = 0(isotropic Tikhonov only). This is numerically stable while leaving a small, well-characterized representation residual. A.4. Empirical Fitting Floor Because OrbFlow predicts GTO coefficients rather than grid densities, learned NMAE cannot fall below the reconstruction error ofc ⋆ . We measure this floor by synthesizing ̃ ρ = sΦ c ⋆ from the stored ridge-fitted coefficients, with no neural network in the loop, and scoring NMAE on the same full probe grid used to evaluate the trained models. On QM9 the fitting residual is0.061%while OrbFlow reaches0.152%, so the remaining gap is due to coefficient prediction rather than incompleteness of the GTO span in Table 6. Relative to SCDP + VO, which uses 21 Equivariant Flow Matching for Electron Density Prediction Table 7: GTO density fitting floor versus OrbFlow NMAE (%). The floor is the NMAE ofρreconstructed from stored ridge ground-truth coefficients with no network. OrbFlow is mean±std over three training seeds; the MD macro mean is unweighted over the six molecules. Datasetn GT floor↓ OrbFlow↓ QM910,0000.0610.152± 0.001 MD benzene5000.3990.421± 0.002 MD ethanol1,0000.8501.078± 0.023 MD phenol1,2000.4690.575± 0.008 MD resorcinol1,0000.5150.663± 0.014 MD ethane2010.8070.871± 0.008 MD malonaldehyde4000.5481.177± 0.012 MD macro mean60.5980.798 the same basis construction and reaches0.178%, OrbFlow removes22%of the error remaining above this floor. On MD the gap varies by molecule: benzene and ethane lie close to the basis limit, whereas ethanol and malonaldehyde remain well above it, so compact rigid molecules are already well represented by this expansion while flexible oxygenates leave a larger residual for the model to close. A.5. Electron Count Conservation Because a spherical Gaussian of angular momentumℓintegrates to zero overR 3 for everyℓ > 0, only the ℓ=0 channel of (1) contributes to the total electron count, N e [c] = ∫︁ R 3 ρ(r) dr = ∑︁ i,j: ℓ i,j =0 c i,j,0 ∫︁ R 3 φ α i,j ,0,0,R i (r) dr,(20) so charge conservation is one linear condition on theℓ=0coefficients rather than a property of the whole field. We do not impose it during training, because a hard constraint would couple the equivariant readout across atoms. When the predicted density is used as an SCF initial guess we renormalize it to the PAW valence countN val e . On the official QM9 test set the closed-form count in(20)for OrbFlow differs fromN val e by0.014 eon average, or0.030%ofN val e . The mean valence count on this set is48.28 e, so the absolute miss is a few hundredths of an electron. The learned density already carries nearly the right charge, and the renormalization used for an SCF initial guess only has to correct a 0.030% discrepancy. B. Proof of Prior Rotational Invariance This section provides the full proof of the proposition stated in Sec. 3.2, that the irrep-isotropic priorp 0 (c 0 ) is invariant under any rotation R∈ SO(3). Proof.Because Wigner D-matricesD (ℓ) (R)are unitary matrices (satisfyingD (ℓ) (R) ⊤ D (ℓ) (R) = I), the quadratic form in the Gaussian exponent of each subspace ℓ remains invariant under rotation, ∥D(R)c (ℓ) 0 ∥ 2 = (︀ D (ℓ) (R)c (ℓ) 0 )︀ ⊤ (︀ D (ℓ) (R)c (ℓ) 0 )︀ = c (ℓ)⊤ 0 D (ℓ) (R) ⊤ D (ℓ) (R)c (ℓ) 0 =∥c (ℓ) 0 ∥ 2 .(21) Thus, the density function is unchanged under any arbitrary 3D rotation, satisfying globalSO(3)-invariance. 22 Equivariant Flow Matching for Electron Density Prediction C. Implementation Details This appendix consolidates the model architecture, target construction, and optimization procedure that support Sec. 3.3. These settings are shared across QM9 and MD except where noted. Per-dataset hyperpa- rameters and training configuration are given separately in Appendix D.1. C.1. Network Architecture and Coefficient Parameterization Graph construction. Each molecule is represented as a graphGwhose nodes comprise nuclear coordinates together with bond-midpoint virtual orbitals (vnode_elem=8). Node features are initialized asℓ=0 spherical-tensor embeddings of atomic numbers, with higher-order channels set to zero. Edges are formed within a cutoff radiusr cut =6 Åand parameterized by Gaussian radial bases (300functions, resolution0.02 Å) together with SO(3) edge frames used by equivariant attention. Equivariant backbone. OrbFlow uses an EquiformerV3 encoder. QM9 employsL=8blocks withℓ max =6, m max =2, and sphere/hidden/edge channel widths128/256/128. Each MD molecule uses a compact variant withL=3,ℓ max =3,m max =2, and widths96/192/96(∼13M parameters). In both settings, attention uses8 heads with attention-hidden width 64. No backbone dropout is applied. Flow conditioning and readout. The scalar flow timetis injected into theℓ=0channel via a sinusoidal embedding. The continuous coefficient statex t is introduced at each layer through an equivariant linear map, using EquiformerV3’s SO(2)-convolution-based (eSCN-style) equivariant readout for coefficient en- tanglement [42], with no separate velocity MLP. This readout mechanism differs from using eSCN as a full alternative backbone. Latent tensors are mapped to per-atom GTO coefficients with multiplicities matched to the underlying basis. Coefficients are then scaled by1/ √ N orb prior to density synthesis so that the training target remains consistent with the normalized coefficient representation. Basis and density synthesis. We use an uncontracteddef2-QZVPPDGTO basis with weighted even- tempered augmentation at ETB scaleβ=1.3and orbital cutoff5 Å. Densities are synthesized in real space under the linear combination ρ(r) = ∑︀ i c i φ i (r). C.2. Flow-Matching Objective Training follows conditional flow matching (CFM) with an endpoint parameterization. The network predicts the terminal coefficient field, and the CFM velocity is implied by the linear bridge. The initial state is sampled from an irrep-structured Gaussian in normalized coefficient space,c norm 0 = σε + μ, whereεis drawn via e3n irrep noise and the scalar centerμis added only to theℓ=0(scalar) coefficient dimension, following the standard convention for additive terms in equivariant architectures, since a nonzero mean is not rotation- invariant for anyℓ>0irrep. Allℓ>0channels remain zero-mean, consistent with the irrep-isotropic prior described in Sec. 3.2 and the invariance result of Proposition 1. On QM9 we use(μ, σ) = (2.5× 10 −3 , 0.85) for the scalar channel. On MD,(μ, σ)is estimated once as the empirical mean and standard deviation of the scalar-channel training ground-truth coefficients for that molecule (Table 8) and held fixed. Electron-count renormalization of the prior is disabled. C.3. Optimization Protocol All production runs use Adam [27] with(β 1 , β 2 ) = (0.9, 0.999),ε=10 −8 , base learning rate1× 10 −3 , and zero weight decay. The learning rate follows a linear warm-up for the first1,000steps (from a factor of0.2to 23 Equivariant Flow Matching for Electron Density Prediction Table 8: MD irrep-Gaussian prior parameters, estimated from each molecule’s training-set ground-truth coefficients and held fixed during training. Moleculeμσ Benzene8.67× 10 −3 1.82 Ethanol5.64× 10 −3 3.04 Phenol5.04× 10 −3 1.98 Resorcinol4.72× 10 −3 2.02 Ethane8.44× 10 −3 3.66 Malonaldehyde1.79× 10 −2 2.79 1.0), then a power decay η t = η 0 α t/β ,α=0.96,β=4433. We maintain an exponential moving average (EMA) of parameters with decay0.995, apply global-norm gradient clipping at0.5, and train in FP32 under PyTorch Lightning DDP. Checkpoints are selected by the lowest validation density error after Euler integration (K=2). All production jobs were run on NVIDIA H200 GPUs. D. Experiment Details This appendix collects the protocols for the four supporting procedures referenced in the main text, the per-dataset training configuration and hyperparameters underlying every reported result, the DFT solver settings used for the SCF acceleration experiments (Sec. 4.3), the inference latency measurement protocol underlying the computational efficiency comparison (Sec. 4.5), and the multipole moment definitions used to evaluate direct physical property prediction (Sec. 4.4). D.1. Training Configuration Table 9 summarizes the primary architecture, coefficient representation, flow-matching, and optimization hyperparameters for both datasets. The training-time sampling ranget ∼ Unif[0.01, 0.99]excludes both endpoints of the bridge to avoid numerical instability during training, since the endpoint parameterization in Eq. 4 is singular att=1where the input state degenerates tox t = c ⋆ , while att=0the input reduces to the pure noise prior and carries no target information. QM9. We train on the standard DeepDFT/SCDP partition (123,835/50/10,000train/val/test) and report both the full test set and Split 2 (last1,600test molecules) as in the main text. Training usesP ′ =60,000 randomly subsampled probe points per molecule (120,000for validation/test), per-GPU batch size4on8 GPUs (effective batch size32), Phase-1 lengthT 1 =7.5× 10 4 steps, and Phase-2 continuation to a total of T=4.5× 10 5 steps. MD. We train a separate model for each of the six trajectories (benzene, ethanol, phenol, resorcinol, ethane, malonaldehyde), using the official per-molecule splits in Table 10. Training, validation, and test all use P ′ =125,000probes per configuration, per-GPU batch size4on4GPUs (effective batch size16), Phase-1 length T 1 =10 4 steps, and Phase-2 continuation to a total of T=1.1× 10 5 steps. 24 Equivariant Flow Matching for Electron Density Prediction Table 9: Primary architecture and optimization hyperparameters. Shared entries are repeated for readability. HyperparameterQM9MD (per molecule) EquiformerV3 backbone # blocks L83 Sphere / hidden / edge channels128 / 256 / 12896 / 192 / 96 Attention hidden / # heads64 / 864 / 8 ℓ max (m max )6 (2)3 (2) Cutoff r cut / # radial bases6 Å / 3006 Å / 300 Approx. # parameters∼99M∼13M Coefficient / density representation GTO basis / ETB β def2-QZVPPD / 1.3 def2-QZVPPD / 1.3 GT projection ridge10 −6 10 −6 Virtual orbitalsbond midpointsbond midpoints Probe subsample P ′ (train / val)6×10 4 / 1.2×10 5 1.25×10 5 / 1.25×10 5 Flow matching Bridge / loss parameterizationCFM / endpointCFM / endpoint Priorirrep-Gaussianirrep-Gaussian Prior (μ, σ)(2.5×10 −3 , 0.85)Table 8 Train t range / integrate grid[0.01, 0.99] / uniform[0.01, 0.99] / uniform Euler steps K (train / eval)2 / 22 / 2 Optimization Optimizer / base LR / weight decayAdam / 10 −3 / 0Adam / 10 −3 / 0 (β 1 , β 2 ) / warm-up steps(0.9, 0.999) / 10 3 (0.9, 0.999) / 10 3 LR decaypower (α=0.96, β=4433) power (α=0.96, β=4433) EMA decay / grad clip / precision0.995 / 0.5 / FP320.995 / 0.5 / FP32 Per-GPU batch / # GPUs (eff. batch)4 / 8 (32)4 / 4 (16) Phase 1→ total steps7.5×10 4 → 4.5×10 5 10 4 → 1.1×10 5 Random seeds0, 42, 1230, 42, 505 Table 10: MD trajectory split sizes used in our experiments. MoleculeTrain Validation Test Benzene1001500500 Ethanol100210001000 Phenol100112001200 Resorcinol100410001000 Ethane2000201201 Malonaldehyde 2000400400 25 Equivariant Flow Matching for Electron Density Prediction D.2. DFT Solver Settings This section provides full plane-wave DFT solver settings for the SCF acceleration experiments in Sec. 4.3. For unseen XC functionals, the evaluation is performed on a fixed test set of 200 molecules. For training- distribution PBE, we expand the test set to 500 molecules because of the lower computational cost. DFT Calculations and Exchange-Correlation Functionals. All calculations use VASP [31] with the projector- augmented wave (PAW) [5] method and the PBE exchange-correlation energy functional[43] (Sec. 4.3 additionally evaluates B3LYP5[3,32] and PBE0 [1,41] for transferability). The kinetic energy for plane-wave basis set is fixed at400eV, and the reciprocal-space sampling is restricted to theΓ-point only, consistent with the isolated-molecule (non-periodic) nature of the QM9 evaluation set. The blocked-Davidson algorithm was used for all SCF calculations except those employing B3LYP5, for which a conjugate gradient-based all-band simultaneous optimization scheme was adopted. This choice was made because blocked-Davidson algorithm exhibited poor SCF convergence for B3LYP5, with pronounced oscillations during electronic iterations that hindered convergence. The all-band optimization scheme provided more robust and stable convergence in these cases. Convergence Criterion. SCF convergence is reached when the change in total energy between consecutive iterations falls below 10 −5 eV, matching the convergence criterion used to report RIC and RTR in Sec. 4.3. Cell Construction. Molecules are placed in an orthorhombic simulation cell with a minimum of 2 Å of vacuum padding on each side of the cell to prevent spurious periodic self-interaction, corresponding to a minimum distance of 4 Å between the atoms from periodic images, same as the original charge density data set [23]. The real-space FFT grid resolution is fixed at approximately 0.15 Å per voxel along each cartesian direction across all systems, ensuring consistent grid density independent of molecule size. Initialization Protocol. For each test molecule, we run two matched SCF calculations that differ only in initial charge density, an atomic superposition baseline (SAD initialization) and initialization from the OrbFlow-predicted density, renormalized to conserve total valence electron count and written directly as the solver’s input charge density file. All other solver settings, includingk-point sampling, energy cutoff, and convergence threshold, are held identical between the two runs to ensure that any difference in required SCF iterations is attributable solely to the quality of the initial density guess. Reported Metrics. For each molecule, we record the number of SCF iterations required to reach convergence from each initialization and compute the Relative Iteration CountRIC = N pred /N SAD × 100%. The Relative Time Ratio additionally accounts for OrbFlow’s own inference latency, ensuring that reported acceleration reflects genuine end-to-end speedup rather than SCF iteration count alone. D.3. Why Iteration Savings Only Partially Translate to Wall-Clock Savings A DFT run costs more than its SCF iterations. Basis and PAW projector construction, FFT grid setup, symmetry analysis, and file I/O are incurred once per calculation and are unaffected by the quality of the initial density, so a reduction in iteration count can never reduce total runtime proportionally. Writing the total time of a run as a fixed overhead plus a per-iteration cost,T = T fix + N t iter , and lettingf = T fix /T SAD denote the non-iterative share of a baseline run, the two reported metrics are related by RTR = f + (1− f ) RIC.(22) Solving(22)forfat each functional in Table 3 givesf = 77.0%for PBE,26.5%for B3LYP5, and10.5%for PBE0. Substituting these back reproduces all three reported RTR values exactly. PBE iterations on molecules 26 Equivariant Flow Matching for Electron Density Prediction of this size are cheap, so fixed setup dominates and even a 46% reduction in iterations buys only about 10% of wall-clock time. Hybrid functionals evaluate exact exchange at every iteration, which makes the iterative part far more expensive, shrinksf, and brings RTR close to RIC. SinceRTR→ RICasf → 0, the wall-clock benefit of a learned initial density grows with the cost of the underlying method, and the PBE row understates what the same iteration savings are worth at higher levels of theory or on larger systems. Note that T pred includes OrbFlow’s own inference latency, so these values of f are mild overestimates. D.4. Multipole Moments This section gives the full definitions of the dipole and quadrupole moments evaluated in Sec. 4.4, and supplements the aggregate accuracy reported there (Table 11) with a per-molecule view. Because the PAW reference densities contain only valence electrons, we pair them with the valence (pseudo- ion) chargesZ val A , so that the total system is charge neutral. The molecular electric dipole moment is then obtained via μ = ∑︁ A Z val A R A − ∫︁ rρ(r) dr,(23) and the traceless electric quadrupole moment tensor is obtained analogously via Q ab = ∑︁ A Z val A (︀ 3R A,a R A,b −|R A | 2 δ ab )︀ − ∫︁ (︀ 3r a r b −|r| 2 δ ab )︀ ρ(r) dr,(24) wherea, b∈x, y, z. BothμandQ ab are therefore fully determined byρonce the nuclear geometry is fixed. Table 11: Dipole and quadrupole MAE (OrbFlow vs. DFT) on held-out QM9. Relative MAE is normalized by the mean DFT-reference magnitude on the same test set. QuantityMAE Rel. MAE [%] Dipole, μ x , μ y , μ z 0.026 D1.7 Dipole magnitude,|μ|0.031 D1.0 Quadrupole, Q ab 0.044 B2.5 Quadrupole norm,∥Q∥ F 0.066 B1.1 Figure 3 shows the per-molecule prediction accuracy for both moments across the held-out QM9 test set, confirming that the aggregate accuracy in Table 11 is uniform across molecules rather than driven by a favorable average over a subset of poorly-predicted outliers. D.5. Efficiency Measurement Protocol This section details the measurement protocol underlying the inference latency comparison in Sec. 4.5. Protocol. All timings are measured on a single NVIDIA H200 GPU with batch size 1. We report the mean over the six MD test molecules (Sec. 4.2) and over three independent timing repeats of the same checkpoints, where each repeat covers the full test split after a short CUDA warm-up of three graphs. The network cost is the wall-clock time to produce final GTO coefficients from the molecular graph, which for OrbFlow includes the fullK-step Euler rollout in(6). The density synthesis cost is the cost of the operatorO(·)in(2)mapping coefficients to the real-space grid. OrbFlow evaluates this operator with a fused Triton reduction of the same GTO contraction used by SCDP. 27 Equivariant Flow Matching for Electron Density Prediction 0246810 Predicted || (D) 0 2 4 6 8 10 DFT | | (D) MAE = 0.0311 Dipole magnitude 64202468 Predicted component (D) 6 4 2 0 2 4 6 8 DFT component (D) MAE = 0.0264 Dipole components Component ideal x y z 10.07.55.02.50.02.55.07.5 Predicted Q component (e·Å²) 10.0 7.5 5.0 2.5 0.0 2.5 5.0 7.5 DFT Q component (e·Å²) MAE = 0.0440 Traceless quadrupole components Component ideal x y z xy xz yz Predicted versus DFT molecular multipoles Figure 3: Predicted (OrbFlow) vs. DFT-reference molecular multipoles on held-out QM9 (no SCF), computed via(23)and(24). Left: dipole magnitude|μ|. Center: individual dipole componentsμ x , μ y , μ z . Right: individual traceless quadrupole componentsQ ab . The dashed line isy = x. Points close to this line indicate accurate direct property prediction, and the reported MAE in each panel matches the corresponding row of Table 11. Table 12: Per-molecule inference latency on the MD benchmark, averaged across the six molecules and three independent H200 timing repeats (mean± std). ModelNetwork Cost [s] Density Synthesis Cost [s]Total [s] ELECTRA0.058± 0.0030.0025± 0.00020.061± 0.003 SCDP0.0265± 0.00040.0150± 0.00010.0414± 0.0005 OrbFlow (ours)0.054± 0.0010.0126± 0.00010.066± 0.001 28 Equivariant Flow Matching for Electron Density Prediction Scaling withK. Two patterns in Table 12 are worth expanding on. First, OrbFlow’s network cost scales close to linearly withK. Its0.054s cost at two integration steps is roughly twice SCDP’s single-shot cost of0.026s, consistent with a per-step cost close to a single regression forward pass, around0.027s per step. OrbFlow’s inference cost is therefore predictable and directly controllable through the choice ofK(Sec. 4.6) rather than compounding unpredictably. Second, the density synthesis cost differs substantially between ELECTRA (0.003s) and the two GTO-based models, SCDP (0.015s) and OrbFlow (0.013s). This gap reflects the underlying density representation, a 3D Gaussian mixture evaluated with a fused KeOps reduction for ELECTRA against the auxiliary GTO basis expansion shared by SCDP and OrbFlow, rather than the flow matching formulation itself. E. Extended Ablation Results This appendix expands on the ablations summarized in Sec. 4.6, covering the deterministic two-step baseline, training compute cost, the training-time rollout horizon, and the inference-time integration budget. E.1. Deterministic Two-Step Baseline The deterministic baseline of Sec. 4.6 uses the same EquiformerV3 configuration as the QM9 models (Table 9), including the equivariant linear map that injects the coefficient state at each layer, and is trained and evaluated on the same 10% QM9 ablation subset as the other ablation rows. The first evaluation receives an all-zero coefficient state and the second receives the output of the first. No noise is sampled, and the continuous time embedding is replaced by a learned two-valued step-index embedding, so that the baseline retains the ability to behave differently on its first and second evaluation. This isolates the stochastic prior and the bridge interpolation, the components specific to the flow formulation, rather than conflating them with the loss of step conditioning. The chain is trained end to end with the real-space MAE of(7)under the same step budget as the corresponding OrbFlow run, with identical optimizer, learning-rate schedule, EMA, probe subsampling, and checkpoint selection. E.2. Curriculum Variants and Training Compute Table 13 collects every configuration discussed in Sec. 4.6 together with its total training cost. The curriculum is compute-efficient in practice because the Phase-1 warm start is shared across any number of downstream rollout horizons. Of the 74 GPU-h total cost for the fullK=2curriculum, only 39 GPU-h come from the Phase-2 fine-tuning stage itself, with the remaining 35 GPU-h of Phase-1 cost shared withK=3. That incremental fine-tuning cost is roughly half that of training the sameK=2objective from scratch (81 GPU-h), and the saving widens once the shared Phase-1 cost is amortized across rollout horizons. E.3. Training-Time Rollout Horizon As shown in Table 13, holding the Phase-1 warm start fixed, we vary the rollout budgetKused during Phase 2.K=1reaches0.296± 0.005%, statistically indistinguishable from Phase-1-only training at0.295± 0.016%. This follows from the endpoint parameterization. A single Euler step starting att 0 = 0evaluatesf θ at exactly the pointL stage1 already supervises, so fine-tuning restricted to this boundary case adds essentially no signal beyond Phase 1. K=2andK=3converge to the same accuracy,0.240± 0.005%and0.240± 0.002%. Total training compute nevertheless grows monotonically withK, since each additional Euler step adds one more sequential network 29 Equivariant Flow Matching for Electron Density Prediction Table 13: All ablation configurations on the 10% QM9 subset, with total training compute. NMAE is mean± std over three seeds. Rows without a compute entry reuse the shared Phase-1 checkpoint and are reported for accuracy only. ConfigurationNMAE (%)↓ GPU-h Direct regression (single pass)0.280± 0.003– Deterministic two-step0.268± 0.003– Phase 1 only0.295± 0.01635 Phase 2 only, K=2 (no warm start)0.250± 0.02781 Joint single-stage, K=20.241± 0.002182 Full curriculum, K=2 (Ours)0.240± 0.00574 Full curriculum, K=10.296± 0.005– Full curriculum, K=30.240± 0.002– evaluation per training example during Phase 2, together with the corresponding backward pass. We therefore adoptK=2as the smallest budget at which rollout supervision is informative and past which additional steps yield no measurable benefit at increasing cost. E.4. Inference-Time Integration Budget Table 14 reports accuracy and relative cost as a function of the number of Euler steps used at inference for ourK=2-trained model, extending to budgets never seen during training. Accuracy is best exactly at the train-matched budget of two steps. Using more steps than trained does not recover accuracy but degrades it, from0.240%to0.539%at four steps, after which the error plateaus in the0.52–0.54%range while relative cost grows to3.98×at sixteen steps. The train-matched setting is therefore both the accuracy-optimal and the most compute-efficient inference choice. Table 14: QM9 NMAE as a function of the number of Euler integration steps used at inference, for a model trained withK=2. NMAE is mean±std over three training seeds (best-val checkpoint). Relative wall-clock inference cost is normalized to the train-matched budget (t K /t 2 ), averaged over seeds. Inference Steps2 (train-matched)4816 NMAE (%)↓0.240± 0.0050.539± 0.0330.519± 0.0440.516± 0.053 Relative Inference Cost1×1.36×2.18×3.98× This lack of generalization beyond the trained horizon follows directly from how Phase 2 supervises the model, and is worth making explicit. Under the endpoint parameterization of(4),v θ is derived algebraically from a predicted clean endpointf θ (G, x t , t), and Phase 2 supervises that endpoint prediction only at the two discretization pointst∈0, 0.5visited by theK=2rollout. The network therefore receives no training signal at the intermediate timesteps an integrator withK > 2would query, nor at the coarser single interpolation stepK=1would require. The plateau in Table 14 is consistent with this account. Rather than learning a smooth field valid under arbitrary discretizations, Phase 2 sharpensf θ into an accurate two-point discrete map, and additional Euler steps subdivide the trajectory into intervals the endpoint predictor was never asked to be accurate on. Removing this restriction, for instance by conditioning on the step size as in shortcut models [13], is the clearest direction for future work. 30