Clustering molecular energy landscapes by adaptive network embedding

Network embedding

Real-world networks, particularly those representing possible molecular structures and other biological and chemical systems, are often large and complex, making them difficult to conceptualize. Network Embedding maps nodes of a given network into a low-dimensional continuous vector space to facilitate downstream machine learning tasks. Using the sparsity of networks, recently developed Network Embedding methods can scale up linearly with regard to the number of edges. Major techniques include factorization of functions of the adjacency matrix, random walk samplings of node neighborhoods, and deep network learning techniques. The idea is that nodes with close proximities in the network should have similar embeddings in the latent space.

The basic setup is an undirected network G (S, E) with node set S and edge set E, while the generalization to directed graphs is straightforward. Let |S| = n and A ∈ R^n×n be the weighted adjacency matrix of the network with weight a_ij ≥ 0 between nodes v_i and v_j. The output will be a map determined by

For methods discussed in this paper, the encoding function (1), referred as direct encoding, is simply a lookup matrix: z_i = f (v_i) = Ze_i, such that Z ∈ R^d×n contains the embedding vectors for all nodes v_i ∈ S, and e_i is an indicator vector; therefore, f (v_i) simply gives the ith column of Z. The set of trainable parameters for direct encoding approaches is the embedding matrix Z, which is to be optimized directly. Vector a_i = {a_ik}ⁿ_k=1 denotes the first-order proximity between node v_i and other nodes. The second-order proximity between v_i and v_j can be determined by the similarity between a_i and a_j, which compares the neighborhood structures of pairs.

In DeepWalk and Node2vec^[7,8], short random walk simulations are used to determine proximity for each pair of nodes. More specifically, random walk runs through nodes of the network, with transition rates determined by the edge weights. Two nodes are close to each other if there is a high probability that a random walk simulation containing one node will also include the other. Similarities between embedded nodes are given by the following conditional probabilities based on the SkipGram model^[9,10], expressed as

Where P (j|i) denotes the conditional probability that a random walk starting at node v_i will include node v_j and z_i is the embedding of v_i. The learning is achieved by minimizing the following cross entropy loss using a Stochastic Gradient Descent (SGD) method^[11]:

Where R (i) represents a k-step random walk trial starting from node i. The efficiency of evaluating (2) can be greatly improved using negative sampling^[10] that randomly selects edges favoring less frequent ones.

In this paper, we will adopt a Network Embedding scheme introduced in^[13], where the embeddings are produced via a sparse approximation of random walks on networks. For a given undirected network G (A) with adjacency matrix A = (a_ij), let D be the diagonal matrix such that D_ii = ∑_ja_ij. The volume of the graph will be given by v = ∑_iD_ii. The Laplacian L = I - D^-1A has eigendecomposition L = ΦΛΦ^T, where Λ represents the diagonal matrix of ordered eigenvalues so that 0 = λ₁ ≤ λ₂ ≤ ... λ_n, and the eigenvectors are given by columns of Φ denoted by $$ \phi _{1} $$, $$ \phi _{2} $$, ..., $$ \phi _{n} $$. Assuming the network is connected, the discrete Green function satisfies

We can further define the commute time ct (i, j) to be the mean time for the Markov process prescribed by transition probability matrix T = D^-1A = I - L on the network to travel from node v_i to node v_j, and back to v_i. It is shown in^[17] that the coordinate matrix for embeddings that preserve commute times has the following form:

To produce an efficient approximation of Θ, we can assume that T is local in the sense that, at least asymptotically, its columns have small support, and high powers of T will be of low rank, which can be justified for potential driven systems by disparate transition rates between different neighboring metastable states. Taking higher powers of T is equivalent to running the Markov chain forward in time, which allows for representations of the random walk, i.e., reaction pathways, at distinct time scales. Making use of this sparsity, we can produce compressed approximations to the dyadic powers of T with its principal components employing fast algorithms such as Lanczos Bidiagonalization for singular value decomposition (SVD) with the complexity depending linearly on the number of nonzero elements.

The following scheme is a modification of the Diffusion Wavelet algorithm^[18]. Starting with T₀ = T, at each iteration, taking U_k and ∑_k to be the top left singular vectors and singular values of T_k ≈ $$ T^{2^{k}} $$, let T_k: = $$ U_{k}^{T} $$T_k-1U_k. From this, we can have a low-rank approximation to the Green function of the random walk using the Schultz method, as given in

The embedding matrix Θ thus satisfies Θ^TΘ= vG, where v is the volume of the network defined above. Denoting the leading singular values and left singular vectors of the matrix vG by ∑_G and U_G, we take Θ: = (∑_G)^1/2U_G^T.

