Predicting ground state configuration of energy landscape ensemble using graph neural network

Finding the ground state configurations of a complex energy landscape is a long standing computational challenge Wales and Bogdan (2006). Short of brute-force enumeration, random search algorithms such as simulated annealing can anneal Markov chains to the global minimum as a simulation temperature approaches zero Kirkpatrick et al. (1983). However, in cases where many interacting degrees of freedom result in a highly rugged energy landscapes, conventional methods suffer from low probability of overcoming energy barriers and the chain may get stuck in local minima Papadimitrou and Steiglitz (1982); Landau and Binder (2014); Wolff (1990).

The classical methods for searching energy landscapes are devised to work for general problems. Yet, many scientific problems often present themselves via an ensemble of energy landscapes with similar underlying patterns, with interactions arising from a single or handful of governing equations. Examples include the energy landscapes of organic molecules built out of chemical building blocks, where potential energies are obtained by solving Schrödinger equation, or the space of protein structures from interactions of individual amino acids. In an ensemble setting, we hypothesize that there exist system-specific sampling rules Wolff (1989); Wang and Swendsen (1990); Houdayer (2001) that make it possible to traverse these particular energy landscapes more efficiently than classical methods. These rules can be learned from examples of energy minima calculated with classical methods for small problems.

Here we demonstrate this approach in the context of a model problem that defines a natural ensemble. We construct Ising spin glasses Mézard et al. (1987); Edwards and Anderson (1975), where the interaction matrix $J$ is a structured random matrix, chosen from protein contact maps. Given the large database of natural proteins Berman et al. (2000) and the distinctive contact pattern of a folded protein Nelson et al. (2008), protein contact map data gives an ideal ensemble for testing whether interaction rules encoded in $J$’s are consistent across varying system sizes.

In recent years, several machine learning techniques have been applied to sample spin configurations of Ising model. The list includes but is not limited to simple regression Liu et al. (2017), restricted Boltzmann machine Huang and Wang (2017), reinforcement learning Bojesen (2018), autoregressive model Wu et al. (2019); McNaughton et al. (2020), and normalizing flow model Hartnett and Mohseni (2020). The goal of these works is to estimate the Boltzmann distribution of a given problem so that the learned model can either completely replace or be used as a proposal distribution for Markov Chain Monte Carlo simulation. However, these schemes do not consider learning with $J$ instances of varying sizes.

Instead we recast spin glass energy minimization, a well known NP-hard problem Barahona (1982), as a node classification problem in graph theory, and employ a graph neural network (GNN) Wu et al. (2020) to parametrize the mapping from a $J$ to the corresponding ground state configuration. We generate the set of most probable configurations from the GNN model to predict low-lying configurations of an energy landscape. If this configuration set misses a ground state configuration, we show that simulated annealing starting from a configuration in this set can search for the ground state configuration more efficiently. The schematic of this strategy is described in Fig. 1.

We further test the utility of the GNN model by constraining the size of $J$—where size refers to the number of amino acid—in a training set and testing the trained model on larger $J$’s. As we increase the size limit of training set $J$ from 30 to 500, model’s test performances quickly reach the level comparable to those obtained with the size limit of 800. We also show that the model trained on $J$ with size less than 800 can predict configurations whose energies are much lower than those found by simulated annealing for $J$ with size around 3000.

We begin by constructing an ensemble of Hamiltonians for which the underlying potential energy landscapes have similar patterns. For simplicity, we consider Ising Hamiltonians of the form

where both coupling and field terms depend on an interaction matrix $J$ and $J_{ij},\sigma_{i}\in\{0,1\}$. The field is chosen to prevent all ground state configurations from collapsing to a trivial ground state of all 1’s. Within this formulation, to obtain an energy landscape ensemble, we need to specify an ensemble of $J$ matrices, whose $J$ is random yet with distinct shared patterns. Due to binary $J_{ij}$, this structured randomness of $J$ shall be encoded in spin connectivity.

