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

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 (10)

• 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 \ref{['eq:delta_transform']}) 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.
• 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 $e r_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.
• Figure 5: MDA dipole and absorption with the OrbEvo-DM-s8 model on test samples. The unit for dipole in the plot is $e r_B$, where $r_B$ is Bohr radius (0.529 Å). The unit for absorption spectra is $0.529e\text{\AA}^2/V$.