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

2026-03-03

AI Summary

This work proposes OrbEvo, a novel model based on an equivariant graph Transformer architecture, to address the high computational cost of conventional real-time time-dependent density functional theory (TDDFT), which requires fine time-step propagation of all occupied states. OrbEvo directly learns the time evolution of electronic wavefunction coefficients in an atomic orbital basis. By incorporating an equivariant external field encoding that reduces SO(3) symmetry to SO(2), and introducing two mechanisms—OrbEvo-WF and the more efficient OrbEvo-DM, which models the time-evolution operator via tensor contraction of the density matrix—the model achieves accurate predictions of field-driven time-dependent wavefunctions, dipole moments, and optical absorption spectra on QM9 and MD17 datasets. Combined with wavefunction pooling, density matrix eigenfeature encoding, and an autoregressive training strategy, OrbEvo maintains high accuracy while significantly reducing computational cost.

Technology Category

• Machine Learning: Quantum Machine Learning
• Search and Optimization: Evolutionary Computation
• Computer Vision: Diffusion Models for Vision

Application Category

• Graph Algorithms and Modeling for the Web: Graph embeddings and representation learning for Web-related graphs
• Economics, Online Markets and Human Computation: Architectures and workflows that use LLMs for crowd work
• Web Mining and Content Analysis: Models for Web evolution

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.

Problem

• time-dependent density functional theory
• wavefunction evolution
• real-time TDDFT
• electronic dynamics
• external excitation

Innovation

• equivariant graph transformer
• time-dependent density functional theory
• wavefunction evolution
• density matrix encoding
• external field conditioning