We can further use Θ as a starting point, introduce parameters by multiplying its jth column by a weight c_j, and optimize the {c_j}’s to minimize cross entropy loss (3) via SGD. Moreover, for robustness of the algorithm, with certain probability, we can reintroduce singular vectors removed in previous truncations, as a residual correction technique.

Metadynamics

Metadynamics was introduced in^[14] as a technique to aid in exploring energy landscapes. The scheme is to create a non-Markovian, and approximately self-avoiding, random walk by adjusting the gradient of the energy landscape after each step with the addition of a derivative of a Gaussian term. Over time, these Gaussians eventually fill up the valleys in the energy potential, which allows the random walk to explore other areas of the landscape and leads to a more complete picture of the system dynamics. Specifically, after each step, the parameter $$ \phi _{i} $$, which represents the derivative of the energy with respect to the ith parameter -$$ \frac{\partial E}{\partial xi} $$, is adjusted according to

where W, δ and $$ x_{i}^{t} $$ are the height, width and center of the Gaussian, respectively, to be chosen based on prior knowledge of the energy landscape.

Transition path theory

The TPT^[15] studies statistical properties of the reactive trajectories such as rates and dominant pathways through probability currents between adjacent states. In the simplest setting, given reactant state A and product state B, any equilibrium path X (t) oscillates infinitely many times between A and B, with each oscillation from A to B being a reaction event. The reactive trajectories are successive pieces of X (t) during which it has left A and is on its way to B next, without coming back to A.

The discrete forward committor $$ q_{i}^{+} $$ is defined as the probability that the process starting in node i will first reach B rather than A, and the discrete backward committor $$ q_{i}^{-} $$ is defined as the probability that the process arriving in node v_i last came from A rather than B. For Markov processes with infinitesimal generator T = (t_ij), the forward committor satisfies the discrete Dirichlet equations:

with the boundary condition $$ q_{i}^{+} $$ = 0, if v_i ∈ A ; $$ q_{i}^{+} $$ = 1, if v_i ∈ B. The backward committor satisfies a similar equation. For time reversible processes, $$ q_{i}^{+} $$ = 1 - $$ q_{i}^{-} $$.

The probability current of reactive trajectories is the average rate at which they flow from one state to another when the process is at statistical equilibrium with distribution π, and can be obtained by

To deal with the fact that transitions between any two states can go forward and backward, the effective current can be introduced as $$ f_{ij}^{+}=\mathrm{max}(f_{ij}^{AB}-f_{ji}^{AB},0) $$.

8-atom LJ cluster

LJ clusters are often used as a model for atomic or molecular dynamics within a fluid, in which the potential energy E of a given configuration of atoms depends only on distances between atoms, as formulated in

where r_ij denotes the distance between atoms i and j, and the parameters σ and ϵ represent pair equilibrium separation and well depth. In the experiments below, we adopt reduced units (e.g., σ = ϵ = 1).

To apply the Network Embedding techniques to the LJ clusters, we first generate a database of local energy minima using the Pele software^[19]. The database was produced using a basin-hopping run of 500 steps and consists of eight local minima connected by ten transition states. It also contains thermodynamic information, including the potential energy at each local minimum. The embeddings here are based on a network constructed with nodes given by the local minima, and edges located between each pair of nodes where a transition state has been identified. The adjacency matrix has entries given by the energy barriers between metastable states.

Initially, every node is embedded into the vector space. Since most points near the global minimum will be close to each other in terms of commute time, this typically results in these nodes being embedded around the global minimum. Then, removing nodes that are further away from the global minimum and re-embedding only those in the residual cluster reveals new, more detailed information about the remaining nodes. We do this by creating a new, smaller adjacency matrix including only re-embedded points, and adjusting the edge weights with Gaussian terms according to Metadynamics

where θ_i is the coordinate of the ith node (energy minimum). Adaptively choosing the center node and removing distant nodes, the process will produce a series of hierarchical embeddings that provide information about the full energy landscape. Each level provides a representation of the energy landscape at a different scale.

Figure 1 gives the disconnectivity tree of the original potential and embeddings after applying the Metadynamics adjustment by equation 11 for the 8-atom LJ cluster. The colors on each minimum on the disconnectivity tree are matched to embeddings of those nodes in Figure 1, with some of the nodes colored red, embedded into the same place. In the embeddings, nodes with closer dynamic relationships are clustered together as a result of the shorter commute time distance between them.

Figure 1. Disconnectivity tree and Metadynamics-based embeddings for the Lennard-Jones cluster with eight atoms. Left: Disconnectivity tree of all local minima; Right: Embeddings for the local minima after applying the Metadynamics adjustment. Color scheme represents the potential energy; e.g., dark blue denotes the lowest, and red indicates the highest. Closely related minima have very similar or identical embeddings; e.g., both yellow minima are embedded at the yellow point on the right.

