Learning the Electronic Hamiltonian of Large Atomic Structures
Authors: Chen Hao Xia
Manasa Kaniselvan
Alexandros Nikolaos Ziogas
Marko Mladenović
Rayen Mahjoub
Alexander Maeder
Mathieu Luisier
Venue: ICML 2025 poster (2025)
Keywords: Density Functional Theory
Equivariant Graph Networks
Materials Science
Primary area: applications->chemistry_physics_and_earth_sciences
TL;DR: We developed an approach combining a local equivariant graph neural network with an augmented partitioning method that enables partitioned training of large atomic structures.

Abstract: Graph neural networks (GNNs) have shown promise in learning the ground-state electronic properties of materials, subverting *ab initio* density functional theory (DFT) calculations when the underlying lattices can be represented as small and/or repeatable unit cells (i.e., molecules and periodic crystals). Realistic systems are, however, non-ideal and generally characterized by higher structural complexity. As such, they require large (10+ Å) unit cells and thousands of atoms to be accurately described. At these scales, DFT becomes computationally prohibitive, making GNNs especially attractive. In this work, we present a strictly local equivariant GNN capable of learning the electronic Hamiltonian (**H**) of realistically extended materials. It incorporates an *augmented partitioning* approach that enables training on arbitrarily large structures while preserving local atomic environments beyond boundaries. We demonstrate its capabilities by predicting the electronic Hamiltonian of various systems with up to 3,000 nodes (atoms), 500,000+ edges, ~ 28 million orbital interactions (nonzero entries of **H**), and $\leq$0.53\% error in the eigenvalue spectra. Our work expands the applicability of current electronic property prediction methods to some of the most challenging cases encountered in computational materials science, namely systems with disorder, interfaces, and defects.

Learning the Electronic Hamiltonian of Large Atomic Structures
Chen Hao Xia
Manasa Kaniselvan
Alexandros Nikolaos Ziogas' Marko Mladenovié Rayen Mahjoub Alexander Maeder
Mathieu Luisier'
Abstract
atomic structures, and thus easily represented by graph- structured data (Veličković al., 2018). They have already enabled new directions of research previously thought com- putationally unaffordable. More recently, GNNs have been applied to predict electronic-level properties. 'he most general of such properties is the ground-state Hamiltonian matrix H, which, if expressed in a localized basis, can be decomposed into sub-matrices H, that encode the cou- pling between the sets of basis elements located on atoms i and j. These coupling terms depend on the nature (atomic type) and relative coordinates of the local environment.
Graph neural networks (GNNs) have shown promise in learning the ground-state electronic properties of materials, subverting ab initio den- sity functional theory (DFT) calculations when the underlying lattices can be represented as small and/or repeatable unit cells (i.e., molecules and periodic crystals). Realistic systems are, how- ever, non-ideal and generally characterized by higher structural complexity. As such, they re- quire large (10+ Å) unit cells and thousands of atoms to be accurately described. At these scales, DFT becomes computationally prohibitive, mak- ing GNNs especially attractive. In this work, we present a strictly local equivariant GNN capable of learning the electronic Hamiltonian (H) of realistically extended materials. incorporates augmented partitioning approach that enables training on arbitrarily large structures while pre- serving local atomic environments beyond bound- aries. We demonstrate its capabilities by predict- ing the electronic Hamiltonian of various systems with up to 3,000 nodes (atoms), 500,000+ edges, 28 million orbital interactions (nonzero entries of H). and <0.53% error in the eigenvalue spec- tra. Our work expands the applicability of current electronic property prediction methods to some of the most challenging cases encountered compu- tational materials science, namely systems with disorder, interfaces, and defects.
Previous works on electronic property prediction focused on the cases of molecules (Zhong et al., 2023; Yu et al., 2023b; Bai et al., 2021) and ordered materials (Li et al., 2022; Gong et al., 2023; Wang et al., 2024a) whose graph representations are fairly small; typical molecules contain only few atoms and all relevant structural information in crystalline materials can be captured by translating a small unit cell. The electronic properties of such systems can be computed "exactly" with density functional theory (DFT) within a few minutes on a single CPU.
Real materials, however, are not made with the repetition of small unit cells. The unavoidable presence of defects, for example, doping atoms, vacancies, strain, compositional variations, or lattice mismatch (Ducry et al., 2020; Kani- selvan et al., 2023; Strand et al., 2018), requires large unit cells composed of to few thousand atoms to capture the associated structural disorder (Repa & Fredin, 2023). Accurate DFT simulations of such non-ideal materials are prohibitively expensive, even on today's 's largest supercom- puters. The prospect of applying machine learning solutions to handle such systems is thus particularly attractive and of high relevance in computational materials science.
1. Introduction
Graph neural networks (GNNs) have proven capable of learning properties that are functions of molecular and
Here, we adapt equivariant GNN approaches to learn the ground-state Hamiltonian H of this more realistic class of atomic structures. As main contributions:
'Equal contribution 'Integrated Systems Laboratory, Depart- ment of Information Technology and Electrical Engineering, ETH Zurich, Zurich, Switzerland. Corresoondence to: Chen Hao Xia <chexia @ iis. ce.te.c.th,,, Manasa Kaniselvan <mkanisel- van @ www.sece.tthc......
We develop strictly local GNN-based architecture tailored for electronic property prediction from small to large scales. Our network leverages efficient SO(2)- convolutions and multi-headed attention, with learn- able embeddings for node and edges that are mapped diagonal (H+,) and off-diagonal (H,,) blocks.
Proceedings of the 42nd International Conference Machine Learning, Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

Learning the Electronic Hamiltonian of Large Atomic Structures
We propose efficient augmented partitioning method that incorporates virtual nodes/edges enable parti- tioning of input graphs while maintaining their exact connectivity. Arbitrarily large graphs can then be de- composed into independent partitions that fit into GPU memory during training without compromising the achievable testing accuracy.
Orbital basis
Atomic graph
Nom2
(3)
We combine our network and augmented partitioning ap- proach to treat custom dataset of three large atomic struc- tures in their amorphous phase. As amorphous materials encompass all aforementioned defect types, they ide- ally suited to test and validate our model. In particular, we achieve 2.17-2.58 meV prediction accuracies on unseen samples, matching the ranges of previous studies for ma- terials with orders of magnitude fewer atoms (Wang et al., 2024b). We also demonstrate how this error translates into downstream applications by matching the eigenvalues of H to within 0.53% relative L1 error on structures with 3,000 atoms, which require several (3.65) hours to compute using DFT. Our work advances applications of equivariant GNNs towards practical use cases in computational materials sci- ence.
2
Node
41
Block
No0.4
Nona
Edge
(6)
Block
Orbital blocks
Hamiltonian matrix
Figure I. Schematic of the mapping between the atomic graph and the blocks of the Hamiltonian matrix H in the localized orbital ba- sis of choice. Each orbital block represents the couplings between atomic orbitals on the same atom (Hss., diagonal) or between two different atoms within Tent (H.,. ditdoiaa.na).
equation at the core of DFT takes the form of generalized eigenvalue problem: H Sit. Here, the Hamiltonian matrix H (NN) has entries H,,, = (pp H (r)l where Á (r) is the so-called Hamiltonian operator. The Overlap matrix (N) is defined as S.j (42: (p3)) They both coarse-grained matrices of size N = Σx N aatoms - where Natoms is the number of atoms, Norb the number of orbitals per atom, and i indices over the different atomic species found in the system. Note that S reduces to the iden- tity matrix in case of an orthogonal basis (sp). Otherwise, can be directly computed from the basis as the problem's physies does not influence it.
2. Background & related work
The electronic properties of material refer to its set of energy levels (8) and wavefunctions (1)) that electrons can occupy. They correspond to the eigenvalues and eigenvec- tors of the Hamiltonian matrix H describing the atomic system of interest. This quantity is a function of the location (relative positions {r: }) and identity (atomic numbers {2,}) of all constituent atoms {i} (Hoheeberg & Kohn, 1964)). Therefore, predicting the electronic properties of materi- als consists of learning the mapping F' {r,, [i} + H between the atomic structure and the elements of the corre- sponding Hamiltonian matrix, before diagonalizing H to obtain E and vb (Fig. 1).
The Hamiltonian matrix can be decomposed into sub- matrices Hi,; of size (Ni No each describing the interactions between all basis elements (orbitals) on atoms i and j. Diagonal blocks (H,,) are the interactions between orbitals on the same atom. When represented on a local basis, the matrix near-sighted; the interactions between orbitals on different atoms decay exponentially with increas- ing interatomic distance. Since atomic orbital basis is used, the sub-matrices are also equivariant under rotation of the atomic bonds, with their transformation properties related by the Wigner-D matrix.
The entries of the ground-star state Hamiltonian matrix are typically computed from first principles with DFT ((Kohn & Sham, 1965)). In several widely used codes, the wave- functions expanded into basis (1) of non-orthogonal atomic orbitals localized around atomic positions, each built, for example, from contracted Gaussian functions (KKühne et al., 2020; Neese, 2)). These orbitals trans- form like spherical harmonics under rotation - /': Y (L((F)) (F') = Σm. D p' (R)),,(()). Here, Y! the spher- {-I....,1) Di D' (R) is the Wigner-D matrix of degree correspond- ing to the rotation R. which transforms the spherical har- monic è and p' are normalized direction vectors.
ical harmonic of degree and order
2.1. Challenges unique to disordered materials
Computing the electronic properties of disordered materials with DFT still requires defining a repeatable 'unit cell' that is translated through space to construct the desired atomic system. Periodic boundaries are applied to the surface of
The localized nature of the basis states leads to finite spa- tial overlaps between them. The resulting Schrödinger-like
2

---

