Orbital Transformers for Predicting Wavefunctions in Time-Dependent Density Functional Theory

Xuan Zhang; Haiyang Yu; Chengdong Wang; Jacob Helwig; Shuiwang Ji; Xiaofeng Qian

cs.LGMar 3, 2026

Orbital Transformers for Predicting Wavefunctions in Time-Dependent Density Functional Theory

Authors: Xuan Zhang, Haiyang Yu, Chengdong Wang, Jacob Helwig, Shuiwang Ji, Xiaofeng Qian

Organizations: Department of Computer Science and Engineering, Texas A&M University · Department of Materials Science and Engineering, Texas A&M University · J. Mike Walker ’66 Department of Mechanical Engineering, Texas A&M University · Department of Electrical and Computer Engineering, Texas A&M University ·

Abstract

We aim to learn wavefunctions simulated by time-dependent density functional theory (TDDFT), which can be efficiently represented as linear combination coefficients of atomic orbitals. In real-time TDDFT, the electronic wavefunctions of a molecule evolve over time in response to an external excitation, enabling first-principles predictions of physical properties such as optical absorption, electron dynamics, and high-order response. However, conventional real-time TDDFT relies on time-consuming propagation of all occupied states with fine time steps. In this work, we propose OrbEvo, which is based on an equivariant graph transformer architecture and learns to evolve the full electronic wavefunction coefficients across time steps. First, to account for external field, we design an equivariant conditioning to encode both strength and direction of external electric field and break the symmetry from SO(3) to SO(2). Furthermore, we design two OrbEvo models, OrbEvo-WF and OrbEvo-DM, using wavefunction pooling and density matrix as interaction method, respectively. Motivated by the central role of the density functional in TDDFT, OrbEvo-DM encodes the density matrix aggregated from all occupied electronic states into feature vectors via tensor contraction, providing a more intuitive approach to learn the time evolution operator. We adopt a training strategy specifically tailored to limit the error accumulation of time-dependent wavefunctions over autoregressive rollout. To evaluate our approach, we generate TDDFT datasets consisting of 5,000 different molecules in the QM9 dataset and 1,500 molecular configurations of the malonaldehyde molecule in the MD17 dataset. Results show that our OrbEvo model accurately captures quantum dynamics of excited states under external field, including time-dependent wavefunctions, time-dependent dipole moment, and optical absorption spectra.

Figures & tables

Figure 1: The framework of RT-TDDFT. (a) Ground state wavefunctions as the initial input. (b) External electric field applied onto the system. (c) Time evolution of wavefunctions under external field. (d) Physical properties calculated from the time-dependent wavefunctions and dipole moments.

Figure 2: (a) Overview of OrbEvo. Top: Given the molecular structure and ground-state wavefunctions, OrbEvo predicts the delta wavefunctions (Equation 3 ) in future steps (one time bundle) autoregressively. Bottom: OrbEvo takes wavefunction coefficients as node features on 3D atom graphs, where each electronic state is represented by one graph. The output node features correspond to the target wavefunction coefficients at the next time bundle. (b, c) OrbEvo architectures. (b) OrbEvo-WF uses layer-wise pooling and global transformer blocks to perform electronic state interactions. (c) OrbEvo-DM computes density matrix features from input wavefunctions via tensor contraction and linear projection. Diagonal block features are added into node features and off-diagonal block features are conditioned in equivariant graph attentions. (d) Embedding layer, where atom type embedding, edge degree embedding and linear projection of input coefficients are added together. (e) EquiformerV2 block with SO(2) equivariance, composed of two SO(2)-LayerNorm layers, one equivariant graph attention layer and one feed forward network. (f) SO(2)-LayerNorm, where the output of the SO(3)-LayerNorm in the original EquiformerV2 is multiplied by a scale vector and added with a bias vector. The scale and bias vectors are computed from the external electric field intensity at current and the next time bundles with an MLP. Scale has different values for different rotation order $\ell$ ’s, which preserves the SO(3) equivariance. Bias has non-zero values only at $m=0$ , which breaks the symmetry from SO(3) to SO(2). (g) Illustration of density matrix featurization via tensor contraction.

OrbEvo Model	Wavefunction			Dipole		Absorption
OrbEvo Model	1-step $\operatorname{\ell_{2}-MAE}$	Rollout $\operatorname{\ell_{2}-MAE}$	Rollout $\operatorname{nRMSE}$	$\operatorname{nRMSE-all}$	$\operatorname{nRMSE-z}$	$\operatorname{nRMSE-\alpha}$
DM-s8	0.0242	0.0947	0.1778	0.3008	0.2326	0.0671
WF-sall	0.0192	0.0853	0.1585	0.3957	0.3066	0.0865

