A fully differentiable framework for training proxy Exchange Correlation Functionals for periodic systems

TL;DR

This paper presents a fully differentiable, ML-integrated framework for training exchange-correlation functionals in density functional theory, enabling end-to-end optimization for periodic systems with promising accuracy.

Contribution

It introduces a novel differentiable framework that seamlessly incorporates neural network models into DFT workflows for solids, facilitating improved functional training.

Findings

Achieved 5-10% relative errors in benchmarks

Implemented in PyTorch for ease of use

Integrated with DeepChem for model reuse

Abstract

Density Functional Theory (DFT) is widely used for first-principles simulations in chemistry and materials science, but its computational cost remains a key limitation for large systems. Motivated by recent advances in ML-based exchange-correlation (XC) functionals, this paper introduces a differentiable framework that integrates machine learning models into density functional theory (DFT) for solids and other periodic systems. The framework defines a clean API for neural network models that can act as drop in replacements for conventional exchange-correlation (XC) functionals and enables gradients to flow through the full self-consistent DFT workflow. The framework is implemented in Python using a PyTorch backend, making it fully differentiable and easy to use with standard deep learning tools. We integrate the implementation with the DeepChem library to promote the reuse of established models and to lower the barrier for experimentation. In initial benchmarks against established electronic structure packages (GPAW and PySCF), our models achieve relative errors on the order of 5-10%.