On a larger scale, we can also see how the higher energy nodes relate to the nodes with the two lowest energies. Specifically, the global minimum (in dark blue) is closer to the nodes in the middle of the tree (in orange), while the second lowest local minimum is much closer to the three highest energy nodes. This allows us to conclude the lowest commute time path through these states: one of the orange-colored states might transition directly to the global minimum, while one of the higher energy states would be more likely to transition to the second lowest energy state first and then either remain there or transition on to the global minimum.

For comparison, we also investigated the inter-node relationships by directly computing the commute times between each pair of nodes, as provided in Table 1. The commute time between nodes v_i and v_j, or mean time for the Markov process on the graph prescribed by the Laplacian L to travel from v_i to v_j and back to v_i, is given by

where λ_k are the eigenvalues of L, and the $$ \phi _{k} $$ are the corresponding eigenvectors as defined in the introduction. The off-diagonal entries for the Laplacian for a LJ cluster with k degrees of freedom are given by

where O and O_k represent point group orders of the local minima and transition states, respectively; similarly, v and v_k indicate mean vibrational frequencies, and E and E_k stand for potential energy levels at each configuration. It can be seen that the commute time diagram is consistent with the Network Embedding output.

Multi-level embeddings of LJ 38 cluster

To scale up this result, we also apply this framework to the 38-atom LJ cluster. Using a 500-step basin-hopping run, we produced a database of 8,797 local minima connected by 8,099 transition states. Figure 2 shows the disconnectivity tree of the 38-atom LJ cluster. For simplicity, we construct the adjacency matrix for this network with entries given by the energy barriers between states. Figure 3 gives the hierarchical embedding of the LJ-38 cluster. The left image shows the initial embeddings, colored by their commute time distances from the global minimum. We re-embedded the points of interest with the Metadynamics adjustment to the potential with two iterations and obtain the right image.

Figure 2. Disconnectivity tree for the 38-atom LJ cluster. The structures of the two lowest-energy configurations are also pictured. LJ: Lennard-Jones.

Figure 3. Hierarchical embeddings for the LJ cluster with 38 atoms. Pictured are the embeddings before (left) and after (right) applying Metadynamics. Color scheme denotes commute time from the global minimum, with dark blue being shortest distances, and red as furthest distances. LJ: Lennard-Jones.

The hierarchical structure of the embedding process is organized as follows. In the first level embeddings [Figure 3, left], the nodes are embedded consistently with their commute time distances from the global minimum, which correspond loosely to the potential energy level. In particular, the nodes with energy E > -170, which are the highest energy clusters in the tree, are embedded further from the global minimum. In the next embedding, these nodes are removed due to their further distances from the global minimum, and more central clusters will be re-embedded. As we consider the second level and later embeddings, this pattern repeats: each re-embedding reveals a new “layer” of nodes embedded closer to the global minimum, which corresponds to nodes on the disconnectivity tree that are of lower energy than nodes removed in the previous level.

In particular, at the third level [Figure 3, right], the embeddings have four small “spokes” originating from a central cluster containing the global minimum. However, these embeddings provide additional context: nodes embedded within a particular “spoke” are more closely related, which means the system is more likely to transition between these states. Since the nodes do not all come from the same group on the disconnectivity tree, these embeddings can also reveal interactions between nodes not indicated on the disconnectivity tree.

Additionally, since the spokes are connected to the cluster containing the global minimum, we can conclude that each spoke represents a potential transition pathway from the outer edge of the cluster to the center. In other words, if the system is at a state represented by the outer point of one of the spokes, its most likely path toward the global minimum will be to travel through the states represented by other nodes in the same spoke.

We can apply similar reasoning to higher level embeddings. Each level reveals a more detailed picture of the dynamics of a different part of the energy landscape of the system. The first levels give a coarse-grained picture, only identifying broad groups of high and low energy nodes, while later levels give a more fine-grained visualization of the nodes most closely related to the global minimum.

Often, we want a more detailed, finer-grained visualization of the energy landscape than the disconnectivity tree in Figure 2 can provide. It is informative, therefore, to re-embed parts of the network of greatest interest to gain further insight. Here, we focus on the lowest energy parts of the LJ energy landscape. We repeated the above process using the subnetwork consisting only of the nodes with potential energy < -170.9, that is, the 163 lowest energy nodes. Figure 4 presents the results of this experiment.

Figure 4. Embeddings of the local minima of the 38-atom LJ cluster with potential energies less than -170.9. The figure shows output of the second level embeddings with Metadynamics adjustment. Color scheme denotes commute time from the global minimum. LJ: Lennard-Jones.

