Policy-guided Monte Carlo: Reinforcement-learning Markov chain dynamics

We introduce policy-guided Monte Carlo (PGMC), a computational framework using reinforcement learning to improve Markov chain Monte Carlo (MCMC) sampling. The methodology is generally applicable, unbiased, and opens up a path to automated discovery of efficient MCMC samplers. After developing a general theory, we demonstrate some of PGMC's prospects on an Ising model on the kagome lattice, including when the model is in its computationally challenging kagome spin ice regime. Here we show that PGMC is able to automatically machine learn efficient MCMC updates without a priori knowledge of the physics at hand.

Figure 1

Sketch of the hierarchy of abstractions in play. A state, e.g. a configuration of a lattice model, constitutes a single point in a high-dimensional state space $\mathcal{S}$. A trajectory of states $s_1\to s_2\to \cdots \to s_n$ of length $n$ corresponds to a single point $c$ in the even higher-dimensional trajectory space $\mathcal{C}$. The darker areas illustrate the regions of higher probability density that contributes the most to the final result.

Figure 2

Trajectory $c$ may be regarded as a number of shorter, time-shifted, and correlated subtrajectories, e.g., $c^{\prime }$, $c^{″}$, and $c^{\prime \prime \prime }$.

Figure 3

Configuration of the Ising model on an $L=3$ kagome lattice. The dark and light disks indicate down ($-1$) and up (+1) spins, respectively.

Figure 4

Generation of two separated monopole excitations from a kagome spin ice ground state by flipping a string of alternating spins. The monopoles are associated with the colored elementary triangles, the only ones that violate the spin ice constraint by not containing two up spins (light) and one down spin (dark).

Figure 5

Tunneling from one kagome spin ice ground state to another. A loop of spins of alternating signs has to be flipped in order to conserve the spin ice rule everywhere. The smallest loop for which this can be accomplished belongs to an elementary hexagon or plaquette.

Figure 6

Learning curves for the nontrivial single-flip policies ${\pi }_2$, ${\pi }_3$, and ${\pi }_4$ [with ${\pi }_x={\pi }_x\left(t\right)]$ showing the expected acceptance rate $\langle \alpha \rangle$ and normalized effective degrees of freedom $d$ as a function of optimization iterations. The quantities displayed are averages over sweeps of $N$ optimization iterations. Each case was initialized with a random configuration and $\mathbf{\theta }=\mathit{0}$.

Figure 7

Snapshots of configurations (top) and the corresponding relative probability weights for selecting the spin to flip in the next iteration (bottom), according to policy ${\pi }_4$. The state on the left is sampled at criticality, while the state on the right is frustrated. The policy, separately optimized to each parameter point, assigns a larger weight to energetic structures like domain walls and monopole excitations, and relatively more so in the case of the frustrated state close to the ground macrostate. The connectivity lines of Fig. 3 and 5 have been omitted and only the values associated with the sites (the small hexagons) are shown.

Figure 8

PGMC sampling performance of Eq. (18) for the parameters listed in Table 1 and all policies investigated in this work, after being optimized. The top plot shows the estimated performance factor (4) calculated from Eq. (26) and (15) and the cost factor $u\left(c\right)$, which is equal to the number of elementary actions used in constructing $c$. The performance factor has been scaled by the number of sites $N$ to make the timescale correspond to a Monte Carlo sweep. The bottom plot shows the acceptance rate. The horizontal lines indicate the optimal performance, i.e., what the results would have been had the configurations been drawn directly following Eq. (18) instead of being generated by Markov chains. The background shading indicate classes of policies sharing a similar structure. Results obtained at the same $(K,B)$ point share a common color and symbol style, with lines as guides to the eye. Results in unfrustrated parameter regimes are shown with closed symbols, while the frustrated ones are shown with open symbols. In some cases, for ${\pi }_1\text{–}{\pi }_3$, ${\pi }_5$, and ${\pi }_7$, finding $\langle {\epsilon }_O\rangle$ for $(K,B)=(-4,2.5)$ was too slow or unreliable, so these points have been omitted. Error bars indicate sample standard deviation. See the main text and the Appendix for discussion and further details.

Figure 9

Example of a three-state trajectory of a chain policy of length 4. Each action, $a_{12}:s_1\mapsto s_2$ and $a_{23}:s_2\mapsto s_3$, consists of four elementary actions. The Metropolis acceptance step is only concerned with the states before and after the actions have been completed (circles with white fill), making room for temporary visits of transit states (open circles) with lower probability weights.

Figure 10

Frequency of updates leading to $n_{\text{flips}}(s\to s^{\prime })\in \{0,...,6\}$ net spin flips in a PGMC sampling of Eq. (18) at $(K,B)=(-4,2.5)$ with policy ${\pi }_8$. Most updates involve no change, either because the proposed action effectively is $a_{⌀}$ or because the proposed update was rejected. The error bars, indicating the sampled standard deviation, are smaller than the symbol size. See the Appendix for technical details.

Figure 11

Worm action $a_{\text{w}}=a_{⌀}{a}_n\cdots a_1:s\mapsto s^{\prime }$ (solid arrows) and its inverse $a_{\text{w}}^{-1}=a_{⌀}{a}_1^{-1}\cdots a_n^{-1}:s^{\prime }\mapsto s$ (dashed arrows).

Figure 12

Snapshots of configurations and most recent actions leading to them, at the critical point (left) and in the kagome spin ice regime (right). In the top panels, the actions have been selected by the worm policy ${\pi }_9$, while in the bottom panels, they have been selected by the short-term self-avoiding policy ${\pi }_{14}$. Stronger intensity indicates that the same spin has been flipped (and flipped back) multiple times within the last action, while the weakly colored background spins have not been flipped at all. Note that while the elementary actions of the worm policy, by construction, form a single connected set, the actions of the short-term self-avoiding policy have no such restrictions imposed. This makes the latter more flexible in adjusting to other parameter regimes, where string or loop updates may no longer be optimal.

Figure 13

(a) Inverse transform and (b) weighted binary tree selection of an action based on some policy $\pi$. Each box represents some action $a$, with the height being given by $\pi \left(a\right)$. In (a), the action space is sorted into feature categories. A category is chosen (colored) according to the total marginal probability of the category by drawing a uniform random number $R_1$. Next, a specific action is selected uniformly at random among the members of the category by random number $R_2$. The total number of operations is constant with respect to the size of the action space $N_{\mathcal{A}}$. In (b), the probability (weights) are arranged in a binary tree structure, with the weight of each parent node being the sum of its children nodes. An action is selected by traversing the tree from root to leaf, each time selecting the left or the right child at random, proportionally to their relative weights. This scheme comes at the (mild) expense of an $\mathcal{O}\left(log{N}_{\mathcal{A}}\right)$ scaling in the number of operations needed. The advantage is that it can be efficiently used even when the number of possible categories is large or even infinite. Naturally, when an action has been applied to a state, both the category groupings in (a) and the weights of the tree in (b) will have to be updated. In cases with short-range feature maps (i.e., when the features only depend on a finite neighborhood of the degree of freedom being updated by the action) this can also be done with $\mathcal{O}\left(1\right)$ and $\mathcal{O}\left(log{N}_{\mathcal{A}}\right)$ numbers of operations, respectively.