SlowGuess/NeurIPS_2023_OCR huggingface dataset:
The solution to this equation is not available in closed form, and in particular  $\rho(t,x) \neq \rho_{\theta(t)}(x) - \rho(t,x)$  is only close to  $\rho_{\theta(t)}(x)$  if we let  $\alpha \to \infty$ , so that the walkers  $X_t$  move much faster than the parameters  $\theta(t)$  but this limit is not easily achievable in practice (as convergence of the FPE solution to its equilibrium is very slow in general if the potential  $U_{\theta(t)}$  is complicated). As a result  $\mathbb{E}[\partial_\theta U_{\theta(t)}(X_t)] \neq \mathbb{E}_{\theta(t)}[\partial_\theta U_{\theta(t)}]$ , implying the necessity to include the weights in the expectation (18).
The iteration rule for  $A_{k}$  in (5) implies that
$$
A _ {k} = \sum_ {q = 1} ^ {k} \left(\alpha_ {q - 1} \left(X _ {q - 1}, X _ {q}\right) - \alpha_ {q} \left(X _ {q}, X _ {q - 1}\right)\right), \quad k \in \mathbb {N}. \tag {29}
$$
For  $k\in \mathbb{N}_0$  , let
$$
\beta_ {k} (x, y) = \left(4 \pi h\right) ^ {- d / 2} \exp \left(- \frac {1}{4 h} | y - x + h \nabla U _ {\theta_ {k}} (x) | ^ {2}\right) \tag {30}
$$
be the transition probability density of the ULA update in (5), i.e.  $\beta_{k}(X_{k},X_{k + 1})$  is the probability density of  $X_{k + 1}$  conditionally on  $X_{k}$ . By the definition of  $A_{k}$ , we have
$$
\begin{array}{l} \exp \left(A _ {k}\right) = \prod_ {q = 1} ^ {k} \exp \left(\alpha_ {q - 1} \left(X _ {q - 1}, X _ {q}\right) - \alpha_ {q} \left(X _ {q}, X _ {q - 1}\right)\right) \tag {31} \\ = e ^ {- U _ {\theta_ {k}} (X _ {k}) + U _ {\theta_ {0}} (X _ {0})} \prod_ {q = 1} ^ {k} \frac {\beta_ {q} (X _ {q} , X _ {q - 1})}{\beta_ {q - 1} (X _ {q - 1} , X _ {q})} \\ \end{array}
$$
where in the second line we added and subtracted  $|X_q - X_{q-1}|^2 / 4h$  and used the definition of  $\alpha_k(x, y)$  given in (6). Since the joint probability density function of the path  $(X_0, X_1, \ldots, X_k)$  at any  $k \in \mathbb{N}$  is
$$
\varrho \left(x _ {0}, x _ {1}, \dots , x _ {k}\right) = Z _ {\theta_ {0}} ^ {- 1} e ^ {- U _ {\theta_ {0}} \left(x _ {0}\right)} \prod_ {q = 1} ^ {k} \beta_ {q - 1} \left(x _ {q - 1}, x _ {q}\right) \tag {32}
$$
we deduce from (31) and (32) that, given an  $f: \mathbb{R}^d \to \mathbb{R}$ , we can express the expectation  $\mathbb{E}[f(X_k)e^{A_k}]$  as the integral
$$
\begin{array}{l} \mathbb {E} \left[ f \left(X _ {k}\right) e ^ {A _ {k}} \right] \\ = \int_ {\mathbb {R} ^ {d k}} f \left(x _ {k}\right) e ^ {- U _ {\theta_ {k}} \left(x _ {k}\right) + U _ {\theta_ {0}} \left(x _ {0}\right)} \prod_ {q = 1} ^ {k} \frac {\beta_ {q} \left(x _ {q} , x _ {q - 1}\right)}{\beta_ {q - 1} \left(x _ {q - 1} , x _ {q}\right)} \varrho \left(x _ {0}, x _ {1}, \dots , x _ {k}\right) d x _ {0} \dots d x _ {k} \tag {33} \\ = Z _ {\theta_ {0}} ^ {- 1} \int_ {\mathbb {R} ^ {d k}} f (x _ {k}) e ^ {- U _ {\theta_ {k}} (x _ {k})} \prod_ {q = 1} ^ {k} \beta_ {q} (x _ {q}, x _ {q - 1}) d x _ {0} \dots d x _ {k} \\ \end{array}
$$
Since  $\int_{\mathbb{R}^d}\beta_k(x,y)dy = 1$  for all  $k\in \mathbb{N}_0$  and all  $x\in \mathbb{R}^d$ , we can perform the integrals over  $x_0$ , then  $x_{1}$ , etc. in this expression to be left with
$$
\mathbb {E} [ f (X _ {k}) e ^ {A _ {k}} ] = Z _ {\theta_ {0}} ^ {- 1} \int_ {\mathbb {R} ^ {d}} f (x _ {k}) e ^ {- U _ {\theta_ {k}} (x _ {k})} d x _ {k} \tag {34}
$$
Setting  $f(x) = 1$  in this expression gives
$$
\mathbb {E} \left[ e ^ {A _ {k}} \right] = Z _ {\theta_ {0}} ^ {- 1} \int_ {\mathbb {R} ^ {d}} e ^ {- U _ {\theta_ {k}} \left(x _ {k}\right)} d x _ {k} = Z _ {\theta_ {0}} ^ {- 1} Z _ {\theta_ {k}} \tag {35}
$$
which implies the second equation in (7); setting  $f(x_{k}) = \partial_{\theta}U_{\theta_{k}}(x_{k})$  in (34) gives
$$
\mathbb {E} \left[ \partial_ {\theta} U _ {\theta_ {k}} \left(X _ {k}\right) e ^ {A _ {k}} \right] = Z _ {\theta_ {0}} ^ {- 1} \int_ {\mathbb {R} ^ {d}} \partial_ {\theta} U _ {\theta_ {k}} \left(x _ {k}\right) e ^ {- U _ {\theta_ {k}} \left(x _ {k}\right)} d x _ {k} = Z _ {\theta_ {0}} ^ {- 1} Z _ {\theta_ {k}} \mathbb {E} _ {\theta_ {k}} \left[ \partial_ {\theta} U _ {\theta_ {k}} \right] \tag {36}
$$
which can be combined with (35) to arrive at the first equation in (7).
A.3 Mini-batched version of Algorithm 1.
In Algorithm 2 we present a mini-batched version of Algorithm 1, where we do not update the positions of every walker at every iteration. Instead, we evolve only a small portion of the walkers, while keeping the other walkers frozen and only updating their weights using the information from the model update (which uses only the model energy and does not require back-propagation). This mini-batched version is more computationally efficient and leads to convergence in training with much fewer steps of ULA. More importantly, the mini-batched version of the algorithm, as compared to the full-batched version, enlarges the total sample size and therefore improves the sample variety with even less computational cost.
Algorithm 2 Mini-batched Sequential Monte-Carlo training with Jarzynski correction
1: Inputs: data points  $\{x_{i}^{*}\}_{i = 1}^{n}$ ; energy model  $U_{\theta}$ ; optimizer step  $\mathrm{opt}(\theta ,\mathcal{D})$  using  $\theta$  and the empirical CE gradient  $\mathcal{D}$ ; initial parameters  $\theta_0$ ; number of walkers  $N\in \mathbb{N}_0$ ; batch size  $N^{\prime}\in \mathbb{N}_0$  with  $N^{\prime} < N$ , set of walkers  $\{X_0^i\}_{i = 1}^N$  sampled from  $\rho_{\theta_0}$ ; total duration  $K\in \mathbb{N}$ ; ULA time step  $h$ ; set of positive constants  $\{c_k\}_{k\in \mathbb{N}}$ .
2:  $A_0^i = 0$  for  $i = 1,\ldots ,N$
3: for  $k = 1: K - 1$  do
4:  $p_k^i = \exp (A_k^i) / \sum_{j = 1}^N\exp (A_k^j)$  ▷ normalized weights
5:  $\tilde{\mathcal{D}}_k = \sum_{i=1}^N p_k^i \partial_\theta U_{\theta_k}(X_k^i) - n^{-1} \sum_{j=1}^n \partial_\theta U_{\theta_k}(x_*^j)$
6:  $\theta_{k + 1} = \mathrm{opt}(\theta_k,\widetilde{\mathcal{D}}_k)$  optimization step
7: Randomly select a mini-batch  $\{X_k^j\}_{j\in B}$  with  $\# B = N'$  from the set of walkers  $\{X_k^i\}_{i = 1}^N$
8: for  $j\in B$  do
9:  $X_{k + 1}^{j} = X_{k}^{j} - h\nabla U_{\theta_{k}}(X_{k}^{j}) + \sqrt{2h}\xi_{k}^{j},\quad \xi_{k}^{j}\sim \mathcal{N}(0_{d},I_{d})$
10:  $A_{k + 1}^{j} = A_{k}^{j} - \alpha_{k + 1}(X_{k + 1}^{j},X_{k}^{j}) + \alpha_{k}(X_{k}^{j},X_{k + 1}^{j})$
11: end for
12: for  $j\notin B$  do
13:  $X_{k + 1}^{j} = X_{k}^{j}$  no update of the walkers
14:  $A_{k + 1}^{j} = A_{k}^{j} - U_{\theta_{k + 1}}(X_{k}^{j}) + U_{\theta_{k}}(X_{k}^{j})$ $\triangleright$  weight update
15: end for
16: Resample the walkers and reset the weights if  $\mathrm{ESS}_{k + 1} < c_{k + 1}$ , see (13).  $\triangleright$  resampling step
17: end for
18: **Outputs:** Optimized energy  $U_{\theta_K}$ ; set of weighted walkers  $\{X_K^i, A_K^i\}_{i=1}^N$  sampling  $\rho_{\theta_K}$ ; partition function estimate  $\tilde{Z}_{\theta_K} = Z_{\theta_0} N^{-1} \sum_{i=1}^N \exp(A_K^i)$ ; CE estimate  $\log \tilde{Z}_{\theta_K} + n^{-1} \sum_{j=1}^n U_{\theta_K}(x_*^j)$
A.4 Resampling Routines
Resampling schemes are necessary to tackle the decay of effective sample size (13). For the sake of completeness, here we recall three of the most widely used routines: multinomial [70], stratified [71], and systematic resampling [72], and refer the reader to the review [73] for more details.
Given a set of normalized scalar weights  $\{p_i\}_{i = 1}^N\in [0,1]$  with  $\sum_{i = 1}^{N}p_{i} = 1$  associated to the  $N$  walkers, we define the cumulative sum
$$
P _ {n} = \sum_ {i = 1} ^ {n} p _ {i}, \quad n = 1, \dots , N \tag {37}
$$
All three methods prescribe a way to choose a set  $\{u_n\}_{n=1}^N \in (0,1]$  used to perform the resampling in the following way: for every  $n,m \in \{1,\dots,N\}$ , the  $m$ -th particle is chosen during the  $n$ -th extraction if
$$
P _ {m - 1} <   u _ {n} <   P _ {m} \tag {38}
$$
Let us now specify how the set of  $u_{n}$  is selected in the cases in study. We denote by  $\mathcal{U}(a,b]$  the uniform probability distribution on the interval  $(a,b]$ .
Multinomial resampling. Sample  $u_{n}^{\mathrm{mult}} \sim \mathcal{U}(0,1]$ , independently for every  $n \in \{1,\dots ,N\}$ . This approach leads to a large number of possible resampled configurations, which is not desirable in practice as it increases the variance of the estimator.
Stratified resampling. Partition the interval  $(0,1]$  into  $N$  sub-intervals, or strata, of size  $1 / N$ ; then, sample  $u_{n}^{\mathrm{str}}\sim \mathcal{U}((n - 1) / N,n / N]$  independently for each  $n = \{1,\dots ,N\}$ . This approach picks a single  $u_{n}$  in each stratus, thereby reducing the number of possible resampled configurations.
Systematic resampling. Partition the interval  $(0,1]$  into  $N$  sub-intervals, or strata, of size  $1 / N$ ; then, sample  $u_{1}^{\mathrm{sys}}\sim \mathcal{U}(0,1 / N]$ , and  $u_{n}^{\mathrm{sys}} = u_{1}^{\mathrm{sys}} + (n - 1) / N$  for  $n > 1$ . This method also reduces the number of possible resampled configurations.
Note that there are various modifications of these three methods (see the review [73]), all of them meeting the so-called unbiasedness condition, that is, the  $i$ -particle is expected to be sampled in average  $Np_{i}$  times. However, these extensions do not lead to a critical lowering of the number of possible resampled configurations compared to systematic resampling. For this reason in our numerical experimentals we only tested the three methods above. Note also that, regardless of the resampling routine one uses, a fundamental role is played by the variance of the weights: if the cumulative sum is dominated by one or very few weights, i.e. resampling is triggered too late, it does not remedy the suppression of population variability.
A.5 Contrastive divergence and persistent contrastive divergence algorithms
For completeness, we give the CD and PCD algorithms we used to compare with our Algorithms 1 and 2. As written these algorithms use the full batch of data points at every optimization step, but they can easily be modified to do mini-batching.
Algorithm 3 Contrastive divergence (CD) algorithm
1: Inputs: data points  $\Omega = \{x_{*}^{i}\}_{i=1}^{n}$  in  $\mathbb{R}^d$ ; energy model  $U_{\theta}$ ; optimizer step  $\operatorname{opt}(\theta, \mathcal{D})$  using  $\theta$  and the empirical gradient  $\mathcal{D}$ ; initial parameters  $\theta_0$ ; number of walkers  $N \in \mathbb{N}_0$  with  $N < n$ ; total duration  $K \in \mathbb{N}$ ; ULA time step  $h$ ;  $P \in \mathbb{N}$ .  
2: for  $k = 1, \dots, K - 1$  do  
3: for  $i = 1, \dots, N$  do  
4:  $X_0^i = \text{SampleMultinomial}(\Omega)$   
5: for  $p = 0, \dots, P - 1$  do  
6:  $X_{p+1}^i = X_p^i - h\nabla U_{\theta_k}(X_p^i) + \sqrt{2h}\xi_p^i$ ,  $\xi_p^i \sim \mathcal{N}(0_d, I_d)$   
7: end for  
8: end for  
9:  $\tilde{\mathcal{D}}_k = N^{-1}\sum_{i=1}^{N}\partial_\theta U_{\theta_k}(X_P^i) - n^{-1}\sum_{i=1}^{n}\partial_\theta U_{\theta_k}(x_*^i)$   
10:  $\theta_{k+1} = \mathrm{opt}(\theta_k, \tilde{\mathcal{D}}_k)$   
11: end for  
12: Outputs: Optimized energy  $U_{\theta_K}$ ; set of walkers  $\{X_P^i\}_{i=1}^N$
Interestingly, we can write down an equation that mimics the evolution of the PDF of the walkers in the CD algorithm, at least in the continuous-time limit: this equation reads
$$
\partial_ {t} \breve {\rho} = \alpha \nabla \cdot (\nabla U _ {\theta (t)} (x) \breve {\rho} + \nabla \breve {\rho}) - \nu (\breve {\rho} - \rho_ {*}), \quad \breve {\rho} (t = 0) = \rho_ {*} \tag {39}
$$
where the parameter  $\nu > 0$  controls the rate at which the walkers are reinitialized at the data points: the last term in (39) is a birth-death term that captures the effect of these reinitializations. The solution
Algorithm 4 Persistent contrastive divergence (PCD) algorithm
1: Inputs: data points  $\Omega = \{x_{i}^{*}\}_{i=1}^{n}$  in  $\mathbb{R}^d$ ; energy model  $U_{\theta}$ ; optimizer step  $\operatorname{opt}(\theta, \mathcal{D})$  using  $\theta$  and the empirical CE gradient  $\mathcal{D}$ ; initial parameters  $\theta_0$ ; number of walkers  $N \in \mathbb{N}_0$  with  $N < n$ ; total duration  $K \in \mathbb{N}$ ; ULA time step  $h$ .
2:  $X_0^i = \text{SampleMultinomial}(\Omega)$  for  $i = 1, \ldots, N$ .
3: for  $k = 1, \dots, K - 1$  do
4:  $\tilde{\mathcal{D}}_k = N^{-1}\sum_{i = 1}^N\partial_\theta U_{\theta_k}(X_k^i) - n^{-1}\sum_{i = 1}^n\partial_\theta U_{\theta_k}(x_*^i)$
5:  $\theta_{k + 1} = \mathrm{opt}(\theta_k,\widetilde{\mathcal{D}}_k)$  optimization step
6: for  $i = 1,\dots,N$  do
7:  $X_{k + 1}^{i} = X_{k}^{i} - h\nabla U_{\theta_{k}}(X_{k}^{i}) + \sqrt{2h}\xi_{k}^{i},\qquad \xi_{k}^{i}\sim \mathcal{N}(0_{d},I_{d})$
8: end for
9: end for
10: Outputs: Optimized energy  $U_{\theta_K}$ ; set of walkers  $\{X_K^i\}_{i=1}^N$ .
to this equation is not available in closed from (and  $\check{\rho}(t,x) \neq \rho_{\theta(t)}(x)$  in general), but in the limit of large  $\nu$  (i.e. with very frequent reinitializations), we can show [55] that
$$
\check {\rho} (t, x) = \rho_ {*} (x) + \nu^ {- 1} \alpha \nabla \cdot (\nabla U _ {\theta (t)} (x) \rho_ {*} (x) + \nabla \rho_ {*} (x)) + O (\nu^ {- 2}). \tag {40}
$$
$$
\begin{array}{l} \int_ {\mathbb {R} ^ {d}} \partial_ {\theta} U _ {\theta (t)} (x) \left(\rho_ {*} (x) - \check {\rho} (t, x)\right) d x \\ = - \nu^ {- 1} \int_ {\mathbb {R} ^ {d}} \partial_ {\theta} U _ {\theta (t)} (x) \nabla \cdot \left(U _ {\theta (t)} (x) \rho_ {*} (x) + \nabla \rho_ {*} (x)\right) d x + O \left(\nu^ {- 2}\right) \tag {41} \\ = \nu^ {- 1} \int_ {\mathbb {R} ^ {d}} \left(\partial_ {\theta} \nabla U _ {\theta (t)} (x) \cdot \nabla U _ {\theta (t)} (x) - \partial_ {\theta} \Delta U _ {\theta (t)} (x)\right) \rho_ {*} (x) d x + O (\nu^ {- 2}) \\ \end{array}
$$
The leading order term at the right hand side is precisely  $\nu^{-1}$  times the gradient with respect to  $\theta$  of the Fisher divergence
$$
\begin{array}{l} \frac {1}{2} \int_ {\mathbb {R} ^ {d}} | \nabla U _ {\theta} (x) + \nabla \log \rho_ {*} (x) | ^ {2} \rho_ {*} (x) d x \tag {42} \\ = \frac {1}{2} \int_ {\mathbb {R} ^ {d}} \left[ | \nabla U _ {\theta} (x) | ^ {2} - 2 \Delta U _ {\theta} (x) + | \nabla \log \rho_ {*} (x) | ^ {2} \right] \rho_ {*} (x) d x \\ \end{array}
$$
where  $\Delta$  denotes the Laplacian and we used  $\int_{\mathbb{R}^d}\nabla U_\theta (x)\cdot \nabla \log \rho_*(x)\rho_*(x)dx = \int_{\mathbb{R}^d}\nabla U_\theta (x)\cdot \nabla \rho_*(x)dx = -\int_{\mathbb{R}^d}\Delta U_\theta (x)\rho_*(x)dx$ . This confirms the known fact that the CD algorithm effectively performs GD on the Fisher divergence rather than the cross-entropy [28].
Here we give additional details about our numerical results. The codes are available at: https://github.com/submissionx12/EBMs_Jarzynski.
B.1 Gaussian mixture distributions
Here we give some additional numerical results for the teacher-student model discussed in Section 4.1 in which the two wells of the teacher are aligned along the first dimension, fixing the first component of the means to be  $a_{*}^{1} = -10$  and  $b_{*}^{1} = 6$ , and  $a_{*}^{\alpha} = b_{*}^{\alpha} = 0$  for any  $\alpha = 2,\ldots ,d$ ; we also set  $z_{*} = -\log (3)$ , corresponding to a mass  $p_* = 1 / (1 + e^{-z_*}) = 0.25$  of the mode centered at  $a$ .
All the simulations are performed in  $d = 50$ , with a time step of  $h = 0.1$  for the ULA update. The number of data points is  $n = 10^4$ . The setup of the teacher is the same for every simulation we display here and the optimization step is performed with full batch gradient descent with learning rate constant in time.
Figure 6: GMM experiments: Evolution of the parameters and the cross entropy during training by Algorithm 1, PCD, and CD. W.r.t. Algorithm 1, we display the results for five different thresholds in the resampling step. 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.
Figure 7: GMM experiments: Evolution of the effective sample size (ESS) for five different choices of threshold associated to  $c$ , constant in time. Bimodal student.
GMM for different choices of  $c_k$ . Results are shown in Figures 6 and 7. As initial conditions we select  $a_0^1 = -10^{-1}$ ,  $b_0^1 = 10^{-1}$ ,  $a_0^\alpha \sim 10^{-2}\mathcal{N}(0,1)$  and  $b_0^\alpha \sim 10^{-2}\mathcal{N}(0,1)$  for any  $\alpha = 2,\dots ,d$ ; this perturbation around  $a = b = 0$  is prescribed to avoid numerical degeneracy. For  $z$ , we fix  $z_0 = 0$ . We run Algorithm 1, as well as the CD and PCD Algorithms 3 and 4 using  $N = 10^{4}$  walkers for  $K = 8\cdot 10^{3}$  iterations. We use a different learning rate for  $z$  and  $a,b$ , namely  $\gamma_{z} = 0.125$  and  $\gamma_{a,b} = 0.2\gamma_z$ . Moreover, these values are multiplied by a factor 10 in CD. With regard to the resampling step in Algorithm 1, we choose a threshold  $c_k = c$  which is fixed in time. We display the result for five possible values of this hyperparameter, namely  $c = [0.1,0.2,0.4,0.8,1.0]$ ; these are related to thresholds in the effective sampling size (ESS) via  $\mathrm{ESS}_{\mathrm{thresh}} = N / (c + 1)$ . With regard to Algorithm 1, we choose  $M = 1$  and  $N' = N$ , that is the full batch version; in the CD Algorithm 4 we choose  $P = 4$  for the number of ULA steps between restarts. Every run is performed five times and the average between them is shown in figures.
Mode collapse in GMM for PCD. Results are shown in Figure 8. In the same setup as above we select as initial conditions  $a_0^1 = -10^{-1}$ ,  $b_0^1 = 10^{-1}$ ,  $a_0^\alpha \sim 10^{-2}\mathcal{N}(0,1)$  and  $b_0^\alpha \sim 10^{-2}\mathcal{N}(0,1)$  for any  $\alpha \in [2,d]$ ; for  $z$ , we fix  $z_0 = 0$ . The learning rate is chosen to be  $\gamma_z = 5$  for  $z$  and  $\gamma_{a,b} = 0.2\gamma_z$  for the means. The time step of ULA is  $h = 0.2$ . Since the objective is to show mode collapse in the PCD algorithm, we run just Algorithm 4 for  $K = 10^{4}$  iterations and  $N = 10^{4}$  walkers.
The need for resampling. In absence of resampling, we observe a dramatic deterioration of ESS (Figure 9). To investigate how this issue is solved by resampling, we compare the three routines presented in Appendix A.4, using three pre-specified lags between the resampling steps. We use the same hyperparameters (learning rate, target distribution, etc.) as in Figure 9. The results are
Figure 8: Mode collapse and oscillations in PCD. Evolution of the probability  $p_k = 1 / (1 + e^{-z_k})$ .
Figure 9: GMM in 50d without resampling step, full batch experiment with  $N = 10^{5}$  walkers, average of 30 runs,  $4.8 \cdot 10^{5}$  iterations, same target density of Subsection 4.1. Left Panel: relative mass of the first mode. Right panel: evolution of ESS
shown in Figure 10: looking at the upper panels, we see that stratified and systematic resampling are more stable than multinomial; moreover, if the resampling step is performed too infrequently (green lines), the method converges to a result where the target relative mass is off. Looking at the lower panels, we see that the minimum statistical error is obtained with systematic resampling; moreover, the behaviour of uncertainty appears to be more stable than stratified.
As a side note, systematic resampling requires just a single random number, contrarily to multinomial and stratified which need  $N$ , making the former also more efficient from a computational point of view. This consideration, plus the experimental results we just discussed, strongly motivate the adoption of such method for the resampling step in our proposed algorithm.
Discussion. GMM in high dimension are challenging for the standard CD and PCD algorithms given in 3 and 4. The experimental results of this section also confirm the theoretical analysis presented in Appendix C below: CD is performing GD but on Fisher divergence rather than cross entropy, and as a result it incorrectly estimates the mass of the modes since Fisher divergence is insensitive to this quantity. On the other hand, PCD causes cycles of mode collapse (see Figure 8 and the left panel of Figure 6), and the KL divergence does not decrease monotonically, since the protocol is not ensured to be gradient descent (right panel in Figure 6).
In this example, our Algorithm 1 outperforms these standard methods as it is an implementation GD on cross-entropy. In particular, the estimation of the relative mass with Algorithm 1 is more
Figure 10: GMM in 50d with three resampling routines, full batch experiment with  $N = 10^5$  walkers, 50 runs per each resampling interval,  $4.8 \cdot 10^5$  iterations. Upper panels: relative mass of the first mode for each resampling method. Lower panels: standard error computed using the empirical standard deviation via  $\sigma_{emp} / \sqrt{50}$ .
accurate than with the CD and PCD algorithms (see Figure 6). Moreover, the computation of KL divergence via Jarzynski weights is fairly precise; in Figure 6 the estimated KL and the exact one overlaps beyond the minimum values reached in with the CD and PCD algorithms. With regard to Figure 7, the choice of the threshold for resampling does not appear to be decisive in this regime of hyperparameters; in fact, looking at the evolution of  $\theta$  and of KL divergence, the overall behavior of Algorithm 1 is not dramatically influenced by the choice of  $c$ .
B.2 MNIST data set
In this section, we provide additional details about our numerical experiments on the MNIST dataset using Algorithm 2. First, as discussed in Section 4.2, we confirm that the Jarzynski weights of the generated samples are directly related to the image quality: this is shown in the left panel of Figure 11, where we display images along with their Jarzynski weights. Moreover, the right panel of Figure 11 indicates that resampling at the end of training using these weights helps improve sample quality. Examples of generated images are shown in Figure 12.
B.3 CIFAR-10
A key component of the success of our method is the resampling step. In Figure 13 we plot the Effective Sample Size (ESS): each peak in Figure 13 indicates a step of resampling of all the samples. We note that training EBMs with mini-batches has the side effect of a rapid loss of the ESS, which measures the sample quality. So, doing resampling with a proper criterion and a reliable resampler is necessary.
C Theoretical analysis of mode collapse in simple 1d-GMM
In this appendix, we explain why and how mode collapse arises when learning multimodal distributions if we do not include the weights prescribed by Jarzynski equality, and why collapse does not arise with these weights included. To show this in the simplest setting, we work with a Gaussian mixture model: a mixture of two one-dimensional Gaussian densities, with the same variance equal
Generated samples during training (w/o resampling)
Generated samples during training (w/ resampling)
Figure 11: MNIST dataset: Left panel: Images randomly chosen during the training from the entire set of generated samples with their associate Jarzynski weights. From left to right, top to bottom, the higher the Jarzynski weight is, the better the image quality is. Right panel: Images obtained under the same training conditions, after resampling and continued training for 120 epochs.
Figure 12: MNIST dataset: Left panel: images generated by the mini-batched version of the Algorithm 2. Right panel: Images generated from the PCD with mini-batches.
to 1, with means  $a, b \in \mathbb{R}$  of the modes: our aim is to learn the probability masses of these modes. The general picture is given in Appendix C.1, with details presented afterwards in Appendix C.2, where we assume that the learning dynamics is that in (47), with no walkers and direct access to the distributions  $\rho_{z(t)}$  and  $\rho_*$ , and Appendix C.3, where we quantify the stochastic fluctuations induced by empirical estimation.
Figure 13: CIFAR-10 dataset: The Effective Sample Size (ESS) during the training with Algorithm 2. Each peak in the plot implies one step of resampling. Notice that the number of resampling in CIFAR-10 experiments is significantly larger than the one in MNIST experiments (4), which suggests the necessity of resampling in the scalability of our method.
C.1 Mode collapse in Gaussian Mixtures