Now, we want to provide a closer inspection on the information flow along the hierarchical sampling with Metadynamics. In the first level without the Metadynamics adjustment, nodes in these embeddings [Figures 3 and 4] are clustered according to their similarity in terms of commute time. In other words, nodes that can be quickly and frequently reached from either the global minimum or second lowest energy node will be grouped near them. As a result, most of the nodes we are most interested in end up in the same cluster, and the embeddings obtained from the first application of the embedding method only give useful clusters for nodes more distant in terms of commute time from the part of the graph of interest. As we progress through additional levels, we pull apart the cluster containing the global minimum, positioned near the origin in the embeddings of the first level, until the embeddings of the final level give details of the dynamics of the process within this cluster.

After the second level of Embedding with Metadynamics adjustment, if two nodes share a cluster or are close in the embedding space, it indicates that the system can easily transition between those nodes, with relatively low energy barriers. As a result, the groupings seen in these embeddings correspond to the groupings in the disconnectivity tree for this system^[3,19]. For instance, one of the clusters in the second level embeddings corresponds to the global minimum and its nearest neighbors, pictured directly right of center in the disconnectivity tree in Figure 2. Clusters can be mapped to the disconnectivity tree by comparing the potential energies of the nodes within the cluster to the tree.

The difference is after the third level, some of the clusters instead represent combinations of multiple disconnectivity tree groups; this is a result of re-embedding the nodes to spread out those previously near the origin. Such nodes ended up being embedded in or near the clusters they are most closely related to, even though they are not actually members of the respective tree groupings. For instance, the global minimum is embedded directly next to a node from a neighboring tree group.

In other words, the embeddings of the first level tell us about higher energy nodes and those that are more distant from the global minima, and further levels reveal information about parts of the network that are closer to the global minimum. Additionally, these embeddings are useful for identifying transition paths. The structure of the third level embeddings, in particular, reveals four transition paths: if the system is initialized from a node near the outside of one of these “spokes”, its lowest energy path to the global minima will involve following the spoke into the center cluster.

We can also use these embeddings to observe the results of entropic changes to the system. In Figure 5, the embeddings for this cluster under two additional temperature conditions are shown. The temperature change results are reasonable according to what was observed previously with the 8-atom cluster. Namely, at lower temperatures [Figure 5, left], closely related nodes are more likely to be embedded much nearer each other, creating the impression that there are fewer embeddings, while at higher temperatures [Figure 5, right], there is greater variation in the node embeddings.

Figure 5. Metadynamics-based embeddings for the 38-atom cluster at the temperatures T = 0.08 (left) and T =1 (right). Color scheme denotes commute time from the global minimum.

Human telomere folding

Next, we want to further develop the Network Embedding technique for organic molecular structures and apply it to a more complex problem: DNA folding in a human telomere. More specifically, we consider a sequence of 22 nucleotide bases A(G₃TTA)₃G₃ which repeats within human telomeres. This sequence is known to form a G-quadruplex [Figure 6], a type of secondary structure formed by groups of four guanine bases called G-tetrads. Its structure and potential energy landscape were previously investigated in^[21], where the potential energy landscape was calculated using the HiRE-RNA model for coarse-grained DNA^[22] with six or seven atoms considered for each of the 22 nucleotides in the telomere. We use this database as a starting point. In particular, we construct a network with nodes given by the 4,000 lowest-energy local minima, and edges between nodes determined by the transition states connecting them.

Figure 6. The four-strand G-quadruplex structure (PDB structure 1KF1), with guanine nucleotides colored green. Image produced with Chimera^[20].

For the initial experiments, we used a random walk based on the energy barriers between states. Figure 7 shows the results of the initial embeddings and fourth levels with Metadynamics adjustments, based on an energy landscape adjusted by a Gaussian term with width 0.75 and height 1 after each successive embedding. As in the LJ cluster experiment, the first embedding includes all nodes in the network, while each successive embedding shows a re-embedding of the subnetwork of nodes most closely related to the global minimum (that is, all nodes whose previous embeddings lie within a small tolerance of the global minimum).

Figure 7. Embeddings of the local minima network for the human telomere sequence, based on the adjacency matrix and the Metadynamics adjustment. Color scheme denotes commute time from the global minimum.

As with the LJ clusters, each level of embedding represents “zooming in” on the part of the network around the global minimum. The first level gives us a global view- the global minimum and its nearest neighbors (with respect to commute times) are clustered at the origin, represented by a dark blue dot. The red and orange dots furthest away from the origin represent local minima which are more distantly related, requiring multiple steps or higher energies to transition to the global minimum. Potential transition paths can be identified by starting at one of these points and moving toward the origin along nearby points. At the second and each of the following levels, the nodes embedded closest to the local minimum are re-embedded to give us a more detailed inspection into the relationships of those nodes. The likely transition paths here can be constructed similarly.