Learning the Hamiltonian of Disordered Materials with Equivariant Graph Networks

Notational clarity: Thank you for pointing this out. We have gone through and standardized all the mathematical notation.

List of contributions: Thank you for the suggestion - we have added a list of contributions to the Introduction. Please also see the ‘general response'.

Generalizability to other disordered systems:

We have now added two additional examples of structural disorder (to appendix I), as detailed in the ‘General response’ (substoichiometric HfOx and amorphous PtGe). Together, these examples demonstrate the capabilities of our approach in tackling different types of structural disorder.

Definition of unique:

We understand the concern - to help illustrate the uniqueness of the structures in the dataset, we have included a plot of the outliers in Hf coordination in Fig. 9 in Appendix F. Amorphous structures are not a completely random distribution of atoms, they typically contain some smaller structural motifs and a range of coordination numbers. However, the spatial distribution of these differently coordinated atoms is different in all samples, indicating that they are all unique.

Near-sightedness:

Yes, the near-sightedness principle holds strictly for electronic Hamiltonians represented in a local basis. Please see our general response for more details and the newly added Fig. 5b, which illustrates this effect for our example system.

Metric for accuracy:

For clarity, we changed the naming from relative MSE/MAE to relative L1 and L2 errors (see Fig. 6). In addition, we computed the same error metrics for the case where only the occupied states are considered (0.84%). The metric is reasonable since the absolute errors shown in Fig. 13 in Appendix H are in the same order of magnitude at band edges.

General

As was pointed out by another reviewer, we did not sufficiently stress the two contributions of our work (the development of a network for Hamiltonian prediction based on SO(2) convolutions by combining concepts from several previous works, and the augmented partitioning approach). We have now added a contributions section to the introduction.

(Q1) Nearsightendness principle + Analysis demonstrating how Hamiltonian matrix elements are affected by atoms at different distances:

We have added (1) an explanation of the near-sightedness principles of orbital interactions to Section 3 (methods), and (2) a sub-figure (Fig. 5 (b)) to illustrate the decay (plotted in log scale) of the Hamiltonian elements as a function of interatomic distance. When using a local basis, near-sightedness strictly holds and is a fundamental physical property of electronic Hamiltonians. Previous studies on Hamiltonian predictions have similarly applied this principle to justify the omission of longer-range interactions [1].

This information can also be interpreted from Fig. 10 in Appendix F.1, where we plot the eigenvalue spectra of the ground-truth Hamiltonian after filtering the matrix elements by different values of $r_{cut}$. While an $r_{cut}$ value of $6Å$ results in a noticeable deviation of the eigenvalues from the unfiltered spectra, $r_{cut}$ = 8 and 10 $Å$ are identical to the unfiltered case. This is consistent with the results in Table 2, where the node and edge prediction errors saturate when determining the graph connectivity with $r_{cut}$ $\ge$ $8Å$.

(Q1) Justification of single message passing layer

There are two more points we would like to stress when it comes to determining the connectivity, and the use of a single layer. The cutoff radius is chosen to capture the relevant blocks of the Hamiltonian. This leads to the radius of connectivity being much larger compared to other networks, which represent materials as graphs and learn the properties of each node/atom, as these networks often define edges according to atomic bonds. With an $r_{cut}$ of $8Å$, the average atom has roughly ~180 edges to other atoms - the aggregation neighborhood is already large. Including a second-hop neighborhood would extend the receptive field to atoms that have blocks of negligibly small (often numerically zero) elements in the Hamiltonian, which do not provide any extra information for training, and also lead to neighborhood explosion considering the densely connected nature of these graphs, and the large cutoff, which would make training unfeasible. Finally, there are also examples of message-passing equivariant neural networks that have achieved success in the prediction of various complex properties of large scale systems while limiting the receptive field and enforcing locality (notably Allegro) [7].

(Q1) Use of smaller systems:

