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

TL;DR

This paper introduces a differentiable workflow for optimizing a deep-learned exchange-correlation functional in DFT and LR-TDDFT, enabling end-to-end training using quantum chemistry data.

Contribution

It develops a JAX-based framework that allows gradient-based optimization of a single functional for both energy and response properties in quantum chemistry.

Findings

Successfully trained a functional on excitation energies of small molecules.

Incorporated self-interaction cancellation as penalty terms.

Assessed transferability to molecular test cases.

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. To enable this training in a computationally efficient and differentiable manner, we developed a JAX-based two-component quantum chemistry framework (IQC), in which the learned functional provides a self-consistent potential and linear-response kernel via automatic differentiation. This construction permits gradient-based optimization through both the self-consistent-field (SCF) fixed-point equations and the Casida eigenvalue problem. We learn an exchange-correlation functional on excitation energies of small molecules while incorporating one-electron self-interaction cancelation as penalty terms, and we assess its possible transfer to molecular test cases.