In this work, we obtain this ensemble using protein contact maps to calculate the $J$ ensemble. Since proteins are characterized by distinct secondary structures, together with non-local contacts, protein contact maps define a set of structured random connectivity matrices. We downloaded fist subunit of all protein structure files deposited in the Protein Data Bank rcs to ensure $J$’s do not have a distinct block diagonal structure due to the presence of multiple domains. Hence the connectivity features in our $J$ ensemble solely originate from the pattern of intra-domain folding, which are referred to as secondary and tertiary protein structures. We excluded protein with missing spatial information or whose protein chain is shorter than 20 residues or longer than 800 residues. We additionally added two largest subunit structures with chain length of 3661 and 2814 for the size generalizability experiments. From these files, we generated contact maps by setting $J_{ij}$ as 1 if the distance between two corresponding amino acid residues is less than 8Å, and 0 otherwise Monastyrskyy et al. (2014). From this procedure, we obtained 64563 different contact maps, excluding the two large cases, to define our $J$ ensemble. We emphasize the resulting spin configurations derived from these energy functions Eq. 1 have no relation to amino sequences; our intent here is not to make predictions about proteins per se, but instead to use the regularity of protein structures to define a natural ensemble.

For all $J$’s except the two largest, we ran simulated annealing for each $J$ starting from 100 random initial configurations, and selected an annealed configuration with the lowest energy as its purported ground state configuration $\sigma_{min}$. The annealing schedule was optimized such that simulated annealing always find the ground state configurations on $J$’s with size smaller than 30, which we identified from brute-force enumeration of all configurations. For the two largest $J$’s, we decreased a cooling rate and increased an equilibration steps at each temperature to account for an enlarged configurational space and ran 30 randomly initialized simulated annealing. Further simulation details are discussed in the Supplemental Material ¹¹1See supplemental material at …. Since finding the global energy minimum of an Ising spin glass in $2^{N}$ configuration space is NP-hard, we settled for this repeated annealing scheme and assume $\sigma_{min}$ closely approximates the actual ground state configuration. From all pairs of $J$ and $\sigma_{min}$, 6400 pairs were randomly selected as validation set, another 6400 pairs as test set, and remaining 51763 pairs as training set. For the first size generalizability experiment, we used the same test set but sub-select from the training set the pairs whose $J$ are smaller than certain size cutoffs to make small-$J$ training sets. For the second size generalizability experiment, we used the entire training set and test on the two large $J$’s.

Our prediction task is to learn the mapping from $J{\to}\sigma_{min}$. This can be cast as a node classification problem in graph theory and hence we parametrize the mapping with a graph neural network. Given a graph, the $L$-layer model generates an expressive feature embedding for node $\sigma_{i}$ by aggregating the features of all $L$-hop neighbor nodes of $\sigma_{i}$ as shown in Fig. 1, and uses this embedding to classify globally whether each node shall be turned on or off. To allow generalization of the mapping across $J$’s with different size and structure, we chose a message passing framework Scarselli et al. (2008); Gilmer et al. (2017) with attention mechanism Vaswani et al. (2017); Veličković et al. (2017), instead of Laplacian-based convolution method Defferrard et al. (2016); Kipf and Welling (2016) which requires a constant graph structure.

The inputs to the graph neural network are the adjacency matrix $J$ and node features, which are initially a node degree the field strength $h$ from Eq. 1. At each layer, the network updates node features by first applying standard nonlinear transformation—expanding feature dimension from 2 to $F$, then calculating attention coefficient $\alpha_{ij}$ to find relative importance of a neighbor node $j$ to node $i$, and taking weighted sum of neighbor nodes’ features using these coefficients. To capture more information from neighbors, this process is repeated $K$ times with different set of weights and newly computed $\alpha_{ij}^{k}$ to produce $K{\times}F$ features for a node. The features of node $i$ are then reduced to a probability $P(\sigma_{i}{=}1)$ in the final layer for node classification. The functional forms of the operations are detailed in the Supplemental Material Note (1). Notably, the final model used in this work consists of six layers.