A crucial point that was not successfully conveyed in our initial submission is that it is impossible to generate small, non-crystalline unit cells for training. The reason why is illustrated in Fig. 2 in the manuscript - only large unit cells are able to capture realistic disorder. To represent large unit cells that maintain the amorphous property of the material, large graphs are unavoidable. However, within these large graphs, we can still exploit the locality of interatomic interactions to avoid computing long-range interactions. If the application is for small molecules, s/p orbitals, or crystals where all information can be captured within a unit cell, then existing equivariant approaches such as DeepH2, DeepH-E3, QHNet, and HamGNN [2-5] will be able to treat them. Our network and training scheme are designed to operate in a regime that is inaccessible to these other works, and nevertheless presents significant scientific value and involves the highest computational burden to treat with conventional approaches. We have included an additional experiment in Table 3 of this response to show that our approach can similarly be applied to predict crystalline HfO₂ in the monoclinic phase, with a much better prediction accuracy compared to the amorphous HfO₂ structure.

In the reply. it was mentioned that: "One could train the model using local environments extracted from large systems. These local environments would contain all the necessary structural information to predict Hij. The trained model should then be transferable to both small and large systems."

Training using local environments extracted from large systems is far from trivial. This is, however, exactly what our augmented partitioning approach was designed for. Let us explain in further detail what the challenges are and why our method addresses them:

To extract a local environment from large systems and use it for training, three relatively “straightforward” naive approaches can be envisioned:

1. Extracting small samples from large amorphous structures, and then performing DFT calculations on these isolated small unit cells. The DFT results from different samples are then used to train the model and predict the large amorphous structures.

2. Same method as 1, but with the addition of periodic boundary conditions to the small samples to account for the influence of the atomic environment.

3. Performing DFT on a full graph and using naive partitioning to extract a local sub-graph from the full graph for training, ignoring connections to other slices.

Approach 1: The main issue with the first approach comes from the boundary conditions that should be applied during the DFT calculations. The small samples we can extract from the large amorphous structure are densely connected to other atoms. Neglecting the latter and considering only the isolated parts, which are treated as molecules, gives rise to numerous dangling bonds at the surface of the small samples and highly inaccurate Hamiltonian matrix elements.

To further substantiate this point, we performed additional DFT calculations on an isolated sample of size 10x10x10 Å extracted from the large amorphous HfO₂ structure, placed in vacuum, i.e., without periodic boundary conditions. We then compared the resulting Hamiltonian matrix elements to those of the same atoms, but taken from the Hamiltonian matrix corresponding to the original large structure. The Mean Absolute Error (MAE) for the onsite elements of the small sample is 752 meV. In other words, the Hamiltonian matrix elements of the small sample are completely different from their counterparts in the large structure. Hence, cannot be used for training the model at all. Furthermore, the offsite Hamiltonian blocks representing the connections between the isolated piece and the rest of the large amorphous structure will also be missing from the training data.

Approach 2: As a second approach, periodic boundary conditions could be added to the small samples to avoid dangling bonds and mimic the atomic environment. However, for small unit cells, periodic boundary conditions typically lead to strong interactions between one atom and its periodic replica and thus very different Hamiltonian matrix entries than those of the original, large structure. We repeated the previous experiment with the inclusion of periodic boundary conditions: the onsite MAE becomes even larger (813 meV). The inaccuracies caused by the consideration of small samples only are the reason (i) why DFT calculations of large disordered structures is imperatively needed despite being computationally expensive, and (ii) why the learning/prediction of such systems has so far remained an open issue in literature.

Approach 3: Alternatively, as a third approach, DFT calculations can be performed on a full graph and, through “naïve" partitioning, a local sub-graph can be extracted from the full graph for training, i.e., all connections between the sub-graph and rest of the graph are ignored. As mentioned in the manuscript, the omission of connections to atoms outside of the partition leads to the wrong/incomplete aggregation of information. To learn the correct aggregation function and be able to make accurate predictions, we must preserve the full connectivity of the partition subgraph. The poor prediction accuracy of this this third approach can be found in the ablation study (Table 3 of the manuscript).

This brings us to our proposed augmented partitioning approach, which addresses this connectivity issue by introducing virtual nodes and edges. Through this method, the local environment is correctly extracted from the large graph, which also contains which also contains accurate data from the DFT computation of the full amorphous structure. As a result, the prediction results using this approach is on par with other literature (refer to previous response on comparison with literature). We hope that we have now addressed this question. Please let us know if there is anything else we missed.

