Learning Equivariant Non-Local Electron Density Functionals

We highly appreciate that the reviewer is impressed by the accuracy of our functional. We would like to address the reviewer's concerns with our revised manuscript and the following answer.

W1: background

We strive to effectively communicate with both the natural science and the machine learning community. Based on the current feedback, we made the following improvements:

1. The introduction includes a proper definition of non-local functionals in the introduction:

   While [semi-local functionals] integrate well into existing quantum chemistry code, non-local interactions like Van der Waals forces exceed the functional class (Kaplan et al., 2023). In contrast, non-local functionals can capture such interactions by depending on multiple points in space simultaneously, e.g., $E_\text{XC}[\rho]=\iint_{R^3} \rho(r)\rho(r')\epsilon(g(r),g(r'))drdr'$.

2. We extended Appendix B with a definition of the Fock matrix $F$ and how it is computed.
3. We added a loss definition and gradient computation to the experimental section.
4. We use atomic-range and molecular-range instead of long-range, see next reply.

We highly appreciate any additional feedback on unclear sections.

W2 + Q1: long-range terminology

We understand the difficulty in terminology as the term long range is differently used in the force field literature. Given the definition of non-locality above, a non-local functional still does not necessarily need to depend on distant points. To increase clarity and avoid confusion with the force field terminology, we introduce the following two ranges:

• atomic-ranged: electronic interactions around a nucleus.
• molecular-ranged: electronic interactions at typical molecular length scales.

While the convolution with our equivariant filters yields non-local features with atomic-range interactions, molecular-range interactions are added through the equivariant message passing. We adjusted the manuscript accordingly. In a new ablation study on 3BPA, we demonstrate the importance of the molecular-ranged interactions:

This demonstrates that the atomic-range interactions through our point cloud embedding and the molecular-range interactions through our equivariant message passing reduce errors.

To further improve our communication, we renamed our method from Global Graph Exchange Correlation (GG-XC) to Equivariant Graph Exchange Correlation (EG-XC), as the interactions may not be global for large structures.

W3 + Q2: Computational complexity

We agree with the reviewer that computational complexity and runtime measurements are important for contexting our work. We added Appendix M and Appendix N to cover these topics. In Appendix M, we break down the computational complexity for each subprocedure of an SCF step. In Appendix N, we measure the runtime of different XC functionals. We summarize the runtime in the following table.

The computational difference between EG-XC and mGGAs is relatively small compared to the runtime of hybrid functionals like B3LYP or $\omega$B97X. For mGGAs, the runtime is dominated by evaluating the density features on the quadrature grid, while for hybrid functionals, the evaluation of the exact exchange term dominates the computation. We refer to the new appendices for more details on the setup and individual runtime complexities.

We highly appreciate the reviewer's detailed feedback and are happy to hear that the reviewer agrees with our data efficiency claims. We want to address the reviewer's concerns with our revised manuscript and the following answer.

W1 + Q1 + Q2: Other DFT observables (density, forces, dipole moments, etc.)

We agree with the reviewer that these are interesting and extremely valuable aspects for evaluating XC functionals. However, we argue that improving energy accuracy is extremely valuable as it distills highly accurate quantum mechanical calculations, e.g., CCSD(T), Møller–Plesset perturbation theory, or neural-network VMC, into DFT. These methods do not necessarily permit computing other observables. To improve DFT's accuracy in such settings, this work augments the design space of ML XC functionals with efficient non-local interactions. For instance, in our new Appendix N, we find that GG-XC offers a 47% runtime reduction compared to range-separated classical hybrids like B97X while being able to learn arbitrary non-local interactions. Nonetheless, we highly encourage future work to investigate other observables.

W2 + Q3: Lack of intermolecular non-covalent interactions and chemical reactions.

While we agree that non-covalent interactions and chemical reactions are interesting applications of learnable XC functionals, such structures typically require unrestricted DFT to handle open-shell problems. We clarify this in the limitation section:

[...] to handle open-shell systems, e.g., in chemical reactions, one would need to extend the equivariant embeddings to include spin information.

W3: The lack of explicit enforcement of physical constraints.

