Differentiable programming for Kohn-Sham DFT: energy derivatives with respect to the parameters of exchange-correlation functional at linear cost

We applied the adjoint method to Kohn-Sham equations taking the energy derivatives with respect to the parameters of the exchange-correlation functional. The results obtained are completely consistent with the Hellmann-Feynman theorem and can serve as a starting point for the study of analytical derivatives of other quantities obtained using the Kohn-Sham method. The prototype we created on the basis of python package PySCF exhibited linear complexity performance behavior and showed a significant speedup in the calculation of the derivative compared to the automatic differentiation approach such as pytorch based DQC. In a more severe sense, the development of efficient methods for computing derivatives with respect to parameters of scientific models is important in applications to machine learning where training is done via gradient-based optimization algorithms or the model is integrated with deep learning and there is a need to speed up the calculation in the backpropagation pass.