Note that the responses to the other two questions (single message passing layer and perturbation) will also be posted soon.

To demonstrate the effect of long and short range perturbations in amorphous structures, we introduced a single perturbation at one chosen location in the sample and measured the mean absolute error of the resulting onsite blocks of the Hamiltonian matrix when compared to that of the unperturbed structure.

The types of perturbation introduced include single atom translation (0.1 Å shift), vacancy (replacing O atom), and substitution (O atom replaced with Hf atom). Their influence is plotted against distance from perturbation in Fig. 15 (a), (b) and (c) respectively (In Appendix K of the updated manuscript). In all cases, the effect of the perturbation rapidly decays with increasing distance.

For the case of a 0.1 Å translation perturbation, the average onsite MAE at a distance of 8 Å away is equal to $0.15$ $mE_h$. Considering the average value of an onsite Hamiltonian block ($63$ $mE_h$), the perturbation only affects the matrix elements by 0.24% overall. Similarly, for vacancy and substitution perturbations, the matrix elements of atoms located 8 Å away only changes by 0.18% and 0.12%, respectively. This implies that for our chosen cutoff of 8 Å, perturbations occurring outside of the selected range around one specific atom have a negligible impact on its Hamiltonian matrix elements. This also means that the electronic structure of that atom can be learned using information from the local atomic environment only.

This feature is also demonstrated through our study of sub-stoichiometric HfO$_x$ with randomly distributed vacancies. Despite training on independent slices, we ensure that every atom within that slice is surrounded by a complete local environment, with the inclusion of all neighbors within a radius of 8 Å. Any vacancies outside of that range have a negligible effect on the atom, and are seen and learned by other partitions. When multiple slices are trained together, the entire distribution of perturbations are captured, allowing the model to generalize well to unseen structures with a completely different distribution of vacancies and local atomic environments.

The key difference between the systems we are investigating and twisted bilayer graphene, is the absence of atomic order/periodicity in our amorphous structures, while periodicity is still present in twisted bilayer graphene.

While it is true that twisted bilayer graphene with varying twist angles relies on large unit cells to be described, periodicity still exists in these structures along the two directions in the graphene plane. The main additional information that must be learned in this case is the interaction between the different graphene layers in the third direction, which can be captured by the small samples. The local atomic environments of the atoms within the large bilayer structure are therefore still highly similar to that of the small training samples with periodic boundary conditions applied in the graphene plane. As a result, a model trained on information from small unit cells can generalize to test structures that are larger along these periodic directions.

In our case, there is no repeated motif/structure along any direction because of the amorphous nature of our samples. Applying periodic boundary conditions to a small unit cell would lead to a crystalline material that does not resemble our large amorphous structures at all. It is essential to note that amorphous HfO2 contains a large variety of different motifs and local environments that, when connected together, exhibit the unique properties that we are trying to predict. Replacing the disordered environments and connections with periodic images leads to the loss of amorphous behavior, and very different Hamiltonian elements. We hope this fully addresses your question.

First, we would like to emphasize that the augmented partitioning approach we propose is a general training method that can be applied to other, more expressive networks. It is not strictly attached to the EquiformerV2/DeepH2 architecture that we implemented in this work. The augmented partitioning approach is mainly applied during graph construction, which is mostly independent of the architecture itself.

Second, we want to point out that similar problems have been successfully treated by networks which only use strictly local interactions. The best example of this is Allegro [1], which has been used to predict energy and forces, and build interatomic potentials. It takes in a similar set of two-body inputs (atomic numbers of atom i and neighbors j, and their distances) as our architecture, requiring no information from atoms beyond the first hop neighborhood. Many body representations can then be constructed through repeated tensor products as more layers are added. It can also be adapted for the prediction of Hamiltonian entries as done by this recent ICLR submission: https://openreview.net/forum?id=kpq3IIjUD3. This is consistent with how local information is sufficient to achieve the prediction accuracies we have shown.