We agree with the reviewer's sentiment that physical constraints may aid in edge cases. However, this is certainly not given due to the loose nature of the constraints. We train the mGGA from Nagai et al. [1] on 3BPA to demonstrate that constraints do not need to be beneficial. While Dick et al. [2] attempts to implement many physical constraints, no such restrictions are put on Nagai's functional. The results are summarized in the following table:

These mixed results present no clear evidence that physical constraints yield meaningful improvements even in this out-of-distribution extrapolation setting. Other functionals like DM21 [3] come to similar conclusions. We added these results to the new Appendix O.

[1] Nagai et al. "Completing density functional theory by machine learning hidden messages from molecules", 2020
[2] Dick et al. "Highly accurate and constrained density functional obtained with differentiable programming", 2021
[3] Kirkpatrick et al. "Pushing the frontiers of density functionals by solving the fractional electron problem", 2021

Thanks for the authors' reply, I would keep my score.

We highly appreciate the reviewer's feedback and hope to resolve the outstanding questions. In addition to this reply, we like to point the reviewer to the corresponding sections of the updated manuscript.

W1: Computational cost

We generally agree that the computational cost between these methods varies significantly and, thus, include $\Delta$-ML methods as a computationally more similar method in all experiments as well as the learnable XC functional from Dick 2021. To better illustrate the runtime difference, we added the new Appendix N, where we benchmark the cost of running DFT. There, we compared several DFT functionals in terms of runtime and found that GG-XC was only 7% slower than mGGAs.

W2: Reliance on mGGA

We acknowledge that the inclusion of a mGGA increases complexity but we disagree that this is a weakness of our approach. Consistent with classical non-local functionals like B3LYP or $\omega$B97X the majority of the XC energy can be captured through such semi-local functionals. In our revised manuscript, we include an ablation study on 3BPA where we remove the mGGA. The following table summarizes the results:

These results support the importance of the mGGA.

Q1: The results on MD17 (Table 1) and QM9 (Figure 3) seem to be good but with mixed results. How to understand these differences in accuracy?

We agree that further interpretation aids the understanding of these results. To better communicate our findings, we elaborate further in our experimental discussion (see updated manuscript). In particular, we hypothesize that the small gap between $\Delta$-ML and XC functionals on MD17 is due to the i.i.d. nature of the dataset. In all extrapolation settings, 3BPA and QM9, we find learnable functionals yielding significant improvements.

Q2: How is non-local defined? Will the re-weighting also make density gradients non local?

We revised the introduction to yield a proper definition of non-local functionals in the introduction:

While [semi-local functionals] integrate well into existing quantum chemistry code, non-local interactions like Van der Waals forces exceed the functional class (Kaplan et al., 2023). In contrast, non-local functionals can capture such interactions by depending on multiple points in space simultaneously, e.g., $E_\text{XC}[\rho]=\iint_{R^3} \rho(r)\rho(r')\epsilon(g(r),g(r'))drdr'$.

We extended Appendix B with a definition of the Fock matrix and how we compute the gradient of the exchange-correlation potential with respect to the density matrix. Since we differentiate through the whole XC model, the gradients of the non-local contributions are part of the XC potential $\frac{\partial E_\text{XC}[\rho]}{\partial P}$.

We highly appreciate that the reviewer finds our work novel and would like to address the reviewer's concerns with our revised manuscript and the following answer.

W1 + Q2: Loss and gradient computation

In our revised manuscript, we added our mean squared error loss definition and a description of how we compute gradients in the experiments section. Concretely, we compute the mean squared error between the energies of the last $I_{loss}\in N_+$ SCF cycles and the target energy:

We follow Dick (2021) and minimize the mean squared error between the converged SCF energy and the target energy over the last $I_\text{loss}\in N_+$ steps [...]

[...] The parameter gradients are obtained through the total derivative of the loss with respect to the parameters of the XC functionals $\frac{dL}{d\theta}$ by backpropagating through the SCF iterations.

We want to clarify that, as stated in the abstract (lines 18–19), introduction (lines 59–60), experiments (lines 349–351), and conclusion (lines 462–463), our training methodology is based solely on energy targets as these are readily available. In contrast, CCSD(T) densities are typically unavailable.

