Finally, we demonstrate how our framework can be used to optimize the adaptive resetting protocol to perform a given task. We do so by representing the state-dependent resetting probability as a neural network and training it to optimize a loss function that can be any observable of the process with adaptive resetting. The loss is calculated based on trajectories without resetting using our reweighing procedure. As a concrete example, we find an optimized adaptive resetting protocol that minimizes the MFPT of conformational transitions in simulations of the mini-protein chignolin in explicit water.
We begin by defining an adaptive, i.e., a state- and time-dependent, resetting rate r(X, t), with X being the state of the system and t representing time. Given a trajectory, , we can define a random variable describing its resetting time R via its cumulative distribution function
Note that R distributes differently for different trajectories under the same functional choice of the resetting rate. We emphasize that any state- and time-dependent resetting strategy can be represented in this way. For example, for an exponential resetting time, one takes a constant resetting rate r(X, t) = r. More generally, to represent a general resetting time distribution whose survival function is Ψ(t) = 1 – Pr(R ≤ t), one takes , which does not depdend on X. Similarly, one may consider an arbitrary spatial dependence of the resetting rate, e.g., r(X, t) = rΘ(X), where Θ(X) is the Heaviside step function. More generally, arbitrary state-time couplings can also be captured as no restrictions are stipulated on r(X, t) other than it being non-negative. For simplicity of demonstration, in concrete examples given in this paper, we take resetting rates r(X) which are time-independent, although the theory developed below is completely general.
We next construct the random variable describing the FPT under the resetting protocol r(X, t). We observe that a first-passage process with resetting can be described in the following way: First, a trajectory is sampled, and this determines T, i.e., the FPT without resetting. The trajectory also sets the distribution of R via Equation (1). We next sample R, and if T≤R, we conclude that the FPT is simply T. If, however, R < T, resetting occurs before first-passage, a new trajectory is sampled, and the procedure is repeated, tallying R. Overall, the random variable describing the FPT under resetting is
where is an independent and identically distributed copy of T. Using the total expectation theorem, and averaging over all the trajectories, and all realizations of R, we find that the MFPT under restart is
From here, it is clear that the MFPT under restart can also be written as . We note that the same result was previously obtained for state-independent resetting using a similar technique. A detailed derivation of Equation (3) is given in section 1 of the Supplementary Information.
We can interpret Equation (3) as illustrated in Fig. 2. For a given trajectory with resetting, the first passage is composed of several failed attempts to complete the process, each ending in resetting. These are followed by one successful attempt, ending in first-passage, and completing the process. The total number of attempts, failed plus successful, is geometrically distributed. Namely, all attempts are statistically independent with a success probability Pr(T ≤ R). Therefore, the average number of attempts is 1/Pr(T ≤ R). The mean duration of an unsuccessful attempt is 〈R∣R j and zero otherwise. Thus, to estimate the propagator under state- and time-dependent resetting, we only need to estimate P and from trajectories without resetting. We do this as follows:
where we recall the definitions of Ψ and p in and above Equation (4), and sum over trajectories. Similarly, we have
The advantage of this approach is that the solution is obtained in real-time, however, for long times, it requires inverting very large matrices.
Alternatively, the long-time behavior and the NESS can be obtained more efficiently by taking the Z-transform of Equation (9) and using the convolution theorem. This gives
where is the Z-transform of the series {f, f, f, . . }. While Equation (13) is given in discrete time, an equivalent equation using the Laplace transform is valid for continuous time. We note that a special case of the continuous-time analog of Equation (13), for an overdamped particle diffusing in a potential and space-dependent resetting rate, was given by ref. . Our work extends this result to a general stochastic process.
The final value theorem for Z-transforms states that . By using it, we get (see section 6 of the Supplementary Information)
where 〈N〉 is the mean number of time steps between consecutive resetting events. Note that the steady-state in Equation (14) is well defined whenever 〈N〉 is finite, regardless of whether or not the process without resetting has a steady-state. This is a generalization of a well-known result in the theory of standard resetting to state- and time-dependent resetting.
To estimate the NESS, we first sample a set of N trajectories without resetting of length MΔt. We stress that M should be large enough such that, had we used resetting, the probability of surviving M steps without resetting would be negligible, i.e., . Then, we use Equations (12) and (14), and the definition of the Z-transform, to obtain
This estimation results in an unnormalized distribution, which should be normalized. The normalization factor provides an estimate for the mean time between consecutive resetting events, 〈N〉. Equation (15) shows that the estimation of the NESS with resetting, from trajectories without resetting, is done by averaging the histogram of positions over time and trajectories, but reweighing each trajectory, at every time step, by its survival probability.
The above results can be used to predict and design NESS of spatially-dependent resetting protocols. We demonstrate this using two examples.
It is well known that for free diffusion with a constant resetting rate, a Laplace distributed NESS emerges. An analytical solution for the NESS of diffusion with a parabolic resetting rate r(x) = rx is also known. Interestingly, in both cases, the tails of the NESS decay as , with α = 1 for the constant resetting rate, and α = 2 for the parabolic resetting rate. This raises a more general question: what is the asymptotics of the NESS for diffusion with a power-law resetting rate r(x) = r∣x∣. While there are currently no known closed-form solutions for the NESS with λ ≠ {0, 2}, we can easily estimate the resulting NESS using the procedure described in the previous section.
Figure 4a presents the NESS for diffusion with a resetting rate r(x) = r∣x∣, for λ = {0, 1, 2, 3}. Red solid lines indicate the analytical solution when it is known. When analytical solutions are not known, we present results of brute-force Langevin dynamics simulations with stochastic resetting (yellow solid lines) instead. The blue dashed lines are obtained using the estimation method developed in the previous section, based on a single set of free diffusion trajectories without resetting. In all cases, we observe a good agreement between our method and the ground truth (theory or brute-force simulations).
In Fig. 4b, we show that the NESS asymptotics obeys and estimate the values of α through linear regression. This procedure retrieves the correct values for λ = {0, 2}, where we estimate α ≃ {1.04, 2.00}, respectively. For λ = {1, 3}, we estimate α ≃ {1.57, 2.53}, respectively. Following these results, we hypothesize that for a resetting rate r(x) ~ ∣x∣, the tails of the NESS decay as . While this hypothesis remains to be proven, it exemplifies the power of our approach, revealing new phenomena and inspiring future work.
Next, we demonstrate how to use our approach to engineer desirable NESS of interest. We do so by designing the arrow-shaped steady-state, that is illustrated in Fig. 1b. We start with a system composed of a diffusing particle within a box with reflecting boundaries at x = {-5, 5} and y = {-5, 5}. Restart takes the particle to a uniformly distributed initial position within the box. The resulting NESS, for a position-independent resetting rate, is uniform. We seek to design a position-dependent resetting rate that will generate the desired arrow-shaped NESS, using only a single set of trajectories without resetting.
A naive approach would be to simply apply a constant resetting rate within the desired shape and zero resetting rate outside (Equation (S15) in the Supplementary Information). However, using our approach to predict the resulting steady-state, we find that this naive guess leads to a very fuzzy distribution, in which the arrow is barely discernible (see Fig. 5a). This problem cannot be solved by increasing the resetting rate. The reason is that, for every reasonably complicated NESS, there will be areas that are very hard to reach without crossing regions with a high resetting rate. In our case, the area engulfed by the arrow can only be approached from the right, which significantly lowers its occupation probability, leading to the fuzzy arrow.
Previously, testing different resetting protocols, in a trial-and-error fashion, would have required running thousands of trajectories with resetting for every protocol until the desired steady-state would have been obtained. Instead, using our approach, we can design an improved adaptive resetting protocol (Equation (S16) in the Supplementary Information) that would lead to the desired shape (Fig. 5b), using the same set of trajectories without resetting that is already available. In the improved resetting strategy, the resetting rate was increased quadratically in the distance from the center of the box, and to the right of the arrow, to prevent the accumulation of density in the box’s corners and to the right of the arrow, respectively. Results from simulations with this resetting protocol are given in Fig. 5c for comparison, showing excellent agreement with our prediction.
To conclude this paper, we show how to systematically optimize an adaptive resetting protocol for an arbitrary task based on our theory. Recently, ref. used reinforcement learning to learn resetting protocols for target search problems. In contrast, we do so by representing the adaptive resetting probability as a neural network. Essentially, the input to the network is the system state and time, and the output is the probability to perform resetting. The state of the system can be represented using a set of descriptors as commonly done in machine learning applications in molecular simulations. Then, we define a loss function for training, which can be any observable of the process with the adaptive resetting protocol given by the neural network. For example, this could be any function of the MFPT, FPT distribution, propagator, and NESS. The training data is the trajectories without resetting, and the value of the loss function at every epoch is given by the predicted value of the desired observable under resetting using our theory. Optimization is carried out using standard techniques, such as stochastic gradient descent (see section 8 of the Supplementary Information for the full details). We demonstrate the usefulness of our machine learning framework to optimize the adaptive resetting probability for minimizing the MFPT of conformational transitions of a protein in molecular dynamics simulations.
Molecular dynamics simulations are a powerful tool, but their accessible timescales are limited to a few microseconds. Therefore, to simulate any physical phenomena on longer timescales, e.g., protein folding and crystal nucleation, one must use enhanced sampling methods. Stochastic resetting recently emerged as a promising technique for that purpose, either as a standalone method or in combination with Metadynamics — a popular enhanced sampling tool. Resetting accelerated rare events in molecular dynamics and Metadynamics simulations by more than an order of magnitude, and provided an accurate inference of the kinetics of the unperturbed process. However, in almost all cases, the resetting rate employed was state-independent. We recently showed that even a very limited functional form of adaptive resetting, designed ad hoc, substantially lowered the MFPT. We now show that representing the rate by a flexible neural network allows automatic optimization and yields new adaptive resetting protocols that result in higher speedups than previously possible.
We will demonstrate our approach by finding the resetting strategy that leads to the highest speedup (lowest MFPT) for the folding of the chignolin mini-protein in explicit water, which consists of 5889 atoms (simulation details are given in section 9 of the Supplementary Information). The system has three metastable states: an unfolded state, a misfolded state, and the folded, native state. Representative configurations of the states are given in Fig. 6a. We identify the states using the C-alpha root-mean-square deviation (RMSD) from the folded configuration. The free energy along this degree of freedom is plotted in Fig. 6b. The blue and yellow stars mark the RMSD values of the unfolded and misfolded configurations of Fig. 6a, respectively. We observe that there is an energy barrier the system has to cross in order to get to the native state (the deep well around RMSD <1.5 Å). There is no substantial barrier between the misfolded and unfolded states, such that multiple misfolding-unfolding events often occur before a successful folding. We consider two first-passage processes leading to the folded state from two initial configurations: the unfolded and misfolded states.
To optimize the resetting strategy, we built a network with a simple architecture: It receives three inputs (the RMSD from the folded configuration, the radius of gyration, and the end-to-end distance of the protein) and has two inner layers with ten nodes each. The output is the probability to reset given the three inputs, restricted to be between zero and one, using a sigmoid activation function in the final layer. The network is trained by going over data of trajectories with no resetting and estimating the MFPT under the resetting probability represented by the network using Eqs. (3-7). The estimated MFPT serves as a loss function that is minimized during training. An example script to train the model is given on GitHub.
Figure 6c shows the training curves of the model as a function of the number of epochs when starting from the unfolded (blue) or misfolded (yellow) states. The plotted value is the estimated speedup, which is defined as the ratio of the MFPT values without and with resetting, respectively. For both curves, there is a plateau around 25 epochs (highlighted with a dotted vertical line), where the algorithm reached a local minimum. The dotted lines in Fig. 6d show the resetting protocol represented by the network at this epoch as a function of the RMSD, averaged over the other two degrees of freedom. Purple and orange colors represent simulations initiated at the unfolded and misfolded states, respectively. In both cases, the networks suggest state-independent resetting probabilities. Remarkably, these are the resetting probabilities we identified as providing the highest acceleration for state-independent resetting in refs. .
In later stages of training, the networks find better, state-dependent resetting protocols, reaching a second plateau around the 100th epoch, highlighted with a dashed vertical line in Fig. 6c. The corresponding adaptive resetting probabilities are plotted with dashed lines in Fig. 6d. For simulations initiated at the misfolded state, the predicted speedup is very close to the one obtained in ref. using an ad hoc, naive functional form of a step function. The resetting probability represented by the network has a similar functional form, starting with a small probability that increases sharply to nearly ~1 for RMSD values larger than ~6 Å. Most importantly, for simulations initiated at the unfolded configuration, the naive step function form did not lead to any speedup, but our machine learning optimization framework finds a protocol that grows gradually with the RMSD and leads to higher speedups.