We can combine this approach with our augmented partitioning method to capture higher body order information through more layers and further improve our prediction accuracy on amorphous materials (which is already comparable to literature), while maintaining the same receptive field. This is a promising direction to explore that is beyond the scope of the paper.

What we have demonstrated in this particular submission using the current architecture, is that the proposed augmented partitioning approach gives a similar level of prediction accuracy compared to full graph training when the receptive field is restricted. We have therefore developed a generally applicable partitioning method that enables the reliable training and prediction of large amorphous structures that were so far not achievable.

References: [1] Musaelian, A., Batzner, S., Johansson, A. et al. Learning local equivariant representations for large-scale atomistic dynamics. Nat Commun 14, 579 (2023).

Graph connectivity/cutoff radius:

The cutoff radius (which determines the connectivity) is a convergence parameter that can be chosen either by looking at the changes in the eigenvalue spectra (as in Fig. 10, where $r_{cut}$ was chosen such that the eigenvalue accuracy is not significantly compromised) or by checking the prediction accuracy of models with different $r_{cut}$ (as in Table 2) and choosing an $r_{cut}$ that achieves sufficient accuracy. However, it is not possible to analytically compute this quantity ahead of the DFT calculations, using only the structural (and basis) information. A higher $r_{cut}$ implies that more interactions are captured (as shown for both HfO2 and PtGe in Table 1 below). Therefore, the general rule for an unknown material is to choose the maximum possible cutoff such that the slice fits into memory, in order to maximize prediction accuracy.

Range of disordered materials:

Please see our general response, which details the changes we made to address this point (as it was brought up by multiple reviewers). In short, we agree that showing only one material is not a convincing demonstration of our application’s ability to handle electronic properties in the case of structural disorder. To fix this, we have added examples of sub-stoichiometric HfOₓ (containing vacancies) and amorphous PtGe. We are exploring the latter for its phase-change-related properties. The results have been added to Appendix I of the revised manuscript.

To add to the response to the previous question on cutoff radius, we conduct a similar investigation on a completely new amorphous PtGe material, with results shown in Table 1 (also added to Appendix I.3). Here, as the PtGe structure occupies less memory, we can choose a very large cutoff radius of 16 Å without exceeding memory limits (unlike HfO2), and can therefore achieve a much higher prediction accuracy of 1.43 meV for all 2688 atoms and 2148055 edges using only one slice for training. The results from Table 3 below (and Table 2 in the manuscript) also show that higher $r_{cut}$ is always preferred over lower $r_{cut}$ due to the significant improvements in overall accuracy. More details on the PtGe structure can be found in Appendix I.

Table 1: Prediction accuracy of the model on amorphous PtGe with different $r_{cut}$

We hope these additional tests and the high prediction accuracy validate the robustness of our approach, especially when applied to new datasets with larger cutoff radius.

Practical limitations to testing larger systems:

The idea behind this approach is that the training structures of a given size should sufficiently sample the range of atomic motifs which can be found in the structure, such that testing/prediction can be done on larger scales at similar prediction accuracy. We therefore do not expect any hard limitations in extending the augmented partitioning method to test on larger systems, which would require more augmented partitions. However in terms of training, there may be a limitation in terms of scaling due to the two issues (virtual node overhead + load imbalance) described in Appendix G (please see Fig. 11). When training larger graphs, these imbalances could affect performance at larger scales.

We realize that we did not sufficiently detail the novelty of our contribution in the initial submission. We added the following ‘contributions’ section to the Introduction:

1. We develop an efficient GNN-based model for electronic property prediction by combining (1) the SO(2) convolution approach detailed in [1], (2) the equivariant attention mechanism proposed in [2], and (3) concepts from [3] and [4] to introduce learnable node/edge embeddings and a basis transformation layer for mapping outputs to the Hamiltonian. We provide the code for this implementation [in the Supporting Materials].

2. As the graphs required to capture amorphous disorder are necessarily large and incur high memory consumption in a full-batch training environment, we propose an efficient augmented partitioning method that partitions the input graphs and corrects atomic environments with masked virtual nodes and edges. Our method allows for arbitrarily large graphs to be broken down into partitions that can fit into GPU memory during training without compromising the achievable testing accuracy.