Learning the Electronic Hamiltonian of Large Atomic Structures
In other words, all network operations / acting on input embedding æf of degree 1 must satisfy: f(D (R) (R)))=))))))))))))))))))))))() D' f(x'). The resulting networks are trained us- ing Message Passing (MP), where each MP layer works as follows: An atom i receives input messages from all its neighboring source atoms j. Each input message goes through convolution operations that combine features with different I while preserving equivariance; a specific out- put embedding z; of degree 17 can be computed through:
2²A Eindy z'j 88/; 11 14, d2.d da..d,daaaaa (f,)). Here, fyy is a normalized vector indicating the direction of the edge con- necting the atoms j and i, and h is a set of trainable weights. The sum runs over tensor products which take æ (a source input embedding of degree 4)) and Y's (a filter spherical harmonic embedding of degree 12) and produce the output embedding:
Figure 2. Illustration of the differences between ideal periodic (left) and compositionally (middle) or structurally (right) disordered materials, whose features can only he described by large unit cells.
In all examples, black box delimits the smallest repeatable unit cell, while the circles/lines correspond to atoms/bonds.
(z)) 8/2 y'a (fri))ing =
E C(al)) Jm dind Ycla (fee), m2.
where C²= Cla.me (1.m)) (1)m))mmm))) are the Clebsch- coef- ficients that are indexed by the order m and degree I of the input, filter, and output embeddings. The combination of feature (x) and geometric (f) information along each edge encodes both the identity and structure of the sys- tem. These 'Tensor Field Networks' (TFNs) (Thomas et al., 2018) achieve state-of-the thetterrrttttt accuracy on small molecule (Yu et al., 2023b) and crystalline (Gong et al., 2023) datasets. However, they are also computationally expensive. The net- work training scales with ((G,o)), where lmaz is the max- imum degree of the angular momentum considered. As a consequence, E(3)-equivariant tensor product networks difficult to apply beyond a few atoms (Zhang et al., 2024).
this cell to avoid dangling bonds. However, the assumed periodicity can alter the disordered nature of the material if the repeating unit cell is too small. If atoms can interact with all their periodie images, non-physical phenomena may arise such as the formation of coherent electronic states across cells, drastically affecting the material properties. Small unit cells thus cannot accurately represent the physics of disordered systems. To avoid such unrealistic scenarios and better approximate disorder, 'large-enough' unit cells should be constructed, with dimensions ranging from 10 Å (Repa & Fredin, 2023) to a few nanometers (Duery et al., 2020). Generating the Hamiltonian matrix H of these systems with DFT involves tens hundreds of self-consistent field (SCF)
=
loops, each requiring the diagonalization of an intermediate H, the solution of Poisson's 's equation, and the creation of new Hamiltonian matrix. As the first operation scales with O O(N analyzing electronic properties for large disordered systems (or different representations of disorder for the same material) is often computationally unaffordable.
Recently, the cost of training equivariant GNNs has been significantly reduced by combining the benefits of data rota- tion and equivariant network operations. These approaches take advantage of the fact that when edges are rotated to align with a fixed axis (y or 2, depending on convention), the only non-zero spherical harmonic components are those of order = O. By keeping track of the bond vectors and performing internal spherical rotations, complex SO(3) con- volutions can thus be reduced to SO(2) linear convolution operations (Passaro & Zitnick, 2023). Furthermore, under these conditions, the Clebsch-Gordan coefficients exhibit a predictable sparsity pattern (non-zer only when my = + m1). Altogether, the scaling reduces to ((2maz) (Wang et al., 2024a), speeding up training, while allowing for higher-order angular momenta (Imaz) and more parameters to capture finer, more complex details of the surrounding environment. An advanced SO(2) convolution network was developed by (Passaro & Zitnick, 2023) (eSCN), and was further expanded by (Liao et al., 2023) (FquiformerV2)
2.2. Hamiltonian prediction models
Only a few studies have attempted to directly predict the Hamiltonian matrix H of a given material with a GNN rather than fitting its invariant quantities such as the total energy. The key is to constrain the solution space by leverag- ing prior knowledge of physical symmetries, e.g., rotational equivariance of orbital blocks.
In such equivariant GNNs, the predicted Hamiltonian ro- tates along with the input (Yu et al., 2023b; Zhang et al., 2024; Batatia et al., 2023; Gong et al., 2023), which re- quires maintaining -O(u)-eqqivariance within the model.
3

---

Learning the Electronic Hamiltonian of Large Atomic Structures
with the inclusion of equivariant attention. An implemen- tation of this architecture on Hamiltonians by (Wang al., 2024a) achieved better performance on custom crystalline 2D-material datasets compared to previous tensor field and invariant networks. However, the prediction of large disor- dered structures still remains an open problem in literature.
components are initially set O.
During the node update phase, each node receives mes- sages from all its neighbors j. consisting of the concatenated embeddings T& R2j. and ezz. To enable fast tensor prod- uet operations in large atomic structures, we adopt eSCN convolutions (Passaro & Zitnick, 2023), with incoming mes- sages rotated to align with the 2-axis during these operations. Non-linearity then introduced through a gate activation layer. On top of this, we also include mechanisms, in EquiformerV2 (Liao et al., 2023), which allows our network to learn better from highly dense local atomic environments with varying node degree. Finally, the resulting output messages are rotated back to their original orientations, aggregated onto the node i, and passed through feedforward network to update its output embedding n. corresponding to onsite (diagonal) Hamiltonian blocks.
3. Methods
attention
We first present an equivariant architecture that is custom- made for large-scale Hamiltonian prediction. A high-level overview is shown in Fig. 3(a)-bb. Relevant implementa- tion details and ablation studies are presented in Section 4 and Appendix
3.1. Orbital block processing
ag Each block H of the Hamiltonian matrix represents the
To learn offsite (non-diagonal) blocks, we use concepts from the Hamiltonian prediction networks in (Gong et al., 2023), (Wang et al., 2024a), and introduce learnable em- beddings for every edge, defined between pairs of atoms located within distance Yout from each other. The updated node embeddings are used update the edges through a similar process, without the attention layer. The predicted outputs are then post-processed back into the uncoupled basis with the reverse transformation: (H H such that they can be used to reconstruct the Hamiltonian matrix block-by-block.
interaction (coupling) between the spherical harmonics of degree la on atom i and that of degree Is on atom j. Mathe- matically, it can be cast into the tensor product la @ la of length (2la + 1) x (21g + 1) between the uncoupled angular momentum eigenstates l,, and (g. Each of these tensor prod- ucts can be decomposed into a direct sum (+) of coupled an- gular momentum eigenstates LL, M), where L ranges from 1/a L6 to Ia tsl: T(a (lg) = 1/4- Lg] @ 2 (la + (g). The transformation T from the uncoupled to the coupled basis is performed using matrix of Clebsch-Gordol coef- ficients. For example, the coefficient for specific [L, .M) coupled angular momentum eigenstate component is given by:
3.3. Training on augmented partitions
A notable difference between our application and that of standard GNNs in computational materials resides in the density of interatomic interactions. Many types of physi- cal interactions are short-range and can be captured with limited receptive field (Batatia al., 2025). However, despite the near-sightedness property of Gaussian functions, orbitals located on different atoms can interact with each other over distances exceeding 100 A (see Appendix F), giving rise to specific nonzero off-diagonal blocks in the Hamiltonian matrix. They must be taken into account by increasing Feut. The graph representation of these structures is thus densely connected, with > 102 edges per node,
Hl Σ Ma = la my
()
M) Dasmalllam (aa, ma) (s,,m))
Applying a rotation R to the resulting H. then consists of independently applying the corresponding Wigner-D trans- formation (R) to each of its subspaces of degree l. To minimize the number of transformations, all orbital blocks are initially transformed into the coupled basis as pre- processing step (Fig. 3(a)).
The combination of large graphs, dense connectivity, and large tensor representations of nodes and edges results in high memory consumption and long computation times per epoch during training (see Appendix H.2). If a distributed computing environment was used, the communication of intermediate values during the forward pass would incur significant overhead, impeding scalability (Wan al., 2022). At the same time, graphs cannot be arbitrarily partitioned by removing inter-atomic connections. Doing so misinforms the network as it tries to fit the target data while aggregating inputs through an incorrect/ncommlete graph structure. The
3.2. Network architecture
The graph's nodes/edges are first initialized with embed- dings of shape (Nv,, (mmax + 1)2, E,, M(Ne, (Imax + 1)2, Ee), where Nn/e is the number of nodes/edges, Enje the dimension of the node/edge embeddings, and Imaz the max- imum degree of the features. The / = 0 channels of the node embeddings are initialized with atomic numbers, while those of the edge embeddings are initialized with the scalar distance between the two connecting nodes, expanded in the chosen basis, here, contracted Gaussian functions. All other
4

---

Learning the Electronic Hamiltonian of Large Atomic Structures
|rl.
Orbital blocks
Basis transformation
Node/edge embeddings
Node update
Edge update
On!
Layer Norm
Layer Norm
1=2
Rotation
Rotation
S0(2) Layer
SO(2) Layer
activation
Gate activation
SO22) Layer
Rotation
Attention
Feedforward
(6)
Augmented partitions
Rolation
pertodic
labeled
Node eggregation
nodes
Feedforward
flos. Sout
naut, 6y
virtual nodes
periodic
Linear Orbital Block Output Head
fsice
Figure 3. (a) Data transformation of the Hamiltonian matrix. Blocks of orbital interactions are first extracted from H and reshaped into input tensors, which are transformed into coupled basis. The tensor corresponding to each node and edge expanded into dimension E,, E.. and this set of initial embeddings for every node and edge is sent into (b) the node update block. The node features (nxx are updated using message- paaesag scheme. The edge features (ex) ) are updated based on the learned nx Z. refers to atomic numbers, while ry is the set of scalar distances between atoms. (e) Illustration of how connectivity beyond the partition boundaries is incorporated through the addition of virtual edges (e) from virtual nodes (j") to the labeled nodes. A set of edges from labeled and virtual nodes is shown for single node (i) within the partition. Solid vertical lines indicate the partition boundaries, and different colors represent different atomic species.
inability divide the graph into batches also slows down the training process overall.
As the set of virtual connections used to augment each graph includes only the 1-hop neighborhood, the network is strictly local. Limiting the receptive field is an approach often used to increase scalability. As demonstrated by previous strictly local architectures, e.g., Allegro ( (Musaelian et al., 2023), information from the local environment is often sufficient to achieve -ftte-of-the-rtt prediction accuracy when inter- actions are strongly localized and the interaction cutoff is sufficiently large. In Appendix F, we present further details on how the Hamiltonian matrix elements satisfy this local- ity. To capture higher body-order information, many-body interactions can flexibly added into the network without increasing the receptive field, as implemented by (Zhouyin et al., 2024).
To efficiently train graphs corresponding to large materials while maintaining correct atomic environments and neigh- boring edge connections, we introduce an augmented par- titioning approach. A visual representation of is pro- vided in Fig. 3()). The graph is first partitioned into sub- graphs where the data can fit into the memory of single GPU. Atoms located outside of given partition, but con- nected to those within, are represented by virtual nodes (Fig. 3(c)- dashed circles). They are attached to the parti- tion through virtual edges (Fig. 3(c) dashed lines). These virtual nodes/edges are initialized similarly to their labeled counterparts with input atomic numbers and distances. How- ever, their outputs are not used. Their only purpose is to inform each partition of its full connectivity so that the network can then learn an accurate and generalizable ag- gregation function during message passing. To leverage the periodic boundaries of the material structures treated here, we partition the imput structures by dividing the graph into 'slices' along the longest dimension (z-axiss, retaining edges across the y- and z-boundaries. Details about the construction of partitions are provided in Appendix B.
3.4. Dataset creation
To generate sufficiently rich training data, existing datasets typically sample molecules various time steps of molec- ular dynamics (MD) trajectories (Yu et al., 2023a; Schütt et al., 2019; Christensen & von Lilienfeld, 2020) or generate multiple small perturbations of the atoms in a crystalline lattice (Li et al., 2022). Here, as a representative subset of realistic (disordered) materials, we consider amorphous crystals and take advantage of the fact that (1) every node
5

---

Learning the Electronic Hamiltonian of Large Atomic Structures
Material
Structure
Purpose
Fout [A]
# atoms
# orbitals
# edges
J: [A]
[Å]
2 [A]
a- a-HIO2
validate
3,000
18,000
527,348
52.876
26,308
26.242
a- 2
train
3,000
18,000
533,364
52.346
26.237
26,293
a-HfO2
test
3,000
18,000
530,920
52.722
26.267
26.191
a-HfO2
3,000
18,000
1,792,760
52.722
26.267
26.191
a-GST a-GST
train/validate (6:1 split)
12
1,008
13,104
230,848
29.541
25,583
41,777
test
1,008
13,104
226,406
25.857
29.857
41.691
a-PlGe
train/validate (10:1 split)
2,688
16,128
319,262
82.283
23.171
25.031
a- PtGe
test
2,688
16.128
319,306
82.283
23.171
25.031
a- a-PiGe
test
10
2,688
16,128
629,790
82,283
23.171
25.031
Table 1. Attributes of the generated dataset for three materials, each with its own training, validation, and test set: The [ar, V, triplet defines the periodic unit cell size. uuzh is the number of non-zero elements in the Hamiltonian, encompassing all orbital interactions. Edges were defined according interaction distance Tent.
components (Brandbyge et al., 2002). Structure examples can be visualized in Fig. 4 for each material. The corre- sponding structural information used for training, validation, and testing is found in Table 1. Details behind the sample generation are given in Appendix D. This dataset will be made publicly available to serve as a reference for large electronic Hamiltonian matrix predictions.
(b)
Ge 56
Te
(c)
4. Results
Pt
For fair comparisons in experiments where the quantity of training data may vary, we used a ReduceLLRonPlatee Plateau scheduler that reduces the learning rate when no further decrease in validation loss is detected. Training is stopped once a minimum learning rate is reached. The values of the hyperparameters and the scheduler settings for different experiments discussed in Appendíx C.
Figure 4. Atomic structure example for the three materials we con- sider. In each case, the black dashed box illustrates the boundaries of the repeating unit cell.
4.1. Ablation studies of the training approach
has a different local atomic environment, and (2) each struc- ture contains a large sampling of different motifs. A wide range of training data can thus be captured within a single sample.
We study the model's s ability to generalize to different con- figurations of large systems by predicting the Hamiltonian matrix H of full a-HfO, sample (structure 3 in Table 1), which remains unseen during the training process. For all subsequent HfO₂ experiments, the augmented partitioning scheme is only applied during training on structure 2, while the H of the full unseen structure is predicted during infer- ence (structure 3). We use an Text of Á. Errors are reported separately for nodes (e,) and edges (ex) distinguish intra- and anter-atomic orbital interactions, which typically have very different magnitudes.
We generate a custom dataset for three different materials in the amorphous (a-) phase (without long-range order): HIO, a-GeSbTe (a-GST), and a-PtGe. Besides exhibiting features that require large unit cells to be described, these materials are also scientifically and technologically rele- vant: a-HfO2 found in almost all transistors and integrated circuits (Chan al., 2008). a-GST and alloys of a-Ge show high resistance contrasts between crystalline and amorphous phases, and be used as switching layers of non-volatile memory cells (Pirovano et al., 2004; Kolobov et al., 2004; Zellweger et al., 2024) that are fully compatible with CMOS fabrica- tion processes. Knowing the electronic Hamiltonian of these materials is essential to understanding the behavior of the corresponding devices and designing better-performing
common high-x dielectric that can be
First, we demonstrate the improvement in accuracy result- ing from the augmented partitioning approach introduced in Section 3.3. In particular, we examine the influence of vir- tual nodes/edges in Table Compared to training with raw partitions, the addition of virtual nodes and edges reduces the node (Cnode) and edge prediction emror (Cedge) by ~60% and 70% respectively. Such an improvement is expected,
6

---

Learning the Electronic Hamiltonian of Large Atomic Structures
cn im Fi.]
6e [m
full unseen structures belonging to the same material. The results are summarized in Table 4.
n 11
1
5.39
0.55 0.16
- 1
2.16
The strictly local architecture trained on augmented parti- tions performs consistently across the different test datasets. In all cases, prediction errors remain relatively insensitive to the partition size (Table 11 in Appendix 1). Partitioning also allows us to extend Neut further (to 12 A for HfO2) while remaining within GPU memory limitations, effectively en- compassing all non-zero matrix elements. As a result, the best achievable error for each material system ranges from 2.17 17 meV to 2.58 me for test structures containing 1000- 3000 atoms and 200,000+ to 1,792,760+ edges. These val- ues are comparable to what a previous study obtained (2.2 meV) using equivariant GNNs for smaller structures with <50 atoms per unit cell (Wang et al., 2024b).
Table 2. Ablation study on the impact of virtual nodes on the pre- diction accuracy. Testing is done on structure 3, using slices of length taiice = 3 Å from structure 2 for training and of length falice =4 A from structure I for validation. The first column indi- cates the connectivity between virtual (xx') and labeled (n) nodes. e,,e, is the error for nodes/edges.
as raw partitions omit a large proportion of boundary edges and thus incorrectly capture atomic neighborhoods. We note that as talice < Teut, the one-hop neighborhood connects all nodes within each partition.
Next, we explore the impact of augmented partitioning on accuracy by training on an increasingly sub-divided graph in Table 3. The total number of labeled atoms used for training remains the same (3,000), while the total number of labeled edges decreases with more partitions. Despite the different divisions ranging from 5 (talice ~1 Å) to 27 (fstice 2 Å) slices, Cn and ce remain very close to the values obtained by training with the full graph (latsce 52.346 A). The prediction error is thus insensitive to partition size. For small slices in this case, the reduced fraction of labeled connections along the z direction does not affect the accuracy, as the remaining data along the y and directions is sufficient to train the network.
4.3. Compositional disorder
a- a-HfO2 often exists in a sub-stoichiometric form (a-HfO,z) due to the presence of oxygen vacancies induced by the growth process. The distribution of these vacancies varies significantly from one sample to the other, introducing a statistical component to computational studies of the HfO,, electronic properties. Accounting for these compositional variations necessitates a distinct DFT simulation for each stoichiometric different structure, which is computation- ally intensive, but an ideal application of machine learning solutions.
To demonstrate our network's 's ability to treat such problems, we train a model on a single sub-stoichiometric a- a-HfO1.s structure where 10% of the oxygen atoms were replaced by randomly distributed vacancies represented as ghost atoms. We then use to predict unseen full structures with a stoi- chiometry of a-HfO,,7, a-HfO 1.8, and a-HfO The en and ea (Table 5) remain within small range (2.34-2.45 m Fn. and 1. mEk, respectively), regardless of the vacancy concentration and distribution, showing that the network generalizes very well to compositional disorder, despite be- ing trained on only one configuration. Models trained on other vacancy concentrations perform similarly. Their re- sults, along with details of the vacancy implementation, can be found in Appendix I.1.
Latioe [A]
N₂
No
Epochs
[mEi.]]
ca [m
2
27
95,398
12,143
2.43
0.23
~3
18
141,512
17,736
2.16
0.16
184,730
17.726
2.24
0.16
8
7
320,324
13,480
2.64
0.20
12
5
381,504
14,251
2.75
0.19
~52
1
533,364
17,528
2.50
0.17
Table 3. Prediction accuracy for HfO, when the network is trained on differently-sizee partitions the same graph (structure 2). using 1 MP layer. N₂ is the number of slices, N₂ is the total number of labeled edges. The total number of labeled nodes remains constant. The number of slices is equal to Ltt, where L = 52.346 Á is the full length of structure 2. The validation set is a slice of length lalice 4 Á starting at x 25 A from structure 1. The models are tested on the full unseen structure 3.
4.4. Eigenvalue spectrum of predicted a-HrO,
Next, we assess whether the prediction accuracy of the trained network is sufficient for practical application. For this, we use HIO2, as it has the highest error in our experi- ments in Table 4 and thus presents an upper bound on the expected accuracy. We assemble the full Hamiltonian of the HfO2 test structure using the network outputs (Hured), trained with 18 partitions of length falice =~3 A. We then compute the eigenvalue spectrum of Hared and its refer- ence HT as well as the error distribution between them
4.2. Structural disorder
We apply the same local architecture to our custom dataset of structurally disordered (amorphous) materials (HfO2, GST, and PtGe from Fig. 4/Table 1), which contain range of different atomic elements (Hf, O, Ge, Sb, Te, Pt), basis sets (SZV, DZVP), and bonding behavior (ionic, covalent). All models are trained using augmented partitions and tested
7

---

Learning the Electronic Hamiltonian of Large Atomic Structures
Material
Estice [A]
N₂
Feut
[A]
N₂
Ne
6n [m)
étot (m Ey.
Ftot [me
V
a-GST
5
6
12
1,008
226,406
0.97
0.08
0.09
2.40
a-PIGe
5
10
8
2,688
319,306
0.77
0,08
0.09
2.49
a-PtGe
5
10
10
2,688
629,790
0.80
0.08
0.08
2.17
a-HfOz
3
8
3,000
530,920
2.16
0.16
0.18
4.84
a-HfO2
3
18
12
3,000
1,792,760
2.13
0.09
0.10
2.58
Table 4. Summary of model performance when trained on amorphous GST, PtGe, and HfO2. N,, and Ne refer to the number of nodes and edges of the test structure. respectively. Key parameters of the structures used for training, validation, and testing are given Table
Stoichiometry (z) Test set
en [m]
m Ex]
4.5. Computational cost
Train
Compared to full-batch training of the graph, our method using 8 augmented slices results in 6.5x speedup per epoch (0.38 vs. 2.5 s) and 7.2x decrease in memory consumption per rank (8.59 61.68 GiB) without af- fecting accuracy. A more complete analysis is provided in Appendix H.2. This scaling behavior is only limited by the overhead introduced to process the virtual nodes/edges and by any computational load imbalance from partitioning. Further computational improvements could be achieved by combining the augmentation approach with optimized graph partitioning algorithms extended to leverage periodicity.
1.8
1.9
2.34
0.15
18
1.8
2.38
0.15
1.8
1.7
2.45
0.16
Table Model trained on a-Hf0.=1.s and tested on full unseen structures with different stoichiometry (1r). 18 slices of each struc- ture, each 3 thick, are used for training. The training method is identical to the from Table 3.
(Fig. 5(a)). The predicted Hored is able to reproduce all eigenvalues of HGT within 0.530% relative L1 error. The error decreases to 0.446% when eigenvalues of unoccupied states (those situated above 0.306 En.) are excluded. The remaining error is carried mostly by the largest eigenvalues and distributed around the edges of energy gaps (see Fig. 9 in Appendix F), which correspond regimes of stronger inter-atomic orbital coupling (Atkins & De Paula, 2009).
The extension of GNN-based predictions to large mate- rial systems could potentially save tremendous amounts of computational time, as the inference phase scales with O (Natomss while DFT calculations are limited to O [ atoms). While DFT calculations to obtain the HGT of small molecules (c.g.. H₂O) take only a few seconds, the same operation for a-HfO, structures made of 3,000 atoms is computationally two orders of magnitude heav- ier (0.04 vs. 3.65 node hours, see Appendix H). We have thus demonstrated the applicability of GNN approaches to a regime where exact solutions are almost entirely prohibitive for downstream applications. The model could also serve as an initial guess DFT packages to reduce the number of self-consistent field iterations that are required obtain converged electron densities (Unke et al., 2021).
Elgenvalue (H) LE,1
1.16%
4/' 0.530
[a=0.2]
-1
H pred
peppped
0.443 %
0.446 %
H er
5. Applications in Device Transport Modeling
61
0
5000
10000
15000
Investigating materials at the device scale is one of the next challenges to tackle for ML-based materials modeling (Miret et al., 2025). Here we demonstrate that Hamiltonian matrices predicted by our approach can serve inputs to ab- initio quantum transport simulations, which aim to compute the electrical current flowing through materials under an applied voltage bias. This is done through the solution of the Schrödinger equation with open boundary conditions, where the Hamiltonian matrix contains the electronic properties of the treated device. From the produced energy-resoll transmission function, the electrical current can be derived via the Landauer- Landauer-Büttiker formula (Brandbyge et al., 2002).
Index
Figure 5. (a) Eigenvalue spectrum of the predicted (Hered) and reference oT Hamiltonian matrices. The alpha value indi- cates the seatter point transparency. (H.,)) prred is symmetrized before diagonalization with H [H Ht). The rel- ative 1.11.2 errors in the eigenvalue spectrum, computed as ([((F))rredd(F) ( ( ([I (5)GT | (( (norm = 1, where E' is the vector of eigenvalues, are shown for all eigenvalues and for the ones corresponding only to occupied states (below 0.305 Et)..
The black dashed box indicates the bandgap.
8

---

Learning the Electronic Hamiltonian of Large Atomic Structures
trical transparency caused by different spatial distributions of oxygen vacancies. Device simulations with ML--aased Hamiltonian can therefore be used to understand the con- nection between atomic structure and electrical current flow in TiN/a- //T/TiN stacks. Under this level of accuracy, the model can now also generate Hamiltonian matrices for systems with sizes which are otherwise computationally infeasible with DFT.
10a
trarsmitsion
10-2
Tranmiision
1
2017
1
-
-0
10
10-1
Predictdd Reference
Prediced Reference
10-
10-3
Device
6n Ea[mEa)]
ca ca[mEa]
Tref [A]
Ipred [A]
9
133 Eneray (eV)
11
15
17
2.3
13
15
17
Energy (elV)
1.71
0.16
8.00 x x10-9
5.54x10-*
1.66
0.16
1.48x
1.06x10
Figure 6. HINTT-IIO/TT/TNN structures with different oxygen cancy distributions (top). which lead to distinct electrical transmis- sion curves (boltom). The positions of vacant oxygen atoms in the former are indicated with red points. The grey boxes in the transmission plots indicate the difference the Fermi levels of the two contacts, created through an applicd bias. Current flow is determined primarily by the transmission through this energy window.
1.59
0.15
6,99x10
x
Table 6. Summary of prediction results and computed currents for devices with vacancy configurations forming different filament shapes. Devices 1-2 are as Fig. 6.
6. Conclusion
We developed an equivariant GNN tailored to learn the elec- tronic properties of large-scale, disordered materials, and introduced an augmented partitioning approach to break down and train the graphs encountered when dealing with realistic materials. More generally, we proposed a method to tackle the training of systems represented by large, highly connected, and near-sighted graphs where strictly local environment is sufficient. Our approach can be straightfor- wardly applied to other complex materials, or adapted to attributes, such as
As test, we designed typical metal-insulator-meet resis- tive switching device by embedding the a- a-HfO, structure in Fig. with two metal contacts and creating a TiN/a- HfO2/Ti/TiN stack (Kaniselvan et al., 2023). We then com- puted its electrical current for a-HfO,, layers with different vacancy configurations. The electrical transparency of HfO, is function of its exact atomic structure well as its local stoichiometric variations. Initially, the stoichiomet- ric a-HfO, blocks current flow due to its finite band gap. The gradual introduction of oxygen vacancies changes this stoichiometry and introduces energy states within the band gap which can carry current (Devices 1-2). This current is further enhanced when defects form a filament bridging the metallic electrodes.
learn their other vibrational properties, e.g., phonon dispersions (Fang et al., 2024). The resulting network captures relevant properties of large structures in sufficient detail to achieve few-meV ac- curacy and reproduce the energy eigenvalues to under 0.6% error. Further data generation, network optimization, and enabling increased expressiveness will be the next steps.
We modeled this process using either a (1) DFT- or (2) MI--gnerraed Hamiltonian. The resulting transmission functions and electrical current values are presented for
A Acknowledgements
two HINNT-TO.TTiTTIN structures in Fig. 6 and Table 6, respectively. Training for the latter occurred on slices of amorphous HfO2 structures with uniform distribution of oxygen vacancies. In the test structures, these oxygen vacan- cies have formed partial or complete conductive filaments. We thus test whether the learned electronic structures of isolated vacancies can be extended to the case of tight inter- actions between them, which are responsible for the current.
We acknowledge funding from the ALMOND project (SNSF Sinergia grant no. 198612) and from the Swiss State Secretariat for Education, Research, and Innovation (SERI) through the SwissChips research project. This work was also supported by the Swiss National Supercomputing Center (CSCS) under project lp16. We would also like to thank Paul Uriarte Vicandi and Luiz Felipe Aguinsky for providing the HfO, and PtGe structures, respectively.
While the quantitative agreement between the transmission functions and electrical currents obtained with DFT and MI. Hamiltonian is not perfect, it is sufficient to (1) track the difference in current across five orders of magnitude, (2) machh physically relevant features such as the band edges of the HfO,O,, layer or localized regions with higher trans- mission, and (3) capture the qualitative difference in elec-
Impact Statement
The goal of this work is to advance research in computa- tional materials science through machine learning. There are many potential societal consequences of our work, none of which we feel must be specifically highlighted here.

---

Learning the Electronic Hamiltonian of Large Atomic Structures
References
ISSN 2516-0230. doi: 10. 10.1039900a00066ff. URL hett::/dx..oi.....
0.9 (0d.///ddn000000m
Anisimov, V. I., Zaanen, J., and Andersen, O. K. Band theory and mott insulators: Hubbard instead of stoner i. Physical Review B, 14:943-954, Jul 1991.
Fang, S., Geiger, M., Checkelsky, J. G., and Smidt, T. Phonon predictions with e33)-equivariant graph neural networks, 2024. URL ptatttss//arxii..rrrrr abs/
Atkins, and De Paula, Atkins' physical chemistry. Oxford University Press, London, England, edition, November 2009.
Gong, X., Li, H., Zou, N., Xu, R., Duan, W., and Xu, Gen- eral framework for e(3)-equivarian neural network repre- sentation of density functional theory hamiltonian. Nature Communicatioons, 14(1), May 2023. ISSN 2041-1723. doi: 10. 00.108886414467-023--86888--8- URL http: org/10.10380 0tt.:-11888/14467--0000 384
Bai, H., Chu, P., Tsai, J.-Y.. Wilson, N., Qian, X., Yan, Q. and Ling, H. Graph neural network for hamiltonian-base bbaseed material property prediction. Neu- ral Computing and Applications, 34(6))4625-4432, November 2021. ISSN 1433-3058. s00521-001-00611600 URL http: dxx doi. 10. 1007 000077/00511-000 0661-6666666 0-
/ /dxx////oo doi.
doi: 10.1007/
Hjorth Larsen, A., Jørgen Mortensen, J., Blomqvist, J., Castelli, E., Christensen, R., Dutak, M., Friis, J., Groves, M. N., Hammer, B., Hargus, C., Hermes, D.. Jennings, P. C.,, Bjerre Jensen, P., Kermode, J., Kitchin, J. R., Leonhard Kolsbjerg, E., Kubal, J., Kaasbjerg, K., Lysgaard, S., Bergmann Maronsson, J., Maxson, T. Olsen, T., Pastewka, L., Peterson, A., Rostgaard, C., Schiotz, J., Schütt, O., Strange, M., Thygesen,
Batatia, I. Schaaf, L. L., Chen, H., Csányi, G., Ortner, C., and Faber, F. A. Equivariant matrix function neural networks, 2023. URL https: /tapps://axxv...//// org/obgg./...ooooooo...e..ooeeee / 2310. 10434.
Batatia, I., Batzner, S., Kováes, D. P., Musaelian, A., Simm, G. C., Drautz, R., Ortner, C., Kozinsky, B., and Csányi, G. The design space c(3)-equivarian atom-centred interatomic potentials. Nature Machine Intelligence, 7(1): 56-67, January 2025. ISSN 2522-5839. doi: 10.1038/
S., Vegge, T., Vilhelmsen, L.,, Walter, M., Zeng. Z., and Ja- cobsen, K. W. The atomic simulation environment-a python library for working with atoms. Journal of Physics: Condensed Matter, 29(27):273002, jun 2017. URL https: dxx aa680e.
s42256-024-009556 URL http: /ttp://dxx /xxx doi. 10 10.1038/s4220 024-00956-1 00956-x.
doi: 10. 1001008/1100 doi Boi.org/100 /00
Brandbyge, M., Mozos, J.-L.,, Ordejón, P., Taylor, J., and Stokbro, K. Density-functional method for nonequi- librium electron transport. Physical Review B, 65(16), March 2002. ISSN 1095-3795. doi: 10.1103/physrevb. 65.165401, URL http::/t///. /xxxxxx/// doi. oow.oorgg/..103///33/0 PhysRevB. 65 6.6..56666666 05.
Hohenberg, P. and Kohn, W. Inhomogeneous electron gas. Physical Review, 136(3B): B871, November 1964. ISSN 0031-899X. doi: 10.11005/hhyseev URL http: // dx www.ddi.orggg....ooooooooooo httpo////.....1//hhaeee... B864
Jain, A., Ong, S. P., Hautier, G., Chen, W., Richards, W. D., Dacek, S., Cholia, S., Gunter, D., Skinner, D., Ceder, G., and Persson, K. A. Commentary: The materials project: A materials genome approach to accelerating materials innovation. APL Materials, (1):001002, 07 2013. ISSN 2166-532X. doi: 10.1063/1.4812325 URL https: /tt.///oiii//o org /00 wttpg//.......///./... 10w./11611111 000.////.....1111.11 4812323.
Chan, M., Zhang. T., Ho, V., and Lee, P. Resistive switching effects of hfo2 high-k dielectric. Microelectronic Engi- neering, 85(12):240-22424, December 2008. ISSN 0167- 9317. doi: 10. mee10066/..mee20000 09.12. URL http:// mee. 2008 .9.021.
dx. o wwwooo/////1 10..g/10... 1 1011111.1.... j/j.
Christensen, A. S. and von Lilienfeld, A.
Re-
vised MD17 dataset (rMD17). 7 2020. doi:
Kaniselvan, M., Luisier, M., and Mladenović, M. An atom- istic model of field-induced resistive switching in valence
10.6084/m9 www.:/if.sssaaeeee22......
URL https:
//figshare. hct./articcess/tateeeeeeee
change memory. ACS Nano, 17(9):8281-827 2023. ISSN 1936-086X. doi: 10. 1el.102aasaaoo2222255777 URL http:///// dow....ddooooooooooo.ooooo org 10. wtt.:..002/aasaaa... 2c12575.
March
Revised MD17 (all__ddtate_t_ttee aatasst__rD17_11221777
Csányi, G., Winfield, S., Kermode, J. R., De Vita, A., Comisso, A., Bernstein, N., and Payne, M. C. Expressive programming for computational physics in fortran 95+ IoP Compul. Phys. Newsletter, pp. Spring 2007, 2007.
Kohn, W. and Sham, I.. J. Self-consistent equations includ- ing exchange and correlation effects. Physical Review, 133-A1138, November 1965. ISSN 0031- 899X. doi: 10. 10.1 140. URL http:// dx dx. 1000 11tt://111110////eeeeoee 140 A1133.
Ducry, F., Aeschlimann, J., and Luisier, M. Electro-thermal transport disordered nanostructures: modeling per- spective. Nanoscale Advances, 2(7):2648 -6678 2020.
10

---

Learning the Electronic Hamiltonian of Large Atomic Structures
Kolobov, A. V., Fons, P., Frenkel, A. I., Ankudinov, A. L., Tominaga, J., and Uruga, T. Understanding the phase- change mechanism of rewritable optical media. Nature Materials, 3(10):703-70) September 2004. ISSN 1476- 4660. doi: 10. 10.108/mmat1215. URL hwtp:httt//... /xx. dof. org/10. 1038/ 1215
Passaro, S. and Zitnick, C. L. Reducing so (3) convolu-
tions so (2) for efficient equivariant gnns. In Interna- tional conference machine learning, pp. 27420-27438. PMLR, 2023.
Perdew, J. P., Burke, K., and Ernzerhof, M. Generalized gradient approximation made simple [phys. rev. lett. 77, 3865 (1996)]. Phys. Rev. Lett., 78:1 18:1396-1399 Feb 1997.
Kühne, T. D., lamnuzzi, M., Ben, M. D., Rybkin, V. V., See- wald, P. Stein, F., Laino, T., Khaliullin, R. Z., Schütt, O., Schiffmann, F., Golze, D., Wilhelm, J., Chulkov, S., Bani- Hashemian, M. H., Weber, V., Borštnik, U., Taillefumier, M.. Jakobovits, A. S., Lazzaro, A., Pabst, H., and et al. CP2k: An electronic structure and molecular dynamies software package- quickstep: Efficient and accurate elec- tronic structure calculations. J. Chem. Phys., 152(19): 194103, May 2020.
Pirovano, A., Lacaita, A., Benvenuti, A., Pellizzer, F,, and Bez, R. Electronic switching in phase-change memo- ries. IEEE Transactions on Electron Devices, 51(3): 452 459, March 2004. ISSN 0018-9383. doi: 10.1109/ ted.2003. 21.2000..23243. URL http: / dx http:/xxxdd.oo//////../ org/10. 1109 P/99/TED.. 2003. 823243.
Repa, G. M. and Fredin, L. A. Predicting electronic structure realistic amorphous surfaces. Advanced Theory and Simulations, 6(11), June 2023. ISSN 2513-0390. doi: 10. 1002/adts. URL http: / www. 10. 15.1002//dde 202300292.
Li, H., Wang, Z., Zou, N., Ye, M., Xu, R., Gong, X., Duan, W., and Xu, Y. Deep-learning density func- tional theory hamiltonian for efficient ab initio electronic- structure calculation. Nature Computational Science, 2(6):367-377, June 2022. ISSN 2662-8457. doi: 10. 00.1038/44388--02--006666- URL http: htpp//xx.dd/////
Schütt, K. T., Gastegger, M.. Tkatchenko, A., Müller, K.-R., and Maurer, R. J. Unifying machine learn- ing and quantum chemistry with a deep neural net- work for molecular wavefunctions. Nature Communi- cations, 10(1), November 2019. ISSN 2041-1723, doi: 10. 10.1038/411467-01--1333333 12--------117722...1 URL http: cxx doi. org/10 1038 /41//441111-1187-----7
Li, S., Zhao, Y., Varma, R., Salpekar, O., Noordhuis, P., Li, T., Paszke, A., Smith, J., Vaughan, B., Damania, P., and Chintala, S. Pytorch distributed: Experiences on accelerating data parallel training, 2020. URL https:
//arxiv. org/ abs/ / 20.. 15704.
Senent, M. and Wilson, S. Intramolecular basis set su- perposition errors. International Journal of Quantum Chemistry, 82(6):282-22) 292, 2001. doi: https://doi.or// 10. 1000qu2 1030. URL https: ///ooomnellbbrrrrrr wiley. htt///oi//aaa 10. 1002..000/00000 http://aaa///...//uuuuaa 1030.
Liao, Y-L.,, Wood, B., Das, A., and Smidt, T. Equiformerv2: Improved equivariant transformer for scaling to higher- degree representations, 2023. URL /tttpp:///rxxiv. org/ org/abs// 2306 12059.
Stillinger, F. H. and Weber, T. A. Computer simulation of
Miret, S., Lee, K. L. K., Gonzales, C., Mannan, S., and Krishnan, N. M. A. Energy force regression on dft trajectories is not enough for universal machine learning interatomic potentials, 2025. URL https: /arxiv. org/abs/2502. 03660.
local order in condensed phases of silicon. Phys. Rev. B. 31:5222-5271, Apr 1985, doi: 10.1103/PhysRevB. 31.5262. URL https: http///innkapppee wtt.appppp.oooo//o////ooo 1103 / PhysRevB. .2262.
Strand, J., Kaviani, M., Gao, D., El-Sayed, A.-M.. Afanas'er ev, V. V., and Shluger, A. L. Intrinsic charge trapping in amorphous oxide films: status and chal- lenges. Journal of Physics: Condensed Matter, 30(23): 233001, May 2018. ISSN 1361-648X. doi: 10.1088/ 1361-648x/aac005. URL http: / /xx. doi. org/10. 0088/1331-648x..0
Musaelian, A., Batzner, S., Johansson, A., Sun, L., Owen, C. J., Korbbluth, M., and Kozinsky, B. Learning local equivariant representations for large-scale atomistic dy- namics. Nature Communications, 14(1):579, 2023.
Neese, F. The orca program system. WIREs Computational Molecular Science, 2(1):73-78, June 2011. ISSN 1759- 0884. doi: 10.1002/wems.81 URL http: /ttp//xxxooo///o org/10. 1002000222 wwms.81 81.
Søren Smidstrup, K. S., Blom, A., Markussen, T., Wellen- dorff, J., Schneider, J., Gunst, T. Verstichel, B.. Khomyakov, P. A., Vej-Hansen, U. G., Brandbyge, M., et al. Quantumatk: An integrated platform of electronic and atomic-scale modelling tools. J. Phys: Condens. Mat- ter (APS), 32:015901, 2020. doi: 10. 10.1088/1361-648X ab4007. URL https: htttps://iossiieeeeeeee http//ppppppcccceneeeeeeeeeeeeee .iw.eeei..o..r........ stticlee/0..108// 1111133311666488////
Nordlund, K., Ghaly, M., Averback, R. S., Caturla, M., Diaz de la Rubia, T., and Tarus, J. Defect production in collision cascades in elemental semiconductors and fee metals. Phys. Rev. B, 57:7556-7555 75756,777,,00 Apr 1998. doi: 10. NOS/PhyyReevB.57.7556 URL https: //tl:::///nn//////// aps. org/doi/10. 1103 1ww.:111///hhhhe.eeeom o 7556.

---

Learning the Electronic Hamiltonian of Large Atomic Structures
Tafen, D. N. and Drabold, D. A. Realistic models of binary glasses from models of tetrahedral amorphous semiconductors. Physical Review B, 68(16), October 2003. ISSN 1095-3795. doi: 10.1103/physrevb.68. 165208. URL http:/ http::/dxxxo....oo www.:/dooooooooooooooooo 1103/ PhysRevB. 68 165208.
Wang, Y., Li, Y., Tang, Z., Li, H., Yuan, Z. Tao, H., Zou, N., Bao, T., Liang, X., Chen, Z., Xu,
S. Bian, C., Xu, Z., Wang, C., Si, C., Duan, W. and Xu, Y. Universal materials model of deep- learning density functional theory hamiltonian. Sci- ence Bulletin, 69(16):2514-2521 2024b. ISSN 2095-9273. doi:
Thomas, N., Smidt, T., Kearnes, S., Yang, L., Li, L.., Kohlhoff, K., and Riley, P. Tensor field networks: Rotation- and translation-cquivarian neural networks for 3d point clouds, 2018. URL https::///// /ttp//axivv./////////aaa abs 100s/180..00 08219.
011. URL httpattsss//w...ceeeeeeeeeeeeee com/ ccctance/arracccleep/////2222222
Wang. Z. Q. and Stroud. D. Monte carlo studies of liq- uid semiconductor surfaces: Si and ge. Phys. Rev. B. 38:1384-1391, Jul 1988, doi: 10.1103/PhysRevB, 38. 38.1384. URL https: //t::k//nnnn//nne aps 1L03/PhysRete 38 1384.
Thompson, A. P., Aktulga, H. M., Berger, R., Bolintineanu, D. S., Brown, W. M., Crozier, P. S., in 't Veld, P. J., Kohlmeyer, A., Moore, G., Nguyen, T. D., Shan, R., Stevens, M. J., Tranchida, J., Trott, C., and Plimpton, LAMMPS a flexible simulation tool for particle-bast based materials modeling at the atomic, meso, and continuum scales. Comp. Phys. Comm. 271:108171, 2022. doi: 10.10161..p::2021108177
wttp.aapp...g///////........ org/doi/10.
Youn, Y., Kang, Y., and Han, An efficient method to gen- erate amorphous structures based on local geometry. Com- putational Materials Science, 95:256-262. December 2014. ISSN 0927-0256. doi: 10.1016ff.commatsii2211 07.053. URL http: / dx doi. 1ttp.o/.//00...1111.... commat 20.4.00440000000... 053.
Unke, O.,. Bogojeski, M., Gastegger, M., Geiger, M., Smidt, T., and Müller, K.-R. Se (3)-equivariant prediction of molecular wavefunctions and electronic densities. Ad- vances in Neural Information Processing Systems, 34: 14434-14447,20 2021.
Yu, H., Liu, M., Luo, Y., Strasser, A., Qian, X., Qian, X., and Ji, S. Qh9: A quantum hamiltonian prediction benchmark for qm9 molecules. Advances in Neural Information Processing Systems, 36:40487-40503, 2023a.
Urquiza, M. L., Islam, M. M., Duin, A. C. T., Car- toixà, X., and Strachan, A. Atomistic insights the full operation cycle of a hfo2-based resistive random access memory cell from molecular dynamics. ACS Nano, 15(8):12945-12944 July 2021. ISSN 1936-
Yu, H., Xu, Z., Qian, X., Qian, X., and Ji, S. Efficient and equivariant graph networks for predicting quantum hamiltonian. In International Conference on Machine Learning, pp. 40412-40424. PMLR, 2023b.
086X. doi: 10. htt.///asmaao..11144... URL http: //dx.dor org 10. http:/1..11111aaaa11
Yuxing Zhou, Wei Zhang, E. M. L. D. Device-scale atomistic modelling ofphase-change memory materials. Nature Electronics, 6 6(10):746-754, 2023.
1601466.
Vande Vondele, J. and Hutter, J. Gaussian basis sets for accurate calculations on molecular systems in gas and condensed phases. J. Chem. Phys., 127(11): 09 2007. ISSN 0021-9606.
Zellweger, T,, F., L., K., P., C.. W., M., L., and Emboras, A. Amorphous germanium multi-functional switching layer for electro- optical memristors. In Book of Abstracts: MEMRISYS 2024, pp- 73, 2024. URL https: /ww. memr memriays2024. org httpllorgoooooooooeeeeeeeeeeeee ffw.es/ug/////////// 1llffc09ldd0445ca8888887
Veličković, P., Cucurull, G., Casanova, A., Romero, A., Liò, P., and Bengio, Y. Graph Attention Networks. In ICLR, 2018.
Zhang, H., Liu, C., Wang. Z., Wei. X., Liu, S., Zheng, N.,
Wan, C., Li, Y., Wolfe, C. R., Kyrillidis, A., Kim, N. S., and Lin, Y. PipeGCN: Efficient full-gne full-graph training of graph convolutional networks with pipelined feature communi- cation. March 2022.
Shao, B., and Liu, T.-Y. Self-consistency training for hamiltonian predictionn 2024. haagg//bb/// /2403.2403.. 09560.
funsity-functional-tee URL https:
Zhong, Y., Yu, H., Su, M., Gong, X., and Xiang, H. Trans- ferable equivariant graph neural networks for the hamil- tonians of molecules and solids. npj Computational Materials, 9(1), October 2023. ISSN 2057-3960. doi: 10. 10.10388441524402--0113338 130-4. URL http: : / // dx doi. org/10. 1 0000113//844554-2-0000003033
Wang, Y., Li, H., Tang, Z., Tao, H., Wang. Y., Yuan, Z.,, Chen, Z., Duan, W., and Xu, Y. Deeph-2: Enhancing deep-learning electronic structure via an equivariant local- coordinate transformer, 2024a. URL https: //axxiv. org/abs/ 2401. 17015.
12

---

Learning the Electronic Hamiltonian of Large Atomic Structures
Zhouyin, Z., Gan, Z.., Pandey, S. K., Zhang, L., and Gu, Q- Learning local equivariant representations for quantum operators, 2024, URL https: / org.v.org// abs/ 2407. 06053.
13

---

Learning the Electronic Hamiltonian of Large Atomic Structures
A. Loss computation during training and inference
For all experiments, a minor difference from procedures reported (Yu et al., 2023b) and (Schütt et al., 2019) that we use the Mean Squared Error (MSE) of the full output and target vectors in the coupled space to compute the fitting and validation loss during training: (T, @4) = # ا (2, ..)², where z and are the flattened targets in the coupled space. These targets padded with zeros to ensure that those of different orbital interactions have the same dimensions. We note that doing so burdens the network to learn the zero-pade zero-padding in addition to the data, which can artificially increase the number of epochs required for training. However, the computational cost of extracting padding- padding-free free data and post-processing it at the end of each epoch is also non-negligible. For this work we therefore use the padded data during training for case of implementation. The final reported loss for all our results uses the Mean Absolute Error (MAE) after converting the output and label tensors back into uncoupled space and extracting the un-padded data: AM #:) Norb
B. Augmenting graph partitions with virtual nodes
In Algorithm 1, we detail the procedure to partition the full graph G, described by the set of vertices V and edges E, into set of slices {g) GN} which are augmented by virtual nodes and edges. The number of labeled nodes (m?)) and the number of labeled edges (n() are collected and passed to the training functions, which then omit the remainder the outputs (the virtual nodes and edge outputs) while computing the loss.
Algorithm 1: Augmented partitioning approach
Graph G(V,E) Slice x-coordinates [21 ... æv] Set of subgraphs {91 ... 9N) Numbers labeled nodes
[n] ... nN] Numbers labeled edges [nf ... nf]
2 for: t 1 to N do
V: []:
n²" = 0;
for u E V do
if' v.z E [r,, 2141) then
Vaappedddv);
n +- 1:
end
10
end
11
ε <- []:
12
mf = 0;
13
for Ui E V:
for 12 E do
15
if
02 E & then ε, appeddd + v2);
16
17
n += 1;
18
end
end
20
end
for U1
E V; do
for 1/ E V V; do
23
if
21 E & then
24
V (pp)
25
E, append(
+ v1);
26
end
end
end
Gi(V,. E,);
end
14

---

Learning the Electronic Hamiltonian of Large Atomic Structures
The augmentation approach be combined with any method of partitioning the input graph. An ideal partitioning scheme would result in the maximal ratio of labeled nodes/edges within each sub-graph, compared to virtual nodes/edges. As consider structures with fully periodic boundaries in all three dimensions, simple heuristic to leverage this periodicity is to partition along only one dimension (z), thus maintaining all the labeled edges across the y- and 2 2- cell boundaries of each periodic image. This motivates the division by slices described in Algorithm 1, which can be very effective when assuming a constant atomic density. Finding more optimal partitions that nevertheless leverage periodicity is subject of future work.
C. Hy pyrraaarsstsess
We use the hyperparameters shown in Table 7 to train on all datasets. The ReduceLLRnnPleeee scheduler decreases the learning rate by the decay factor when it does not detect further decrease in validation loss within the decay patience Epatience- The threshold refers to the sensitivity of the scheduler to changes in validation loss. Once the minimum learning rate is reached, the training stops.
Hyper-parameters
a-HHoo//PGGe/GST dataset
Optimizer
Adam
Precision
single (f32)
Scheduler
ReduceLLROnPlataau
Initial learning rate
x 10
Minimum learning rate
10-5
Decay patience pppttonce
500
Decay factor
0.5
Threshold
x 10-3
Maximum degree L Lmar
4
Maximum order M Maa
Embedding size Number of attention heads N₂
16
Feedforward Network Dimension
64
Table 7. Hyper-parameters used for a-HfO2, a-PIGe and a-GST data.
D. Dataset generation
Atomic structures corresponding materials in the amorphous phase are not straightforward to generate since they must accurately capture the structural motifs underlying this phase and realistic range of atomic coordination. To accurately reproduce long-range structural disorder, the structures used must also be large enough to avoid the creation of wavefunctions that repeat over periodic boundaries. Existing methods to do so include melt-quench (Urquiza et al., 2021), cecd-and-coordinate (Youn et al., 2014), or 'decorate and relax' (Tafen & Drabold, 2003) approaches.
In this work, we use melt-quench processes with molecular dynamics (MD) to evolve each the three materials considered from their crystalline phases, following similar procedure as the ones described in (Kaniselvan et al., 2023) and (Urquiza al., 2021). We then perform structural relaxation with CP2K code (Kühne al., 2020) correct for any discrepancies between the relaxed bond lengths attained with the force field used for MD and those obtained with DFT. Due the large cell sizes of the a-HfO, structures, all necessary information is contained within the I point (where the wavevectors k, k, = k, = 0). The energies at this location can be computed by directly diagonalizing H. The datasets are publicly available at https: //bggg www.ungginfcaceeo face.e....... co/ //ddaassees chexias a8 Amorpphus Hami ltonians.
Further details specific to each material are provided in the sections below.
D.I. a-HfO
We generate 3 independent structures of a-HfO, using the QuantumATK toolkit (Søren Smidstrup et al., 2020). As a first step, we run an MD NVT simulation at 3000K for 50 ps with a step size of 1 fs. We use the MTP-HIO,-2022 potential provided by the software. Next, we run an NPT simulation for 300 ps (and the same fs step size), with an initial reservoir
15

---

Learning the Electronic Hamiltonian of Large Atomic Structures
temperature of 3000K and final temperature of 300K, for cooling rate of 9K/ps. Finally, we anneal the structure 300K for 50 ps.
Due to the computational cost of using more complete Double- Valence Polarized (DZVP) basis set (Vande Vondele & Hutter, 2007), we use simpler Single- Valence (SZV) basis (Vande & Hutter, 2007), which uses 4 basis functions per Oxygen atom and 10 basis functions per Hafnium atom. The plane-wave cutoff is set to 500 Ry, while a cutoff of 60 Ry is used for mapping the Gaussian-type orbitals onto the grid. We use the PBE functional for the exchange-correlation energy (Perdew et al., 1997). To accurately capture the band gap of a 2-HfO we apply the Hubbard correction (Anisimov et al., 1991) of U 7 eV to the 3d orbital of Tii and the Hubbard correction of U 10 eV to the 2p orbital of O.
D.2. Substoichiometric HfO,,
We create a dataset for sub-stoichiometric HfO,, structures by introducing randomly distributed oxygen vacancies into the original, pristine HfO22 structures. The sub-stoichiometric structures are generated for x 1.9, 1.8, and 1.7 (corresponding vacancy concentrations of 5%, 10%, and 15 %, respectively). Vacancies are treated as ghost atoms (atoms with no orbitals, but with basis set defined at their locations), to mitigate the basis set superposition error (Senent & Wilson, 2001), known problem related to localized basis sets. More precisely, by treating vacancies as ghost atoms, one prevents the excessive borrowing of the basis sets from neighboring atoms by the vacancy, which improves the accuracy of the predicted electronic properties. These ghost atoms are assigned an atomic number 0f
D.3. a-GST
We use two amorphous GST-124 (Ge(SbTe2)2) structures containing 1008 atoms for training and validation. The first structure is extracted from the MD simulation of GST- crystallization, provided by (Yuxing Zhou, 2023) (Supplementary Material). The initial geometry for the second structure is a perfectly crystalline GST-124 structure, taken from Materials Project database (Jain al., 2013). The amorphous structure is then generated through a standard melt-quench procedure, consisting of atomic position randomization at 3000K for 20 ps, cooling the melting point of 600K at rate of 10 holding for 30 ps, quenching to 300K at a rate of 2.5 x 10 K.s and finally, annealing at 300K for 50 ps. Both amorphous structures are obtained via MD simulations in LAMMPS (Thompson et al., 2022), equipped with the QUIP library for Gaussian Approximation Potential (GAP) (Csányi et al., 2007). The corresponding Hamiltonian matrices are obtained using CP2K, where we run calculations with the DZVP basis, the plane-wave cutoff of 300 Ry, the Gaussian-type orbitals mapping cutoff of 50 Ry, and the PBE functional.
D.4. a-PtGe
To generate the PtGe structures, germanium structures are taken from the Materials Project database (Jain et al., 2013). This is followed by an NVT melt-quench process using LAMMPS and Stillinger- Weber parameters (Stillinger & Weber, 1985; Nordlund et al., 1998; Wang & Stroud, 1988). The structures are heated to a melting temperature of 5000K at rate of 0.47 10 kept the melting temperature for 20000 ps (structure 1) or 22000 ps (structure 2), quenched at rate of 4.7 K.s-1 and finally annealed at 300K for 100 ps. Using the Atomic Simulation Environment (ASE) tool (Hjorth Larsen et al., 2017). 1/3 of the Ge atoms are replaced by Pt atoms. The cell of the alloy is then stretched to match the cell of a PiGe, structure (taken from the Materials Project and optimized using CP2K). Fixed-volume geometry relaxation is then performed on the PtGe alloy. For the structural optimization, as well as for the H and S generation, SZV basis set and PBE exchange-correlation functionals are used. We apply a plane- wave cutoff of 1000 Ry and a cutoff for Gaussian-type orbitals mapping of 70 Ry.
E. Atomic bonding environments in the amorphous phase
We use the example of a-HfO, to investigate the structural differences between samples generated from different starting melt-quench processes. In Fig. 7 we plot the O-coordination of each Hf atom and the radial distribution function g(r) (where is the inter-atomic distance) for each of the three structures. The distribution in the coordination and dispersion of the peaks in g(r) indicates the amorphous nature of the three structures. Variations between them appear as perturbations in these two quantities. To gain more insights into how different the structures are, we additionally plot the spatially resolved O-coordination of Hf atoms along the longest, x coordinate for the three structures, as well as the distributions of outliers (Hf atoms with very low and very high O O-coordination) in three-dimensional space. These outliers are situated at different
16

---

Learning the Electronic Hamiltonian of Large Atomic Structures
locations in different structures, demonstrating a significant degree of dissimilarity among the structures.
(b)
Structure
Structure
Structure $
Structure
3-
Prouunnyy
Frequency
62
0.2
0.2
0.2
1
0.0
0.0
0.0
Coortination
Coardination
Coordination
[[A]
Coordination outlier
Figure O-coordination of Hf atoms (number of O bonding Hf atom) for each the a-HfO2 structures, showing distribution around a coordination number of 6, and variation between the structures. h) The radial distribution function (g(r) dn(r) Vigmuin )
de Arr Natoms
where rr) is the number of atoms with distance between them for the structures 1-3. (c) Spatial distribution coordination outliers (Hf atoms with O-coordination equal to 8 for the three structures, which indicator of the uniqueness of the three structures.
F. Near-sightedness of the Hamiltonian matrix
In Fig. 8, we visualize the matrix elements corresponding to the Hamiltonian of one sample of each material from the dataset. Specifically, show the interaction ranges which non-zero matrix elements exist, well their decay with increasing interatomic distance. When using local basis, this near-sightedness holds strictly.
a-GST
a-PtGe
a-HfOz
Frequency (nomallzacc
Frequency (normalized
frequency (normalized)
0.20
0.15
0.15
0.2
0.10
0.1
0.05
0.05
0.00
0.0
0.00
5
10
10
15
Inyl (A)
Iny (A)
|ry (A)
10-1
10-1
maxiH [[A
10
maxiH
10-3
10-3
10-3
10-3
10-5
10-3
10-7
10-7
10-7
10
10
Inyl (A)
[nyl (A)
Ingl (A)
Figure 8. Summary of the Hamiltonian matrix properties for each of the materials our custom dataset. Shown are (top row) atomic structures, (middle row) frequency of matrix elements function of interatomic distances, and (bottom row) distribution of the maximum element of cach block of H as a function of interatomic distance Irsyl-
17

---

Learning the Electronic Hamiltonian of Large Atomic Structures
Next, using the three a-HfO2 structures in Table 1 example, we show the distribution of energy eigenvalues in Fig. different values of Text- In the second row,
into the range of eigenvalues around the energy bandgap, which is defined by the transition between occupied and unoccupied electronic states (the Fermi level EF 0.3.% in all cases). Values of T 2 8A create no noticeable difference on the eigenvalue spectra. Note that the value of Teut = oc corresponds to the case where no nonzero values were filtered from HGT
Structure
Structure
Structure 3
Eigenvalue (H) (Ew)
Eigenvalue (H) [E.]
Elgenvalue (H) (E.)
Four (.00.05)
0
Fou 0.05
Fou (..0000)
-1
-1
-1
5000
10000
15000
5000
10000
15000
5000
10000
15000
0.6
0.6
0.6
Eigenvaive (E.)
Structure
Eigenvalue (M) (Ex)
Structure
Eigenvalue (H) [E.]
Structure 3
0.5
0.4
[e=0.2]
0.4
(a-0.2)
0.4
[a=0.2)
Egg
6
Eg
Eg
8
8
0.3
10
0.3
10
12000 Index
14000
10000
12000
14000
10000
12000
14000
Index
Figure 9. Figenvalues of the ground-truth Hamiltonian matrix, showing (top) the full eigenvalue spectrum and (bottom) the spectrum around the handgap, for the three structures, the order of their appearance Table 1.
To demonstrate the locality of long- and short-range perturbations, we introduce a single perturbation at one chosen location in the structure and measure the mean absolute error of onsite Hamiltonian blocks when compared to that of unperturbed structure. The types of perturbations introduced include translation of a hafnium atom, oxygen vacancy, and substitution an oxygen atom with hafnium one, and are plotted against the distance from perturbation in Fig. 10 (a), (b). and (c), respectively.
10
10
10-3
MAE Value (EN)
10-3
MAE Value (En)
MAE Value (En)
10-3
10
10-4
10-4
10-5
10-5
10-5
10-6
10-6
10-6
10
20
30
10
20
30
20 Distance from Perturbation (A)
10
30
Distance from Perturbation (A)
Distance from Perturbation (A)
Figure Scatter plots showing the decay of MAE with increasing distance from different perturbations, including (a) 0.1 translation of Hf atom, (b) replacement of O atom with vacancy, and (c) replacement of atom with Hf atom.
In all cases, the effect of the perturbation rapidly decays with increasing distance. For the case of the 0.1 À translation, the average onsite MAE at a distance Hamiltonian block (63 EA), the perturbation only affects the matrix elements by 0.24% overall. Similarly, for vacancy and substitution perturbations, the matrix elements of atoms located Á away only changed by 0.18% and 0.12% respectively.
away given by 0.15 m, Considering the average value of an onsite
18

---

Learning the Electronic Hamiltonian of Large Atomic Structures
This implies that for our chosen cutoff of 8 A, perturbations occurring outside of the radius surrounding the atom have negligible effect its Hamiltonian matrix elements. This also indicates that the electronic structure of that atom be learned using information from the local atomic environment.
G. Cutoff radius and connectivity
now explore the minimal graph connectivity that can be used by the network to accurately learn relevant features of HfO2. To do this we use the slice partition approach introduced in Section 3.3, using 18 slices of length &alice 3 Á train the network. Results reported in (Table. 8). Reducing the value of Four below 8 Å noticeably increases the error (6n/e), thus demonstrating the sensitivity of H to the exact connectivity of the graph. Going from Feut 8 Å to 12 A, the edge error improves, while the node error begins to plateau, but the node degree (which is proportional to the memory consumption of the network) grows by 1.7x. This is consistent with the negligible changes to the eigenvalue spectra with Teut of 8 (Fig. 9).
Fout [Å] deg(n)
deg(n)'
én Em
6o [m
étot [mEx]
étot [meV]
4
21.02
10.60
2.84
0.52
0.64
17.50
72.41
32.01
2.27
0.26
0.29
7.98
8
177.61
49.89
2.16
0.16
0.18
4.84
10
347.16
78.85
2.18
0.12
0.12
3.38
12
590.25
138.59
2.13
0.09
0.10
2.58
Table Prediction accuracy of the network with different Teue- Training was done with a single slice of length fslice = A. The edge connectivity of the matrix is set by Text. deg(n) is average node degree, and deg(n)' the reduced node degree omitting virtual node neighbors. Note that for this value of takce. the majority of neighbors for the average node are virtual. and 6, are the Mean Average Error (MAE) nodes/edges, respectively. All units are in [x - E. The validation loss of the model computed from slice similar length extracted from structure 1. The networks are tested an unseen full graph (structure 3) constructed with the same Tout.
We perform similar study the cutoff radius of PtGe in Table 9. Increasing the cutoff radius once again increases the overall prediction accuracy of the trained model, mostly due to improvements in teee-prrdiitoon accuracy.
Cutoff [Å] deg(n)
deg(n)'
6n [m
6o [m
Cbot [m.E..
Étot [me V
4
14.28
10.00
1.08
0.26
0.33
8.96
6
50.17
27.37
0.80
0.14
0.16
4.28
8
118.71
51.57
0.77
0.08
0.09
2.49
10
234.28
84.07
0.80
0.08
0.08
2.17
Table 9, Prediction accuracy of model on amorphous PtGe material with different Peue
H. Compute environment and runtime comparisons
The training is performed with PyTorch Distributed Data Parallel (Li et al., 2020), where the graph partitions (slices) can be distributed between GPUs.
H.1. Memory consumption of full-graph training
During the training of the full graph model, the peak memory consumption observed was 61.68 GiB on a single NVIDIA A100 GPU. Most of the consumption does not stem from the network and the structure but from the additional memory needed for the convolution operations.
H.2. Scalability of augmented partitioning
In Fig. 11, we show the decrease in time per epoch and resulting speedup when using the augmented partitioning approach introduced in Section 3.3. Since the partitions are independent, the only communication involved every epoch is a collective to inform each GPU/rank of the loss of each other rank. The time per epoch thus decreases uniformly with the number of slices (N₂) used.
19

---

Learning the Electronic Hamiltonian of Large Atomic Structures
Despite the independence of each batch and the minimal communication per epoch, the scaling perfectly linear. The deviation from ideal speedup be attributed to factors:
Load imbalance: The partitioning approach was designed to leverage the periodicity in the yy- and 2- direction within straightforward implementation. However, it is not ideal in terms of the number of cuts/number of virtual nodes/edges required, resulting in slightly different amount of work per rank which leads observable load imbalance higher Nr- This effect can be seen in the allocated memory per partition (Fig. 11(c)). We note that the augmented partitiioingg method can be used with any standard graph- sraph-partitionin algorithm.
Computational overhead of the virtual nodes and edges: Individual nodes and edges of the graph can be repeated in labeled and virtual node lists. Treating the replicas introduces additional computational cost while training the network, which increases with N₁. This overhead is maximum with the use of very small slices (large N2), thus introducing trade-off between parallelism and time epoch.
(a)
(b)
(c)
8
2.5
Time per epoch (s)
7
2.0
6-
Memory [GB]
Speedup
40
1.5
1.0
20
0.5
1
4
5
6
7
8
1
2
3
4 5
4
8
# Slices (N+)
# Slices (N₂)
#Slices (N₂)
Figure 11. (a) Time per epoch and (b) speedup resulting from the use of increasing numbers of slices N.. Median values are shown, while the eиror bands are one standard deviation. Experiments were run on NVIDIA A A100 GPUs with # ranks set to N.. Measurements are only shown up 8 slices/8 GPUs due to limitations in available compute resources at the time of submission. The fill-between indicates the range runtime over the first 30 minutes training. The dashed black line corresponds to the ideal speedup. in which case the use of N; slices would enable Ni x speedup in the runtime per epoch.
(c) Measured peak memory consumption partitions, where each bar corresponds to different GPU. Variation in memory consumption between GPUs at each individual value of N₁ translates load imbalance, which correlates with the deviation from ideal scaling shown in (b).
function of the number of
H.3. H2O vs HfO, runtimes
In Section 4.5, we make a comparison between the computational cost of computing the Hamiltonian for H₂O molecule and the HfO₂ structure. To approximate the cost of generating them under the same computational conditions, we set up CP2K simulations with DZVP basis for H2O. The computation time per H₂O molecule was 7s, when run on 12 nodes with 12-core Intel Xeon E5-2680 CPUs and NVIDIA P100 GPU, resulting in a total of 0.04 node hours. The HfO, structures require 3.65 node hours the same compute environment (but distributed to 27 nodes). The difference, omitting scaling behavior, is roughly ~000x
Additional Tests
L.I. Sub-stoichiometrie hafnium oxide
Here, we provide a more detailed study on the prediction of substoichiometric HfO2, where the train/test structures contain different vacancy fractions. 18 slices (tatice = Å) were used all cases. The results are summarized in Table 10. The 6, and & values across different experiments lie within a small range (2.24-2.73 m.,, and 0.13-10.15 mE respectively), showing that the network generalizes well structures of different vacancies, regardless of which vacancy configuration was trained
20

---

Learning the Electronic Hamiltonian of Large Atomic Structures
Stoichiometry (æ)
En [m]
te [m
Training
Testing
1.9
1.9
2.24
0.13
1.9
1.8
2.44
0.14
1.9
1.7
2.73
0.15
1.8
1.9
2.34
0.15
2.38
0.15
1.8
1.7
2,45
0.16
1.7
1.9
2.32
0.14
1.7
1.8
2.37
0.15
2,51
0.18
Table 10. HtO,, models trained and tested with different stoichiometry (x)) using augmented partitioning. slices of each structure, each 3 Á thick, was used for training. The training method is identical the one used to obtain Table 3. Models trained HfO,= (5%). =I 1.8 (10 %), and 1107 (15% vacancies) are tested on test structures with vacancies ranging from 5% to 15%.
on. To demonstrate that the augmented partitioning approach similarly does not affect accuracy for sub-stoichiometric HIO,, we also perform full graph training using structure 2 with 15% vacancies, and compare with the augmented partitioning approach in Table 12. The comparison of en and e. values between full and partitioned approaches indicates that both approaches generalize well to different stoichiometry. These values are also close that of stoichiometric HfO, in Table 3.
1.2. Partitioning tests for other materials
For each material, train the model on partitions of different thicknesses containing the same number of atoms. The results summarized in Table 11. The same test also repeated on substoichiometrie HfO1.7 and tested structures with different stoichiometry (shown in Table 12. It can be seen that augmented partitioning preserves the accuracy of prediction for all material cases studied. For HfO,, it also preserves the model's s ability to generalize to structures of different stoichiometry.
Material
tstseelA [A]
Nr
[ew
.n.
Ne
<n <n/mbh)
6o [m[
écot [mb]].
Ctot [mel
a-GST
1
12
710
138532
1.15
0.08
0.08
2.28
a-GST
5
6
12
710
48356
0.97
0.08
0.09
2.40
a-PtGe
1
1633
182372
0.77
0.08
0.08
2.26
a-PtGe 5
10
8
1633
84416
0.77
0.08
0.09
2.49
Table 11. Summary of model performance when trained amorphous GST and PtGe structure with slices of different thicknesses. The structures used training, validation, testing are displayed in Table N, and N, refers total number labeled nodes edges in the training respectively.
Training method
Stoichiometry (x) Testing set
[mFF]]
[m Fi..]
partitioned partitioned
1.9
2.32
0.14
1.8
2.37
0.15
partitioned
1.7
2.51
0.18
full full
1.9
3.02
0.12
1.8
2.54
0.11
full
1.7
2.52
0.11
Table 12, Comparison between graph training and the augmented partitioning training using the same HfO structure with 15% vacancies. Models are tested on structures with vacancies ranging from 5% to 15%.
I.3. Small Molecule Benchmarks
We test our backbone architecture setup on part of the MD17 benchmark (Schütt et al., 2019), which consists of small- molecules, each with (water), 9 (malondialdehyde) or (uracíl) atoms. The hyperparameters used are listed in Table. 14.
21

---

Learning the Electronic Hamiltonian of Large Atomic Structures
Dataset
Model
Train
Validate
Test
Batch size
Ftot [(
QHNet
3900
10.79
water
This work
500
500
3900
5.60
QHNeet 25,000
1,478
21.52
malondialdehyde
This work
25,000
50
1,478
16.60
QHNet 25,000
4,500
20.12
uracil
This work
25,000
50
4,500
5
13.80
Table 13. Comparison with QHNet MD17 benehmark datasct containing water, malondialdehyde and uracil
The number of training steps for each experiment is fixed 200,000 match that of QHNet in (Yu al., 2023b). The results in Table 13 show that the performance of our setup comparable that of the baseline for all measured cases.
MD 17 dataset
Optimizer
Adam
Precision
double (f64)
Scheduler
Reducel.ROnnlltteal
Criterion
RMSE MAE
Initial learning rate
10
Minimum learning rate
1 10-10
Number of MP Layers
2
Decay patience /aationee
50
Decay factor
0.5
Threshold
1 x 10-5
Maximum degree Lmaz
4
Maximum order Mmax
Embedding size
128
Number of attention heads N₁.
2
Feedforward Network Dimension
64
Table Hyper-parameters used MD17 dataset.
22


Reviews:
[{"review_id": "QGOZW7m1KK", "reviewer": "Reviewer_Kuqs", "summary": "The authors propose a method based on SO(2) graph neural networks to learn the electronic Hamiltonian matrix in the structural relaxation process of inorganic materials. To address computational challenges associated with defective or large-scale crystalline structures, they introduce an efficient partitioning strategy that enables the model to handle arbitrarily large structures while preserving local atomic environments.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "**1. Application performance**\n\nUsing 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. \n\n**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. \n\nWe 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.\n\n| Device      | Node Error [mEₕ] | Edge Error [mEₕ] | Current (Label H) [A] | Current (Predicted H) [A] |\n|------------|------------|------------|----------------|----------------------|\n| Filament 0  | 1.81       | 0.20       | 1.98E-09      | 1.47E-09           |\n| Filament 1 | 1.74       | 0.20       | 1.12E-05      | 9.99E-06           |\n| Filament 2 | 1.66       | 0.20       | 1.45E-06      | 1.96E-06           |\n\n**Table 1:** Summary of prediction results and computed currents for devices with vacancy configurations forming different filament shapes.\n\nThe 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). \n\nFirstly, 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.**\n\n**2. Model Comparison:**\n\n**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. \n\n\n**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**.\n\n\n**Note that the list of references is found in the response to Reviewer 3 (nuFS)**", "claims_and_evidence": "The paper makes two primary contributions: (1) the development of a strictly local graph neural network (GNN)-based architecture for efficient prediction of the electronic Hamiltonian matrix; and (2) the introduction of an augmented partitioning method that enables large crystalline structures to be decomposed into smaller subgraphs while maintaining prediction accuracy comparable to global computation, thereby significantly reducing memory consumption.\nFor (1), the authors provide an extensive review of existing GNN-based Hamiltonian prediction methods in Section 2.2, but lacks a systematic comparison between the proposed architecture and previous methods, such as its advantages in accuracy, computational efficiency, and scalability.\nFor (2), I think the augmented partitioning method is well-supported with convincing evidence.", "methods_and_evaluation_criteria": "The authors construct three large-scale datasets of disordered crystalline materials and conduct a comprehensive analysis of prediction. Additionally, the authors employ three data augmentation strategies—atomic position perturbation, oxygen vacancy filling, and atomic substitution, to evaluate the model’s ability to capture both short-range and long-range interactions. The methodological design is robust.", "theoretical_claims": "I have reviewed the theoretical claims in the paper and have not identified any apparent issues.", "experimental_designs_or_analyses": "The authors’ experimental design and analysis are comprehensive and align well with general methodologies in Hamiltonian matrix prediction research. The experiments involve multiple large-scale disordered crystalline datasets, and the partitioning-based prediction method is rigorously evaluated for both model performance and effectiveness. \nHowever, the paper provides little discussion on the practical effectiveness of the method or its impact on computational efficiency. In the final sentence of Section 4.5, the authors mention that the proposed method could serve as an initial guess for DFT packages to reduce computational costs. Given that one of the key applications of Hamiltonian matrix prediction is to expedite electronic structure calculations, it would be beneficial for the authors to include a more detailed quantitative analysis.", "relation_to_broader_scientific_literature": "In computational materials science, reducing computational time complexity while maintaining prediction accuracy has long been a key objective. In this work, the authors propose a partitioning-based method that addresses the challenge of directly handling large-scale disordered systems, offering promising applications in other disordered material systems.", "essential_references_not_discussed": "I think the authors have adequately cited and discussed prior work.", "other_strengths_and_weaknesses": "Strengths: The proposed partitioning algorithm achieves good computational performance without sacrificing accuracy, particularly for large-scale disordered material systems.\nWeaknesses: The paper lacks comparative analysis with other existing Hamiltonian prediction methods and does not provide quantitative analysis, such as how many SCF iterations can be reduced or how much speedup can be achieved in actual calculations.", "other_comments_or_suggestions": "No.", "questions_for_authors": "1. Application performance, see \"Experimental Designs Or Analyses\".\n2. Model comparasion, see \"Claims And Evidence\".\n3. How would this method perform when applying to small crystal structure (maybe a slice of the unit cell), is it possible to provide the prediction result analysis of a small slice of one structure from the proposed dataset (compared with previous works)?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "S8Z4VveHiZ", "reviewer": "Reviewer_9oL5", "summary": "The paper proposes a new method for applying GNNs to learn the electronic Hamiltonian for atomic structures beyond the unit cell. The paper starts by introducing the use of GNNs for property prediction on fairly small size unit cell materials and molecules and motivates the need to capture more complex materials behaviors, such as defects and strains, at larger size scales. To enable GNNs to be applied for larger scale atomic systems, the paper proposes two contributions: 1. Learning the ground-state Hamiltonian using a SO(2)-equivariant GNN that learns local embeddings for the diagonal and off-diagonal blocks of the Hamiltonian. 2. A graphs partitioning method that enables partitioning larger input graphs, which represent larger atomic systems, to more effectively deploy the proposed GNN. The final paragraph of the introduction highlights that the authors deploy their method on amorphous systems, which include various defect types thereby providing a reasonable test case.\n\nSection 2 describes relevant background for electronic structure modeling related to energy levels and wavefunctions that make up the Hamiltonian matrix. The Hamiltonian matrix is computed based on a basis of atomic orbitals that transform under spherical harmonics. The Hamiltonian matrix itself can be decomposed into sub-matrices that describe the interactions of orbitals on the same atom (diagonals) and between different atoms (off-diagonal). The main challenge with calculating Hamiltonian matrices is the computational cost, which has lead to approximations that rely on creating a repeatable unit cell of atoms. That unit cell, however, often fails to properly capture atomic behavior that require larger size, such as defects. Section 2 also describes related work on using machine learning for directly predicting the Hamiltonian matrix, focusing mostly on equivariant GNNs that operate in a spherical harmonic basis.\n\nSection 3 introduces the primary method of the paper starting with the representation for the block of the Hamiltonian matrix and the proposed network architecture. The network architecture mainly relies on a combination of eSCN equivariant convolution to learn over the atomic graph and multi-headed attention to learn in embeddings for the local atomic environments. Section 3.3 describes the graph partitioning method that enables the proposed method to model interactions beyond local atomic interactions while still maintaining a local receptive field for scalability. The proposed partitioning approach leverages virtual nodes and edges that augment the graph for a given partition but are not used for the message passing. As such, the receptive field includes only a 1-hop neighborhood ensuring scalability to larger systems while maintaining sufficient accuracy. Section 3.4 describes the dataset generation focusing on moelcular dynamics trajectories of amorphous materials, namely: HfO2, GeSbTe, and PtGe. The systems span thousands of atoms with >10k orbitals and >100k edges. \n\nSection 4 presents the main results. The first set of experiments presents ablations related to the proposed graph partition mechanism and generally find that partitioned graphs achieve competitive results when compared to full-graph training. The results in Section 4.2 claim that the proposed method can effectively capture the structural disorder contained in the amorphous materials data. Section 4.3 claims that the proposed method can capture compositional disorder in HfO2 and Section 4.4 outlines a case study comparing predictions of the full Hamiltonian using the trained network and various partition slices. Section 4.5 describes computational savings achieved by the graph partitioning approach outlining both compute and memory savings. Section 5 provides a brief conclusion.\n\n## update after rebuttal\n\nThe authors answered my questions to my satisfaction and agreed to include relevant details that will make the paper stronger, including clarifying related work for their paper and providing a comparison on scalability. As such, I increased my score and support of the paper.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "**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. \n\n**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. \n\n**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.\n\n**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.**\n\n**How our work exceeds previous work:** \n\nFirstly, 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:\n\n(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) \n\n(2) H predictions require higher degree tensors to capture all orbital interactions.\n\n\n**Scalability comparisons can, however, be made with previous Hamiltonian models. We provide a summary in Table 2. There are two main categories:**\n\n**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. \n \n**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.\n\n\n **Architecture** | **Approach** | **Max # Atoms per Unit Cell** |\n|-----------------|-------------|-----------------------------|\n| QHNet [6]      | SO(3)       | ~30                         |\n| WANet [7]      | **SO(2)**   | <200                        |\n| DeepH2 [3]     | **SO(2)**   | ≤150                        |\n| SLEM [8]       | **SO(2)**   | <100                        |\n| **This work**  | **SO(2)**   | 3000+                        |\n\n**Table 2:** Comparison of our model to state-of-the-art architectures on task, approach, and complexity of their respective training/testing datasets.\n\n**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.**\n\n**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. \n\n**Note that the list of references is found in the response to Reviewer 4 (nuFS)**", "claims_and_evidence": "The primary claims are backed up with evidence from targeted experiments on large-scale atomic systems. The method descriptions are through and provide relevant details for both the physics and machine learning considerations.", "methods_and_evaluation_criteria": "The methods and evaluation criteria used are created by the authors themselves and are generally well chosen to demonstrate their proposed method. The experiments could be strengthened by describing why relevant baselines (e.g. Allegro) may not be able to scale to the systems used.\n\nThe paper could be strengthened by showing additional experiments on smaller systems - ideally these results would show that minimal performance drops occur with smaller (unit cell scale) systems, providing more evidence for the capabilities of the proposed method.", "theoretical_claims": "The theoretical claims presented are generally well supported.", "experimental_designs_or_analyses": "The experiment design for the three amorphous materials, described in the main paper and supplementary material, are generally sound and detailed.", "relation_to_broader_scientific_literature": "The paper generally covers a good amount of the relevant literature. It could be further strengthened by discussing references related to the challenges of applying GNNs directly in MD simulation [1][2][3], providing broader context for GNN-based methods for atomistic modeling [4] and tying the challenges to broader applications [5].\n\n[1] Fu, X., Wu, Z., Wang, W., Xie, T., Keten, S., Gomez-Bombarelli, R. and Jaakkola, T., Forces are not Enough: Benchmark and Critical Evaluation for Machine Learning Force Fields with Molecular Simulations. Transactions on Machine Learning Research.\n\n[2] Bihani, V., Mannan, S., Pratiush, U., Du, T., Chen, Z., Miret, S., Micoulaut, M., Smedskjaer, M.M., Ranu, S. and Krishnan, N.A., 2024. EGraFFBench: evaluation of equivariant graph neural network force fields for atomistic simulations. Digital Discovery, 3(4), pp.759-768.\n\n[3] Gonzales, C., Fuemmeler, E., Tadmor, E.B., Martiniani, S. and Miret, S., 2024. Benchmarking of Universal Machine Learning Interatomic Potentials for Structural Relaxation. In AI for Accelerated Materials Design-NeurIPS 2024.\n\n[4] Duval, A., Mathis, S.V., Joshi, C.K., Schmidt, V., Miret, S., Malliaros, F.D., Cohen, T., Liò, P., Bengio, Y. and Bronstein, M., 2023. A hitchhiker's guide to geometric gnns for 3d atomic systems. arXiv preprint arXiv:2312.07511.\n\n[5] Miret, S., Lee, K.L.K., Gonzales, C., Mannan, S. and Krishnan, N.M., 2025. Energy & Force Regression on DFT Trajectories is Not Enough for Universal Machine Learning Interatomic Potentials. arXiv preprint arXiv:2502.03660.", "essential_references_not_discussed": "I encourage the authors to discuss the references in the above box. The papers related to challenges for GNNs in simulations appear particularly relevant. The paper could also benefit by providing more details on how their work exceeds the capabilities of prior work, such as Allegro [1] and Allegro-Legato [2] which is mentioned in the paper, that have tried scaling GNNs to large scale simulation.\n\nThe paper should also discuss relevant context for Hamiltonian Neural Networks [3] as a general method.\n\n[1] Musaelian, A., Batzner, S., Johansson, A., Sun, L., Owen, C.J., Kornbluth, M. and Kozinsky, B., 2023. Learning local equivariant representations for large-scale atomistic dynamics. Nature Communications, 14(1), p.579.\n\n[2] Ibayashi, H., Razakh, T.M., Yang, L., Linker, T., Olguin, M., Hattori, S., Luo, Y., Kalia, R.K., Nakano, A., Nomura, K.I. and Vashishta, P., 2023, May. Allegro-legato: Scalable, fast, and robust neural-network quantum molecular dynamics via sharpness-aware minimization. In International Conference on High Performance Computing (pp. 223-239). Cham: Springer Nature Switzerland.\n\n[3] Greydanus, S., Dzamba, M. and Yosinski, J., 2019. Hamiltonian neural networks. Advances in neural information processing systems, 32.\n\nThe paper would also benefit from discussing prior work on graph partitioning for GNNs [4] in the context of the implemented partition algorithm. The algorithm in [4] seems to have a similar approach and claimed benefits.\n\n[4] Mostafa, H., 2022. Sequential aggregation and rematerialization: Distributed full-batch training of graph neural networks on large graphs. Proceedings of Machine Learning and Systems, 4, pp.265-275.", "other_strengths_and_weaknesses": "Overall, the paper presents an interesting and potentially compelling method for scaling GNNs to larger materials systems. The paper could be strengthened by:\n\n* Providing more details on whether the models evaluated are one for each system or a generalized model.\n* Whether their proposed method would be applicable to smaller systems without performance drops. As mentioned before, this would help strengthen the reasons for using the method.\n* Discussing whether the graph partitioning approach can be applied to any model architecture and what limitations may exist.", "other_comments_or_suggestions": "N/A", "questions_for_authors": "1. Did you train one model for each material system?\n2. Do you think your approach can generally across different structures and systems? What would be needed for that?\n3. Can you clarify if there is a different in the cut-off distance between slices and the cut-off distance used for graph creation?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "VJF0BBWPX3", "reviewer": "Reviewer_zXDf", "summary": "The author propose a graph partitioning strategies to localize message passing and new networks which leverages SO(2) convolution for predicting the Hamiltonian for large atomic structures. The author demonstrates its capabilities by predicting the electronic Hamiltonian of various systems with up to 3,000 nodes with and ≤0.55% error in the eigenvalue spectra.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "**Non-diagonal part of Hamiltonian beyond cutoffs:** \n\n**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:\n\n1. **In Fig. 7**, we illustrated the decay in values of Hamiltonian elements with increasing inter-atomic distance for all three datasets.\n\n2. **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. \n\n**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.  \n\n**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). \n\n**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. \n\n**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. \n\n**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.\n\n**Note that the list of references is found in the response to Reviewer 4 (nuFS)**", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "Yes", "theoretical_claims": "N/A", "experimental_designs_or_analyses": "The derived Hamiltonian should also be evaluated with SCF iterations speedup or the MAE on system energy deriving from the Hamiltonian.", "relation_to_broader_scientific_literature": "It is somewhat relavent, it could be broadly applicable to distributed graph computing.", "essential_references_not_discussed": "Some other paper also leverages the SO(2) convolution such as [1].\n\n[1] Enhancing the scalability and applicability of Kohn-Sham Hamiltonian for scalable molecular systems. ICLR 2025.", "other_strengths_and_weaknesses": "1. Strength: The significance of this work is high, it could be a promising avenue for large-scale simulation or all-atom protein simulations.\n2. Weakness: The evaluation on the real-world applications is very limited (such as SCF speedups, system energy prediction, dipole moment prediction). The graph partition strategy is not that novel, and the author did not addresses the non-diagonal blocks beyond cutoffs. Missing of the long-range interactions could potentially impact the real-world applications. The graph-partitioning could also potentially lead to non-smoothness in real world simulations.", "other_comments_or_suggestions": "N/A", "questions_for_authors": "1. How does the network handle the non-diagonal part of the Hamiltonian beyond cutoffs? Methods like QHNet leverages the pair construction layer which performs a tensor product between every pair of the atoms. I think it's not doable here.\n2. What is the absolute error for the eigen-spectrum, is it within the chemical accuracy? (< 1 kcal/mol)\n3. The graph would generate multiple graph slices, how do you combine them to generate the final H (a diagram or algorithm would suffice)  ? And for the embedding of virtual nodes, will their embedding be updated by the embedding from other graph slices the next time you encounter the graph?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "NipcRUzmD4", "reviewer": "Reviewer_nuFS", "summary": "This work focuses on scaling Hamiltonian prediction to large periodic structures. While a common ML problem, Hamiltonian prediction has been challenging in large structures due to the quadratic scaling of the Hamiltonian. The authors propose a partitioning scheme to enable distributed computing and batching of large unit cells. Further, the authors present a new dataset consisting of 3 periodic large compounds with a total of 7 structures. Due to the size of the structures, training on a single sample suffices.\n\n## After the rebuttal\n\nI appreciate the authors’ response but remain with my many points of my initial critisim. Primarily, comparisons to previous works are possible by evaluating on smaller datasets. While new datasets are always welcome this does not rule out the evaluation on preivous established ones. This undermines their second contribution as listed by the authors. Overall, I do not strongly object an acceptance but also do not recommend it.", "questions": "", "limitations": "", "rating": 2, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "**Isolation of Contributions:** Our two main contributions not found in other works are:\n\n1. 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.  \n\n2. 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. \n\n**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**.\n\n**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.\n\n**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.\n\n**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].\n\n**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.\n\n**References**\n\n[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)\n\n[2] Kaniselvan, M., Luisier, M., and Mladenovic, M. An atomistic model of field-induced resistive switching in valence change memory. ACS Nano, March 2023.\n\n[3] Wang, Y., Li, Y., et. al. Universal materials model of deep-learning density functional theory hamiltonian. Science Bulletin, 69(16):2514–2521, 2024b.\n\n[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.\n\n[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\n\n[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.\n\n[7] Y Li, et. al Enhancing the scalability and applicability of Kohn-Sham Hamiltonian for scalable molecular systems. ICLR 2025.\n\n[8] Zhouyin, Z., et. al. Learning local equivariant representations for quantum operators, ICLR 2025\n\n[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", "claims_and_evidence": "The following claims are made:\n1. A single message passing GNN achieves accurate Hamiltonian prediction on large structures.\n2. The authors demonstrate that Hamiltonian prediction scales to thousands of atoms.\n3. The proposed partitioning scheme enables efficient training on large compounds.\n\nIn general, these claims also seem supported but the limited evaluation makes the specific judgements hard.", "methods_and_evaluation_criteria": "The author propose a new GNN for predicting Hamiltonians that is similar to previous works but only uses a single message passing step. They evaluate it with their partitioning scheme.", "theoretical_claims": "The paper does not introduce theoretical claims.", "experimental_designs_or_analyses": "The authors evaluate only on their newly proposed dataset their newly propose method. This combination poses a challenge for the reader to identity meaningful transferable improvements. Evaluation on standard datasets for their proposed GNN and previous GNNs on their new dataset would strengthen the work significantly and would improve its contextualization. Additional downstream application, e.g., by running DFT, or computing observables would strengthen the paper's position.", "relation_to_broader_scientific_literature": "The authors should consider contextualizing their work with the batching works done in graph neural networks where similar issues for large graphs arise.", "essential_references_not_discussed": "-", "other_strengths_and_weaknesses": "Strengths:\n* The paper is well written and easy to understand if one has the necessary background. \n* The problem is interesting and the results are encouraging for large-scale Hamiltonian prediction.\n\nWeaknesses:\n* The works presents limited theoretical and technical contributions and heavily relies on its empirical results.\n* The empirical evaluation is quite limited and does not allow the reader to identify the downstream performance or comparisons to previous works.", "other_comments_or_suggestions": "* Please use consistent units, using eV and Eh in different places makes comparisons more difficult.\n* l. 338 should be Table 3.", "questions_for_authors": "I am willing to increase my score if the authors manage to clearly isolate their contributions and show the evaluate in each contribution via ablation studies, comparisons with previous works, and evaluation on established datasets.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: The paper presents a mean-field electronic Hamiltonian prediction model that can be applied to large systems. This arises from the need to process disordered materials, in which a huge unit cell should be assumed. The model implements this by assuming diagonal sparsity of the Hamiltonian matrix, restricting interatomic interaction to be within each segment, while keeping a few neighboring atoms ("virtual nodes") for each segment to account for inter-segment interaction. The authors applied the model on some real-world workloads with unit cells larger than before, and showed useful prediction results in terms of derived physical quantities.

All the reviewers and I myself appreciate the contribution of a useful Hamiltonian prediction model that is tractable on large atomic systems, based on a graph partitioning strategy and implementation. Reviewers also raised a few insufficiencies include seemingly minor technical novelty (using SO(2) convolution and assuming locality for scalability in Hamiltonian prediction are not new), and the lack of comparison with existing Hamiltonian prediction methods, even if only on a small-scale system.

Personally, I tend to value the contribution of connecting and adjusting prevailing machine learning techniques to a new (disordered material system) challenging (need of handling a large unit cell) practical need. The ablation studies of their own method provide a reasonable understanding on how various techniques and components work, which could also provide valuable information to the community. I won't judge too much on missing some evaluation settings/metrics as the selected ones seem more directly relevant to the real application, but I would encourage the authors to also provide such information to complete the analysis on the work (even if it shows a limitation).

As the rating of this submission appear a bit diverse, I posted the comment to the reviewers before decision. No other opinions were raised.
Decision: Accept (poster)