Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

End-to-End Differentiable Learning of a Single Functional for DFT and Linear-Response TDDFT

Xiaoyu Zhang

TL;DR

The paper tackles the challenge of creating a single exchange–correlation functional that consistently governs ground-state energies, self-consistent potentials, and LR-TDDFT kernels. It introduces an end-to-end differentiable workflow in a two-component quantum chemistry framework (IQC using JAX) to train a neural xc functional by jointly optimizing KS-DFT and adiabatic LR-TDDFT targets, with gradients propagated through the SCF fixed point and the Casida-like eigenproblem. A neural spectral descriptor of the density matrix is learned via per-eigenvalue embeddings to produce a scalar energy correction, while penalties enforce one-electron self-interaction cancellation and adherence to the Lieb–Oxford bound. Demonstrated on a helium-based proof-of-concept in a fixed cc-pVDZ basis, the approach achieves rapid convergence (ten iterations) and favorable accuracy for S1/T1 excitations and SIE compared with standard functionals, suggesting potential transfer to molecular systems and a path toward more transferable, end-to-end trained functionals.

Abstract

Density functional theory (DFT) and linear-response time-dependent density functional theory (LR-TDDFT) rely on an exchange-correlation (xc) approximation that provides not only energy but also its functional derivatives that enter the self-consistent potential and the response kernel. Here we present an end-to-end differentiable workflow to optimize a single deep-learned energy functional using targets from both Kohn-Sham DFT and adiabatic LR-TDDFT within the Tamm-Dancoff approximation. Implemented in a JAX-based two-component quantum chemistry code (IQC), the learned functional yields a consistent potential and LR kernel via automatic differentiation, enabling gradient-based training through the SCF fixed point and the Casida equation. As a proof of concept in a fixed finite basis (cc-pVDZ), we learn an exchange-correlation functional on the helium spectrum while incorporating one-electron self-interaction cancelation and the Lieb-Oxford inequality as penalty terms, and we assess its possible transfer to molecular test cases.

End-to-End Differentiable Learning of a Single Functional for DFT and Linear-Response TDDFT

TL;DR

The paper tackles the challenge of creating a single exchange–correlation functional that consistently governs ground-state energies, self-consistent potentials, and LR-TDDFT kernels. It introduces an end-to-end differentiable workflow in a two-component quantum chemistry framework (IQC using JAX) to train a neural xc functional by jointly optimizing KS-DFT and adiabatic LR-TDDFT targets, with gradients propagated through the SCF fixed point and the Casida-like eigenproblem. A neural spectral descriptor of the density matrix is learned via per-eigenvalue embeddings to produce a scalar energy correction, while penalties enforce one-electron self-interaction cancellation and adherence to the Lieb–Oxford bound. Demonstrated on a helium-based proof-of-concept in a fixed cc-pVDZ basis, the approach achieves rapid convergence (ten iterations) and favorable accuracy for S1/T1 excitations and SIE compared with standard functionals, suggesting potential transfer to molecular systems and a path toward more transferable, end-to-end trained functionals.

Abstract

Density functional theory (DFT) and linear-response time-dependent density functional theory (LR-TDDFT) rely on an exchange-correlation (xc) approximation that provides not only energy but also its functional derivatives that enter the self-consistent potential and the response kernel. Here we present an end-to-end differentiable workflow to optimize a single deep-learned energy functional using targets from both Kohn-Sham DFT and adiabatic LR-TDDFT within the Tamm-Dancoff approximation. Implemented in a JAX-based two-component quantum chemistry code (IQC), the learned functional yields a consistent potential and LR kernel via automatic differentiation, enabling gradient-based training through the SCF fixed point and the Casida equation. As a proof of concept in a fixed finite basis (cc-pVDZ), we learn an exchange-correlation functional on the helium spectrum while incorporating one-electron self-interaction cancelation and the Lieb-Oxford inequality as penalty terms, and we assess its possible transfer to molecular test cases.
Paper Structure (3 sections, 18 equations, 2 figures, 1 table)

This paper contains 3 sections, 18 equations, 2 figures, 1 table.

Table of Contents

  1. Supplemental Materials: End-to-End Differentiable Learning of a Single Functional for DFT and Linear-Response TDDFT
  2. Implicit differentiation through the SCF fixed point
  3. Geometries of $\mathrm{H}_2$ and $\mathrm{H_2O}$

Figures (2)

  • Figure 1: End-to-end DFT and LR-TDDFT differentiable workflow. A learned xc energy functional is used self-consistently in SCF; its first and second derivatives provide a consistent potential contribution and an adiabatic LR kernel via automatic differentiation. Gradients are propagated through the SCF fixed point by implicit differentiation and through the TDDFT eigenproblem by differentiating the LR eigenvalue solution.
  • Figure 2: The evolution of four quantities during the training on He: (1) the total training loss; (2) the deviation of the $\mathrm{S}_1$ excitation energy; (3) the deviation of the $\mathrm{T}_1$ excitation energy; and (4) the self-interaction error for He$^+$. Compared data acquired by HF, SVWN, PBE, and B3LYP is presented as dashed lines.