Table 1: Results on the MDA dataset.

OrbEvo Model	Wavefunction			Dipole		Absorption
OrbEvo Model	1-step $\operatorname{\ell_{2}-MAE}$	Rollout $\operatorname{\ell_{2}-MAE}$	Rollout $\operatorname{nRMSE}$	$\operatorname{nRMSE-all}$	$\operatorname{nRMSE-z}$	$\operatorname{nRMSE-\alpha}$
DM-s8	0.0190	0.0797	0.1885	0.1946	0.1459	0.0752
WF-sall	0.0164	0.0874	0.2071	0.6045	0.4629	0.1270

Table 2: Results on the QM9 dataset.

Figure 3: QM9 dipole and absorption with the OrbEvo-DM-s8 model on test samples 0, 10, 20, 30, 40. Note that the test samples are randomly shuffled during dataset generation. The unit for dipole in the plot is $er_{B}$ , where $r_{B}$ is Bohr radius (0.529 Å). The unit for absorption spectra is $0.529e\text{\AA}^{2}/V$ . We highlight that there is no explicit supervision on dipole or absorption during training and validation.

Figure 4: Wavefunction rollout using the OrbEvo-DM-s8 model compared with the ground truth.

Appendix figures & tables17 assets

Supplementary material from the paper’s appendix.

Appendix

OrbEvo Model	Wavefunction			Dipole		Absorption
OrbEvo Model	1-step $\operatorname{\ell_{2}-MAE}$	Rollout $\operatorname{\ell_{2}-MAE}$	Rollout $\operatorname{nRMSE}$	$\operatorname{nRMSE-all}$	$\operatorname{nRMSE-z}$	$\operatorname{nRMSE-\alpha}$
DM-sall	0.0244	0.0997	0.1888	0.3203	0.2494	0.0729
DM-s8	0.0242	0.0947	0.1778	0.3008	0.2326	0.0671
DM-s4	0.0257	0.1010	0.1902	0.3096	0.2396	0.0734
DM-sall-cat	0.1269	0.4429	0.7875	2.063	1.6345	0.8040
DM-s8-no-dm(t)	0.0508	0.2788	0.5457	0.8738	0.6768	0.1758

Appendix

Table 3: Ablation studies on the MDA dataset.

Dataset	Model	# iterations	GPU	Wall Clock Time	GPU Memory (MB)
MDA	OrbEvo-DM-s8	300k	2 $\times$ 11GB 2080Ti	3.475 days	13,848
MDA	OrbEvo-WF	300k	2 $\times$ 11GB 2080Ti	3.345 days	14,248
QM9	OrbEvo-DM-s8	395k	4 $\times$ 48GB A6000	3.118 days	49,434 - 54,700
QM9	OrbEvo-WF	395k	2 $\times$ 80GB A100	5.003 days	46,652 - 69,662

Appendix

Table 4: Training cost of OrbEvo models. MDA models are trained with a batch size 32. QM9 models use a batch size of 16. All models are trained with Pytorch distributed data parallel ( torch.ddp ) for multi-gpu training and with num_workers=16 in dataloader for MDA and num_workers=32 for QM9. As a rough estimation, 2

\times

2080Ti is roughly equivalent to 1

\times

A6000 in terms of speed. The GPU memory usage is tested by running training on 1 single A100 GPU for 10 minutes. For QM9, The GPU memory can vary depending on the molecule sizes in a batch. We note that with a slightly optimized push-forward training implementation, we are able to fit the training within 2

\times

A6000 GPUs for both OrbEvo models on the QM9 dataset with similar training time.

Dataset	Model	GPU	Batch Size	Wall Clock Time / Batch		GPU Memory
Dataset	Model	GPU	Batch Size	Wavefunction	Wavefunction + Property	(MB)
MDA	OrbEvo-DM	1 $\times$ A6000	20	3.67 seconds	5.23 seconds	5742
MDA	OrbEvo-WF	1 $\times$ A6000	20	2.84 seconds	4.60 seconds	2032
QM9	OrbEvo-DM	1 $\times$ A6000	20	18.00 seconds	26.74 seconds	34,164 - 42,842
QM9	OrbEvo-WF	1 $\times$ A6000	20	11.86 seconds	20.31 seconds	17,204

