AI Learns Quantum Chemistry Functions, Boosting Material Design and Discovery

Scientists are continually seeking to improve the accuracy and efficiency of density functional theory (DFT) and linear-response time-dependent density functional theory (LR-TDDFT) calculations. Xiaoyu Zhang from Peking University, alongside colleagues, present a novel end-to-end differentiable workflow for optimising a single deep-learned energy functional using data from both Kohn-Sham DFT and adiabatic LR-TDDFT. This research is significant because it allows for gradient-based training through the self-consistent field (SCF) fixed point and the Casida equation, yielding a consistent potential and response kernel via automatic differentiation. By learning an exchange-correlation functional on the helium spectrum, incorporating self-interaction cancellation and the Lieb-Oxford inequality, the team demonstrate a proof of concept and assess its potential for broader molecular applications.

This breakthrough addresses a long-standing limitation in computational chemistry, where the accuracy of these methods is constrained by approximations in the exchange-correlation functional.

The research introduces a novel approach to simultaneously control energies, self-consistent potentials, and linear-response kernels within a unified framework. Implemented within a JAX-based two-component quantum chemistry code named IQC, the learned functional facilitates gradient-based training through both the self-consistent field fixed point and the Casida equation.
A key innovation lies in the end-to-end differentiability of the entire process, enabling the automatic calculation of gradients necessary for optimising the deep-learned functional. Conventional quantum chemistry packages lack this inherent differentiability, requiring manual derivation of analytical gradients which is both complex and prone to error.

The team overcame this challenge by constructing a fully differentiable system, allowing for robust and stable training of the functional using gradient-based optimisation methods. This approach moves beyond iterative training schemes and direct potential training, offering a more comprehensive and efficient solution.

As a proof of concept, the researchers trained the exchange-correlation functional on the helium spectrum, incorporating one-electron self-interaction cancellation and the Lieb-Oxford inequality as penalty terms. The resulting functional yields a consistent potential and linear-response kernel via automatic differentiation, demonstrating its ability to accurately predict both ground-state and excited-state properties.

Furthermore, initial assessments suggest the potential for transferring this learned functional to more complex molecular test cases. This work represents a significant step towards data-driven functional development, explicitly including linear-response excitation information in the training objective.

By optimising the functional end-to-end, the study bypasses limitations of traditional parameterisation methods and opens new avenues for improving the accuracy and transferability of DFT and LR-TDDFT calculations. The development of IQC, a fully differentiable quantum chemistry package, provides a powerful platform for future research in this area and promises to accelerate the discovery of improved exchange-correlation functionals for a wide range of applications.

Implementation of a differentiable two-component density functional theory and linear-response framework

A JAX-based two-component code, IQC (intelligent quantum chemistry), underpins the research, enabling a fully differentiable workflow for optimizing a deep-learned energy functional. This approach circumvents the limitations of conventional, non-differentiable quantum chemistry packages by facilitating gradient-based training directly through the self-consistent field (SCF) fixed point and the Casida equation.

The study learns an exchange-correlation functional using targets from both Kohn-Sham density functional theory (DFT) and adiabatic linear-response time-dependent density functional theory (LR-TDDFT) within the Tamm-Dancoff approximation. The methodology employs two-component formulations of DFT and LR-TDDFT, utilising Γ, Λ, Θ for two-component atomic orbital basis functions and P, Q, R for two-component molecular orbitals.

Molecular orbitals are expressed using expansion coefficients, defined as P = CΛPΛ, and the first-order reduced density matrix is calculated as DΛΓ = C∗ΛICΓI. The matrix representation of the Fock operator for two-component DFT is determined by hΓΛ + (ΓΛ|ΠΘ)DΘΠ −cHF(ΓΘ|ΠΛ)DΘΠ + ∂Exc ∂DΓΛ, with subsequent partial differentiation yielding the kernel operator KΓΛΘΠ = (ΓΛ|ΠΘ) −cHF(ΓΘ|ΠΛ) + ∂2Exc ∂DΓΛ∂DΠΘ.

To address challenges with backward differentiation in SCF calculations, the work treats the converged SCF solution as the root of a fixed-point equation in Fock space. An SCF update, denoted as Sθ, maps an input Hermitian Fock matrix to an output matrix by solving the generalized eigenproblem and rebuilding the Fock matrix, incorporating the learned functional parameters θ.

This avoids the excessive memory consumption associated with unrolling iterative SCF procedures, allowing for application to larger systems than previous methods. The Casida equation for TDDFT calculations, simplified using the Tamm-Dancoff approximation, is solved after an initial DFT calculation determines the wavefunctions of the reference state.

Helium excited state and ionisation potential accuracy via JAX-trained IXC functional optimisation

The developed IXC functional, trained using a JAX-based two-component code, achieved deviations below 0.01 atomic units for the first singlet and triplet excited states of helium. Specifically, for the S1 state, deviations were smaller than 0.01 au, while for the T1 state, both the IXC functional and B3LYP yielded deviations below the same threshold.

Furthermore, the self-interaction error in the He+ system was minimized, with both B3LYP and IXC producing deviations below 0.005 au. During the training process, the total loss decreased monotonically within only ten iterations, demonstrating substantially enhanced efficiency compared to alternative training strategies.

The ∆LO value remained positive throughout training, indicating stable and reliable optimization. This successful training protocol accurately fitted a diverse set of properties, as evidenced by the consistent potential and linear response kernel derived via automatic differentiation. Benchmarking the trained IXC functional on H2, Li+, and H2O revealed a mean absolute error (MAE) below 0.1 au for all evaluated properties: Ω(S1), Ω(T1), self-interaction error, and ∆LO.

Notably, IXC exhibited a significantly smaller deviation in the self-interaction error, below 0.01 au, surpassing the performance of the commonly used B3LYP hybrid functional. For H2O, IXC achieved values of 0.342822 au for Ω(S1) and 0.308343 au for Ω(T1), demonstrating its potential transferability despite being trained on a single atom and with a fixed basis set.

Optimised IXC functional demonstrates transferable accuracy across diverse chemical systems

Researchers have developed a fully differentiable workflow for training a deep-learned energy functional using both ground-state and excited-state density functional theory calculations. This approach utilizes a JAX-based two-component code, termed IQC, to optimise the functional via automatic differentiation, ensuring consistency between the energy, self-consistent field (SCF) procedure, and linear response calculations.

The functional is trained on the helium spectrum, incorporating penalties to address self-interaction error and uphold the Lieb-Oxford inequality. Evaluations on hydrogen, lithium cation, and water demonstrate the potential transferability of this training strategy, with the learned functional, IXC, achieving mean absolute errors comparable to, and in some cases surpassing, those of established functionals like B3LYP.

Specifically, IXC exhibits a notably small deviation in the self-interaction error for hydrogen, below 0.01 atomic units, and superior accuracy for the mean absolute error compared to B3LYP. The authors acknowledge limitations including the use of a fixed finite basis set and a minimal spectral descriptor, and suggest future work focusing on expanding the descriptor set, enforcing size-extensivity, employing more diverse training data, and moving beyond the adiabatic approximation. These advancements represent a step towards more accurate and efficient electronic structure calculations.