Learning the Electronic Hamiltonian of Large Atomic Structures

Isolation of Contributions: Our two main contributions not found in other works are:

1. An augmented partitioning approach that allows arbitrarily large graphs to be broken down into independent partitions that (1) maintain the connectivity of the full structure through virtual nodes/edges, (2) fit into GPU memory during training, and (3) ensure high scalability without compromising achievable test accuracy.

2. A GNN model that is custom-made for scalable Hamiltonian predictions of large, complex atomic structures. Unlike previous models [3-8], it combines strict locality with efficient SO(2) convolutions plus multi-head equivariant attention to better distinguish between complex atomic environments.

Ablation studies: An ablation study of our proposed augmented partitioning approach is shown in Table 2 of the results section, where the omission of virtual components led to a 50% increase in node error and an 88% increase in edge error, as the graph tries to aggregate information through an incomplete graph structure. The strict locality is also motivated and justified by previous work like Allegro [5,8], plus we have also conducted studies on their effects on Hamiltonian matrices in Appendix F.

Discussion of other batching methods for large graphs: Previous popular works on large graph partitioning [9] employ neighborhood sampling approaches to reduce communication volume between partitions. This cannot be applied to Hamiltonian predictions, as omitting any graph connections leads to the wrong atomic neighborhood. Our method is a novel way to (1) batch the data, and (2) removes the need for communication between different partitions while maintaining graph connectivity.

Evaluation on established datasets: Established Hamiltonian datasets (e.g., QH9 [6]) consist of small molecules, which cannot be used to evaluate a model's performance on the large-scale problem we are trying to tackle. The use of custom datasets is also encouraged in the applications track, where ML is applied to a wide variety of problems that have yet to be tackled, and a valid comparison does not yet exist.

For evaluating previous methods on our custom dataset, the closest accuracy comparison (within memory limits) is shown in Table 3 in the paper for HfO2 (and also in Table 11 in Appendix I.3 for GST and PtGe). Full graph training is the default baseline approach used by previous SOTA. We showed that for all three datasets, the application of augmented partitioning approach maintains accuracy, while significantly increasing scalability (shown in Appendix H.2). The values we obtained (0.99-5.16 meV) are also within range of what SOTA models obtained (1.5-3.23 meV) for simpler structures [3-4].

Downstream applications: We applied our approach to the simulation of valence change memory in the response to Reviewer Kuqs. The obtained accuracy allows us to observe trends in resistance contrast changes as a function of applied voltage and resulting ion movements. By subverting expensive DFT calculations entirely, we enable a range of previously unfeasible large-scale device simulations. The true meaning of our work, therefore, lies in its ability to address a crucial bottleneck that previous methods (including DFT) could not. It also sets a precedent for the prediction of complex structures at the scale of ~10^3 atoms per unit cell and can be used as an initial benchmark for future models that also aim to do so for various applications.

References

[1] Jónsson, H., Mills, G. & Jacobsen, K. W. in Classical and Quantum Dynamics in Condensed Phase Simulations (eds Berne, B. J. et al.) 385–404 (World Scientific, 1998)

[2] Kaniselvan, M., Luisier, M., and Mladenovic, M. An atomistic model of field-induced resistive switching in valence change memory. ACS Nano, March 2023.

[3] Wang, Y., Li, Y., et. al. Universal materials model of deep-learning density functional theory hamiltonian. Science Bulletin, 69(16):2514–2521, 2024b.

[4] Zhong, Y., Yu, H., Su, M., Gong, X., and Xiang, H. Transferable equivariant graph neural networks for the hamiltonians of molecules and solids. npj Computational Materials, 2023.

[5] Musaelian, A., Batzner, S., Johansson, A., Sun, L., Owen, C. J., Kornbluth, M., and Kozinsky, B. Learning local equivariant representations for large-scale atomistic dynamics. Nature Communications, 14(1):579, 2023

[6] Yu, H., Liu, M., Luo, Y., Strasser, A., Qian, X., Qian, X., and Ji, S. Qh9: A quantum hamiltonian prediction benchmark for qm9 molecules, 2023a.

[7] Y Li, et. al Enhancing the scalability and applicability of Kohn-Sham Hamiltonian for scalable molecular systems. ICLR 2025.

[8] Zhouyin, Z., et. al. Learning local equivariant representations for quantum operators, ICLR 2025

[9] M. Besta and T. Hoefler, "Parallel and Distributed Graph Neural Networks: An In-Depth Concurrency Analysis," in IEEE Transactions on Pattern Analysis and Machine Intelligence

Non-diagonal part of Hamiltonian beyond cutoffs:

The use of a cut-off radius is standard practice in Hamiltonian prediction literature of bulk structures [3-4], which exploits the nearsightedness of the Hamiltonian matrix. Appendix F include two studies on the nearsightedness of Hamiltonians for our dataset:

1. In Fig. 7, we illustrated the decay in values of Hamiltonian elements with increasing inter-atomic distance for all three datasets.

2. In Fig. 8 we quantify the effect of discarding orbital blocks beyond a predefined cutoff on the prediction accuracy. Removing those with rcut > 4 Angstroms (the chosen rcut for HfO2 in the study) from the DFT reference data showed a negligible difference in the eigenvalue spectra of H.

