Tangent Space Least Adaptive Clustering
Abstract
The biasing of dynamical simulations along collective variables uncovered by unsupervised learning has become a standard approach in analysis of molecular systems.
However, despite parallels with reinforcement learning (RL), state of the art
RL methods have yet to reach the molecular dynamics community.
The interaction between unsupervised learning, dynamical simulations, and
RL is therefore a promising area of research.
We introduce a method for enhanced sampling that uses non-linear geometry estimated by an unsupervised learning algorithm in a reinforcementlearning enhanced sampler.
We give theoretical background justifying this method, and show results on data.
1. Introduction
Molecular dynamics (MD) simulations are an empirical method for understanding the properties of materials from proteins to solar panels (15, 22, 9, 11).
In these simulations, positions of individual atoms or sub-molecular agglomerations are sampled through time according to interatomic forces, but despite knowledge of these forces, the overall conﬁgurational behavior is unknown.
Simulations thus provide information about the low-dimensional dynamics and properties of the simulated system (16).
Unfortunately, the computational burden of these simulations is high (25).
The potential energy and its gradient are often challenging to compute and low-probability but important regions like transitions between energy wells require huge numbers of samples to access.
Satisfactory exploration of chemical compound behavior in simulation is therefore a challenging and unresolved problem.
A variety of machine learning methods have been proposed to reduce this computational cost (33, 20).
For example, potentials generated using supervised learning have been used to sample from hard-to-compute (6) or
hard-to-ﬁnd (26) conﬁgurations.
Unsupervised learning is also used to assist simulations through discovery of collective variables (CVs) that parametrize the slow manifold generated by the underlying dynamics (24, 21, 10, 28).
Once this data manifold is learned, it can be used to bias future similar simulations towards interesting high-energy or transitional regions.
Many methods for biasing simulations based on CVs have been proposed (31, 19, 5).
There are many similarities between reinforcement learning and this collective-variable based sampling paradigm.
Simulators balance exploration versus exploitation, and can use tools like differentiable deep networks to identify good actions for progressing the simulation to interesting states (3, 23).
The variation of experimental conditions like solvents or temperatures to create a more robust understanding of molecular behavior resembles domain randomization (32).
As in simulated robotics (34), a complicated heirarchy of simulation ﬁdelity relates the simulated system to the real world phenomena of ultimate interest (12, 30).
Perhaps due to the relatively recent success of reinforcement learning methods, formal application of reinforcement learning principals is still limited.
However, it is reasonable to expect that application of these principals should integrate with existing unsupervised-learning based approaches to enhanced sampling using CVs.
One approach that does perform enhanced sampling through a combination of CVs and RL is the REinforcement learning based
Adaptive samPling
(R EAP ) method of (27).
However, this approach does not use any non-trivial unsupervised learning.
Our idea therefore is to adapt this approach from the perspective of non-linear geometry.
We therefore introduce Tangent Space Least Adaptive Clustering (T SLC ) that favors exploration of variables that parametrize the data manifold.
We use local tangent space estimation as a simple and straightforward method for learning the manifold, introduce a new reward based on projection onto tangent spaces, and give an efﬁcient clustering algorithm for determination of sets on which to compute the tangent spaces.
2. Background
We motivate this method in the context of unsupervised learning for dynamical systems.
Dynamical systems Dynamical models allow simulators to progress forward in time t through computation of the derivative of
the state vector x ( t ) ∈
R D w.r.t.
t .
A simple model of dynamics is given by the Langevin Equation (27)
[Equation]
where is E ( x ( t )) is the potential energy, η ( t ) is a random force and γ is a constant that, together with the temperature T and Boltzmann constant k B , determines the degree of randomness of the random walk through the conﬁguration space.
Langevin dynamics for even small systems are chaotic, and even though we can compute E and its derivatives, we do not in general know the properties of the slow-manifold around which the random walk progresses (1).
Additionally, there is no guarantee that this walk will traverse the slow-manifold in a reasonable amount of time.
Therefore, these dynamics are a motivating application of unsupervised learning for acceleration of dynamical simulations (17, 2).
Unsupervised learning and tangent spaces Suppose, in the ideal case, that we are able to simulate from the Boltzmann distribution until the density of x ( t ) is exactly proportional to E .
Theoretical analyses often guarantees that these data are distributed on a manifold M containing a central slow-manifold S within the D dimensional feature space (13).
Unsupervised learning then estimates a map φ such that φ ( M )
=
φ | M ( S ) is 1) diffeomorphic to S and 2) in a low-dimensional space R m ( m << D ) .
Given such a φ , the typical approaches for accelerating simulation are 1) to ﬁnd a geometric feature of the conﬁguration parameterizing φ ( S ) and move the simulation in its direction (25), or 2) if φ is differentiable, to move the simulation in a direction ∇ φ k where φ k is a coordinate of the learned latent space (5).
In either case, the geometric feature or coordinate functions φ k are a CV that parameterize S .
One issue with the above is that we may not have enough data to learn a good φ .
In this case, tangent space estimation is a less data-intensive alternative to learning φ .
The reason is that, if φ is suitably well behaved, then the singular value decomposition of its differential in coordinates of R D and
R m gives Dφ ( x ) =
( T x S )
Σ x
( T φ ( x ) φ ( S )) T (18), where
T x S is the tangent space of the slow-manifold.
Biasing simulations by the method of collective variables is essentially moving along the tangent space T x S .
In our approach, we will substitute in a tangent space estimator for T x S without explicit calculation of φ .
REinforcement learning based Adaptive samPling RE-
inforcement learning based
Adaptive samPling
(R EAP ) is an application of RL principals to accelerating traversal of conﬁguration spaces in dynamical simulations (27).
It is an extension of least counts adaptive sampling (L C ), a simple iterative paradigm for enhanced sampling that uses clustering to identify undersampled regions.
Our proposed method for sampling is largely motivated by R EAP .
In L C , every iteration consists of simulating dynamics from given starting states, clustering the trajectory, and then selecting the next starting states from the smallest, or least count, clusters.
Clusters with fewer members correspond to more sparsely sampled regions of the conﬁguration space, and will tend to be towards the boundaries of the explored region.
R EAP also proceeds iteratively by simulating, clustering, and selecting new initial states to start sampling from, but the way the next initial states are chosen from the least count clusters is more involved.
R EAP supposes preselection of a set of possible functions of the state, θ 1 ( x ) . . .
θ K ( x ) :
R D → R , also known as collective variables.
With no a priori knowledge of S , this set will be larger than ideal, and the R EAP algorithm must decide their relative importance on-the-ﬂy.
Every iteration, after clustering to ﬁnd least count clusters, a weight for each CV is recomputed.
These weights w ∈ R K are determined by maximizing
[Equation]
where C p is the set of states in the least count clusters, θ i ( s m ) is the value of the i th CV at state s m ∈ C p , and (cid:104)
θ i ( T ) (cid:105) and σ i ( T ) are the mean and standard deviation of θ i over the full trajectory T thus far, respectively.
The w i are constrained to sum to one, be nonnegative, and each w i may not change by more than a parameter δ from its value in the previous round.
These weights are then used to select initial states for the next round of simulation, namely a ﬁxed number of states in C p which maximize
[Equation]
Comparison with the rich universe of RL methods indicates that there are many areas of potential improvement to this approach (8).
We focus in particular on adapting Equation 2 with the geometrical underpinnings of the method of collective variables.
The R EAP objective only considers population statistics like mean and standard deviation - much less information than can be captured through unsupervised learning (29).
3. Methods
We thereby propose a method for determining weights for
CVs which differs from R EAP in two key ways.
First, we select weights which capture global geometry of the underlying data manifold from only local geometry.
Second, we replace k-means clustering with our own clustering method
C LUST , which is amenable to the sampling problem.
Our method, Tangent Space Least Adaptive Clustering (T SLC ), like R EAP , consists of timestepping, clustering, reweighting, and reinitialization.
CVs are known, and weights calculated for each CV at each iteration are used to select points from least count clusters.
However, unlike R EAP , we assume, as is typical in unsupervised learning for dynamical systems, that the explored potential energy surface lies near a low-dimensional manifold.
Algorithm
1: Input: Differentiable energy E , collective variables θ 1 . . .
θ K , number of least count clusters to sample from
N , number of parallel runs M , number of iterations I , timesteps per iteration s , intrinsic dimension d 2: for t ∈ 1 . . .
I do 3:
Compute clusters, centers, and least count clusters:
C , q, C p ← C LUST ( x (1 : tsM )) , N ) 4:
Compute V i ← PCA ( x ( C i )) , d ) for i ∈ 1 . . .
|C| 5:
Compute ∇ θ k ( q i ) for i ∈ 1 . . .
|C| , k ∈ 1 . . .
K 6:
Normalize ∇ θ k ( q i ) ← n ∇ θ k ( q i ) (cid:80) |C| i =1 ||∇ θ k ( q i ) || for i ∈ 1 . . .
|C| , k ∈ 1 . . .
K 7:
Compute A as in Equation (5) 8:
Compute ˜ w , the largest eigenvector of A 9:
Compute w k
=
˜
w 2 k for k ∈ 1 . . .
K 10:
Select S M , the M states from C p with the highest values of Equation (3) 11:
For each state s m in S M , timestep Equation (1) s times starting from s m 12: end for
Algorithm description There are three important observations in the development of our choice of weights.
One is that the collective variables which characterize the potential and encourage exploration should have gradients which are parallel to the tangent spaces of S .
In particular, the projection of gradients onto local tangent spaces should be large.
Two is that local tangent spaces can be computed on each cluster, using Principal Component Analysis.
Since our data are distributed near S , we can estimate its tangent spaces as long as our clusters are small enough to reﬂect local information about the manifold, while also large enough to not be dominated by small-scale ﬂuctuations.
These requirements on the clusters for our purpose are what motivate our clustering algorithm and ﬁnal observation, that the volume of the clusters chosen should scale with the volume explored.
Since T SLC only differs from R EAP in the computation of weights and choice of clustering, we only explain those steps in detail.
Weight Computation (1)
We compute a clustering of the full trajectory thus far C LUST ( x full ) =
C .
We calculate for each cluster i a d dimensional local tangent space using standard Principal Component Analysis.
Denote a suitable choice of d D -dimensional orthogonal vectors which span this space as columns of the matrix V i ∈ R D × d .
(2) The center of the cluster i , q i , is used as a representative of that cluster.
We compute the values of ∇ θ j ( q i ) for each
CV θ j on each cluster i , which we can can calculate with automatic differentiation since the CVs are analytic.
In order to avoid favoring CVs that are scaled larger, we normalize the gradients over the clusters (see Algorithm 1).
Let G i ∈
R D
× K be a matrix so that the j -th column is the normalized gradient ∇ θ j ( q i ) .
(3) To ﬁnd the weights, we solve the following optimization problem: ﬁnd the unit 2-norm column vector w ∈ R K so that
[Equation]
is maximized.
This selects a linear combination of the gradients of the CVs that on average makes the projection onto the local tangent spaces largest.
Since the columns of each
V i are orthogonal,
||
V
i V
Ti G
[Equation]
V Ti G i w || 2 .
It follows that this objective is
[Equation]
to which the solution is the eigenvector of A ∈ R K × K associated with the maximum eigenvalue of A , which exists because A is symmetric.
Finally, let the weights w i be the squares of the entries of w .
These are nonnegative and sum to one.
The states to start sampling from in the next iteration are the M states in C p which maximize Equation (3).
Clustering
C LUST ( x, N ) takes as input x full , the full trajectory so far and N , the number of least count clusters we want, and returns the clustering, centers of each cluster, and least count clusters.
It is inspired by the initialization method of (7), extended to K-means by (4), and simpliﬁed to take into account the need for speed and for small clusters.
In T SLC , we replace k-means clustering with a clustering procedure that recomputes the number of clusters n clu to use at every iteration.
We calculate the number of clusters starting from the following premise.
In a uniform random walk, the distance from the start of the data grows asymptotically as √ n , where n is the total number of sampled points.
We expect our sampling to explore faster than a random walk, and the distance to grow as n γ for a parameter γ ∈ (1 / 2 , 1) .
Hence, the volume explored will grow as ( n γ ) d and to maintain clusters of approximately
[Figure]
equal volume, the number of centers should grow proportionally with this number.
Therefore n clu is taken to be floor ( b ( n γ ) d ) , where b is a parameter.
Then, a number of centers n (cid:48) clu = floor ( n clu log n clu )
+ l are sampled uniformly from x full where l is a parameter.
To these n (cid:48) clu centers, we apply the Fastest First Traversal algorithm (14), pruning the centers until only n clu are left.
Finally, we perform a single step of k-means with these initialized centers to assign points to clusters and to recalculate the centers.
The result is n clu clusters which retain approximately the same volume every iteration of T SLC and whose centers are spread uniformly over the explored region.
4. Experiments Comparison of T SLC with R EAP
We compare
T SLC with R EAP on a toy circular potential function embedded in D = 10 dimensions.
The potential is lowest on the circle(in the x 1 - x 2 plane) and its magnitude decays exponentially with the square of the distance away from the circle.
It is given by V = − c ∗ exp ( − a ∗ dist 2 ) , where a = 10 , c = 250 , and dist is the distance from the circle in
Euclidean space.
Both methods are given the same CVs: the coordinate projections functions of the state x , θ i
= x i for i = 1 ...D , and θ 0 = arctan
( x 2 /x 1 ) , the angle on the circle.
The deﬁnition of θ 0 must be adjusted so that it is correctly deﬁned on every quadrant of the x 1 - x 2 plane.
To compare, we run both sampling methods I = 700 . iterations, and look at the percentage of the circle explored.
R EAP was implemented using code provided by (27), available online.
Timestepping was done the same way as in (27).
Both timestep from two initial states s = 5 times each every iteration, and sample M = 2 new points from the N = 2 least count clusters found by their respective clustering algorithms.
R EAP uses δ = .
02 and k =for k-means for all iterations, while ours uses parameters b = .
07 , l = 2 and γ = .
We repeat this proceduretimes each and plot histograms of the percentages of the circle explored in Figure 1(e).
On average, R EAP explored approximately 56 % of the circle in the given time, while
T SLC explored 85 % , indicating strong performance of
T SLC .
In Figure 1(a) we show what one round ofiterations of R EAP typically looks like using the same methodology as above.
The plot shown is a contour plot of the potential V only as a function of x 1 and x 2 , with the other 8 state variables set to 0.
In 1(b) we do the same for
T SLC .
In Figures 1(c) and 1(d) we plot weights of θ 1 , θ 2 , and θ 0 that R EAP and T SLC select over typical runs, as a function of iteration number.
T SLC is able to identify θas the CV which characterizes the potential, and gradually becomes more conﬁdent in this, while R EAP tends to select one variable at a time which it believes is the most important.
5. Conclusion
We have described a novel reinforcement learning-based approach for simulating from dynamical systems.
This method is inspired by unsupervised learning methods currently used in molecular dynamics, and is a step towards a synthesis of unsupervised-learning CV based methods with RL.
6. Acknowledgements
This material is based upon work supported by the National
Science Foundation under Grant DMS-1810975, the ARCS
Foundation Award, and the GO-MAP Graduate Excellence
Award.