The model’s predicted configuration, $\hat{\sigma}$, are then obtained by choosing the greater of the two node classification probabilities, $\texttt{argmax}[P(\sigma_{i}{=}0)\ P(\sigma_{i}{=}1)]$. However, this point estimation does not take a full advantage of the learned embedding. This scheme is especially problematic for nodes with probabilities around 0.5 because the non-argmax configurations would have been just as likely. Therefore, we generate a set of top most probable configurations from the configuration probability output of the model, giving a broader coverage of the low-lying region in the energy landscape. To obtain $M$ such configurations, we pick top $\log_{2}{M}$ nodes whose $P(\sigma{=}1)$ are close to 0.5, and order all permuted configurations according to their corresponding sum combination of probabilities.

A GNN model trained on the entire training set correctly predicted $\sigma_{min}$ for 1700 of 6400 $J$’s in the held-out test set. To further quantify the model’s performance, we investigate following two metrics. Define accuracy as the ratio between the number of correctly predicted nodes in $\hat{\sigma}$ and total number of nodes and energy difference $\Delta E$ as the energy gap between a predicted configuration and the true ground state. Fig. 2(a) shows the prediction accuracy decreases, while the energy difference increases with increasing $J$. The average accuracy and $\Delta E$ across the entire ensemble are $0.978$ and 2.79 respectively, due to the size distribution of $J$ skewed towards small $J$’s (Supplemental Material Fig. S1 Note (1)). Since energy histograms of small $J$’s obtained via complete configuration enumeration are peaked at positive energy and negative energy configurations occur in far-left tail region (Supplemental Material Fig. S2 Note (1)), predicting configurations with such small energy differences is surprising. We emphasize again that the model does not evaluate the energy function of Eq. 1 to optimize a configuration. This suggests the GNN model has learned a generalizable node feature transformation for this particular class of energy landscape simply by comparing its predicted configurations to known ground state configurations.

Fig. 2(b) shows the averaged histogram of node classification probability $P(\sigma{=}1)$ from high accuracy configurations in blue, and that of low accuracy configurations in orange. A striking feature is that most nodes in both cases are predicted with high certainty as evinced by the peaks at both ends. In addition, the histogram of low accuracy configurations shows more nodes in the middle, indicating that the model’s prediction accuracy may be directly related to the node classification probability $P(\sigma)$. We thus set a threshold probability, $P_{thr}$, to select nodes with low uncertainty where $P(\sigma_{i}{=}1)\geq P_{thr}$ or $P(\sigma_{i}{=}1)<1{-}P_{thr}$ and calculated an error rate among these nodes as $P_{thr}$ is varied. As shown in the inset of Fig. 2, the number of misclassified nodes among such nodes goes down as we increase the threshold. This result in turn confirms that most misclassifications indeed occur among uncertain nodes in the middle region of the histogram.

Given there are only a handful of uncertain nodes, the set of top most probable configurations can account for most of permutations of their node configurations because the first few nodes to be changed are those with $P(\sigma_{i}{=}1){\approx}0.5$. This enumerated set allows for coverage of configuration space around the model’s initial prediction. We enumerated top 1000 most probable configurations for each $J$ in the test set to cover 10 most uncertain nodes since $1000{\approx}2^{10}$. We then calculated the energy of these configurations and picked the lowest energy configuration as an improved prediction of the model, $\hat{\sigma}_{top}$. From this procedure, we additionally found the ground state configurations in 1673 $J$’s. This improvement of $\hat{\sigma}$ by configuration enumeration suggests that the uncertain nodes contain frustrated nodes to which the configuration energy is highly sensitive and, thus, that the GNN model has an implicit representation of energy in the node embedding.