W2 + Q3: Design choices We agree that ablation studies provide additional insights into our contributions. Thus, we added a new ablation study on the 3BPA dataset where we disable several components of our model. For convenience, we repeat the table here.

It is apparent that each component is a significant contribution to the accuracy of GG-XC. For a detailed interpretation, we refer to the revised manuscript's experiments section.

W3 + Q4: Fock matrix definition and implementation

In our updated manuscript, we extended our SCF explanation in Appendix B to cover the construction and the computation of the Fock matrix:

[...] the Fock matrix is given by

To evaluate the $E_{XC}$ and its derivative, we discretize $R^3$ with spherical integration grids around the nuclei. Specifically, for a set of quadrature points $r_i\in R^3$ and weights $w_i\in R_+$, we first compute the density and mGGA features $g(r_i)=\left[\rho(r_i), \Vert\nabla\rho(r_i)\Vert, \tau(r_i)\right]$ on these points:

[where $\chi:R^3\to R^B$ are the basis functions.] From this discretized input, we compute the scalar output $E_{XC}$, e.g., $E_{XC}[\rho] = \sum_{i=1}^{N} w_i \rho(r_i) \epsilon(g(r_i))$ for an mGGA, and rely on automatic differentiation via backpropagation to compute the derivative w.r.t. $P$.

W4 + Q5: Computational cost

We acknowledge that runtime and computational complexity are valuable insights. To improve our coverage of these topics, we added Appendix M, where we analyze the computational complexity of SCF and GG-XC steps, and Appendix N, where we measure runtimes for different XC functionals. To measure runtime, we implemented all functionals in JAX and measured the time of a SCF calculation. We ignore the precomputations like the core Hamiltonian, density fitting, and the evaluation of the atomic orbitals. Here, we summarize the new Figure 7 in a Table:

Compared to hybrid functionals, our GG-XC yields a computationally efficient approach to learning non-local interactions. In contrast to hybrid functionals, GG-XC can learn arbitrary non-local interactions beyond exchange-related quantities.

In computational complexity, the evaluation of mGGA-like functionals (including GG-XC) is primarily dominated by the evaluation of the density features on the quadrature grid, which requires $O(N_\text{quad}N_\text{bas}^2)$ operations. For hybrid functionals, the dominating computation is the evaluation of the exact exchange term requiring $O(N_\text{aux} N_\text{bas}^3)$ operations. For further complexity analysis, we refer the reviewer to Appendix M.

Thank you very much for your detailed response and for conducting additional experiments to address the questions raised in my review. As a computational chemistry researcher, albeit not specializing in DFT method development, I took long time to fully evaluate and appreciate the significance of your work. Your thorough answers and experiments have been very helpful in clarifying key points, and I would like to express my gratitude for your efforts. I have updated my rating to 6 in recognition of the quality and potential of your work. On the other hand, I have a few additional questions that I am curious about. These are not requests for further revisions but rather points of personal interest to deepen my understanding of your methodology.

1. Generalization to Element Types and Charged Species

Your method seems to demonstrate strong generalization capabilities beyond the limitations of delta learning, which is very impressive. Have you tested the GG-XC functional on systems with diverse element types or charged species (e.g., ionic systems)? If so, could you share any relevant results or insights? I view this methodology as fundamentally different from delta learning or potential learning, targeting distinct aspects of DFT method that set it apart from those approaches.

2. Position on Jacob’s Ladder

The cost efficiency of your method, particularly when compared to range-separated hybrid functionals, is remarkable. I am curious to know how you would position the GG-XC functional on Jacob’s Ladder of DFT accuracy. Specifically: Do you believe GG-XC can achieve accuracy comparable to or exceeding B3LYP? Is there a reason why comparisons with established functionals like B3LYP or ωB97X were not included in the main manuscript?

3. Analysis of TS Calculations and Optimization

The energy landscapes shown in Figures 4, 5, and 6 are particularly striking. Based on these: Have you performed transition state (TS) calculations or structure optimizations using GG-XC? If so, have you analyzed the differences in optimized structures or energy values compared to B3LYP or ωB97X? Any results or observations on this would be very interesting.

4. System size