Appendix

Table 5: Inference cost of OrbEvo models. All models are tested one a single A6000 GPU using num_workers=10 in dataloader. The reported times are wall clock time per batch. We report both the time for producing the wavefunction trajectory (Wavefunction), as well as the time for producing the wavefunction trajectories and computing the dipoles and absorptions (Wavefunction + Property). Note that the properties are not parallelized with batch processing and are computed on CPUs. We note that electronic state sampling is not enabled during inference, which leads to increased GPU memory usage for OrbEvo-DM. In comparison, during training OrbEvo-DM is able to use electronic state sampling to reduce GPU usage.

Dataset	# CPU cores	Wall Clock Time / Molecule
Dataset	# CPU cores	Ground-state DFT	Total
MDA	24	34.3 seconds	1.5 hours
QM9	24	73.1 seconds	3.2 hours

Appendix

Table 6: Simulation time per molecule. The simulation time is averaged over 40 simulations. Ground-state DFT is the time to compute the initial wavefunction coefficients from molecular structures. The initial wavefunction coefficients are used as input to OrbEvo models.

Figure 5: MDA dipole and absorption with the OrbEvo-DM-s8 model on test samples. The unit for dipole in the plot is $er_{B}$ , where $r_{B}$ is Bohr radius (0.529 Å). The unit for absorption spectra is $0.529e\text{\AA}^{2}/V$ .

OrbEvo Model	Wavefunction			Dipole		Absorption
OrbEvo Model	1-step $\operatorname{\ell_{2}-MAE}$	Rollout $\operatorname{\ell_{2}-MAE}$	Rollout $\operatorname{nRMSE}$	$\operatorname{nRMSE-all}$	$\operatorname{nRMSE-z}$	$\operatorname{nRMSE-\alpha}$
DM-s8	0.0224	0.0863	0.1613	0.1997	0.1499	0.0539
WF-sall	0.0225	0.1080	0.2008	0.3758	0.2881	0.0822

Appendix

Table 7: Results on the MDA dataset with the new training.

Figure 6: Tensor product visualization produced by the e3nn library.

Panel of Figure 7 — Figure 7: Equivariance error of TDDFT data. Left: real part of the wavefunction coefficients of an unrotated MDA molecule at one time step. Right: the difference between the wavefunctions at the same time step in a second simulation produced from a rotated version of the same molecule, and the coefficients manually rotation-transformed from the left plot. In the second simulation the molecule is rotated by 35 degree around the electric field direction.

Panel of Figure 8 — Figure 8: Equivariance error of OrbEvo-DM. Left: real part of the model’s predicted wavefunction coefficients for a MDA molecule using the ground-truth wavefunctions at one time step as input. Right: the difference between the model’s predicted wavefunctions using the rotated structure and manually rotation-transformed ground-truth wavefunctions, and the coefficients manually rotation-transformed from the left plot. The molecule’s rotation and the rotation transformation is 35 degree around the electric field direction.

Model	Dataset	1-step $\ell_{2}$ -MAE	Rollout nRMSE	Dipole $z$ nRMSE	Absorption nRMSE
OrbEvo-DM-s8-ood	QM9_ood - val	0.0142	0.1498	0.1113	0.0615
OrbEvo-DM-s8-ood	QM9_ood - test	0.0132	0.1482	0.1175	0.0697

Appendix

Table 8: Performance on QM9_ood validation and test sets.

Model	Dataset	1-step $\ell_{2}$ -MAE	Rollout nRMSE	Dipole $z$ nRMSE	Absorption nRMSE
OrbEvo-DM-s8-ood	QM9_id_ood_intersect - val	0.0141	0.1500	0.1126	0.0633
OrbEvo-DM-s8	QM9_id_ood_intersect - val	0.0146	0.1579	0.1153	0.0628
OrbEvo-DM-s8-ood	QM9_id_ood_intersect - test	0.0137	0.1496	0.1142	0.0666
OrbEvo-DM-s8	QM9_id_ood_intersect - test	0.0139	0.1521	0.1074	0.0636

Appendix

Table 9: Performance comparison on QM9_id_ood_intersect validation and test sets.