Problem specificity:

We agree that our methodology is specifically designed for electronic property prediction, from the construction of equivariant networks for Hamiltonian blocks to the augmented partitioning approach designed to tackle full-batch training of large, densely-connected graphs. Both of these techniques can find other applications that will be explored in future work.

However, we do not want to stray from our application focus in this paper because this problem is of high scientific relevance. Predicting the electronic properties of large disordered structures from first-principles using density functional theory is a key research activity in materials science. It consumes a large portion of the computing time of several supercomputers in the world [5]. We have tried to stress the potential scientific impact of research in this area in the Outlook section of the manuscript. More generally, several previous papers at top venues have been dedicated specifically to predicting electronically-resolved information, several of which are cited in our work [6,7]. We strongly believe that this application is of significant interest to ICLR.

Our approach is similarly straightforward when extending to crystalline materials. This is because crystalline materials are composed of a periodic tiling of a unit cell in space, with atomic environments varying only within the unit cell. When represented on a local basis, the corresponding Hamiltonian matrix is composed of identical sub-blocks. This also implies that the difference between training and testing data is almost negligible since the same/very similar local atomic environments are seen in both cases. After developing a similarly-sized structure of monoclinic HfO₂ (a crystalline phase), for example, we can achieve near sub-meV accuracy (more information can be found in Appendix I). On the other hand, amorphous structures contain a large range of different local atomic configurations that vary across different samples, necessitating the use of large unit cells to be fully captured. The accurate prediction of electronic properties of amorphous materials therefore remains an open problem in literature [8]. The approach introduced in our paper tackles this challenging problem.

It is not straightforward to compare to our work to other studies, as no existing equivariant network that we are aware of is computationally capable of treating our datasets without extensive additional computational optimization. However, we can show generalizability to different disordered systems. In Appendix I, we have added examples of sub-stoichiometry HfOₓ ( with $x < 2$ ) and amorphous Pt-doped Ge, as mentioned in the ‘general response’ above. The code we provided in the Supporting Information also includes several examples applied to small-molecule datasets made from MD17 trajectories, including H₂O, uracil, and malondialdehyde. Results for one example - H₂O - are shown in Appendix B of our paper.

We are glad to have addressed some of your concerns. In response to your question, we have recently added the theoretical justification of the augmented partitioning approach to Appendix J of the manuscript. There, we added Equations 1-4 and explained them in detail. Hope it addresses your question. Please let us know if there are further clarifications needed. (e.g. notations etc.)

Thank you for the constructive feedback. Appendix J.1 is indeed a reformulation of the update rule behind augmented partitioning. When prompted to provide a more rigorous mathematical analysis, specifically in ‘how the augmented partitioning approach guarantees accurate node and edge predictions’, we concluded that the best approach would be to show that ‘the individual feature update rule is unchanged by the augmented partitioning approach’. Each embedding thus undergoes a standard message passing procedure. The hypothesis statement in this case, is that for node and edges within the partition, the feature update using augmented partitioning is equivalent to the feature update for full graph training for the first layer.

We would also like to clarify that this paper is submitted to the applications track, which in our experience typically involves adapting ML architectural concepts for downstream scientific research use. During this process, we do not make contributions to mathematical foundations of machine learning, or specifically to the expressiveness of GNNs. We follow the example of previously accepted ICLR papers in the area of materials science and chemistry that have also proposed new approaches (https://openreview.net/forum?id=CUfSCwcgqm, https://openreview.net/forum?id=KwmPfARgOTD, https://openreview.net/forum?id=mCOBKZmrzD ), where rigorous mathematical proofs of approaches and lemmas are much less common.

I appreciate the authors' detailed explanation, which has clarified the focus and methods of this paper. I do not have any questions regarding the methodology. However, since the methods appear relatively straightforward, I wondered if there are any underlying theoretical foundations for this approach, particularly concerning its effectiveness and efficiency. Incorporating such insights could significantly enhance the impact of this work. I have increased my rating previously; however, I hope my suggestion has highlighted some potential aspects that might be missing in this paper for both the authors and the broader audience.