Infusing Self-Consistency into Density Functional Theory Hamiltonian Prediction via Deep Equilibrium Models

TL;DR

A versatile framework that combines DEQ with off-the-shelf machine learning models for predicting Hamiltonians, and uses the fixed-point solving capability of the deep equilibrium model to iteratively solve for the Hamiltonian.

Abstract

In this study, we introduce a unified neural network architecture, the Deep Equilibrium Density Functional Theory Hamiltonian (DEQH) model, which incorporates Deep Equilibrium Models (DEQs) for predicting Density Functional Theory (DFT) Hamiltonians. The DEQH model inherently captures the self-consistency nature of Hamiltonian, a critical aspect often overlooked by traditional machine learning approaches for Hamiltonian prediction. By employing DEQ within our model architecture, we circumvent the need for DFT calculations during the training phase to introduce the Hamiltonian's self-consistency, thus addressing computational bottlenecks associated with large or complex systems. We propose a versatile framework that combines DEQ with off-the-shelf machine learning models for predicting Hamiltonians. When benchmarked on the MD17 and QH9 datasets, DEQHNet, an instantiation of the DEQH framework, has demonstrated a significant improvement in prediction accuracy. Beyond a predictor, the DEQH model is a Hamiltonian solver, in the sense that it uses the fixed-point solving capability of the deep equilibrium model to iteratively solve for the Hamiltonian. Ablation studies of DEQHNet further elucidate the network's effectiveness, offering insights into the potential of DEQ-integrated networks for Hamiltonian learning. We open source our implementation at https://github.com/Zun-Wang/DEQHNet.

Figures (6)

• Figure 1: Schematic of the DEQH model. A Hamiltonian solver must be engineered to concurrently process structural information, the Hamiltonian, and the overlap matrix, and to output the subsequent Hamiltonian iteration for DEQ convergence. Within the Hamiltonian solver, the module handling structural information is termed injection, while the remaining components are collectively referred to as filter.
• Figure 2: A schematic diagram of DEQHNet. (a) The overall architecture of DEQHNet, with the injection mechanism encapsulated within the yellow frame and the remainder designated as the filter. (b) TransBlock. Integrates the Hamiltonian and overlap matrix with structural information (relative position vectors). (c) ConvNetLayer. Responsible for processing structural information. (d) Diagonal Reduction. Transforms the Hamiltonian and overlap matrix into equivariant node features. (e)-(g) Couple node features to form diagonal and off-diagonal blocks of the Hamiltonian.
• Figure 3: (a) The network predicts the Hamiltonian using atomic number $Z$ and coordinates $R$ as inputs, outputting the molecular Hamiltonian $H$. (b) The DEQH model also includes an input for $H$ and iteratively refines the Hamiltonian until reaching a fixed-point solution $H^*$, hence we refer to this block as the Hamiltonian solver.
• Figure 4: (a) The variation in the number of DEQ iterations within DEQHNet as a function of training steps. For this analysis, 50 random configurations were selected from the uracil test set, and the iteration counts were tallied after performing inference with DEQHNet checkpoints saved during the training process. (b) The change in Hamiltonian MAE with respect to iteration count for a randomly chosen molecule from the QH9-dynamics-geo test set, comparing the self-consistent field (SCF) iterations using PySCF, DEQHNet inference, and DEQHNet inference initialized with Hamiltonians guessed by PySCF.
• Figure 5: Ablation study for MAE of Hamiltonian $H$, orbital energy $\varepsilon$, and orbital coefficients $\mathbf{C}$ for DEQHNet, DEQHNet without the overlap matrix as an input, QHNet, and QHNet with the addition of the overlap matrix as an input, respectively.