Time bundle	1-step $\ell_{2}$ -	8-step Rollout	16-step Rollout	32-step Rollout	64-step Rollout	100-step Rollout	Dipole $z$	Absorption
size	MAE	nRMSE	nRMSE	nRMSE	nRMSE	nRMSE	nRMSE	nRMSE
1	0.0093	0.0780	0.0340	0.1363	0.4433	0.9032	0.9526	0.1684
2	0.0130	0.0668	0.0340	0.1106	0.3087	0.5765	0.5669	0.1228
4	0.0139	0.0693	0.0289	0.0572	0.1235	0.1979	0.2670	0.0758
8	0.0242	0.0720	0.0378	0.0677	0.1245	0.1778	0.2328	0.0672
16	0.0588	0.1026	0.0544	0.1075	0.2040	0.2872	0.2641	0.0922

Appendix

Table 10: Time bundling analysis on the MDA dataset.

Figure 9: Global phase rollout on the QM9 sample.

Figure 10: Global phase rollout on the MDA sample.

OrbEvo Model	Wavefunction			Dipole		Absorption
OrbEvo Model	1-step $\operatorname{\ell_{2}-MAE}$	Rollout $\operatorname{\ell_{2}-MAE}$	Rollout $\operatorname{nRMSE}$	$\operatorname{nRMSE-all}$	$\operatorname{nRMSE-z}$	$\operatorname{nRMSE-\alpha}$
DM-s8	0.0242	0.0947	0.1778	0.3012	0.2329	0.0672
DM-s8-w/-quadratic-dm	0.0290	0.1110	0.2088	0.3538	0.2744	0.0784

Appendix

Table 11: Density matrix analysis on the MDA dataset.

OrbEvo Model	Wavefunction			Dipole		Absorption
OrbEvo Model	1-step $\operatorname{\ell_{2}-MAE}$	Rollout $\operatorname{\ell_{2}-MAE}$	Rollout $\operatorname{nRMSE}$	$\operatorname{nRMSE-all}$	$\operatorname{nRMSE-z}$	$\operatorname{nRMSE-\alpha}$
DM-s8-noise	0.0204	0.1262	0.2423	0.3868	0.3036	0.0815
WF-sall-noise	0.0155	0.0866	0.1617	0.4045	0.3157	0.0788

Appendix

Table 12: Noise injection results on the MDA dataset.

Hyperparameters	Value
Optimizer	AdamW
Learning rate scheduling	Cosine Annealing
Maximum learning rate	$1\times 10^{-3}$
Weight decay	$1\times 10^{-3}$
Number of epochs for MDA	129 (300k iterations)
Number of epochs for QM9	17 (395k iterations)

Appendix

Table 13: OrbEvo model hyperparameters.

Quantum Chemistry

Quantum chemistry models molecular electrons and nuclei to predict energies, structures, spectra, and reactivity using classical and quantum methods.

See all 64 papers

Equivariant Flow Matching for Electron Density Prediction
Oct 2, 2026Chenxing Liang, Chengdong Wang, Yuchao Lin +2Computational Quantum Chemistry Quantum Chemistry
Department of Computer Science and Engineering, Texas A&M University · Department of Materials Science and Engineering, Texas A&M University · Department of Electrical and Computer Engineering, Texas A&M University +2
Surrogate Functionals for Machine-Learned Orbital-Free Density Functional Theory
Apr 22, 2026Roman Remme, Fred A. HamprechtSurrogate Modeling Quantum Chemistry
Interdisciplinary Center for Scientific Computing (IWR), Heidelberg University
A Fixed-Point Neural Operator for Size- and Functional-Transferable Hamiltonian Prediction
Jun 12, 2026Yunhong Lou, Xihang Yue, Xinran Wei +2Quantum Chemistry Molecular Property Prediction
1Zhejiang University, Hangzhou, China. · 2Zhongguancun Academy, Beijing, China. · 3Zhongguancun Institute of Artificial Intelligence, Beijing, China.

Oct 2, 2026J/K move · Enter open · S save

Equivariant Flow Matching for Electron Density Prediction

Chenxing Liang, Chengdong Wang, Yuchao Lin, Xiaofeng Qian+1

Department of Computer Science and Engineering, Texas A&M University · Department of Materials Science and Engineering, Texas A&M University · Department of Electrical and Computer Engineering, Texas A&M University+2

Computational Quantum Chemistry Quantum Chemistry Density Functional Theory Flow-Based Generative Modeling Molecular Property Prediction Flow Matching

Open paper AskPDF arXiv

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.

Abstract

Quantum Chemistry

Explore similar work

Equivariant Flow Matching for Electron Density Prediction

Surrogate Functionals for Machine-Learned Orbital-Free Density Functional Theory

A Fixed-Point Neural Operator for Size- and Functional-Transferable Hamiltonian Prediction