Density functional theory (DFT) stands as a cornerstone method in computational quantum chemistry and materials science due to its remarkable versatility and scalability. Yet, it suffers from limitations in accuracy, particularly when dealing with strongly correlated systems. To address these shortcomings, recent work has begun to explore how machine learning can expand the capabilities of DFT; an endeavor with many open questions and technical challenges. In this work, we present Grad DFT: a fully differentiable JAX-based DFT library, enabling quick prototyping and experimentation with machine learning-enhanced exchange-correlation energy functionals. Grad DFT employs a pioneering parametrization of exchange-correlation functionals constructed using a weighted sum of energy densities, where the weights are determined using neural networks. Moreover, Grad DFT encompasses a comprehensive suite of auxiliary functions, notably featuring a just-in-time compilable and fully differentiable self-consistent iterative procedure. To support training and benchmarking efforts, we additionally compile a curated dataset of experimental dissociation energies of dimers, half of which contain transition metal atoms characterized by strong electronic correlations. The software library is tested against experimental results to study the generalization capabilities of a neural functional across potential energy surfaces and atomic species, as well as the effect of training data noise on the resulting model accuracy.

翻译：密度泛函理论（DFT）因其卓越的通用性和可扩展性，成为计算量子化学与材料科学领域的基石方法。然而，该理论在精度方面存在局限，特别是在处理强关联系统时。为弥补这些不足，近期研究开始探索如何利用机器学习拓展DFT的能力——这一探索仍面临诸多开放性问题与技术挑战。本文提出Grad DFT：一个基于JAX的全可微DFT库，能够快速实现机器学习增强型交换关联能量泛函的原型设计与实验。Grad DFT采用开创性的交换关联泛函参数化方法，通过神经网络确定的权重对能量密度进行加权求和来构建泛函。此外，Grad DFT还包含一套完整的辅助函数，特别值得一提的是其具备即时编译能力的全可微自洽迭代程序。为支撑训练与基准测试工作，我们额外整理了一套经实验验证的双原子分子解离能数据集，其中半数分子含有强电子关联的过渡金属原子。通过将该软件库的实验结果进行对比测试，我们研究了神经泛函在势能面与原子种类间的泛化能力，以及训练数据噪声对模型精度的影响。