Long range interactions beyond cutoff: In Appendix F, Fig. 9 we studied the effects of various types of perturbations (shift, vacancy, substitution) beyond the cutoff radius and showed that they all have a negligible effect on the Hamiltonian matrix elements (less than 0.24%). This indicates that the Hamiltonian of the atoms in our datasets can be learnt from their local atomic environment. The use of strict locality is also well-established in literature, and previous architectures (e.g., Allegro for force fields [5] and [8]) have used this to improve scalability while still achieving state-of-the-art prediction accuracy.

Absolute Error for Eigenspectrum: The threshold of < 1 kcal/mol is normally used for the prediction of thermochemical quantities (e.g., interaction, ionization energy) and is not commonly applied to eigenspectra of bulk structures with several atoms. Still, we understand the need for a quantitative measure with units, so we computed the mean absolute error for all eigenvalues, obtaining 3.08 mHartrees (2.53 mHartrees for occupied energies).

Real World Applications: Unlike the molecular Hamiltonians predicted by models like WANet, whose applications include the extraction of molecular quantities (e.g., orbital energies), the Hamiltonians of large structures (> 10^3 atoms) targeted by our approach have their own set of downstream applications (e.g., semiconductor physics). To truly assess its practical relevance in this domain, we applied it to the simulation of valence change memory cells, with details outlined in the response to Reviewer Kuqs. Using predicted Hamiltonians for simulations, we achieved levels of accuracy that allow us to observe key trends in current vs. voltage characteristics of VCM cells and reveal the dependence of the current on the atomic geometry of these devices. In general, it enables large-scale simulation workflows involving repeated DFT updates over 100-1000 time steps that were previously unfeasible, and also allows large structures beyond DFT capabilities (10k+ atoms) to be simulated.

Multiple graph slices and Hamiltonian construction: In this paper, slices are only used to train the model. The final trained model was then used to predict the full structure/graph at once (e.g., all 3000 atoms), with only one final predicted Hamiltonian matrix obtained, so there is no need to assemble any slices. For very large structures, the Hamiltonian can also be predicted separately as different slices with multiple GPUs/CPUs. Since the nodes and edges in each slice define an independent sub-block, the final matrix can then be reconstructed by populating it piece by piece with the components of each partition. Both approaches lead to the same result.

Role of virtual nodes/edges (whether they are updated by other graph slices): We do not update the embeddings (outputs) of the virtual nodes and edges in any case. In the paper, we mentioned that their output embeddings have no meaning/do not exist (hence the term virtual), and are not involved in training and predictions. Their only purpose lies in their initialized inputs (atomic numbers and distances), which are used to inform the labelled (non-virtual) atoms within the partition of their correct one-hop atomic neighborhood. The use of only the inputs of virtual nodes and edges to maintain connectivity without communication is one of the key novel contributions of our approach.

Note that the list of references is found in the response to Reviewer 4 (nuFS)

1. One model for each material system: Yes, the models are trained for a specific material system, similar to most other literature ([4-8]). The current work does not demonstrate a universal Hamiltonian prediction model, though this can be achieved with a larger library of datasets, as shown by [3]. No public database for structures with 1000+ atoms per unit cell currently exists, so the necessary data must be manually generated. Data creation is computationally expensive due to the melt-and-quench and relaxation processes required. With enough collective effort and time, however, a universal model for all materials, including large, amorphous structures, can be feasibly trained with our approach.

2. Generalization to different structures and systems: Our approach (strictly local model + augmented partitioning) can be straightforwardly applied to other materials systems, including periodic structures (crystals, 2D materials) and molecules. We also applied it to the prediction of crystalline HfO2 structure and obtained far more accurate predictions (1.1 meV MAE) relative to amorphous HfO2 (5.16 meV) without any further optimization. Molecular examples are also included in the code attached in supplementary materials (though augmented partitioning and strict locality are excessive in these contexts due to their small, isolated systems and information-poor atomic neighborhoods). Here, the small number of atoms per unit cell, plus the larger similarity between training and test datasets, makes higher accuracies much more achievable.

3. Cutoff for graph and cutoff for slice: The cutoff (rcut) used to define the graph connectivity is not related to the slice length (t) used to define the partitions. They can be independently chosen.

Applicability of partitioning to other architectures: Augmented partitioning can be applied to other strictly local architectures (e.g., Allegro [5]), since it is implemented during graph construction.**

How our work exceeds previous work:

Firstly, for Allegro, the prediction task is for force and energies. The computational challenge they encounter is fundamentally different from ours, and comparisons are not possible as:

(1) Energy and force predictions involve only node embeddings, while Hamiltonian predictions also include edge embeddings to capture off diagonal blocks, and edges far outnumber the nodes (e.g. 3000 atoms, 500,000+ edges)

(2) H predictions require higher degree tensors to capture all orbital interactions.

Scalability comparisons can, however, be made with previous Hamiltonian models. We provide a summary in Table 2. There are two main categories:

