Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

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

Rakshit Kumar Singh, Aryan Amit Barsainyan, Bharath Ramsundar

TL;DR

This work tackles the computational bottleneck of DFT for periodic systems by introducing a fully differentiable framework that integrates neural network exchange–correlation functionals into KS-DFT for solids, with gradients flowing through the SCF loop. The approach is implemented in Python with a PyTorch backend and integrated into DeepChem, enabling end-to-end learning via a drop-in XC functional API and a HybridXC formulation that blends ML with conventional GGA physics. Benchmarks on Al, Ni, and Ca show 5–10% relative errors relative to established codes, while graphene highlights current limitations and the need for improved basis representations. The framework advances the ability to train ML-based XC functionals in periodic systems, offering a scalable path toward accurate, differentiable, and physically grounded solids simulations, with plane-wave extensions proposed for future work.

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%.

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

TL;DR

This work tackles the computational bottleneck of DFT for periodic systems by introducing a fully differentiable framework that integrates neural network exchange–correlation functionals into KS-DFT for solids, with gradients flowing through the SCF loop. The approach is implemented in Python with a PyTorch backend and integrated into DeepChem, enabling end-to-end learning via a drop-in XC functional API and a HybridXC formulation that blends ML with conventional GGA physics. Benchmarks on Al, Ni, and Ca show 5–10% relative errors relative to established codes, while graphene highlights current limitations and the need for improved basis representations. The framework advances the ability to train ML-based XC functionals in periodic systems, offering a scalable path toward accurate, differentiable, and physically grounded solids simulations, with plane-wave extensions proposed for future work.

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%.
Paper Structure (16 sections, 5 equations, 3 figures, 1 table)

This paper contains 16 sections, 5 equations, 3 figures, 1 table.

Table of Contents

  1. Introduction
  2. Related Work
  3. Methods
  4. Dataset
  5. Lattice Parameter Optimization Workflow
  6. Results and Observations
  7. Conclusion
  8. Future Work
  9. Appendix
  10. Data Generation
  11. Training Objective
  12. Data Representation
  13. K-Point Sampling
  14. Limitations Observed in Graphene Calculations
  15. Reference Training Data
  16. ...and 1 more sections

Figures (3)

  • Figure 1: This figure represents the Lattice constant optimization flow using a trained NN XC functional used alongside a DFT XC functional. The framework was tested using lattice constants from experimental data as initial guess of the system.
  • Figure 2: This figure presents the lattice constant scans for Al, Ni, and Ca computed using PySCF, GPAW, and DeepChem. The lattice parameters were varied from $3.40 \AA$ to $4.45 \AA$ for Al, $3.00 \AA$ to $4.00 \AA$ for Ni, and $4.50 \AA$ to $6.50 \AA$ for Ca, with a step size of $0.05 \AA$.
  • Figure 3: Original PySCF graph