For the half of test set where our model missed the ground state configurations, all predicted configurations still have small energy differences relative to the ground state configurations. We exploited this by running 5 simulated annealings with $\hat{\sigma}_{top}$ as a starting configuration for each remaining $J$. Since we are now annealing from a low-lying point in energy landscape, the starting temperature of the annealing should concurrently decreased to prevent the chain from sampling arbitrarily high energy states. We used the temperature value at which the energy trajectory of sampled states drifts up to a bit higher energy at the beginning to allow for initial exploration of energy landscape Note (1). The $\hat{\sigma}_{top}$-seeded simulated annealing found the ground state configuration for additional 1417 $J$’s with about two orders of magnitude reduction in the averaged number of sampling steps as shown in Fig. 3. In about 20% cases, the minimum number of sampling steps from 5 trials were only few hundreds as only one or two node were misclassified in $\hat{\sigma}_{top}$. In total, we found ground state configurations for 75% of test set $J$’s. This seeded simulated annealing result shows that the predicted configuration falls in the vicinity of $\sigma_{min}$, which is often close enough that simulated annealing can locate $\sigma_{min}$. Given the top most probable configurations, we could also run simulated annealing from other configurations or perform parallel tempering Swendsen and Wang (1986); Earl and Deem (2005) with multiple configurations to account for the possibility of $\hat{\sigma}_{top}$ falling in a basin that is too far away from the one containing $\sigma_{min}$.

To test for size generalizability of GNN model, we first trained models on eight small-$J$ training sets with increasing size cutoffs and test them on the existing test set. As shown in Table 1, test set accuracy and energy difference roughly reach those of the original model trained on the entire training set when the size cutoff for $J$ is above 300. Note that the number of the match between the annealed configurations, $\hat{\sigma}_{anneal}$, and $\sigma_{min}$ decreases as the size constraint increases because most of $\sigma$’s are already recovered through model predictions. In all cases, the configuration enumeration and seeded simulated annealing improve upon initial model predictions. If $\hat{\sigma}$, $\hat{\sigma}_{top}$, and $\hat{\sigma}_{anneal}$ are considered together, GNN model provides comparable performance even at the size cutoff of 200. The local interaction pattern of 6-hop neighbor networks in a protein shorter than 100 amino acids should be similar enough to that of much longer chain. It is thus likely that relatively poor performances with training sets with size cutoff less than 100 are simply due to a limited amount of available data.

To test this hypothesis in more practical use case setting, we tested the original model on $J$’s of size 3661 and 2814. The model predicted $\hat{\sigma}$ with energy -443 and configuration enumeration further improved the energy to -447 for $J$ of size 3661, whereas the lowest energy configuration found from 30 randomly initialized simulated annealing runs was -432. On $J$ with size 2814, we obtained energy values of -366 for $\hat{\sigma}$ and -370 for $\hat{\sigma}_{top}$ whereas randomly initialized simulated annealing only reached down to -362. As in previous analysis, we launched simulated annealing from $\hat{\sigma}_{top}$ and obtain annealed configurations with energy -463 and -391 for $J$ with size 3661 and 2814, respectively. Fig. 4 highlights the efficiency of the GNN model over randomly initialized simulated annealing.

Our work shows that it is indeed possible to use an ensemble of energy landscapes with known ground state configurations to train a neural network to deduce the ground state configuration of similar energy landscapes. On our model problem, we deterministically found the ground state configurations on 50% of the held out test set $J$’s and stochastically on additional 25% through a graph neural network, top configuration enumeration, and seeded simulated annealing. Although this number may appear modest, we emphasize that all configurations predicted by the model were extremely low-lying configurations, often in the vicinity of the ground state configurations. Since the loss function does not include other local minima—or, for that matter, the energy function itself, we believe that such an informed prediction is possible only if the learned node feature embedding of GNN correctly captures the local interaction rule encoded in $J$ interaction matrices, and hence the topological undulation in configuration space.

In addition, we showcased the practical utility of the GNN model with size generalizability experiments. The GNN model predicted the configurations that could not be reached by naive simulated annealing with random initial guess, and we were able to further improve it by combining the enumeration scheme and a seeded simulated annealing. Therefore, the GNN model presents an appealing method to produce extremely good initial guesses for a class of energy landscape problem where governing physics is local.

In future work, we will apply this framework to a variety of problems where the discovery of global minima would have technological consequences.

Supplemental Material for “Predicting ground state configuration of energy landscape ensemble using graph neural network”