How large of a system have you tested with GG-XC? For example, when calculating a molecule like $C_{38}$ with 232 electrons, how much memory does the computation require?

Your work is truly inspiring, and I look forward to seeing how it impacts the computational chemistry community. Thank you again for your efforts and your time in addressing these questions! Best regards.

Thank you. It's still early stage of functional learning, but the direction of the authors' research is very interesting. I will keep watching the follow-up study. I will raise the rating to 8.

We are excited to hear that our research has piqued your interest and are happy to discuss and expand upon your questions.

1. Inquiry: How does EG-XC perform on datasets that include more diverse nuclei and charged systems?

In principle, EG-XC is not limited to organic chemistry, but in our implementation of the differentiable SCF-solver, we have been focused on stable organic chemistry and used a spin-restricted approach to reduce the computational cost. As such, the solver is presently limited to spin-unpolarized systems. While evaluating charged closed-shell systems, such as salt ions, is straightforward for EG-XC, force fields are known to struggle with long-range (e.g., Coulomb interactions). We postpone such inquiries to future work, but we strongly agree that this is an interesting line of questioning, especially when focussing more on the comparison of different XC density functional approximations.

2. Inquiry: a) Where is EG-XC placed in Jacob’s Ladder?

Our method is clearly non-local (above the third rung on Jacob’s ladder [1]), but it does not fall into any of the preexisting categories to the best of our knowledge. Even so, we appreciate the sentiment of your question: Typically, both accuracy and computational cost increase as we ascend the ladder. Here, our evaluation of EG-XC indicates a promising shift, as its computational cost is only marginally higher than that of meta-GGAs, but its accuracy could surpass both that of hybrid functionals (e.g., B3LYP) and range-separated hybrids (like $\omega$B97X) which come at a significantly higher cost. Further work is required to make any definitive statements in this direction.

b) How does it compare to B3LYP and $\omega$B97X?

In general, our aim here is to develop a new functional form, rendering a comparison to general functionals (fitted to different data) difficult. We want to illustrate this point using the MD17 dataset on the CCSD(T) level of theory (unlike the DFT data of QM9 (B3LYP) and 3BPA ($\omega$B97X)) for the SCAN functional, as we already have the data. We note that even this physically constrained functional has free parameters fitted to reference data[2]. In the evaluation on MD17 with a fixed basis set, learnable functionals have an advantage, as their model parameters can adapt to the specific compound at hand. Additionally, the model can learn to correct for systematic errors, e.g., errors due to the finite basis set and density fitting.

As we can see, the MAE of SCAN is an order of magnitude larger than Dick 2021 and EG-XC. Even when performing an optimal constant shift (the median of the training set), which does not affect the forces, we observe an, on average, significantly worse performance compared to the learnable functionals and EG-XC in particular, despite its approximately equal computational cost.

3. Inquiry: How does EG-XC perform on structure optimization tasks and in TS calculations

We have not yet performed such tests, but we encourage the investigation of these aspects in future work. We would be very interested in suggestions for datasets or other suitable test cases and would highly value your perspective as a computational chemist.

4. Inquiry: What is the memory requirement of an SCF calculation with EG-XC for large systems?

The largest system we have trained on is 3bpa, as shown in the main document. The largest system we evaluated our method on was $C_{36}H_4$, which just fitted on a GPU with 40GB memory. As you correctly point out yourself, memory (specifically GPU memory) is indeed the key limiting factor. During training, the caching of all intermediates tensors for backpropagation is the memory bottleneck. At inference, the density-fitted ERI-tensor $B$ is the memory bottleneck. For the $C_{38}$ calculation this tensor alone is 20GB.

Closing Remarks

We are encouraged by your interest in EG-XC and appreciate your insights. Your questions provide clear directions for further development and exploration, and we look forward to seeing how this work might inspire the computational chemistry community.

[1] Perdew, J. P. “Jacob’s Ladder of Density Functional Approximations for the Exchange-Correlation Energy”, 2001 https://doi.org/10.1063/1.1390175.

[2] Sun, J.; Ruzsinszky, A.; Perdew, J. P. “Strongly Constrained and Appropriately Normed Semi-local Density Functional”, 2015 https://doi.org/10.1103/PhysRevLett.115.036402.