Orbital Transformers for Predicting Wavefunctions in Time-Dependent Density Functional Theory
Authors: Xuan Zhang
Haiyang Yu
Chengdong Wang
Jacob Helwig
Shuiwang Ji
Xiaofeng Qian
Venue: ICLR 2026 Poster (2026)
Keywords: Machine learning density functional theory
Time dependent neural PDE solver
Primary area: applications to physical sciences (physics, chemistry, biology, etc.)
TL;DR: We developed deep learning models to efficiently evolve electronic wavefunctions in real time.

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 characterized by dipole oscillator strength. It also shows strong generalization capability on the diverse molecules in the QM9 dataset. Our dataset is available at https://huggingface.co/divelab, and our code is available as part of the AIRS library https://github.com/divelab/AIRS/.

Published as a conference paper at ICLR 2026
ORBITAL TRANSFORMERS FOR PREDICTING WAVEFUNCTIONS IN TIME-DEPENDENT DENSITY FUNCTIONAL THEORY
Xuan Zhang'
Haiyang Yu'
Chengdong Wang
Jacob Helwig'
Shuiwang 1i Xiaofeng Qian 24,5
'Department of Computer Science and Engineering, Texas A&M University
*Department of Materials Science and Engincering, 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
$ Department of Physics and Astronomy, Texas A&M University
{ [xuan. zhang, sji,
ABSTRACT
We aim to learn wavefunctions simulated by time-dependent density functional theory (TDDFT), which can be efficiently represented as linear combination co- efficients of atomic orbitals. In real-time TDDFT, the electronic wavefunctions of 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 re- lies time-consuming propagation of all occupied states with fine time steps. In this work, we propose OrbEvo, which is based on an equivariant graph trans- former architecture and learns to evolve the full electronic wavefunction coeffi- cients across time steps. First, to account for external field, we design an equivari- ant conditioning 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, pro- viding a more intuitive approach to learn the time evolution operator. We adopt 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 the malonaldehyde molecule in the MD17 dataset. Results show that OrbEvo model accurately captures quantum dynamics of excited states under external field, including time-dependent wave- functions, time-dependent dipole moment, and optical absorption spectra char- acterized by dipole oscillator strength. It also shows strong generalization capa- bility on the diverse molecules in the QM9 dataset. Our dataset is available httpss//hugginggiice.oodvvvve and our code is available as part of the AIRS li- brary https///tthhbcco//dveee//AAsA
INTRODUCTION
Density functional theory (DFT) (Hohenberg & Kohn 1964 Kohn & Sham 1965) provides efficient way to solve time-independent many-body Schrödinger equation using a variational prin- ciple and has been widely applied to compute the properties of the ground state of molecules and solids. However, many important physical and chemical phenomena involve the excited states and the dynamic responses of the systems external perturbations. such cases, time-dependent
"qqull senior authorship.

---

Published as a conference paper at ICLR 2026
density functional theory (TDDFT) Runge & Gross 1984) provides natural extension of the time-dependent many-body Schrödinger equation. It can be formulated and solved in frequency space in linear-esponnee TDDFT (Casida 1995), or in the time domain via real-time TDDFT (RT- TDDFT) (Runge & Gross 1984 Yabana & Bertsch 1999. Qian et al. [006] Ullrich 201T), enabling the investigation of excited state properties such as excitation spectra, optical absorption, charge transfer, and electron dynamics under time-dependent external fields such as electromagnetic fields. Starting from the static electronic wavefunctions obtained within ground-state DFT, RT-TDDFT propagates these wavefunctions in the time domain under the influence of an external field, allowing direct investigation of both linear and nonlinear physical properties.
However, RT-TDDFT is computationally demanding due to the temporal and spatial discretization of Kohn-Sham wavefunctions, long-time propagation, repeated evaluations of the Kohn-Sham Hamil- tonian, and the increasing number of Kohn-Sham wavefunctions with system size. To accelerate this procedure, machine learning (ML) provides a promising way to replace or approximate the costly propagation steps, thereby accelerating quantum dynamical simulations while retaining ac- curacy (Zhang et al. 2025). In this work, we propose a new model, OrbEvo, designed to learn the full wavefunction evolution while incorporating the underlying physical symmetries of the TDDFT problem. In particular, we consider the SO(2) equivariance induced by the presence of an exter- nal field, and we demonstrate how ML-based partial differential equation (ML-PDE) frameworks can be adapted to capture quantum dynamics effectively. We extend PDE learning to the setting of wavefunction coefficient evolution on atom graphs, while enforcing SO(2) equivariance to respect the system's 's symmetry constraints. Furthermore, we propose effective methods to handle multiple electronic states, which remain agnostic the choice of backbone neural architecture. Together, these innovations allow our approach to bridge the gap between ab initio quantum dynamics and scalable ML-based approximations.
2 PRELIMINARIES
In this section, we will provide a formulation of the RT-TDDFT problem. At the same time, the constraints inherent to this physical problem will be elaborated on, serving the motivation for the techniques developed. Our method is built upon and enabled by existing literature. We review them in Appendix A
DFT with predefined localized atomic orbital basis set. DFT provides a practical approximation solve the many-body Schrödinger equation of a molecular or material system. Instead of explicitly modeling the many-body wavefunctions, DFT represents the system using a set of single-partice Kohn-Sham wavefunctions (vn: 3 C)==1....N.. where Noce denotes the number of occupied electronic states. Each electronic state can be occupied by up to two electrons according to the Pauli exclusion principle. To construct these Kohn-Sham states, DFT often employs basis set, such as the localized atomic orbitals in this work, { : R3 - C)o= Iu.,,N., with Norb the total number of orbitals in the system. These atomic orbitals are spatially localized around atoms and describe the electronic states of isolated atoms, forming the Hilbert space of the system. In the linear combination of atomic orbitals (LCAO) method, each electronic wavefunction th can be expressed as a linear combination of atomic orbitals, Un = ENun Cno 6o. where C E N XN is the coefficient matrix defining the contribution of each orbital. At the ground state, the coefficients are determined by solving the Kohn-Sham equation (Kohn & Sham 1965) in the matrix form, denoted as
Nar.
HC,, = SC,.
(1)
where H E Nan is the Kohn-Sham Hamiltonian matrix, S E X. is the overlap matrix, and en E R are the eigen energies for the Kohn-Sham eigen states. This formulation highlights the central role of the Hamiltonian and overlap matrices in determining the electronic structure.
TDDFT under external electric field. For the TDDFT problem in this paper, the input consists of atom types and 3D atomic positions of the molecule, denoted as E NN and R E N.Nx3 x3 respec- tively, where N₂ is the number of atoms in the system, together with an applied dime-dependent uni- form external electronic field E(t) E R3. as well the initial ground state wavefunction coefficients C(0). The goal is to predict the temporal evolution of the electronic wavefunction, represented by a sequence of coefficient matrices (C that reconstruct the wavefunctions each time step.
2

---

Published as a conference paper at ICLR 2026
(n Ground Sirie
(c) Time Evolution
(d) Physical Progertes
60.0
6010
Occupiod Elestooic
stooptee
States
0.0
(6)
Blotric Field
2000
Figure 1: The framework of RT-TDDFT. (a) Ground state wavefunctions 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 mo- ments.
In the absence of external electronic field, the dynamics reduces to simple unitary evolution over time, (t) exp(-ic,,tt (/) CO,, electronic wavefunction. However, under a time-dependent electronic field E(t), the perturbation couples these electronic wavefunctions, leading nontrivial transitions that must be captured by the time-dependent Kohn-Sham equations in the LCAO basis as follows,
1,..., .ooe, corresponding to phase rotations of the
d dt C.(t) (c..t)==is--(((
S-
where Hoo (t) (t)(t) (t))o> ()))) hi the Planck constant, and (t) is the Kohn- Sham Hamil- tonian operator time given by (1) Ta + Hn(p(r, 1)] Vxclp(r, 6)] + Vox(!). Within LCAO, the time-dependent electron density is p(r, t) Nob ENer = Doo (r)) (r), where D is the density matrix given by D (1) = >n'r
(C) (4). Th is the occupation number in electronic state th The central task of RT-TDDFT is therefore to integrate Equation 2 over time, compute wavefunction coefficients C(t) in the local orbital basis, subsequently calculate electron density p(r, t). update the density-dependent operators the Kohn-Sham Hamiltonian and compute Hod (t), and repeat this process iteratively for many time steps. In RT-TDDFT, each Kohn-Sham wavefunction un evolves in time under the time-ordered evolution operator Ü(t, to), starting from the initial time to: (t) Üt, (to), where (t, tu) Texp ( -- f( (t')t'' and
T time-ordering operator. More details about time evolution of wavefunctions can be found in Appendix G For the machine learning model, the objective is to learn the time evolution of the Kohn-Sham wavefunctions v/, (f), or equivalently (t,, in order to accelerate TDDFT calcula- tions.
SO(2) equivariance in TDDFT. While property prediction, force field prediction, and Hamiltonian matrix prediction are typically formulated under SO(3) equivariance, meaning that when the input geometry is rotated, the corresponding predicted properties transform consistently under the same rotation. this full rotational symmetry can be broken the presence of an external field. In partic- ular, when a uniform external electronic field along a specific direction is introduced. it defines preferred spatial direction. As result, rotations that modify the angle between the field direction and the molecular orientation will alter the system, whereas rotations around the field axis preserve SO(2) equivariance. Consequently, the overall symmetry of the system reduced. In this work, we focus on the case of uniform external electronic field applied along specific direction, where the molecular system is SO(2) under rotations around its axis, thereby reducing the symmetry require- ment for predicted properties from SO(3) to SO(2) equivariance to consider the effect of uniform external electronic field. The squiaa-que of TDDFT data tested Appendix H.2 Fig- ure
3

---

Published as a conference paper at ICLR 2026
3 METHOD
3.1 OVERALL FRAMEWORK
The overall problem framework for TDDFT is illustrated in Figure | We describe the inputs and targets of this framework, along with the multi time step outputs strategy used during both training and inference of our machine learning model.
Delta transformation for capturing small changes in wavefunction coefficients. One particular challenge in our data is how to define the prediction target. Due to the small magnitude of external electric field, the coefficients at future time steps differ only by small amount compared to the initial step by the factor of a global phase. Directly learning the wavefunction coefficient will make the model only learn the global phase changes. To correctly model the delta wavefunction, we define a global phase factor and delta coefficients for each electronic state as
(t) = Co) (t) (O,(0)) E EC, (n(()
(0---1CCL C,. B To
1
C2((0) E C.n.
(3)
with 3 = 1 10-3 amplify the delta, in which case we have Ct (0) (t) (t)). Note that (0) is vealevalued and the conjugation takes no effect. In the absence of external electric field, i.e., when C., (t) is real-valued, and (1) 0. This highlights that the proposed delta transformation is able to extract the delta wavefunctions induced by the external electric field E(I). Since the A() carries the most information related to physical properties, we focus on learning A() in the main text, while the learning of the (t) can be found in Appendix K
C.) we will obtain (t) = exp(-icnt/h since C,, (0)
Time bundling. Time bundling (Brandstetter et al., 2022) is technique in PDE surrogate models. Instead of advancing time by one at each prediction step, we predict multiple future time steps at once so that the total number of auto-regressive steps will be reduced to produce the same number of total time steps. Formally, our model learns the mapping
M(0) C(O), A(t - i)..................... A(t - 1) + A()....A(t++fD)
(4)
where M is the neural network with parameters 8, is the number of conditioning steps, and f is the number of future steps. We use h f = Nn 8 in our implementation.
By using neural networks to approximate the time propagation process, the simulation time can be greatly reduced compared to classical numerical solvers. For example, the simulation time for one molecule using TDDFT solver would take hours, compared to ~ second for neural network inference. Given the predicted wavefunction coefficients, we can then calculate the properties of the molecule, including dipole moments and absorption spectra.
3.2 MODEL
3.2.1
EQUIVARIANT GRAPH TRANSFORMER
Our model is based on EquiformerV2 V (iao et al. 2024), which is an SO(3)-equivariant graph trans- former, and we use SO(()-equivariant electric field conditioning to break the symmetry SO(2). In EquiformerV2. each node of the graph has an equivariant feature f, E IR X where dugh is the number of spherical channels and demb is the embedding dimension. The spherical channels are partitioned into different segments where each segment has a different rotation order 2 O. The rotation order É defines the equivariance property of each segment when the global reference frame of input space undergoes 3D rotation, and an order-f segment has 2 + spherical channels, in- dexed by m E [-e, é). For example, when the input reference frame is rotated by rotation matrix 0 features will transform as scalars and remain unchanged, l = -1 features
R 1x33 then will transform as 3D vectors and will be rotated by the same matrix R, and = 2 features will transform as order-2 spherical harmonics and will be rotated by the corresponding wigner-D matrix (R) E R Although EquiformerV2 has the possibility of reducing the range of m to be smaller than f.. we always use the full 2 + 1 spherical channels in implementation.
Equivariant graph transformers are composed of equivariant transformer blocks, which process the features with equivariant graph attention and node-wis feedforward networks. The key operation
4

---

Published as a conference paper at ICLR 2026
DF7
A.R
Cllft
Immeendeng data
Dy C
2424102064
Skot Piy
Dlccronc Stacee It
Dkamadi Saz M.5
Conractan
D.s
Your mily
Vour 20l
RCd clutte
Oskiro
So Contioan
A Findee
Pu
DrbFoww-FFF
(c) Frh
(g Dimiity Mariis Computation
Kifi
Lon
Craghh
(4) Ovexiew
()) Embeelling
(el Eq Block
(0 SOC3LLayerNem
Figure 2: (a) Overview of OrbEvo. Top: Given the molecular structure and ground-state wave- functions, 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 corre- spond to the target wavefunction coefficients at the next time bundle. (b, e) OrbEvo architectures. (b) OrbEvo-VV WF uses layer-wise pooling and global transformer blocks to perform electronic state interactions. (c) OrbEvo-DAy 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 LOy2)-LayerNorm layers, one equivariant graph attention layer and one feed forward network. (1) SO(2)-LayeNormm where the output of the LayerrLm 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 current and the next time bundles with an MLP. Scale has different values for different rotation order f's, which preserves the SO(3) equivariance. Bias has non-zero values only at 0, which breaks the symmetry from SO(3) SO(2). (g) Illustration of density matrix fcaturization via tensor contraction.
in equivariant graph attention is to compute a rotation invariant attention score ai and rotation equivariant message may between node and its neighbor node j. aiy and m:j are computed using tensor products between the concatenated node features [f,, fy] and the spherical harmonics projec- tion of their relative vector ray as asj. m.y TPp ([f, f,), r.y), where TP, contains parameters that encode the distance information and mix different rotation orders. Node i's feature is then updated as the weighted sum of messages f = Ejen(a) Ciyy where (i) denotes the node i's neighbors.
3.2.2 WAVEFUNCTION GRAPHS WITH SHARED GEOMETRY
We model the wavefunctions atom graphs where each atom has feature its atom type E N and its coordinates r, E R².
Wavefunction as node features. The wavefunction coefficients for atomic orbitals of the same atom are grouped together to form initial wavefunction features. The coefficients are further grouped according their rotation orders f. The resulting wavefunction feature for electronic state and atom i is WWF , E [d x xduan, where d/2 = 9 corresponds to the concatenation of rotation orders up to l 2, and the daand 2(2N + 1) corresponds to the concatenation of real and imaginary
5

---

Published as a conference paper at ICLR 2026
parts of the conditioning Nb steps and the initial state C(O), which is real-valued. The additional multiplicative factor of is the multiplicity of rotation orders, which accounts for the fact that each atom has two orbitals and up to two p orbitals. Since each atom has zero data, we use zero padding to fill the second multiplicity channel of rotation order = 2. We also zero-pad atoms with fewer orbitals to the same maximum rotation order and multiplicity, which practically only affects hydrogen atoms, which have orbitals 1s, 2s and 1p.
orbital in our
Electronic states set of graphs. As the wavefunctions of all occupied electronic states jointly decide the electron density, and consequently the propagation operator, it is important to consider the interaction between electronic states when evolving each individual electronic state. One straightfor- ward option would be ordering the electronic states according to their energy levels {e, he...... and concatenating
electronic states together into a global feature vector. However, as shown in Appendix B we find such an approach fails to learn the propagation. We attribute this failure to the fact that the electronic states are eigen vectors of the initial Hamiltonian matrix and are better interpreted set, thus mixing them as separate feature channels would make learning difficult.
Instead, we propose to model each electronic state as individual graphs (n=I............. where Gn. (FWF, 2, R). FWF (f. (fWF) Ji=1..., the node features of electronic state n. z and are atom types and coordinates shared by all electronic states.
Wavefunction encoding. We apply linear layer to FWF and increase its number of channels from deond demb, where different weights are used for different rotation order f's, and bias is added to We also add the atom type embedding and the edge degree embedding from EquiformerV2 (Liao et al. 2024) to the projected wavefunction features.
3.2.3 LEARNING INTERACTION OVER ELECTRONIC STATES
We introduce two ways to interact among electronic states.
Interaction via wavefunction pooling. Following set learning methods (Qi et al. 2017 Maron et al. 2020), we do average pooling after each graph transformer block over electronic states. The pooled feature is processed with another graph transformer block and is subsequently broadcasted back to each individual electronic states, Formally,
ppool GT
1
fr,,
(5)
Noce n=1
f, = fn + poool
(6)
where GT is graph transformer block (as in EquiformerV2).
Interaction via density matrix. We use tensor product contraction to extract features from di- agonal and off-diagonal blocks of the density matrix. The density matrix is defined as D(t) Noa Th (t) @ (t) E Na x Nut, where @ is the outer product between vectors. We divide the density matrix into matrix blocks Dij according which atom pairs the left and right coefficients in the outer product belong to. We then use tensor contraction to re-organize each Dij matrix into a of equivariant features with rotation orders up / 4. In practice, we use the linearity tensor contraction to implement this process by first computing the atom pair features for each electronic state as
D jjn TC (tn,at) & (t)(
The tensor contraction operation TC flattens the matrices into equivariant feature vectors via a change of basis using Clebsch-Gorde coefficients. We then aggregate over electronic states and compute the density matrix feature as
N Σ
DR
Th D
D.g
T/ D.j...
(8)
==
The resulting high-order features Dis and D.j describe the density matrix blocks for the self- interaction of each atom and the interactions between pairs of different atoms, respectively. An illustration for the density matrix feature computation shown in Figure 2gg). Additional infor- mation on tensor product contraction can be found in Appendix H.T Due delta transform, the

---

Published as a conference paper at ICLR 2026
density matrix will contain both linear term and quadratic term on delta wavefunctions. We find that including the quadratic term will hurt the performance (as shown in Appendix [ potentially due to its small contribution in the density matrix which may be more sensitive to noise, thus only keep the linear term in our model. The diagonal pairs of the density matrix are linearly projected and added the initial node features and the off-diagonal density matrix features are projected into the same channels using linear layers and are used in computing the graph attention, denoted as Orjj maj = ([G, fj, linear (Di)),, riy).
3.2.4 ORBEVO MODELS
We design two OrbEvo models based on the above two interaction methods. The model architectures are shown Figure
The model uses pooling as electronic state interaction. It has 6 local graph transformer blocks, each followed by pooling and global graph transformer block except for the last layer, resulting in 5 global blocks in total. The model is called full wavefunction model as it makes use of wavefunction features from all electronic states at each layer.
OrbEvo-WF:
OrbEvo-DM.
The model uses density matrix interaction. The density matrix is computed from the input coefficients. The model has 6 layers in total, the off-diagonal blocks are feed into the first two layers of the model. We use L = 4 for the first two layers and & 2 for the later 4 layers since the computational cost associated with higher-order features much higher. The feature conversion from ( = 4 to & 2 is done by only keeping the lower é features.
Electronic state sampling. Since OrbEvo process different electronic states parallel, and the number of electronic states grows linearly with the number of atoms, the computational cost of OrbEvo will be the cost of processing one molecular graph using the backbone equivariant graph transformer multiplied by the number of electronic states. This increase the training cost sig- nificantly, particularly for larger systems. To mitigate this, we do sampling on the electronic states the sampled electronic states. As result, only subset of
during training and only supervise electronic states will be processed by the network layers during training. We indicate the electronic state sampling using suffix -s. For example, WF-sall means we use all electronic states when training OrbEvo-VV WF, and DM-88 means we randomly sample 8 electronic states when training OrbEvo-DM, We find that sampling will degrade the performance of the full wavefunction model significantly while it will not affect the density matrix model. This is because the density matrix model aggre- gates information from electronic states at the input of the model and thus sampling will not affect the interaction between electronic states, while the full wavefunction model will have less information when using sampling.
OrbEvo-DM and OrbEvo-WF have 27,977,056 and 26.963,360 parameters, respectively. We opti- mize the implementation by sharing the radial function computation for different electronic states. We use automatic mixed precision for acceleration.
3.2.5 SO(2) EQUIVARAAT ELECTRIC FIELD CONDITIONING
Following Gupta & Brandstetter (2023) Herde et al. (202)); Helwig et al. (2025).
FiLM (Perez et al. 2018)) like method to insert the conditioning information by computing a scaling factor and a shifting factor from the conditioning, and apply them the feature map. We apply the conditioning after each layer norm layer in the graph transformer blocks.
Since the feature maps are equivariant features, the conditioning features must also satisfy the equiv- ariant constraints. Specifically, we apply a different scaling factor to different l's and compute the bias according the direction of the electric field. In our case, the electric field is always along the z-axis, so the spherical harmonics encoding of it is vector with non-zero entries at m = 0 positions and zero otherwise. Mathematically,
Ul 88 LN + be,
where LN(x), E N E R IXIXC be E (2 is nonzero for m = 0 and LayerNorm in EquiformerV2, st and by
zero otherwise. Here LN is SO(3) computed using a MLP from the electric field intensities at current next time bundles, and O is multiplication with broadcasting. Since the scale term 8 is the same for each f, it preserves the

---

Published as a conference paper at ICLR 2026
SO(3) equivariance. On the other hand, the bias term by adds predefined directional information into the features and consequently breaks the SO(3) equivariance to SO(2).
We show in the ablation studies in Appendix B that breaking the symmetry is essential to conrectly learn the mapping from ground state to the first evolution step. The -quia-equivariance of the OrbEvo model is tested in Appendix [H.2] Figure
Wavefunction readout. We apply an additional equivariant graph attention block to readout the wavefunctions, which is the same the force prediction in EquiformerV2 but keep the order up to 8 2.
3.3 TRAINING STRATEGY
Loss. We use the per-atom 22-MAE loss
Chanussot et al. 2021
Liao et al. 2024). defined as
Aplotch
lg 2g-MAE(Cred) C
c c 122
(10)
N
Nbatch
1=1
where Ctred and Charget are the predicted and ground-truth wavefumction coefficients, respectively, Cped and Cf denote the predicted and ground-truth coefficients for the i-th atom in the batch, where different orbitals are concatenated into one vector, and || |/2 denotes the &2- 22-nomm. The atom index runs over all sampled electronic states and all molecules in batch. The loss is averaged over all time steps in the time bundle.
1) uncorrupted by error, a distribution shift occurs during regressee rollout, where errors made in previous predictions leads to inputs A(L h) + =(L h), ..... A(I 1) s(1 (((-1, 1). Previous works have attempted to mitigate this misalignment by intentionally corrupting training inputs with errors E(i) sampled from distribution approximating the rollout error distribution. Pushforward training (Brandstetter al., 2022) samples these errors directly from the onestep eπror distribution of the model as
Push-forward training. Although training inputs A(t - h). ...
A(t .,)
6(1 ..), El 1) StopGrad (M (C(0), A(L 2h
h 1)) A(L h: 1)) (11)
Practically, this amounts to letting the model unroll once and then use the unrolled prediction as the new input. However, the onestep error distribution at the outset of training produces noise that dominates the signal at the beginning of training. Thus, in addition to maintaining uncorrupted inputs A; or adding A, 4 & from the pushforward distribution with equal probability, we multiply the &, with a warm-up up factor which increases linearly from 0 to maximum value of 1 according to the training step Finally, because the first target ((1), A(2). ..... A(h) cannot be modeled with -oshed-forward inputs, we double the weight for its loss in any batch that it appears in to balance its utilization relative to other targets, which can all be modeled using either pushed-forward forward or uncorrupted targets.
4 EXPERIMENTS
4.1 DATASET DESCRIPTION
We randomly selected 5, 000 diverse molecules from the QM9 (Ramakrishnan et al. 2014) dataset to demonstrate the generalization capability our model, and use 1, 500 molecular lar configurations of the malonaldehyde (MDA) molecule from the MD17 dataset (Chmiela et al.) 2018 for the ab- lation study. Both QM9 and MD17 are widely used in machine learning for materials science and computational chemistry. We then performed self-consistent field (SCF) DFT calculation for each molecule to obtain their -tatt Kohn-Shin Sham wavefunctions using the open-source ABACUS software package (Chen al. 2010 Li al., 2016; Lin et al.) (022).. Subsequently, we carried out RT-TDDFT calculations to propagate all occupied electronic states for 5 fs in total of 1, 000 steps with a time step of 0.005 under a spatially uniform, time-dependent electric field. During each time step, wavefunction coefficient matrices were extracted and uniformly downsampled for every 10 steps. After downsampling, each time-dependent wavefunction trajectory contained 101 steps including the first step, which were used as input data for the training, validation, and testing of our OrbEvo model. More details about dataset generation and description can be found in Appendix F.
We that the forh-forward warm-up factor may not always be helpful, as shown in Appendix
8

---

Published as a conference paper at ICLR 2026
Table Results on the MDA dataset.
Wavefunction
Dipole
Absorption
OrbEvo Model
1-step 6--MAE
Rollout
Rollout nRMSE
nRMSE-all
nRMSE-Z
nRMSE-a
2--MAE
DM-s8
0.0242 0.0192
0.0947 0.0853
0.1778 0.1585
0.3008 0.3957
0.2326 0,3066
0.0671
WF-sall
0.0865
Table 2: Results on the QM9 dataset.
Wavefunction
Dipolc
Absorption
OrbEvo Model
1-step
Rollout MAE
Rollout nRMSE
nRMSE-all
nRMSE-Z
nRMSE-av
f2-MAE
DM-s8
0.0190
0.0797 0.0874
0.1885 0.2071
0.1946 0.6045
0.1459
0.0752
WF-sall
0.0164
0.4629
0.1270
4.2 SETUP
Dataset split and normalization. For QM9, we use 4,000 molecules for training, 500 molecules for validation, and 500 molecules for testing. For MDA, we use 800 configurations for training, 200 configurations for validation, and 500 configurations for testing. We normalize the initial and delta wavefunction coefficients by dividing their respective orbital-wise rooted mean square (RMS) across all orbitals in training dataset. For delta wavefunction, we also average across time. We normalize the electric field by scaling the maximum intensity to 1.
Evaluation metries. We evaluate the performance of our OrbEvo model on three key physical properties: time-dependent wavefunction coefficients, time-dependent dipole moments, and opti- cal absorption spectra characterized by dipole oscillator strengths. These properties are crucial for downstream tasks in TDDFT, and thus provides comprehensive evaluation of the model's outputs. The detailed information about these three metrics are provided in Appendix
4.3 RESULTS
4.3.1
QUANTITATIVE RESULTS
The results on MDA and QM9 datasets are summarized in Table and Table 2 respectively. The wavefunction coefficients do not have unit. The nRMSE errors also do not have units since they are relative errors. Hence all metrics in the tables are unitless.
Overall, the results on the QM9 dataset shown in Table 2 suggest that the OrbEvo-Doe DM model ing density matrix as interaction between occupied electronic states outperforms the OrbEvo-Wi WF model which employs layer-wise pooling of the features of occupied electronic states. This may be because the density matrix in the OrbEvo-DM model is inherently consistent with the mathematical formulation of TDDFT: the density functional used to evaluate the time-dependent Kohn-Sham Hamiltonian in RT-TDDFT. Consequently, it is more straightforward for the OrbEvo-DM model to learn the time evolution operator which depends directly on the density matrix (t).
We conduct ablation studies on the MDA dataset to verify the model design choices and training strategies. A lower wavefunction error shows model's ability evolve the wavefunctions time while lower error in dipole and absorption shows a model's ''s ability in capturing the underlying physics. The results are summarized in Appendix B We also note that the results of OrbEvo-DM- s8 can be further improved with minor changes in training, we report the training and inference cost, as well the simulation time using the classical solver in Appendix D. out-of-distributio analysis in Appendix [ and time bundling analysis in Appendix J
shown in Appendix E Additionally,
4.3.2 QUALITATIVE RESULTS
We show the computed dipole and absorption spectra produced by OrbEvo-DM-s8 in Figure The plots show that the wavefunctions produced by OrbEvo-DM-s8 starting from ground states reproduce the per-time-step dipole moment with high correlation. The optical absorption produced by the dipole prediction faithfully locate the peaks in the spectra, which provides insightful

---

Published as a conference paper at ICLR 2026
Moleculc
Dipole
Dipole y
Dipolc Z
Absorption
Wwwm
snaud sup
wwwmm
www
sniued much
VE V3 us to IEE S0
Www
wwwwmmm
wwwmmom
Gcosn Tack
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 erB, where rn is Bohr radius (0.529 Å). The unit for absorption spectra 0.529eA* /V. We highlight that there is explicit supervision on dipole absorption during training and validation.
information into the molecular excited states. We also show the wavefunction rollout using OrbEvo- DM-s8 in Figure | which demonstrates the good match against the ground truth wavefunctions. Finally, show plots for MDA dipole and absorption predictions in Appendix E
QM9
MDA
41
61
11
t=
t=
61
Prediction
Ground truth
Difference
Figure 4: Wavefunction rollout using the OrbEvo-DM-s8 model compared with the ground truth.
5 CONCLUSION
In this paper, we propose OrbEvo, which is built upon an equivariant graph transformer architecture. We identify the key issues in modeling inter-electronic-saal interaction and propose to model elec- tronic states separate graphs. We further propose models based on density matrix featurization and full wavefunction pooling interaction. Together with pushforward training, our models can accu- rately learn the wavefunction evolution accurately. Moreover, we show that the density- matrix-based model is able to learn the underlying physical properties without providing explicit supervising sig- nal to the model. However, standard TDDFT faces limitations, such as difficulties in dealing with conical intersection, and its performance is limited by the accuracy exchange-correlanne energy functionals, which remains important direction for future development.
10

---

Published as a conference paper at ICLR 2026
ACKNOWLEDGMENTS
This work was supported in part by the U.S. Department of Energy Office of Basic Energy Sciences under grant DE-SC002386, National Science Foundation under grant MOMS-2331036, National Institutes of Health under grant UOIAG0701 12, and Advanced Research Projects Agency for Health under IAY1AX0000S: We thank Texas A&M HPRC for providing CPU resources.
REFERENCES
Nicholas Boyer, Christopher Shepard, Ruiyi Zhou, Jianhang Xu, and Yosuke Kanai. Machine- Learning Electron Dynamics with Moment Propagation Theory: Application Optical Absorp- tion Spectrum Computation Using Real-Time TDDFT. Journal of Chemical Theory and Compu- tation, 24-1):114-123, 2024.
Johannes Brandstetter, Daniel E. Worrall, and Max Welling. Message Passing Neural PDE Solvers. In International Conference on Learning Representations, 2022. URL https://openreviewwneel
forum""d=wSSxx3PXKS
Mark E Casida. Time-dependent density functional response theory for molecules. In Recent Ad- vances Density Functional Methods: (Part 1), pp. 155-192. World Scientific, 1995.
Lowik Chanussot, Abhishek Das, Siddharth Goyal, Thibaut Lavril, Muhammed Shuaibi, Morgane Riviere, Kevin Tran, Javier Heras-Domingo, Caleb Ho, Weihua Hu, et al. Open Catalyst 2020 (OC20) Dataset and Community Challenges, ACS Catalysis, 11(10)6659966772 2021.
Mohan Chen, G-C Guo, and Lixin He. Systematically improvable optimized atomic basis sets for ab initio calculations. Journal of Physics: Condensed Matter, 22(44):445501, oct 2010, doi: 10.1088/0953- 8984/22/44/445501. URL
Stefan Chmiela, Huziel E. Sauceda, Klaus-Rober Müller, and Alexandre Tkatchenko. Towards ex- act molecular dynamics simulations with machine- learned force fields. Nature Communications, 9(1), 2018. ISSN 2041-1723 doi: 10. (0338//144--018-066666 URL http:://dxxdxoooo.rrr 10.38/ 441447-018-06169-2
Jayesh K Gupta and Johannes Brandstetter. Towards multi-spatittepporal-ccall generalized PDE modeling. Transactions on Machine Learning Research, 2023. ISSN 2835-8856. URL https:
Adrián Gómez Pueyo, Miguel A. L. Marques, Angel Rubio, and Alberto Castro. Propagators for the time-dependent kohn-sham equations: Multistep, runge-kutta, exponential runge-kutta, and commutator free magnus methods. Journal of Chemical Theory and Computation, 14(6):
3040-3052, 2018. ISSN 1549-9626.
doi: 10. 11002//acs.cccc.b0001110 URL http://dx. doi.oog//ooooooo0o
1021/acssjetc.8500199
D. R. Hamann. Optimized norm-conserving vanderbilt pseudopotentials. Phys. Rev. B, 88:085117, Aug 2013. doi: N0.1103/PhyseevB.88.888 URL
88.0851171
Jacob Helwig, Sai Srecharsha Adavi, Xuan Zhang, Yuchao Lin, Felix S Chim. Luke Takeshi Vizzini, Haiyang Yu, Muhammad Hasnain, Saykat Kumar Biswas, John J Holloway, et A two-phase deep learning framework for adaptive time-stepping in high-speed flow modeling. arXiv preprint arXiv:2506.0799 2025.
Maximilian Herde, Bogdan Raonic, Tobias Rohner, Roger Käppeli, Roberto Molinaro, Emmanuel de Bezenac, and Siddhartha Mishra. Poseidon: Efficient foundation models for PDEs. In The Thirty-eightl Annual Conference on Neural Informatinn Processing Systems, 2024. URL https: Snmpenrreeew.n//foouu//kddCC/lC
Hohenberg and W. Kohn. Inhomogeneous Electron Gas. Phys, Rev., -3611184-B 1964. doi: 10.1 103/PhysRev. 136. B864. http:////wiikkaaoooo//////////// B86

---

Published as a conference paper at ICLR 2026
W. Kohn and L. J. Sham. Ceff-Consistent Equations Including Exchange and Correlation Effects.
Phys. Rev. 140:A1133-A111 1965. doi: 10.113/PMyyRee.44..A115.. URL htttph:s///mkkaappp
org//oo/10. 10w.....ss//yyseeee.... 100.1111
Pengfei Li, Xiaohui Liu, Mohan Chen, Peize Lin, Xinguo Ren, Lin Lin, Chao Yang, and Lixin He. Large-scale ab initio simulations based on systematically improvable atomic ba- sis. Computational Materials Science, 112:503-517, 2016. ISSN 0997-0256. doi: https: vidoi.org/10.1011 http://0...111/..oomaasees htt.ccmmmaaaci......00... 0044. URL 0100770055555044444
wttp://wiiieeeeee...ooo/eceeeeee Computational Materials Science in China.
Zongyi Li, Nikola Borislavov Kovachki, Kamyar Azizzadenesheli, Burigede liu, Kaushik Bhat- tacharya, Andrew Stuart, and Anima Anandkumar. Fourier Neural Operator for Parametrie Par- tial Differential Equations. In International Conference on Learning Representations, 2021. URL
Yi-Lun Liao, Brandon M Wood, Abhishek Das, and Tess Smidt. EquiformerV2: Improved Equivariant Transformer for Scaling to Higher-Degree Representations. In The Twelfth Inter- national Conference on Learning Representations, 2024. URL httpss//oonntrrii.nn//foomm
mCOBKZmrZD
Peize Lin, Xinguo Ren, and Lixin He. Strategy for constructing compact numerical atomic orbital basis sets by incorporating the gradients of reference wavefunctions. Phys. Rev. B, 103:235131, Jun 2021. doi: 10. 10.1103/physeev..10003 103.235131. Ut/////ts//inkkaasoo////////1111 103.235131
Peize Lin, Xinguo Ren, Xiaohui Liu, and Lixin He. Ab initio electronic structure calculations based on numerical atomic orbitals: Basic fomalisms and recent progresses. WIREs Computational Molecular Science, 14(1)=e1687, 2024. doi: https://doi.org// http://o.orr11111000////.. 1687. URL https:
//wires.onlineline
www.aarrywwlll......o wwwpwec.oooooooooo//o/////////// wttp:////.aa//....000/////aaamma 1002//cms.s...cc 1687
Phillip Lippe, Bastiaan S. Veeling, Paris Perdikaris, Richard E Turner, and Johannes Brandstetter. PDE-refiner: Achieving accurate long rollouts with neural PDE solvers In Thirty-seventh Con-
ference on Neural Information Prucessing Systems, 2023. URL /tLLlttts//oppeerrvvveeee//oomo
Qv64681IWS
Haggai Maron, Or Litany, Gal Chechik, and Ethan Fetaya. On learning sets of symmetric elements.
In International Conference on Machine Learning, pp. 6734-6744. PMLR, 2020.
Saro Passaro and C Lawrence Zitnick. Reducing SO(3) Convolutions to SO(2) for Efficient Equiv- ariant GNNs. In International Conference Machine Learning (ICML), 2023.
Ethan Perez, Florian Strub, Harm De Vries, Vincent Dumoulin, and Aaron Courville. FiLM: visual reasoning with a general conditioning layer. In Proceedings of the AAAI conference on Artificial Intelligence, volume 32, 2018.
Charles R Qi, Hao Su, Kaichun Mo, and Leonidas J Guibas. PointNet: Deep Learning on Point Sets for 3D Classification and Segmentation. In Proceedings of IEEE Conference Computer Vision and Pattern Recognition, pp. 652-660, 2017.
Xiaofeng Qian, Ju Li, Xi Lin, and Sidney Yip. Time-dependent density functional theory with ultrasoft pseudopotentials: Real-time electron propagation across a molecular junction. Phys. Rev. B, 73:035408, Jan 2006. doi: 10.110:/Physsee1.7...55544 URL htttps://inkk.as..rrg/// 10. 1L11003/hyysee8..3.333344
Raghunathan Ramakrishnan, Pavlo O. Dral, Matthias Rupp, and O. Anatole von Lilienfeld. Quan- tum chemistry structures and properties of 134 kilo molecules. Scientific Data, 1(1), 2014. ISSN 2052-4463. doi: 10. (033//ddaa..2014.20 URL http://dddbbb.or///11100////saa
Erich Runge and E. K. U. Gross. Rev. Lett. 52:997-1000, Mar 1984, doi: 10. 10.110%/PhysseeeLett...9999 URL https://iik..aps.orgg
Theory for Time-Dependen Systems. Phys.
Abl/700.100/PhyyReelee....
12

---

Published as a conference paper at ICLR 2026
Kristof T Schütt, Michael Gastegger, Alexandre Tkatchenko, K-R Müller, and Reinhard J Maurer. Unifying machine learning and quantum chemistry with a deep neural network for molecular wavefunctions. Nature Communications, 10(1):5024, 2019.
Karan Shah and Attila Cangi. Accelerating electron dynamics simulations through machine learned time propagators. In ICML 2024 AI for Science Workshop, 2024. URL https://ppenreviewnnee
forum"td-IsddSIIgHHA
Karan Shah and Attila Cangi. Machine learning time propagators for time-dependent density func- tional theory simulations. arXiv preprint arXiv:2508. /6554, 2025.
Yasumitsu Suzuki, Ryo Nagai, and Jun Haruyama. Machine learning exchange-correlation potential in time-dependent ceenstt--function theory. Physical Review A, 101(5)050501, 2020.
Alasdair Tran, Alexander Mathews, Lexing Xie, and Cheng Soon Ong. Factorized Fourier Neural
Operators. In The Eleventh International Conference on Learning Representations, 2023. URL
http://oppenrevie.nn//ooomm http://wwwwwwwwemmoooooooomm/inn
Carsten A. Ullrich. Time-Dependent Density-Functiooul Theory: Concepts and Applications. Ox-
ford University Press, 12 2011. ISBN 9780199563029. doi: 10. 10033ecpottooo/77779999933333333 001.0001.
Oliver Unke, Mihail Bogojeski, Michael Gastegger, Mario Geiger, Tess Smidt, and Klaus-Robert Müller. SE(3)-equivariant prediction of molecular wavefunctions and electronic densities. Ad- vances in Neural Information Processing Systems, 34:14444-11441 2021.
K. Yabana and G. F. Bertsch. Time-dependent lecal-density approximation in real time: Application to conjugated molecules. International Journal of Quantum Chemistry, (5(1):55666, 1999. doi: https://doi.org/100 1002/((1 X(119946XXX199955555 (55:AID--
Haiyang Yu, Zhao Xu, Xiaofeng Qian, Xiaoning Qian, and Shuiwang Ji. Efficient and Equivariant Graph Networks for Predicting Quantum Hamiltonian. In International Conference Machine Learning. pp. 10412-40422 PMLR, 2023.
Xuan Zhang, Jacob Helwig, Yuchao Lin, Yaochen Xie, Cong Fu, Stephan Wojtowytsch, and Shui- wang Ji. Sinenet: Learning temporal dynamics in time-dependent partial differential equa- tions. In The Twelfth International Conference on Learning Representations, 2024a. URL
Xuan Zhang, Jacob Helwig, Haiyang Yu, Xiaofeng Qian, and Shuiwang Ji. Learning time-dependent density functional theory via geometry and physics aware latent evolution. 2024b. URL https:
Xuan Zhang, Limei Wang. Jacob Helwig, Youzhi Luo, Cong Fu, Yaochen Xie, Meng Liu, Yuchao Lin, Zhao Xu, Keqiang Yan, Keir Adams, Maurice Weiler, Xiner Li, Tianfan Fu, Yucheng Wang. Alex Strasser, Haiyang Yu, YuQing Xie, Xiang Fu, Shenglong Xu, Yi Liu, Yuanqi Du, Alexan- dra Saxton, Hongyi Ling, Hannah Lawrence, Hannes Stärk. Shurui Gui. Carl Edwards, Nicholas Gao, Adriana Ladera, Tailin Wu, Elyssa F. Hofgard, Aria Mansouri Tehrani, Rui Wang, Ameya Daigavane, Montgomery Bohde. Jerry Kurtin, Qian Huang. Tuong Phung, Minkai Xu, Chai- tanya K. Joshi, Simon V. Mathis, Kamyar Azizzadenesheli, Ada Fang, Alán Aspuru-Guzik, Erik Bekkers, Michael Bronstein, Marinka Zitnik, Anima Anandkumar, Stefano Ermon, Pietro Liò, Rose Yu, Stephan Güinnemann, Jure Leskovec, Heng Ji. Jimeng Sun, Regina Barzilay, Tommi Jaakkola, Connor W. Coley, Xiaoning Qian, Xiaofeng Qian, Tess Smidt, and Shui- wang Ji. Artificial Intelligence for Science in Quantum, Atomistic, and Continuum Systems. Foundations and Trends@ in Machine Learning, 18(4):385 912, 2025. ISSN 1935-8245. doi: 10.1561/2200000155 15. URL 1ttp//////////////.oooo//////0/0 httpp//dd..bii.og////.////000000
13

---

Published as a conference paper at ICLR 2026
APPENDIX
A
Related Works
14
Ablation studies
15
Evaluation Metries
16
C.1
Density conservation
16
C.2
Wavefunction Metric
16
C.3
Dipole Moment.
16
C,4
Optical Absorption
17
Computational Cost & Comparison
17
E
Qualitative Results on MDA and Efficient Training
17
F
Dataset Description
18
G
Time Evolution of Kohn-Sham Wavefunctions in RI IDDFT
19
Implementation Details
20
H.l
Tensor Product
20
H.2
Equivariance Test
20
I
Out-of-distribution Analysis
20
J
Time Bundling Analysis
22
K
Prediction of Global Phase
22
Additional Ablations
23
M Model Hyperparameters
24
Large Language Model Usage
24
A RELATED WORKS
DFT surrogate models aim to bypass the expensive self-consistency calculation by directly mapping from inputs to the converged DFT outputs. Hamiltonian prediction models (Schütt al., 2019 Unke et al. 2021 Yu et al. 2023) learn to map from atom types and their 3D) coordinates to the converged Hamiltonian matrix. Equivariant 3D graph neural networks enable effective learning with spherical basis through tensor products, albeit the increased computational complexity. For efficiency, eSCN (Passaro & Zitnick 2023) reduces SO(3) tensor products to SO(2) operations by rotating the relative direction. EquiformerV2 (Liao et al.) (024) incorporates the eSCN convolution into a graph transformer architecture. These models take atom types and coordinates as input. We extend it setting where the input features also high-order equivariant features.
Besides molecules, machine learning has enabled surrogate models for time-dependent PDEs (Li al. 2021 Tran et al. [2023 Gupta & Brandstetter 2023 Zhang al. (024a) for applications such as modeling fluid dynamics. These surrogate models frequently require conditioning on external information, such as force magnitude or time steps Gupta & Brandstetter 2023); Herde al. (2024);
Helwig et al. (2025). PDE surrogate models have also been developed for graph data (Brandstetter al. 2022). where the pushforward trick and temporal bundling were proposed enhance stability over long time-integration periods. We adopt the temporal bundling and apply pushforward training on more realistic 3D graph. While|Lippe al. (202)) showed that the pushforward training may not be helpful in general settings, show that it can be indeed helpful for realistic graph data.
Machine learning TDDFT is relatively under-explored. Suzuki et al. (2020) use neural networks to improve the exchange-correlatio potential in TDDFT. Boyer et al. (2024) learns dipole mo- ments using ridge regression. For time propagation within the ML-PDE paradigm, Shah Cangil 2024 2025) study the evolution of charge density in nne-dimnnsionll diatomic systems. TDDFT- Net (Zhang et al. 2024b) learns the density evolution starting from the ground-state density for complex molecules. To the best of our knowledge, existing work directly addresses the learning
14

---

Published as a conference paper at ICLR 2026
Table Ablation studies on the MDA dataset.
Wavefunction Rollout
Dipole
Absorption
OrbEvo Model
1-step
Rollout nRMSE
nRMSF-all
nRMSF-2
nRMSF-o
f2-MAE
2--MAE
DM-sall
0.0244
0.0997
0. 181888
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()
0.0508
0.2788
0.5457
0,8738
0.6768
0718
DM-s8-onestep
0.0200
0.1501
0.2851
0,4369
0.3386
0.1211
WF-sall
0.0192
0.0853
0.1585
0.3957
0.3066
0,0865
WF-s8
0.0334
0.2074
0.4054
0.6579
0.5218
0.1338
WF-s4
0.0414
0.2527
0.4961
0.7762
0,6104
0.1582
-alssll-onestep
0.0205
0.1978
0.3708
0,7400
0.5754
0,1590
WF--sllliny--ool
0.0224
0.6773
1.1564
1.3405
1.2632
0.1667
of time-dependent wavefunctions, representing critical gap in the field. Here study TDDFT di- rectly in the wavefunction space, which captures the underlying physical process and enables more accurate predictions. The orbital-based representation that adopted also allows for more efficient data encoding.
B ABLATION STUDIES
We conduct ablation studies on the MDA dataset verify the model design choices and training strategies. A lower wavefunction error shows a model's ability to evolve the wavefunctions time while lower error in dipole and absorption shows model's ability in capturing the underlying physics. The results are summarized Table
Electronic states sampling. Models with suffix "-all" use all electronic states during training. Mod- els end with "-s8" and "-s4" randomly sample 8 and 4 electronic states during training, respectively. The results show that the sampling does not affect OrbEvo-DM's performance while it degrades the performance of OrbEvo-WF significantly. It shows that by aggregating the electronics state infor- mation early via density matrix can effectively capture the inter-eleetronic-state interaction. The OrbEvo-WF results show the importance of considering all electronic states' information.
Electronic state graph construction. In DM-sall-cat, we concatenate wavefunctions from all elec- tronic states along the channel dimension at model's input instead of considering them as individual graphs. The result shows that the model cannot learn the wavefunction mapping correctly, demon- strating the importance of our graph modeling method.
Density matrix ablation. In DDM-s-no-dmm(), we remove the dependency on the time-evolving density matrix. The results show that the model cannot learn correctly, showing the importance of time-evolving density learning the propagation.
Training strategy. We show the results without using pushforward for DM-s8-onester and WF- sall-onestep. The results show that although the models are able learn the one-step mapping more accurately, the rollout error is significantly worse, showing importance of pushforward training for learning error accumulation during rollout.
Equivariant conditioning, In WF-sall-in-condd, we disable the equivariant electric field condition- ing and add the bias term into the invariant = part instead. The results show that although the one-step error go down normally, the rollout does not work. We observe that the model cannot learn the mapping from the initial ground state to the first step correctly, although it is able to evolve the subsequent steps given the ground truth.
15

---

Published as a conference paper at ICLR 2026
C EVALUATION METRICS
C.I DENSITY CONSERVATION
The propagation of electronic wavefunctions conserves the total density. The square of electron density for electronic state n is computed as
C (t) R,
(12)
which is equal to where S is the overlap matrix. We normalize the predicted coefficients C, (t) as
(t
(t) = cl,ttsc,,(t
(13)
We note that the global phase Yn (f) in Equation 3 cancels out in the product.
C.2 WAVEFUNCTION METRIC
We report the f2- error (Equation (00) for the time-dependent wavefunctions. For a more interpretable metric, we also report the normalized rooted mean square (nRMSE) error, defined for each molecule as
NNos
Σ
Nun o=1
Cpeed
barget
nRMSE Cored
Charget)
1=
(14)
Σ
Noc
ΣE=
T
Nab
Caarge
where Nov and Nor denote the number of occupied electronic states and local atomic orbital bases in the molecule.
C.3 DIPOLE MOMENT
Dipole moment describes the density distribution over spatial directions and are defined as (w))). where f' is the position operator along m E {x, y. direction. With the local atomic orbital basis, given the position matrices for three spatial directions Um.ij = (0x (Aly)) Non the dipole moment of each molecule can be computed as
(t) Σ The (t)) m f2,
(15)
where density conservation is applied to the unrolled wavefunctions post-processing step prior to computing the dipoles. We are interested the dipole difference against time Apm (t) (m.(t) P.. (0), =). We report the nRMSE of the dipole moment for all directions.
defined as
V ET= Eme(umme) (m.me) Apmmrd (t) Ap
target pm
(t)
nRMSF-all pApd
Ap target)
(16)
V Eme(xan.))
(Apir
(t)) mm
larget
as well for z direction, defined
V Et= Ap pred (t)
App
targer (t)
aRMSE-z (Apred Ap Larget)
(17)
DT (Ap.
targer (t))
16

---

Published as a conference paper at ICLR 2026
C.4 OPTICAL ABSORPTION
Optical absorption is an important physical property which reflects the ability of molecule to absorb light specific frequencies. is characterized by dipole oscillator strength which be calculated from the time-dependent dipole moment in response to the applied external electric field as follows:
[ Pt) de [ (t)eftt
Ogg (w: Im
(18)
We report the nRMSE for the dipole oscillator strength along the direction, defined as
V EEw
a0 pred (1) Oz
target (1)
2
nRMSE-a (a!td) a²
pred Farget)
(19)
VEw (al targoo (cl2r)((ww)) (w))
D COMPUTATIONAL COST & COMPARISON
In this section we report the training (Table 4 and inference cost of OrbEvo (Table 5). report the simulation time with the classical solver ABACUS (Table
We also
Dataset
Model
# iterations
GPU
Wall Clock Time
GPU Memory (MB)
MDA
OrbEvo-DMr
300k
2x 1IGB 2080Ti
3,475 days
13,848
MDA
OrbEvo-WFF WF
300k
2x 11GB 2080Ti
3.345 days
14,248
QM9
OrbEvo-DM-s8
395k
4 48GB A6000
3.118 days
49,434 54,700
QM9
OrbEvo-WF WF
395k
x 80GB A100
5.003 days
46,652 69,662
Table 4: Training cost of OrbEvo models. MDA models are trained with batch size 32. QM9 models use batch size of 16. All models are trained with Pytorch distributed data parallel (torch. ddp) for multi-gpu training and with num workers=1e 6 in dataloader for MDA and num workers=32 for QM9, As rough estimation, 2x 2080Ti is roughly equivalent in terms of speed. The GPU memory usage is tested by running training on 1 single A 100 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-forwane training implementation, we able to fit the training within x A6000 GPUs for both OrbEvo models on the QM9 dataset with similar training time.
A6000
Wall Clock Time Batch Wavefunction
GPU Memory (MB)
Dataset
Model
GPU
Batch Size
Wavefunction
Property
MDA
OrbEvo-DMr
1x A6000
3.67 seconds
5.23 seconds
5742
MDA
OrbEvo WF
A6000
seconds
4.60 seconds
2032
QM9
OrbEvo-DM
A6000
20
18.00 scconds
26.74 seconds
34,164 -2,,8442
QM9
OrbEyo-WF
A6000
11.86 seconds
20.31 seconds
17,204
Table 5: Inference cost of OrbEvo models. All models are tested one single A6000 GPU using nam_workers=11 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 increased GPU memory usage for OrbEvo-DMr In comparison, during training OrbEvo-DM is able to use electronic state sampling to reduce GPU usage.
E QUALITATIVE RESULTS ON MDA AND EFFICIENT TRAINING
We show the dipole and absorption produced using the predicted wavefunctions on MDA samples in Figure where train the OrbEvo-DM-s8 model using fosw-forward training with some minor
17

---

Published as a conference paper at ICLR 2026
Wall Clock Time Molecule Ground-state DFT
Dataset
# CPU cores
Total
MDA QM9
24
34.3 seconds
1.5 hours
73.1 seconds
3.2 hours
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 struc- tures. The initial wavefunction coefficients are used as input to OrbEvo models.
changes compared to the models in the main text: (1) We disable the linear warm-up factor and always enable pushforward with 50% probability. (2) We switch the model evaluation mode during push-forward unrolling. This could make the push-forward noise closer to the real rollout during test. (3) We fix the number of push-forward samples be half of the per-GPU batch size to have a more stable GPU usage. The result of models trained in this way is summarized in Table 7 On the other hand, we observe that such changes in training may be helpful for OrbEvo-WF.
Molecule
Dipole
Dipole
Dipolez
Absorption
Grand
Siarid Buar
Gecule
Porduation
wwwmmmn
wwwmmmmn
Mm Ammmmm
www.mmm
Grand Nue
Rec brct
M w mwdgmma
Figure 5: MDA dipole and absorption with the OrbEvo-DMr model on test samples. The unit for dipole in the plot is erb, where TB is Bohr radius (0.529 Å). The unit for absorption spectra is 0.529eA* /V.
Table 7: Results on the MDA dataset with the new training.
Wavefunction
Dipole
Absorption
OrbEvo Model
1-step f2-MAE
Rollout
Rollout nRMSE
nRMSE-al all
nRMSE-2
nRMSE-o
2gMMAE
DM-s8 WF-sall
0.0224 0.0225
0.0863 0.1080
0.1613 0.2008
0.1997 0.3758
0.1499 0.2881
0.0539
0.0822
F DATASET DESCRIPTION
The molecules and their configurations used in this work were sourced from the QM9 (Ramakrish- nan et al.) 2014) and MD17 databases (Chmiela et al,) 2018). The QM9 dataset contains a large number of chemically diverse molecules. This combination allows our model to cover wide range
18

---

Published as a conference paper at ICLR 2026
of potential molecular behaviors and properties. The MD17 dataset provides high- regh-resolution molec- ular dynamics trajectories for small number of molecules with many different conformations. Both QM9 and MD17 are widely used in machine learning for materials science and computational chem- istry. For this work, we randomly chose 5, 000 different molecules from the QM9 dataset consisting of C, H, O, and N elements demonstrate the generalization capability of our model, and randomly selected 1, 500 molecular configurations of the malonaldehyde (MDA) molecule from the MD17 dataset for the ablation study.
To generate the RT-TDDFT datasets for the above QM9 and MDA molecules, we utilized the open- source ABACUS software package (Chen et al. 2010 Li et al. 2016 Lin et al. 2024) to perform the DFT and RT-TDDFT calculations. Consistent imput parameters were used to ensure comparabil- ity between datasets. Specifically we employed the SG15 Optimized
Norm-Conserving Vanderbilt (ONCV) pseudopotentials (SG15-V1,, (Hamann 2013), standard atomic orbitals basis set hierar- chically optimized for the SG15-V1.0 pseudopotentials (Lin et al. (02T), and kinetic energy cutoff' of 100 Rydberg. The ground-state Kohn-Sham wavefunctions were obtained by self-consistent field (SCF) calculations of DFT with a dimensionless convergence threshold 10-6
For RT-TDDFT calculations, we used ground-state Kohn-Sham wavefunctions as the initial states at = and performed time propagation for 5 fs in a total of 1, 000 steps with a time step of 0,005 fs. To simulate the quantum dynamics of the system under an external field, time-dependent uniform electric field E3 (f) was applied along the z direction:
(t 10)² 2o
(1) Eo (cos(2m (n) cos (20 /2(1 (o))) exp
It consists of two frequencies of fi 3.66 fs and f2 = 1.22 fs with a Gaussian width 0.2 fs, a field amplitude Eg = 0.01 V/A, and a central time of to = 0.75 fs. During each time step, wavefunction coefficient matrices saved and then extracted, serving input data for model training, validation and testing.
To enhance computational efficiency and accuracy, modified the ABACUS source code caleu- late the overlap matrix only once at O. Furthermore, we ensured that the output matrix retained 16 significant digits of precision. This modification allowed to generate reliable data with greater efficiency, making it well suited for model training and testing. The DFT and RT-TDDFT calcula- tions were performed using 24 parallel CPU cores.
G TIME EVOLUTION OF KOHN-SHAM WAVEFUNCTIONS IN RT-TDDFT
In RT-TDDFT, each Kohn-Sham wavefunction th evolves in time under the time-ordered evolution operator (t, to). starting from the initial time ty: vb,(t) = (t, to to)) where
(-iS S-1 S. ((()))t))
to) = -tt,to)=Tex
Í is time-ordering operator. In RT-TDDFT, total simulation time Tex is discretized into Nox steps with each time step of At Tia/No /Nw, and (t, fo) is approximated by the product of evolution operators over the discretized time grid (Gómez Pueyo et al.) 2018).
No
Ú(1,10) = II m=1
Ü/6o + mAl, + (m 1)
In general, U[t to (m-1)At) should satisfy the unitary condition conserve the density:
Ü' + m&t, to + (m - 1) = Ú-1 [to solids under external electric field, it should satisfy time-reversal symmetry: Ütto +mtt, to + +mm- 1)At] = Ütt + (m -1) to + mdt). Such time evolution needs to be applied to all occupied electronic states for N, time steps, making it computationally demanding.
m&t, + (m - 1) Moreover, for molecules and
19

---

Published as a conference paper at ICLR 2026
Table Performance on QM9 .odd validation and test sets.
Model
Datasel
1-step Ê2-MAE
Rollout nRMSE
Dipole a nRMSE
Absorption nRMSE
OrhEvo-DMr DM-s8-ood
QM9 ood - val
0.0142
0.1498 0,1482
0.1113 0.1175
0.0615
OrbEvo-DMi DM-s8-ood
QM9 pod test
0.0132
0.0697
H
IMPLEMENTATION DETAILS
H.
TENSOR PRODUCT
In 6 we visualize the tensor product for computing density matrix feature from wavefunction, which is implemented using enn. o3. FullTensorProduct.
28
x De
2 00
2e
WWSASBSABDOCCCOOO XX XXXXX XXXX XX
Figure 6: Tensor product visualization produced by the e 3nn library.
H.2 EQUIVARTANCE TEST
In this section we test the SO(())equivariaace error for both the TDDFT numerical simulation and the OrbEvo model.
In Figure 7 we run two simulations using ABACUS with original rotated molecule. Figure R we use the model to make predictions using inputs before and after rotation. In both cases we rotate around the electric field direction by 35 degree and we conduct manual rotation-transform to align the resulting coefficients or to produce rotation-transformmed input. When applying the rotation transformation to the coefficients, s orbitals and m = components in p and orbitals remain unchanged, +1 components in p and d orbitals are rotated by 35 degree around the electric field direction, and m = +2 components in orbitals are rotated by 70 degree around the electric field direction.
I OUT-OF-DISTRIBUTION ANALYSIS
Using our generated 5000 QM9 data, we created an OOD split for it based on the number atoms in molecule. In particular, use the number of atoms from 8 to 20 for training (3955 samples), number atoms = 21 for validation (518 samples), and number of atoms from 23 to 29 for testing (527 samples). We train the OrbEvo-DM-sS model using this OOD split, the trained model is dubbed OrbEvo-DM-s --oodo In the Table we evaluate the trained model on the validation and test sets of the OOD split.
20

---

Published as a conference paper at ICLR 2026
0.001
65
Data
Equivariance Error
0.001
10
20
30
40
50
00
80
10
20
30
40
50
00
70
80
00.001
-10
-0.003
Figure 7: Equivariance enror of TDDFT data. Left: real part of the wavefunction coefficients of an unrotated MDA molecule at time step. Right: the difference between the wavefunctions the same time step in a second simulation produced from a rotated version of the same molecule, and the coefficients manually ination-transfommen from the left plot. In the second simulation the molecule is rotated by 35 degree around the electric field direction.
1.0
0.0000
Prediction
Equivariance Error
0.0002
0.0000
10
10
-0.0002
20
50
20 50
00
000000
-10
0..000
Figure 8: Equivariance еггоr of OrbEvo-DM. Left: real part of the model's predicted wavefunction coefficients for MDA molecule using the ground-truth wavefunctions one time step input. Right: the difference between the model's predicted wavefunctions using the rotated structure and manually rotation-transfomen 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.
Rather surprisingly, the OOD split gives better average test accuracy than the random split reported in our main paper (Rollout nRMSE 081555, Dipole nRMSE 0. 1459, Absorption nRMSE 0.0752). However, the numbers are not directly comparable since different test data are used. To fairly compare the ID and OOD performance, data that are both in the OOD split and the random split. Specifically, there are 67 molecules in the intersection of the OOD validation set and the random validation and 43 molecules in the intersection of the OOD test set and random test set (QM9 id ood intersect test). We report the results of both the model trained on the OOD split (OrbEvo-Day Sar--v---------o-- and the model trained using the random split (OrbEvo-DMe in the main paper) in Table
evaluate the models the validation and test
(QM9 id ood intersect val),
Despite the randomness due to the small amount of common validation test data, we can see that overall the model trained the random split performs closer to the model trained the OOD split. Overall, the above results show that the OrbEvo model able to generalize on larger systems than those in the training data. The results also suggest that larger systems may not necessarily be more challenging for learning the dynamics of wavefunctions, potentially due to the fact that smaller systems exhibit dynamic patterns with larger magnitude
complex behaviors.
21

---

Published as a conference paper at ICLR 2026
Table 9: Performance comparison on QM9 id ood intersect validation and test sets.
Model
Datasct
1-step 12-MAE
Rollout nRMSE
Dipole : nRMSE
Absorption nRMSE
OrbEvo- OrbEvo-DM-sS-cod
jd__od_intersect- QM9 id ood intersect - val
QM9_id 1199iid_oo_iii. Lintersect-
0.0141
0.1500
0.1126
0.0633
OrbEvo-DM-s8
0.0146
0.1579
0.1153
0.0628
ShhEvo-DM-s8-oon
QM9
intersect test
0.1496
0.1142
0.0666
OrbEvo-DM-s8
QM9 iddoodd.dddoooooooooooo.olo. intersect test
0.0139
0.1521
0.1074
0.0636
J TIME BUNDLING ANALYSIS
We additionally conduct an experiment on the MDA dataset with various time bundle sizes 1, 2, 16 and with 8 (used in the main paper) where time bundle size 1 corresponds without time bundling. We report 1-step error (average error for the time bundle), rollout errors of different trajectory lengths (start from time 0, produce trajectories with lengths 8, 16, 32, 64, 100, larger bundle needs less autoregressive steps), as well as relative errors for dipole- and absorption on the full rollout. The test results are summarized in Table
Table 10: Time bundling analysis on the MDA dataset.
Time bundle
I-step Fo- MAF
8-step Rollout aRMSK
16-step Rollout aRMSE
32-step Rollout nRMSE
64-step Rollout nRMSE
100-step Rollout nRMSE
Dipole
Absorption
size
nRMSE
nRMSE
0.0780
0,9032
2
0.0130
0.0668
0.0340
0.11006
0,3087
0.5765
0.5669
0.1228
0.0139
0.0693
0.0289
0,0572
0.1235
0.1979
0.2670
0.0758
0.0242
0.0720
0.0378
0.0677
0.1245
0.1778
0.2328
0.0672
0.0588
0.1026
0.0544
0.1075
0,2040
0.2872
0.2641
0.0922
We observe that using a smaller time bundle size results in smaller 1-step error. This is because the model needs to predict fewer steps, and closer steps time are easier to predict. For rollout errors, time bundle sizes of and can produce correlated rollouts at 16 steps, but start to diverge for 32 or more steps. Time bundle size of 4 performs well for 32 steps, but becomes less effective than time bundle size 8 for longer rollout. Time bundle size produces the best wavefunction in the overall rollout. Time bundle size of 16 stays stable but the accuracy not as good as or 4.
In terms of training, we observe that time bundle size 1 and 2 start to overfit to onestep mapping and the validation rollout errors start to overfit during the training, with the best validation 100-step rollout error of around 0.4 occurs around 100k iteration for time bundle size 1, and around 0.6 at around I 20k iterations for time bundle size 2 (total training iteration is 300k). Moreover, we observe some oscillations in the validation rollout curves for time bundle size 4 during training.
Note that during training we randomly sample from all 100 steps for feasible starting time for on- estep training. So the total number of training pairs are similar for different time bundle sizes. Overall, a time bundle size of 8 remains a reasonable choice, in which case the model needs to unroll 13 steps produce the entire 100-step trajectory.
K PREDICTION OF GLOBAL PHASE
We re- re-purpose the OrbEvo-DM- model to predict the global phase y(t). In particular, we replace the wavefunetion readout block with a feedforward network. During training, the model takes the ground truth wavefunction coefficients and the global phases in the current time bundle input, and predicts the global phases at the next time bundle. The real and imaginary parts are predicted as 2D vector for each electronic state. We use the MAE as loss function. For MDA, the model is first trained for 100k iterations and subsequently resumed training for 200k iterations. For QM9, the model is first trained for 140k iterations and subsequently resumed training for 186k iterations. We show examples of predicted rollout for QM9 in Figure 9 and for MDA in Figure I0 We observe that the predicted global phases are in good agreement with the ground truth.
22

---

Published as a conference paper at ICLR 2026
Wevefunction #1
Wovelunction #1
1,00
1.40
0.75
0.75
0.50
0.50
0.25
0.25
0.00
0.00
0.25
0.25
0.50
-4.50
-0.25
-0.75
20
40
60
100
50
60
100
Wavefunction v112
Wevefunction #112
1,00
1.00
0.75
0.75
0.5
0.50
0.25
0.25
0.00
0.00
-025
-0.25
-0.50
-0.50
-0.75
-0.75
-1.00
20
40
1.00
90
100
Wavvffanction 4223
Wavefuunction #233
1.00
1.40
0.75
0.75
0.50
0.50
0.25
0.25
0.00
0.00
0.25
0.25
0.50
0.50
0.75
0.75
-100
1.00
20
40
1.00
60
80
100
time step
time step
Figure 9; Global phase rollout on the QM9 sample.
L ADDITIONAL ABLATIONS
We study the effect of keeping the quadratic term delta wavefunctions the density matrix calculation in Table II well the effect of replacing push-forward training with noise injection Table
Table 1I: Density matrix analysis on the MDA dataset.
Wavefunction Rollout f-MMAE
Dipole
Absorption
OrbEvo Model
1-step 2--MAE
Rollout nRMSF
nRMSF-all
nRMSF-2
nRMSF-a
DM-s8
0.0242
0.0947
0.1778 0.2088
0.3012 0.3538
0.2329 0.2744
0.0672
DM-s8-w/ w/---uudrrrtti-dm
0.0290
0.1110
0.0784
Table 12: Noise injection results on the MDA dataset.
Wavefunction Rollout f2-MAE
Dipole
Absorption
OrbEvo Model
1-step f2-MAE
Rollout nRMSE
uRMSE-all
nRMSE-2
uRMSE-a
DM-s8-noise Pool-sall-noise
0.0204 0.0155
0.1262
0.2423 0.1617
0.3868 0.4045
0.3036 0.3157
0.0815
0.0866
0.0788
23

---

Published as a conference paper at ICLR 2026
Wevefunction #1
Wovelunction #1
1,00
1.40
0.75
0.75
0.50
0.50
0.25
0.25
0.00
0.00
0.25
0.25
-0.50
-4.50
-0.25
-0.75
20
40
60
60
100
20.
40
50
100.
Wevefunction #1
Wevefunction #7
1,00
1.00
0.75
0.75
0.0
0.50
0.25
0.25
0.00
0.00
-025
-0.25
-0.50
-0.50
-0.75
-1.00
20
40
1.00
90
100
Waveofanction 4014
Waveffunction di
1,00
1.00
0.75
0.75
0.50
0.50
0.25
0.25
0.00
0.00
0.25
-.25
0.50
0.50
0.75
0.75
-100
1.00
20
40
1.00
60
300
time step
time step
Figure 10: Global phase rollout on the MDA sample.
M MODEL HYPERPARAMETERS
We summarize OrbEvo's hyperparameters in Table 13 Most of them hyperparameters for the EquiformerV2 (Lioo et al. 2024) backbone.
N LARGE LANGUAGE MODEL USAGE
We use large language models to aid or polish writing sparsely. LLMs are also used lightly to help write data processing scripts.
24

---

Published as a conference paper at ICLR 2026
Hyperparameters
Value
Optimizer
AdamW Cosine Anncaling
Learning rate scheduling
Maximum learning rate
1 10-3
Weight decay
1 x 10-3
Number of epochs for MDA
129 (300k iterations)
Number of epochs for QM9
17 (395k iterations)
Maximum cutoff radius
5.0
Number of layers
6
Number of sphere channels
128
Number of attention hidden channels
128
Number of attention heads
8
Number of attention alpha channels
32
Number of attention value channels
16
Number of FFN hidden channels
512
mmxx list
[4]. [2]
Mmax list
[4]. [2]
Grid resolution
eSCN default
Number of sphere samples
128
Number of edge channels
128
Number of distance basis
250
Alpha drop rate
0.1
Drop path rate
0.05
Projection drop rate
0.0
Number of future time steps
8
Number of conditioning time steps
8
Table 13: OrbEvo model hyperparameters.
25


Reviews:
[{"review_id": "BXaxWyNQ81", "reviewer": "Reviewer_GR5K", "summary": "This paper introduces OrbEvo, an equivariant graph transformer framework for learning the time evolution of Kohn–Sham wavefunctions in real-time time-dependent density functional theory (RT-TDDFT). Unlike prior works such as OrbFormer, which focus on static ground-state properties, OrbEvo aims to learn the dynamics of electronic states under external electric fields.\nThe authors propose two model variants: OrbEvo-FullWF, which aggregates wavefunction features through pooling across occupied states, and OrbEvo-DM, which computes density-matrix-based interactions between states via tensor contraction. The model employs SO(2)-equivariant conditioning to represent field-induced symmetry breaking and a pushforward training scheme to stabilize long-horizon rollout. Experiments on QM9 and MD17 demonstrate that OrbEvo-DM outperforms the pooling-based variant, capturing physically consistent time-dependent dipole moments and absorption spectra.", "questions": "1. **Applicability to static DFT**: Can the OrbEvo architecture be used for static ground-state DFT tasks, such as predicting stationary wavefunction coefficients or density matrices? Or is it strictly limited to time-dependent TDDFT propagation?\n\n2. **Ablation on time bundling**: How much does the time bundling technique contribute to the model’s performance and stability? It would be helpful to see results comparing models with and without time bundling.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "**Easy to follow and clear motivation**: The paper provides an intuitive and well-structured explanation of the challenge in modeling time-dependent DFT, with smooth transitions from motivation to formulation.\n\n**Tackling an impactful and general problem**: The work addresses a scientifically meaningful and practically impactful challenge: learning the time evolution of Kohn–Sham wavefunctions to accelerate quantum dynamics simulations. Importantly, the model is trained across thousands of diverse molecules (QM9) and demonstrates generalization to unseen molecular systems, showing potential as a shared, cross-molecular surrogate model for electronic dynamics. This highlights the scalability and versatility of the approach beyond single-molecule modeling.\n\n**Sound model design**: The use of SO(2)-equivariant conditioning, density-matrix features, and push forward training demonstrates strong physical insight and solid engineering.", "weaknesses": "Dataset scale is limited, making it difficult to assess generalization to larger molecules or different field conditions.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "0HKETKMwSW", "reviewer": "Reviewer_F1P1", "summary": "The paper proposed a new model and method that learns the time-dependent DFT's properties, and has shown that the new proposed method, combined with a serious method improvement, can predict nicely the properties from TDDFT.", "questions": "1. In real applications, we often focus on a specific kind of system, where we often do not have that much data. In terms of this, how is this method transferable to similar structures but with limited training data? How is the data efficiency in that?\n\n2. How to transfer the model to different external field conditions?\n\n3. How is the SO3-SO2 mapping achieved in layer norm? I haven't found it in the paper.", "limitations": "", "rating": 4, "confidence": 5, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The network is well designed with very good respect for the physics.\n2. The increase in accuracy is very impressive.", "weaknesses": "1. The writing needs more clarification, for example, the model design, more words in the caption and the method part can be helpful to understand it clearly.\n2. The experimental work is comparably weak. More comparison with other models, and specific studies on in-distribution, out-of-distribution systems can make this study more solid.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "bgy26jEHMo", "reviewer": "Reviewer_6czi", "summary": "This paper proposes *Orbital Transformers*, an equivariant graph Transformer designed to directly predict the *time evolution of Kohn–Sham wavefunctions* in real-time time-dependent density functional theory (RT-TDDFT). Unlike prior approaches that predict energies, Hamiltonians, or spectral observables, this model learns the mapping $C(t) \\to C(t+\\Delta t)$ (or $\\Delta C_t$) directly, effectively learning the quantum propagation operator. The authors introduce an SO(2)-equivariant attention mechanism that takes the external electric field direction as the reference axis, and use FiLM-style conditioning to inject both the field’s direction and time-dependent amplitude. A local autoregressive temporal modeling scheme, along with pushforward training, enables the model to track the dynamic evolution of the system stably over several femtoseconds. Experiments on RT-TDDFT trajectories of QM9 and MD17 molecules under external fields show that the model accurately reproduces dipole dynamics and orbital evolution.", "questions": "The authors generate and RT-TDDFT trajectories with external fields for standard datasets like QM9 and MD17, which may enrich the data resources for time-dependent electronic-structure ML. Will the generated trajectories and other data be publicly released, providing valuable data for future time-dependent quantum ML research?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. This is the first work to *directly predict the time evolution of wavefunctions* in RT-TDDFT rather than energies or Hamiltonians. The idea of learning an implicit propagation operator $f_\\theta: C(t)\\mapsto C(t+\\Delta t)$ is conceptually novel and impactful.  \n2. Modeling system evolution under an external field is *physically meaningful* and directly corresponds to realistic nonequilibrium dynamics.  \n3. The paper presents an *innovative autoregressive temporal modeling* strategy that allows the neural network to continuously track and predict the system’s electronic evolution, combining local ΔC prediction with pushforward training to reduce error accumulation.  \n4. The *integration of field information into the atomic graph network* is well designed: the electric-field direction defines the SO(2) equivariant reference axis, and the field information enters through FiLM-style modulation. This construction is both physically grounded and computationally efficient.", "weaknesses": "1. The network predicts the next-step evolution solely from the current state, **and although the authors adopt certain stabilization strategies** — such as local autoregressive temporal modeling, pushforward training, and ΔC prediction — these mechanisms only address short-term error accumulation. **There remains no explicit architectural component (e.g., temporal attention or recurrent memory) to model long-range temporal dependencies**, which consequently limits robustness and stability during long-horizon propagation beyond several femtoseconds.\n\n2. The SO(2)-equivariant design relies on the external-field direction as the rotational reference axis. Once the external field vanishes, this reference loses physical meaning, and the network no longer has a well-defined axis for the equivariant operations. Moreover, if the external-field direction varies over time or differs across samples, the reference frame for the SO(2) operations also changes with time. It is unclear how well the model would perform under such cases.\n\n3. The paper appears to **lack systematic efficiency evaluation experiments**, for example, comparisons on runtime and computational resource usage, which would be important to assess the practical utility of the proposed model.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "nfuK4RDRJ7", "reviewer": "Reviewer_ce9Z", "summary": "This paper tackles the important and challenging problem of accelerating Real-Time TDDFT (RT-TDDFT) computations using deep learning.  \nSpecifically, it adopts an autoregressive framework to accelerate the propagations of RT-TDDFT, where the wavefunctions of previous steps are input into the network for the prediction of the next steps' wavefunctions. The paper proposes two model architectures (OrbEvo-FullWF and OrbEvo-DM) with different electronic state interacting strategies and compares their performance on their self-generated TDDFT dataset.\n\nI think the paper is in a good shape, with nontrivial contributions for a novel application (RT-TDDFT) and specifically designed models (OrbEvo-FullWF and OrbEvo-DM). Nonetheless, there exist several concerns, which should be addressed before acceptance.", "questions": "* About motivation: Could you provide a clearer picture of the framework? For example, how is the learned model used in RT-TDDFT to accelerate the computation?\n* About performance: How much acceleration is achieved? How should one interpret the metrics in Table 1 and Table 2?\n* How does the model compare to direct property prediction methods?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 4, "strengths": "* Machine learning TDDFT is an important and relatively under-explored research field. This paper approaches the problem with a novel setting (directly learning the wavefunctions), representing a pioneering attempt in this direction. The handling of the high-dimensional orbital coefficient object (the time, the electronic state and the atomic orbital dimensions are all different from the number of atoms) is especially interesting.\n* The design choices of the prediction target and the model architectures seem to be thoroughly considered and physics-grounded.\n* The practices of improving model performance, such as the delta transformation of wavefunctions, the time bundling setting and the push-forward training, are detailed and provide a good impression of adopting state-of-the-art techniques in related fields.\n* The paper is written in a clear and easy-to-understand language, with moderate background knowledge included.", "weaknesses": "* Although the paper has provided some background information for RT-TDDFT, I find myself unclear about the big picture and lost in the implementation details. It is not clear how the learned model is used in the RT-TDDFT framework.\n* The units are missing from all the performance numbers in the paper.\n* Neither the source code nor the dataset are provided. This renders the paper hard to follow.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)