SO(3) models (e.g. QHNet) use direct tensor products, and scale poorly with the lmax (maximum spherical degree) dimension. Because of this, they cannot be scaled to a large number of atoms.

SO(2) convolution models (e.g. DeepH2, SLEM) are more scalable with respect to lmax, but their examples are still relatively simple (< 100-200 atoms per unit cell) as the size of the graph used during training is limited by memory. So far, only our method has managed to train on and predict H for unit cells containing thousands of atoms due to efficient parallelization through strictly locality + augmented partitioning.

Table 2: Comparison of our model to state-of-the-art architectures on task, approach, and complexity of their respective training/testing datasets.

For performance comparison, we achieved errors (0.99-5 meV) within the range of the values obtained by DeepH2 (2.2 meV) and HamGNN (1.5-3.23 meV), despite our more challenging dataset. To further demonstrate what we can do with H predictions at such scale,we provide a concrete application example in the response to Reviewer Kuqs.

Graph partitioning literature: The paper brought up by the reviewer (Mostafa) focuses on reducing the size of the computational graph during fully-distributed training. However, in our case the memory limits can already be overcome when processing the input graph embeddings, even without tracking gradients. The approaches are thus entirely different. To the best of our knowledge, no previous work on large graph partitioning(e.g., GraphSAGE, ClusterGCN [9]) has ever targeted quantitative predictions of materials properties, for which the graph connectivity must absolutely be preserved. Our partitioning approach is therefore also original in this regard.

Note that the list of references is found in the response to Reviewer 4 (nuFS)

1. Application performance

Using the predicted H as an initial guess for SCF iterations is not our main focus, and the reduction in iterations also depends on the convergence threshold. We instead focus on key downstream applications where the Hamiltonian matrix H is the final product needed. Examples include ab-initio quantum transport (QT) simulations of nanoscale devices where electrical current is computed or the investigation of the solubility of doping atoms in semiconductors with the nudged elastic band (NEB) method [1,2]. In these cases, obtaining the required H from DFT is a crucial step that can consume more than 90% of the computational time, particularly in structures made of 1000 atoms or more.

We provide here a concrete demonstration of our approach on a valence change memory (VCM) cell made of a TiN/HfO2/Ti/TiN stack and 5,268 atoms. Normally, stoichiometric HfO2 is an insulator that blocks current flow. Applying a voltage introduces point defects (oxygen vacancies) that assemble into a conductive filament, changing the current by several orders of magnitude along the way.

We first train a model on slices of two HfO2 structures with randomly distributed vacancies, and use it to predict the H of structures where vacancies were arranged into filaments. As these filament shapes were not in the training data, this is an out-of-distribution generalization task. The predicted H can then replace its DFT counterpart in ab-initio quantum transport simulations to compute the electrical current, as summarized in Table 1.

Table 1: Summary of prediction results and computed currents for devices with vacancy configurations forming different filament shapes.

The currents computed with predicted H are well within acceptable error for downstream applications, given that the current values of the three filaments are on different orders of magnitude (1e-9, 1e-6 and 1e-5 A). They provide qualitative insight on the spatial localization of current flow, and allow us to study the reversible dielectric-breakdown process in HfO2 oxides. This has important implications, considering that direct ML inference is much faster than DFT (~2 seconds for forward pass vs 3.94 node hours on GH200 superchips).

Firstly, we can enable previously unfeasible workflows involving repeated DFT updates. By replacing the DFT step with model inference, we can rapidly obtain the Hamiltonian matrices of different structures along, for example, an Molecular Dynamics (MD) trajectory, thus amortizing the training cost (<40 node hours) across 100s to 1000s of time steps. Besides filament formation in VCM, this can also be used to study electrical behavior during phase transition of GST in memory devices. Secondly, we can perform inference on structures far larger than what DFT can handle (10k+ atoms). It is therefore not only an efficient substitute for DFT, but also an essential step to unlocking new research areas.

2. Model Comparison:

Accuracy: Firstly, slices cannot be isolated from a large structure to test other models due to periodic boundary conditions imposed during DFT. Interactions with periodic images cause the H elements of the isolated slices to be entirely different from those of the same atoms in bulk structures (MAE of 837 meV). The use of slices for training/testing is uniquely enabled by our strictly local model + augmented partitioning approach. Regardless, we can still compare our values with the accuracy achieved by SOTA architectures on smaller structures. Despite the complexity of our dataset comprising of disordered structures with thousands of atoms, we achieved errors of 0.99 to 5.16 meV that are within the range of the values obtained by DeepH2 (2.2 meV) [3] and HamGNN (1.5-3.23 meV) [4] for simpler periodic structures.

Scalability: Note that many previous works use the term ‘scalable’ to refer to the l-max dimension, which can be tackled with SO(2) convolutions [3]. On top of this approach, our augmented partitioning method also enhances scalability with respect to the number of atoms/edges in the graph (1k+ nodes ~ 100k+ edges), by allowing large graphs to be broken down and trained independently. To the best of our knowledge, no other work has trained on graphs with 1k+ atoms. We also provide analysis of the speedup offered by our approach when compared to conventional full graph training used by other models in Appendix H.

Note that the list of references is found in the response to Reviewer 3 (nuFS)