SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Generative modeling. During the training stage, i.e. as the weights  $A_{k}^{i}$  evolve, the algorithm produces weighted samples  $X_{k}^{i}$ . At any iteration, however, equal-weight samples can be generated by resampling. Notice that, even if we no longer evolve the model parameters  $\theta_{k}$ , the algorithm is such that it removes the bias from ULA coming from  $h > 0$  – this bias removal is not perfect, again because  $N$  is finite, but this can be controlled by increasing  $N$  at the stage when the EBM is used as a generative model.
4.1 Gaussian Mixtures
In this section, we use a synthetic model to illustrate the advantages of our approach. Specifically, we assume that the data is drawn from the Gaussian mixture density with two modes given by
$$
\rho_ {*} (x) = Z _ {*} ^ {- 1} \left(e ^ {- \frac {1}{2} | x - a _ {*} | ^ {2}} + e ^ {- \frac {1}{2} | x - b _ {*} | ^ {2} - z _ {*}}\right), \quad Z _ {*} = (2 \pi) ^ {d / 2} \left(1 + e ^ {- z _ {*}}\right) \tag {14}
$$
where  $a_{*}, b_{*} \in \mathbb{R}^{d}$  specify the means of the two modes and  $z_{*} \in \mathbb{R}$  controls their relative weights  $p_{*} = 1 / (1 + e^{-z^{*}})$  and  $q_{*} = 1 - p_{*} = e^{-z_{*}} / (1 + e^{-z^{*}})$ . The values of  $a_{*}, b_{*}, z_{*}$  are carefully chosen such that the modes are well separated and the energy barrier between the modes is high enough such that jumps of the walkers between the modes are not observed during the simulation with ULA. Consistent with (14) we use an EBM with
$$
U _ {\theta} (x) = - \log \left(e ^ {- \frac {1}{2} | x - a | ^ {2}} + e ^ {- \frac {1}{2} | x - b | ^ {2} - z}\right), \tag {15}
$$
where  $\theta = (a,b,z)$  are the parameters to be optimized. We choose this model as it allows us to calculate the partition function of the model at any value of the parameters,  $Z_{\theta} = (2\pi)^{d / 2}(1 + e^{-z})$ . We use this information as a benchmark to compare the prediction with those produced by our method.
Figure 1: GMM experiments: Evolution of the parameters and the cross entropy during training by Algorithm 2, PCD, and CD. Average of 20 runs. Left panels: evolution of  $p_k = 1 / (1 + e^{-z_k})$ ; middle panel: evolution of  $a_k$  and  $b_k$ ; right panel: evolution of the Kullback-Leibler divergence. All three methods capture the location of the modes accurately, but only ours get the relative weights of these modes accurately (whereas PCD leads to mode collapse, and CD to an inaccurate estimate). Our method is also the only one that allows for direct estimation of the cross-entropy during training, and the only one performing GD on this cross-entropy-for better visualization we subtract the entropy of the target  $H(\rho_{*})$  and plot the Kullback-Leibler divergence instead of the cross-entropy.
In our numerical experiments, we set  $d = 50$ , use  $N = 10^{5}$  walkers with a mini-batch of  $N' = 10^4$  and  $n = 10^5$  data points. We initialize the model at  $\theta_0 = (a_0, b_0, z_0)$  with  $a_0$  and  $b_0$  drawn from an  $N(0, \epsilon^2 I_d)$  with  $\epsilon = 0.1$  and  $z_0 = 0$ , meaning that the initial  $\rho_{\theta_0}$  is close to the PDF of an  $N(0, I_d)$ . The training is performed using Algorithm 2 with  $h = 0.1$  and fixed learning rates  $\gamma_k = 0.2$  for  $a_k$  and  $b_k$  and  $\gamma_k = 1$  for  $z_k$ . We perform the resampling step by monitoring  $ESS_k$  defined in (13) with constant  $1/c_k = 1.05$  and using the systematic method. We also compare our results to those obtained using ULA with these same parameters (which is akin to training with the PCD algorithm) and with those obtained with the CD algorithm: in the latter case, we evolve the walkers by ULA with  $h = 0.1$  for 4 steps between resets at the data points, and we adjust the learning rates by multiplying them by a factor 10. In all cases, we use the full batches of walkers, weights, and data points to estimate the empirical averages. We also use (12) to estimate the cross-entropy  $H(\rho_{\theta_k}, \rho_*)$  during training by our method (CD and PCD do not provide estimates for these quantities), and in all cases compare the result with the estimate
$$
\tilde {H} _ {k} = \log \left((2 \pi) ^ {d / 2} (1 + e ^ {- z _ {*}})\right) - \frac {1}{n} \sum_ {j = 1} ^ {n} \log \left(e ^ {- \frac {1}{2} | x _ {*} ^ {j} - a _ {k} | ^ {2}} + e ^ {- \frac {1}{2} | x _ {*} ^ {j} - b _ {k} | ^ {2} - z _ {k}}\right) \tag {16}
$$
The results are shown in Figure 1. As can be seen, all three methods learn well the values of  $a_*$  and  $b_*$  specifying the positions of the modes. However, only our approach learns the value of  $z$  specifying their relative weights. In contrast, the PCD algorithm leads to mode collapse, consistent with the theoretical explanation given in Appendix C.1, and the CD algorithm returns a biased value of  $z$ , consistent with the fact that it effectively uses the Fisher divergence as the objective. The results also show that the cross-entropy decreases with our approach, but bounces back up with the PCD algorithm and stalls with the CD algorithms: this is consistent with the fact that only our approach actually performs the GD on the cross-entropy, which, unlike the other algorithms, our approach estimates accurately during the training.
4.2 MNIST
Next, we perform empirical experiments on the MNIST dataset to answer the following question: when it comes to high-dimensional datasets with multiple modes, can our method produce an EBM that generates high-quality samples and captures the relative weights of the modes accurately?
To this end, we select a subset of MNIST consisting of only three digits: 2, 3, and 6. Then, we choose 5600 images of label 2, 2800 images of label 3, and 1400 images of label 6 from the training set (for a total of  $n = 9800$  data points), so that in this manufactured dataset the digits are in have respective weights  $4/7$ ,  $2/7$ , and  $1/7$ .
Generated samples (right after resampling)
Images from MNIST test set
Generated samples (without Jarzynski)
Figure 2: MNIST: Left panel: Examples of images generated by our method right after resampling in the last epoch. Middle panel: Images randomly selected from the test dataset of MNIST. Right panel: Examples of images generated by training using the persistent contrastive divergence (PCD) algorithm.
Figure 3: MNIST: Relative error of the weight estimation of the three modes (i.e. three digits). Our method outperforms the PCD algorithm in terms of recovery of the weight of each mode.
We train two EBMs to represent this data set, the first using our Algorithm 2, and the second using the PCD Algorithm 4. We represent the energy using a simple six-layer convolutional neural network with the swish activation and about  $77K$  parameters. We use the ADAM optimizer for the training with a learning rate starting from  $10^{-4}$  and linearly decaying to  $10^{-10}$  until the final training step. The sample size of the walkers is set to  $N = 1024$  and it is fixed throughout the training.
The results are shown in Figure 2. Both our Algorithm 1 and the PCD Algorithm 4 generate images of reasonable quality; as discussed in Appendix B.2, we observe that the Jarzynski weights can help us track the quality of generated images, and the resampling step is also helpful for improving this quality as well as the learned energy function.
However, the real difference between the methods comes when we look at the relative proportion of the digits generated in comparison to those in the data set. In Figure 3, we show the relative error on the weight estimation of each mode obtained by using a classifier pre-trained on the data set (see Appendix B.2 for details). Although the EBM trained with the PCD algorithm can discover all the modes and generate images of all three digits present in the training set, it cannot accurately recover the relative proportions of each digit. In contrast, our method is successful in both mode exploration and weight recovery.
Unlike in the numerical experiments done on Gaussian mixtures with teacher-student models, for the MNIST dataset, we cannot estimate the KL divergence throughout the training as we do not know the normalization constant of the true data distribution. Nevertheless, we can still use the formula (16) to track the cross-entropy between the data distribution and the walker distribution a plot of which is given in the right panel Figure 4 for the mini-batched training algorithm (Algorithm 2).
In these experiments with MNIST, we used the adaptive resampling scheme described after equation 13. In practice, we observed that few resamplings are needed during training, and they can often be avoided altogether if the learning rate is sufficiently small. For example, in the experiment reported in the left panel of Figure 4) a single resampling step was made. We also found empirically that the results are insensitive to the choice of the parameters  $c_{k}$  used for resampling: the rule of
Figure 4: MNIST dataset: Left Panel: Convergence of the cross-entropy estimated by using formulae (12) with the mini-batched algorithm (Algorithm 2). Right Panel: Evolution of the effective sample size (ESS) defined in (13) - here resampling was started after 240 epochs (with  $c_k = 0$  before and  $c_k = 0.5$  afterwards), and occurred only once immediately after being switched on.
thumb we found is to not resample at the beginning of training, to avoid possible mode collapse, and use resampling towards the end, to improve the image quality.
4.3 CIFAR-10
We perform an empirical evaluation of our method on the full CIFAR-10  $(32 \times 32)$  image dataset. We use the setup of [47], with the same neural architecture with  $n_f = 128$  features, and compare the results obtained with our approach with mini-batching (Algorithm 2) to those obtained with PCD and PCD with data augmentation of [33] (which consists of a combination of color distortion, horizontal flip, rescaling, and Gaussian blur augmentations, to help the mixing of the MCMC sampling and stabilize training).
Generated CIFAR-10 samples with our approach
Figure 5: CIFAR-10 dataset: Left panel: images generated by training with Algorithm 2. Right panel: Images generated from the PCD with mini-batches. Up to date images available at https://github.com/submissionx12/EBMs_Jarzynski.
Generated CIFAR-10 samples with PCD
The code used to perform these new experiments is available in the anonymized GitHub referenced in our paper. The hyperparameters are the same in all cases: we take  $N = 4096$  Langevin walkers with a mini-batch size  $N' = 256$ . We use the Adam optimizer with learning rate  $\eta = 10^{-4}$  and inject a Gaussian noise of standard deviation  $\sigma = 3 \times 10^{-2}$  to the dataset while performing gradient clipping in Langevin sampling for better performance. All the experiments were performed on a single A100 GPU. Training for 600 epochs took about 34 hours with the PCD algorithm (w/ and w/o data augmentation) and about 36 hours with our method.
Some of the images generated by our method are shown in Figure 5. We also quantitatively evaluate the performance of our models with the commonly used metrics (e.g. FID and Inception Score): the results are given in Table 1. They indicate that our method can achieve slightly better performance than the PCD algorithms w/ and w/o data augmentation at a similar computational cost. Furthermore, these results on CIFAR-10 suggest that our method scales well to complicated training tasks on more realistic data sets.
Method	FID	Inception Score (IS)
PCD with mini-batches	38.25	5.96
PCD with mini-batches and data augmentation	36.43	6.54
Algorithm 2 with multinomial resampling	32.18	6.88
Algorithm 2 with systematic resampling	30.24	6.97
Table 1: textitCIFAR-10 dataset: Comparison of FID and Inception Score (IS) for PCD and Algorithm 2. Experiments performed using the neural architecture in [47] to model the energy.
It is worth stressing that a key component of the success of our method is the resampling step, as discussed in Appendix B.3 where we plot the Effective Sample Size (ESS) in Figure 13.
In this work, we proposed a simple modification of the persistent contrastive divergence algorithm which corrects its biases and provably allows one to minimize the cross-entropy between the model and the target densities to train EBMs. Our approach rests on results from non-equilibrium thermodynamics (in particular, Jarzynski equality) that show how to take exact expectations over a probability density function (PDF) that is evolving in time—in the present context the PDF is that of the EBM, and it is changing as we train the model parameters. These formulas naturally lend themselves to practical implementation using sequential Monte Carlo sampling methods. The only difference with respect to training methods based on ULA is that the proposed approach maintains a set of weights associated with every walker. These weights correct for the bias induced by the evolution of the energy, and they also give a way to assess the quality of the samples generated by the method. On synthetic examples using Gaussian mixture densities as well as on some simple data sets involving MNIST, we showed that our method dodges some common drawbacks of other EBM training methods. In particular, it is able to learn the relative weights of various high-density regions in a multimodal distribution, thus ensuring a higher level of fairness than common EBM training techniques based on contrastive divergence.
To focus on the core idea, in the present study we put aside tuning considerations and did not sweep over the various hyperparameters and design choices, like learning rates, number of walkers, or resampling criteria and strategies. These considerations could greatly enhance the performance of our method. We believe that our results show that the method deserves deeper practical investigation on more realistic datasets and with more complex neural architectures.
D.C. has worked under the auspices of Italian National Group of Mathematical Physics (GNFM). M.H. is supported by NYU McCracken doctoral fellowship. EVE is supported by the National Science Foundation under awards DMR-1420073, DMS-2012510, and DMS-2134216, by the Simons Collaboration on Wave Turbulence, Grant No. 617006, and by a Vannevar Bush Faculty Fellowship.
[1] Diederik P Kingma and Max Welling. An introduction to variational autoencoders. Foundations and Trends in Machine Learning, 12(4):307-392, 2019.
[2] Carl Doersch. Tutorial on variational autoencoders. arXiv preprint arXiv:1606.05908, 2016.
[3] Shengjia Zhao, Jiaming Song, and Stefano Ermon. Infovae: Balancing learning and inference in variational autoencoders. In Proceedings of the aaai conference on artificial intelligence, volume 33, pages 5885-5892, 2019.
[4] Jonathan Ho and Stefano Ermon. Generative adversarial imitation learning. In D. Lee, M. Sugiyama, U. Luxburg, I. Guyon, and R. Garnett, editors, Advances in Neural Information Processing Systems, volume 29, 2016.
[5] Ian Goodfellow, Jean Pouget-Abadie, Mehdi Mirza, Bing Xu, David Warde-Farley, Sherjil Ozair, Aaron Courville, and Yoshua Bengio. Generative adversarial networks. Communications of the ACM, 63(11):139–144, 2020.
[6] Esteban G. Tabak and Eric Vanden-Eijnden. Density estimation by dual ascent of the log-likelihood. Communications in Mathematical Sciences, 8(1):217 - 233, 2010.
[7] E. G. Tabak and Cristina V. Turner. A family of nonparametric density estimation algorithms. Communications on Pure and Applied Mathematics, 66(2):145-164, 2013.
[8] Danilo Rezende and Shakir Mohamed. Variational inference with normalizing flows. In International conference on machine learning, pages 1530-1538. PMLR, 2015.
[9] George Papamakarios, Eric Nalisnick, Danilo Jimenez Rezende, Shakir Mohamed, and Balaji Lakshminarayanan. Normalizing flows for probabilistic modeling and inference. The Journal of Machine Learning Research, 22(1):2617-2680, 2021.
[10] Emile Mathieu and Maximilian Nickel. Riemannian continuous normalizing flows. Advances in Neural Information Processing Systems, 33:2503-2515, 2020.
[11] Yang Song, Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. In International Conference on Learning Representations, 2021.
[12] Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. In Advances in Neural Information Processing Systems, volume 33, pages 6840-6851, 2020.
[13] Yang Song, Conor Durkan, Iain Murray, and Stefano Ermon. Maximum likelihood training of score-based diffusion models. Advances in Neural Information Processing Systems, 34: 1415-1428, 2021.
[14] Geoffrey E Hinton. A practical guide to training restricted boltzmann machines. Neural Networks: Tricks of the Trade: Second Edition, pages 599-619, 2012.
[15] Ruslan Salakhutdinov and Hugo Larochelle. Efficient learning of deep boltzmann machines. In Proceedings of the thirteenth international conference on artificial intelligence and statistics, pages 693-700. JMLR Workshop and Conference Proceedings, 2010.
[16] Hannes Schulz, Andreas Müller, Sven Behnke, et al. Investigating convergence of restricted boltzmann machine learning. In NIPS 2010 Workshop on Deep Learning and Unsupervised Feature Learning, volume 1, pages 6-1, 2010.
[71] Genshiro Kitagawa. Monte carlo filter and smoother for non-gaussian nonlinear state space models. Journal of computational and graphical statistics, 5(1):1-25, 1996.
[72] James Carpenter, Peter Clifford, and Paul Fearnhead. Improved particle filter for nonlinear problems. IEE Proceedings-Radar, Sonar and Navigation, 146(1):2-7, 1999.
[73] Tiancheng Li, Miodrag Bolic, and Petar M Djuric. Resampling methods for particle filtering: classification, implementation, and strategies. IEEE Signal processing magazine, 32(3):70-86, 2015.
[74] Michael S Albergo, Nicholas M Boffi, and Eric Vanden-Eijnden. Stochastic interpolants: A unifying framework for flows and diffusions. arXiv preprint arXiv:2303.08797, 2023.
A.1 Jarzynski equality in continuous-time
Before proving Proposition 1, we give a continuous-time version of this result whose proof helps guide the intuition.
Proposition 2. Assume that the parameters  $\theta$  are evolved according to some time-differentiable protocol  $\theta(t)$  such that  $\theta(0) = \theta_0$ . Given any  $\alpha > 0$ , let  $X_t \in \mathbb{R}^d$  and  $A_t \in \mathbb{R}$  be the solutions of
$$
\left\{ \begin{array}{l l} d X _ {t} = - \alpha \nabla U _ {\theta (t)} \left(X _ {t}\right) d t + \sqrt {2 \alpha} d W _ {t}, & X _ {0} \sim \rho_ {\theta_ {0}}, \\ \dot {A} _ {t} = - \partial_ {\theta} U _ {\theta (t)} \left(X _ {t}\right) \cdot \dot {\theta} (t), & A _ {0} = 0. \end{array} \right. \tag {17}
$$
where  $U_{\theta}(x)$  is the model energy and  $W_{t} \in \mathbb{R}^{d}$  is a standard Wiener process. Then, for any  $t \geq 0$ ,
$$
\mathbb {E} _ {\theta (t)} \left[ \partial_ {\theta} U _ {\theta (t)} \right] = \frac {\mathbb {E} \left[ \partial_ {\theta} U _ {\theta (t)} \left(X _ {t}\right) e ^ {A _ {t}} \right]}{\mathbb {E} \left[ e ^ {A _ {t}} \right]}, \quad Z _ {\theta (t)} = Z _ {\theta_ {0}} \mathbb {E} \left[ e ^ {A _ {t}} \right], \tag {18}
$$
where the expectations on the right-hand side are over the law of the joint process  $(X_{t},A_{t})$ .
The second equation in (18) can also be written in terms of the free energy  $F_{\theta} = -\log Z_{\theta}$  as  $F_{\theta(t)} = F_{\theta_0} - \log \mathbb{E}[e^{A_t}]$ : this is Jarzynski's equality [34]. We stress that it is key to include the weights in (18) and, in particular,  $\mathbb{E}[\partial_{\theta}U_{\theta(t)}(X_t)] \neq \mathbb{E}_{\theta(t)}[\partial_{\theta}U_{\theta(t)}]$ . This is because the PDF of  $X_t$  alone lags behind the model PDF  $\rho_{\theta(t)}$ : the larger  $\alpha$ , the smaller this lag, but it is always there if  $\alpha < \infty$ , see the remark at the end of this section for more discussion on this point. The inclusion of the weights in (18) corrects exactly for the bias induced by this lag.
An immediate consequence of Proposition 2 is that we can evolve  $\theta(t)$  by the gradient descent flow over the cross-entropy by solving (17) concurrently with
$$
\dot {\theta} (t) = \frac {\mathbb {E} \left[ \partial_ {\theta} U _ {\theta (t)} \left(X _ {t}\right) e ^ {A _ {t}} \right]}{\mathbb {E} \left[ e ^ {A _ {t}} \right]} - \mathbb {E} _ {*} \left[ \partial_ {\theta} U _ {\theta (t)} \right], \quad \theta (0) = \theta_ {0} \tag {19}
$$
since by (18) the right hand side of (19) is precisely  $-\partial_{\theta}H(\rho_{\theta(t)}, \rho_{*}) = \mathbb{E}_{\theta(t)}[\partial_{\theta}U_{\theta(t)}] - \mathbb{E}_{*}[\partial_{\theta}U_{\theta(t)}]$ . Assuming that we know  $Z_{\theta_0}$  we can also track the evolution of the cross-entropy (3) via
$$
H \left(\rho_ {\theta (t)}, \rho_ {*}\right) = \log \mathbb {E} \left[ e ^ {A _ {t}} \right] + \log Z _ {\theta_ {0}} + \mathbb {E} _ {*} \left[ U _ {\theta (t)} \right]. \tag {20}
$$
Proof. The joint PDF  $f(t,x,a)$  of the process  $(X_{t},A_{t})$  satisfying (17) solves the Fokker-Planck equation (FPE)
$$
\partial_ {t} f = \alpha \nabla_ {x} \cdot (\nabla_ {x} U _ {\theta (t)} f + \nabla_ {x} f) + \partial_ {\theta} U _ {\theta (t)} \cdot \dot {\theta} (t) \partial_ {a} f, \quad f (0, x, a) = Z _ {\theta_ {0}} ^ {- 1} e ^ {- U _ {\theta_ {0}} (x)} \delta (a). \tag {21}
$$
Let us derive an equation for
$$
\hat {\rho} (t, x) = \int_ {- \infty} ^ {\infty} e ^ {a} f (t, x, a) d a \tag {22}
$$
To this end, multiply (21) by  $e^a$ , integrate the result over  $a \in (-\infty, \infty)$ , and use integration by parts for the last term at the right-hand side to obtain:
$$
\partial_ {t} \hat {\rho} = \alpha \nabla_ {x} \cdot \left(\nabla_ {x} U _ {\theta (t)} \hat {\rho} + \nabla_ {x} \hat {\rho}\right) - \partial_ {\theta} U _ {\theta (t)} \cdot \dot {\theta} (t) \hat {\rho}, \quad \hat {\rho} (0, x) = Z _ {\theta_ {0}} ^ {- 1} e ^ {- U _ {\theta_ {0}} (x)} \tag {23}
$$
By general results for the solutions of parabolic PDEs such as (23) (see [69], Chapter 7), we know that the solution to this equation is unique, and we can check by direct substitution that it is given by
$$
\hat {\rho} (t, x) = Z _ {\theta_ {0}} ^ {- 1} e ^ {- U _ {\theta (t)} (x)}. \tag {24}
$$
$$
\int_ {\mathbb {R} ^ {d}} \hat {\rho} (t, x) d x = Z _ {\theta_ {0}} ^ {- 1} Z _ {\theta (t)}. \tag {25}
$$
Since by definition  $\mathbb{E}[e^{A_t}] = \int_{\mathbb{R}^d}\int_{-\infty}^{\infty}e^a f(t,x,a)dadbx = \int_{\mathbb{R}^d}\hat{\rho} (t,x)dx$  this establishes the second equation in (18). To establish the first notice that
$$
\begin{array}{l} \frac {\mathbb {E} \left[ \partial_ {\theta} U _ {\theta (t)} \left(X _ {t}\right) e ^ {A _ {t}} \right]}{\mathbb {E} \left[ e ^ {A _ {t}} \right]} = \frac {\int_ {\mathbb {R} ^ {d}} \int_ {- \infty} ^ {\infty} \partial_ {\theta} U _ {\theta (t)} (x) e ^ {a} f (t , x , a) d a d x}{\int_ {\mathbb {R} ^ {d}} \int_ {- \infty} ^ {\infty} e ^ {a} f (t , x , a) d a d x} = \frac {\int_ {\mathbb {R} ^ {d}} \partial_ {\theta} U _ {\theta (t)} (x) \hat {\rho} (t , x) d x}{\int_ {\mathbb {R} ^ {d}} \hat {\rho} (t , x) d x} \tag {26} \\ = \frac {Z _ {\theta_ {0}} ^ {- 1} \int_ {\mathbb {R} ^ {d}} \partial_ {\theta} U _ {\theta (t)} (x) e ^ {- U _ {\theta (t)} (x)} d x}{Z _ {\theta_ {0}} ^ {- 1} Z _ {\theta (t)}} = \mathbb {E} _ {\theta (t)} [ \partial_ {\theta} U _ {\theta (t)} ] \\ \end{array}
$$
The need for Jarzynski's correction. Suppose that the walkers satisfy the Langevin equation (first equation in (17)):
$$
d X _ {t} = - \alpha \nabla U _ {\theta (t)} \left(X _ {t}\right) d t + \sqrt {2 \alpha} d W _ {t}, \quad X _ {0} \sim \rho_ {\theta_ {0}}, \tag {27}
$$
where  $\theta (t)$  is evolving according to some protocol. The probability density function  $\rho (t,x)$  of  $X_{t}$  then satisfies the Fokker-Planck equation (compare (23))
$$
\partial_ {t} \rho = \alpha \nabla \cdot (\nabla U _ {\theta (t)} (x) \rho + \nabla \rho), \quad \rho (t = 0) = \rho_ {\theta_ {0}} \tag {28}
$$