SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$$
\begin{array}{l} 0 \leq \mathrm {K L} \left(p _ {t, \text {d a t a}} \| p _ {t, \text {e b m}}\right) \\ = c (t) + \int U _ {t} (\mathbf {y}) p _ {t, \mathrm {d a t a}} (\mathbf {y}) \mathrm {d} \mathbf {y} - \int U (\mathbf {x}) p _ {\mathrm {d a t a}} (\mathbf {x}) - \int \log p _ {\mathrm {e b m}} (\mathbf {x}) p _ {\mathrm {d a t a}} (\mathbf {x}) \\ = c (t) - \operatorname {E D} _ {\gamma_ {t}} (p _ {\mathrm {d a t a}}, U) - \mathbb {E} _ {p _ {\mathrm {d a t a}} (\mathbf {x})} [ \log p _ {\mathrm {e b m}} (\mathbf {x}) ] \leq \frac {1}{2 t} \mathbb {W} _ {2} ^ {2} (p _ {\mathrm {d a t a}}, p _ {\mathrm {e b m}}) \\ \end{array}
$$
A.4 Representing Energy Discrepancy as Multi-Scale SM for General Diffusion Processes
We now prove the connection between energy discrepancy and multi-noise scale score matching in a general context. For all following results we will assume that  $\mathbf{x}_t$  is some stochastic diffusion process which satisfies the SDE  $\mathrm{d}\mathbf{x}_t = a(\mathbf{x}_t)\mathrm{d}t + b(\mathbf{x}_t)\mathrm{d}\mathbf{w}_t$  and assume that  $\mathbf{x}_0\sim p_{\mathrm{data}}$ . Let  $q_{t}$  denote the associated transition probability density. To make the exposition cleaner we write  $U_{t}\coloneqq U_{q_{t}}$
The main idea will be the following observation:
Proposition 4. The diffusion-based energy discrepancy is given by the expectation of the Ito integral
$$
\mathrm {E D} _ {q _ {t}} (p _ {\mathrm {d a t a}}, U) = - \mathbb {E} \left[ \int_ {0} ^ {t} \mathrm {d} U _ {s} (\mathbf {x} _ {s}) \right]
$$
Proof. The stochastic integral with respect to the differential  $\mathrm{d}U_s(\mathbf{x}_s)$  is defined to satisfy the following generalisation of the fundamental theorem of calculus:
$$
U _ {t} (\mathbf {x} _ {t}) = U _ {0} (\mathbf {x} _ {0}) + \int_ {0} ^ {t} \mathrm {d} U _ {s} (\mathbf {x} _ {s})
$$
We obtain the desired result by taking expectations on both sides in  $(\mathbf{x}_s)_{0\leq s\leq t}$ .
Notice that the law of the random variable  $\mathbf{x}_s$  is fixed by the initial distribution of the diffusion  $\mathbf{x}_0 \sim p_{\mathrm{data}}$ . These distributions are implied when taking the expectation. We will now explore this connection further. For this we make some basic assumptions which allow us to connect stochastic differential equations with partial differential equations.
Assumption 2. Consider the stochastic differential equation  $\mathrm{d}\mathbf{x}_t = a(\mathbf{x}_t)\mathrm{d}t + b(\mathbf{x}_t)\mathrm{d}\mathbf{w}_t$  for drift  $a:\mathbb{R}^d\to \mathbb{R}^d$  and  $b:\mathbb{R}^d\to \mathbb{R}^k$ . Further, define the diffusion matrix  $\Sigma (\mathbf{x}) = b(\mathbf{x})b(\mathbf{x})^T\in \mathbb{R}^{d\times d}$ . We make the following assumptions:
1. There exists a  $\mu > 0$  such that for all  $\pmb{\xi}, \mathbf{x} \in \mathbb{R}^d\langle \pmb{\xi}, \Sigma(\mathbf{x})\pmb{\xi} \rangle \geq \mu \| \pmb{\xi} \|^2$
2.  $\Sigma$  and a are bounded and uniformly Lipschitz-continuous in  $\mathbf{x}$  on every compact subset of  $\mathbb{R}^d$
3.  $\Sigma$  is uniformly Hölder-continuous in  $x$
Theorem 4 (Fokker-Planck equation). Under Assumption 2,  $\mathbf{x}_t$  has a transition density function given by
$$
\mathbb {P} \left(\mathbf {x} _ {t} \in A | \mathbf {x} _ {0} = \mathbf {x}\right) = \int_ {A} q _ {t} (\mathbf {y} | \mathbf {x}) \mathrm {d} \mathbf {y}.
$$
Furthermore,  $q_{t}$  satisfies the Fokker-Planck partial differential equation
$$
\partial_ {t} q _ {t} (\mathbf {y} | \mathbf {x}) = \sum_ {i = 1} ^ {d} \partial_ {\mathbf {y} _ {i}} \left(- a _ {i} (\mathbf {y}) q _ {t} (\mathbf {y} | \mathbf {x}) + \frac {1}{2} \sum_ {j = 1} ^ {d} \partial_ {\mathbf {y} _ {j}} \left(\Sigma_ {i j} (\mathbf {y}) q _ {t} (\mathbf {y} | \mathbf {x})\right)\right) \tag {10}
$$
$$
q _ {0} (\mathbf {y} | \mathbf {x}) = \delta (\mathbf {y} - \mathbf {x})
$$
For a reference, see (Friedman, 2012, Theorem 5.4)
The Fokker-Planck equation yields the following important differential equation for the contrastive potential  $U_{t}$ :
Proposition 5. Consider the stochastic differential equation  $\mathrm{d}\mathbf{x}_t = a(\mathbf{x}_t)\mathrm{d}t + b(\mathbf{x}_t)\mathrm{d}\mathbf{w}_t$  for drift  $a:\mathbb{R}^d\to \mathbb{R}^d$  and  $b:\mathbb{R}^d\to \mathbb{R}^k$ , and diffusion matrix  $\Sigma (\mathbf{x}) = b(\mathbf{x})b(\mathbf{x})^{T}\in \mathbb{R}^{d\times d}$  that satisfies assumptions 2. Let  $q_{t}$  be the associated transition density and define the contrastive potential  $U_{t}(\mathbf{y}):= -\log \int q_{t}(\mathbf{y}|\mathbf{x})\exp (-U(\mathbf{x}))\mathrm{d}\mathbf{x}$ . Furthermore, we define the scalar field
$$
c (a, \Sigma) (\mathbf {y}) := \sum_ {i = 1} ^ {d} \left(\partial_ {\mathbf {y} _ {i}} a _ {i} (\mathbf {y}) - \frac {1}{2} \sum_ {j = 1} ^ {d} \partial_ {\mathbf {y} _ {i}} \partial_ {\mathbf {y} _ {j}} \Sigma_ {i j}\right)
$$
and the linear operator
$$
\mathcal {L} (a, \Sigma) := \sum_ {i = 1} ^ {d} - a _ {i} \frac {\partial}{\partial \mathbf {y} _ {i}} + \frac {1}{2} \sum_ {i, j = 1} ^ {d} \left(2 \partial_ {\mathbf {y} _ {j}} \Sigma_ {i j} \frac {\partial}{\partial \mathbf {y} _ {i}} + \Sigma_ {i j} \frac {\partial^ {2}}{\partial \mathbf {y} _ {i} \partial \mathbf {y} _ {j}}\right).
$$
Then, the contrastive potential satisfies the non-linear partial differential equation
$$
\partial_ {t} U _ {t} (\mathbf {y}) = \mathcal {L} (a, \Sigma) U _ {t} (\mathbf {y}) + \frac {1}{2} \| b ^ {T} (\mathbf {y}) \nabla U _ {t} (\mathbf {y}) \| ^ {2} + c (a, \Sigma) (\mathbf {y})
$$
Proof. We commute the linear operator of the Fokker-Planck equation to see that  $e^{-U_t}$  satisfies the Fokker-Planck equation in Theorem 4, too, i.e.
$$
\partial_ {t} e ^ {- U _ {t} (\mathbf {y})} = \sum_ {i = 1} ^ {d} \partial_ {\mathbf {y} _ {i}} \left(- a _ {i} (\mathbf {y}) e ^ {- U _ {t} (\mathbf {y})} + \frac {1}{2} \sum_ {j = 1} ^ {d} \partial_ {\mathbf {y} _ {j}} \left(\Sigma_ {i j} (\mathbf {y}) e ^ {- U _ {t} (\mathbf {y})}\right)\right)
$$
$$
e ^ {- U _ {0} (\mathbf {y})} = e ^ {- U (\mathbf {x})}.
$$
We now expand the term corresponding to the drift term:
$$
\sum_ {i = 1} ^ {d} \partial_ {\mathbf {y} _ {i}} \left(- a _ {i} (\mathbf {y}) e ^ {- U _ {t} (\mathbf {y})}\right) = \sum_ {i = 1} ^ {d} \left(- \partial_ {\mathbf {y} _ {i}} a _ {i} (\mathbf {y}) + a _ {i} (\mathbf {y}) \partial_ {\mathbf {y} _ {i}} U _ {t} (\mathbf {y})\right) e ^ {- U _ {t} (\mathbf {y})}
$$
Similarly, we treat the diffusion term:
$$
\begin{array}{l} \frac {1}{2} \sum_ {i, j = 1} ^ {d} \partial_ {\mathbf {y} _ {i}} \partial_ {\mathbf {y} _ {j}} \left(\Sigma_ {i j} (\mathbf {y}) e ^ {- U _ {t} (\mathbf {y})}\right) = \frac {1}{2} \sum_ {i, j = 1} ^ {d} \left(\partial_ {\mathbf {y} _ {i}} \partial_ {\mathbf {y} _ {j}} \Sigma_ {i j} (\mathbf {y}) + \Sigma_ {i j} (\mathbf {y}) \partial_ {\mathbf {y} _ {j}} U _ {t} (\mathbf {y}) \partial_ {\mathbf {y} _ {i}} U _ {t} (\mathbf {y}) \right. \\ \left. - 2 \partial_ {\mathbf {y} _ {j}} \Sigma_ {i j} (\mathbf {y}) \partial_ {\mathbf {y} _ {i}} U _ {t} (\mathbf {y}) - \Sigma_ {i j} (\mathbf {y}) \partial_ {\mathbf {y} _ {j}} \partial_ {\mathbf {y} _ {i}} U _ {t} (\mathbf {y})\right) e ^ {- U _ {t} (\mathbf {y})} \\ \end{array}
$$
Finally, the time derivative simply becomes  $\partial_t e^{-U_t(\mathbf{y})} = -\partial_t U_t(\mathbf{y})e^{-U_t(\mathbf{y})}$ . We can now collect all terms independent of  $U$  and identify
$$
c (a, \Sigma) (\mathbf {y}) = \sum_ {i = 1} ^ {d} \left(\partial_ {\mathbf {y} _ {i}} a _ {i} (\mathbf {y}) - \frac {1}{2} \sum_ {j = 1} ^ {d} \partial_ {\mathbf {y} _ {i}} \partial_ {\mathbf {y} _ {j}} \Sigma_ {i j} (\mathbf {y})\right)
$$
as well as the linear operator term
$$
\mathcal {L} (a, \Sigma) := \sum_ {i = 1} ^ {d} - a _ {i} \frac {\partial}{\partial \mathbf {y} _ {i}} + \frac {1}{2} \sum_ {i, j = 1} ^ {d} \left(2 \partial_ {\mathbf {y} _ {j}} \Sigma_ {i j} \frac {\partial}{\partial \mathbf {y} _ {i}} + \Sigma_ {i j} \frac {\partial^ {2}}{\partial \mathbf {y} _ {i} \partial \mathbf {y} _ {j}}\right)
$$
$$
\begin{array}{l} \sum_ {i, j = 1} ^ {d} \Sigma_ {i j} (\mathbf {y}) \partial_ {\mathbf {y} _ {j}} U _ {t} (\mathbf {y}) \partial_ {\mathbf {y} _ {i}} U _ {t} (\mathbf {y}) = \sum_ {i, j = 1} ^ {d} \sum_ {l = 1} ^ {k} b _ {i l} (\mathbf {y}) b _ {j l} (\mathbf {y}) \partial_ {\mathbf {y} _ {j}} U _ {t} (\mathbf {y}) \partial_ {\mathbf {y} _ {i}} \\ = \sum_ {l = 1} ^ {k} \sum_ {i = 1} ^ {d} \left(b _ {l, i} ^ {T} (\mathbf {y}) \partial_ {\mathbf {y} _ {i}} U _ {t} (\mathbf {y})\right) ^ {2} = \| b ^ {T} (\mathbf {y}) \nabla U _ {t} (\mathbf {y}) \| ^ {2} \\ \end{array}
$$
This gives us the partial differential equation
$$
- \partial_ {t} U _ {t} (\mathbf {y}) e ^ {- U _ {t} (\mathbf {y})} = \left(- \mathcal {L} (a, \Sigma) U _ {t} (\mathbf {y}) + \frac {1}{2} \| b ^ {T} (\mathbf {y}) \nabla U _ {t} (\mathbf {y}) \| ^ {2} - c (a, \Sigma) (\mathbf {y})\right) e ^ {- U _ {t} (\mathbf {y})}
$$
Cancelling all exponentials from both sides of the equation yields the desired result.
Theorem 5. The energy discrepancy takes the form of a generalised multi-noise scale score matching loss:
$$
\mathrm {E D} _ {q _ {t}} (p _ {\mathrm {d a t a}}, U) = \mathbb {E} \left[ \int_ {0} ^ {t} - \sum_ {i, j = 1} ^ {d} \partial_ {\mathbf {x} _ {j}} (\Sigma_ {i j} (\mathbf {x} _ {s}) \partial_ {\mathbf {x} _ {i}} U (\mathbf {x} _ {s})) + \frac {1}{2} \| b ^ {T} (\mathbf {x} _ {s}) \nabla U _ {s} (\mathbf {x} _ {s}) \| ^ {2} \mathrm {d} s \right] + \mathrm {c o n s t .}
$$
Proof. For this proof we return to the stochastic process  $U_{s}(\mathbf{x}_{s})$  from Proposition 4. By Ito's formula,  $U_{s}(\mathbf{x}_{s})$  satisfies the stochastic differential equation
$$
\begin{array}{l} \mathrm {d} U _ {s} (\mathbf {x} _ {s}) = \left(\partial_ {s} U _ {s} (\mathbf {x} _ {s}) + \sum_ {i = 1} ^ {d} a _ {i} (\mathbf {x} _ {s}) \partial_ {\mathbf {x} _ {i}} U _ {s} (\mathbf {x} _ {s}) + \frac {1}{2} \sum_ {i, j = 1} ^ {d} \Sigma_ {i j} (\mathbf {y}) \partial_ {\mathbf {x} _ {i}} \partial_ {\mathbf {x} _ {j}} U (\mathbf {x} _ {s})\right) \mathrm {d} s \\ + \sum_ {i = 1} ^ {d} \sum_ {l = 1} ^ {k} \partial_ {\mathbf {x} _ {i}} U _ {s} (\mathbf {x} _ {s}) b _ {i, l} (\mathbf {x} _ {s}) d \mathbf {w} _ {s} ^ {l} \\ \end{array}
$$
Under the additional integrability condition that  $\mathbb{E}\int_0^t\| b^T (\mathbf{x}_s)\nabla U_s(\mathbf{x}_s)\| ^2\mathrm{d}s < \infty$ , the stochastic integral with respect to Brownian motion  $d\mathbf{w}_s$  has expectation zero. Furthermore, we can replace  $\partial_sU_s(\mathbf{x}_s)$  with our previously obtained non-linear partial differential equation
$$
\partial_ {s} U _ {s} (\mathbf {x} _ {s}) = \mathcal {L} (a, \Sigma) U _ {s} (\mathbf {x} _ {s}) + \frac {1}{2} \left\| b ^ {T} (\mathbf {x} _ {s}) \nabla U _ {s} (\mathbf {x} _ {s}) \right\| ^ {2} + c (a, \Sigma) (\mathbf {x} _ {s}).
$$
Due to opposing signs, the drift  $a$  cancels, i.e.
$$
\begin{array}{l} \mathcal {L} (a, \Sigma) U _ {s} (\mathbf {x} _ {s}) + \sum_ {i = 1} ^ {d} a _ {i} (\mathbf {x} _ {s}) \partial_ {\mathbf {x} _ {i}} U _ {s} (\mathbf {x} _ {s}) + \frac {1}{2} \sum_ {i, j = 1} ^ {d} \Sigma_ {i j} (\mathbf {y}) \partial_ {\mathbf {x} _ {i}} \partial_ {\mathbf {x} _ {j}} U (\mathbf {x} _ {s}) \\ = \sum_ {i, j = 1} ^ {d} \left(\partial_ {\mathbf {y} _ {j}} \Sigma_ {i j} (\mathbf {y}) \frac {\partial}{\partial \mathbf {y} _ {i}} + \Sigma_ {i j} (\mathbf {y}) \frac {\partial^ {2}}{\partial \mathbf {x} _ {i} \partial \mathbf {x} _ {j}}\right) U _ {s} (\mathbf {x} _ {s}) \\ = \sum_ {i, j = 1} ^ {d} \partial_ {\mathbf {x} _ {j}} \left(\Sigma_ {i j} \partial_ {\mathbf {x} _ {i}} U (\mathbf {x} _ {s})\right) \\ \end{array}
$$
Consequently, we obtain the final energy discrepancy expression using Proposition 4
$$
\begin{array}{l} \mathrm {E D} _ {q _ {t}} \left(p _ {\text {d a t a}}, U\right) = - \mathbb {E} \left[ \int_ {0} ^ {t} d U _ {s} \left(\mathbf {x} _ {s}\right) \right] \\ = - \mathbb {E} \left[ \int_ {0} ^ {t} \sum_ {i, j = 1} ^ {d} \partial_ {\mathbf {x} _ {j}} \left(\Sigma_ {i j} (\mathbf {x} _ {s}) \partial_ {\mathbf {x} _ {i}} U (\mathbf {x} _ {s})\right) - \frac {1}{2} \| b ^ {T} (\mathbf {x} _ {s}) \nabla U _ {s} (\mathbf {x} _ {s}) \| ^ {2} \mathrm {d} s \right] + \text {c o n s t .} \\ \end{array}
$$
with  $U$ -independent constant  $\int_0^t c(a,\Sigma)(\mathbf{x}_s)\mathrm{d}s$ . This completes the proof.
As a corollary we obtain the proof of the first statement in Theorem 2: Assume that  $q_{t}$  is defined through the stochastic differential equation  $\mathrm{d}\mathbf{x}_t = a(\mathbf{x}_t)\mathrm{d}t + \mathrm{d}\mathbf{w}_t$ . In this case,  $\Sigma = \mathbf{I}$  and  $\sum_{i,j=1}^{d} \partial_{\mathbf{x}_j} (\Sigma_{ij}(\mathbf{x}_s) \partial_{\mathbf{x}_i} U(\mathbf{x}_s)) = \Delta U(\mathbf{x}_s)$ . Consequently, we obtain from Theorem 5 the score matching representation of  $\mathrm{ED}_{q_t}$  in Theorem 2. In the special case that  $\Sigma = bb^T$  is independent of  $\mathbf{x}$  we obtain an integrated sliced score-matching loss (Song et al., 2020).
A.5 Connections of Energy Discrepancy with Contrastive Divergence
The contrastive divergence update can be derived from an energy discrepancy when, for  $E_{\theta}$  fixed,  $q$  satisfies the detailed balance relation
$$
q (\mathbf {y} | \mathbf {x}) \exp (- E _ {\theta} (\mathbf {x})) = q (\mathbf {x} | \mathbf {y}) \exp (- E _ {\theta} (\mathbf {y})).
$$
To see this, we calculate the contrastive potential induced by  $q$ : We have
$$
- \log \int q (\mathbf {y} | \mathbf {x}) \exp (- E _ {\theta} (\mathbf {x})) d \mathbf {x} = - \log \int q (\mathbf {x} | \mathbf {y}) \exp (- E _ {\theta} (\mathbf {y})) d \mathbf {x} = E _ {\theta} (\mathbf {y}).
$$
Consequently, the energy discrepancy induced by  $q$  is given by
$$
\operatorname {E D} _ {q} (p _ {\mathrm {d a t a}}, E _ {\theta}) = \mathbb {E} _ {p _ {\mathrm {d a t a}} (\mathbf {x})} [ E _ {\theta} (\mathbf {x}) ] - \mathbb {E} _ {p _ {\mathrm {d a t a}} (\mathbf {x})} \mathbb {E} _ {q (\mathbf {y} | \mathbf {x})} [ E _ {\theta} (\mathbf {y}) ].
$$
Updating  $\theta$  based on a sample approximation of this loss leads to the contrastive divergence update
$$
\Delta \theta \propto \frac {1}{N} \sum_ {i = 1} ^ {N} \nabla_ {\theta} E _ {\theta} (\mathbf {x} ^ {i}) - \frac {1}{N} \sum_ {i = 1} ^ {N} \nabla_ {\theta} E _ {\theta} (\mathbf {y} ^ {i}) \quad \mathbf {y} ^ {i} \sim q (\cdot | \mathbf {x} ^ {i})
$$
Three things are important to notice:
1. Implicitly, the distribution  $q$  depends on  $E_{\theta}$  and needs to be adjusted in each step of the algorithm
2. For fixed  $q$ ,  $\mathrm{ED}_q(p_{\mathrm{data}}, E_\theta)$  satisfies Theorem 1. This means that each step of contrastive divergence optimises a loss with minimiser  $E_\theta^* = -\log p_{\mathrm{data}} + c$ . However,  $q$  needs to be adjusted in each step as otherwise the contrastive potential is not given by the energy function  $E_\theta$  itself.
3. This result highlights the importance to use Metropolis-Hastings adjusted Langevin-samplers to implement CD to ensure that the implied  $q$  distribution satisfies the detailed balance relation. This matches the observations found by Yair & Michaeli (2021).
A.6 Derivation of Energy Discrepancy from KL Contractions
A Kullback-Leibler contraction is the divergence function  $\mathrm{KL}(p_{\mathrm{data}}\parallel p_{\mathrm{ebm}}) - \mathrm{KL}(Qp_{\mathrm{data}}\parallel Qp_{\mathrm{ebm}})$  (Lyu, 2011) for the convolution operator  $Qp(\mathbf{y}) = \int q(\mathbf{y}|\mathbf{x})p(\mathbf{x})\mathrm{d}\mathbf{x}$ . The linearity of the convolution operator retains the normalisation of the measure, i.e. for the energy-based distribution  $p_{\mathrm{ebm}}$  we have
$$
Q p _ {\mathrm {e b m}} = \frac {1}{Z _ {U}} \int q (\mathbf {y} | \mathbf {x}) \exp (- U (\mathbf {x}) \quad \text {w i t h} \quad Z _ {U} = \int \exp (- U (\mathbf {x})) \mathrm {d} \mathbf {x}.
$$
The KL divergences then become with  $U_{q} \coloneqq -\log Q \exp (-U(\mathbf{x}))$
$$
\operatorname {K L} \left(p _ {\text {d a t a}} \| p _ {\text {e b m}}\right) = \mathbb {E} _ {p _ {\text {d a t a}} (\mathbf {x})} \left[ \log p _ {\text {d a t a}} (\mathbf {x}) \right] + \mathbb {E} _ {p _ {\text {d a t a}} (\mathbf {x})} [ U (\mathbf {x}) ] + \log Z _ {U}
$$
$$
\operatorname {K L} \left(Q p _ {\text {d a t a}} \| Q p _ {\text {e b m}}\right) = \mathbb {E} _ {Q p _ {\text {d a t a}} (\mathbf {y})} \left[ \log Q p _ {\text {d a t a}} (\mathbf {y}) \right] + \mathbb {E} _ {Q p _ {\text {d a t a}} (\mathbf {y})} \left[ U _ {q} (\mathbf {y}) \right] + \log Z _ {U}
$$
Since the normalisation cancels when subtracting the two terms we find
$$
\mathrm {K L} \left(p _ {\text {d a t a}} \| p _ {\text {e b m}}\right) - \mathrm {K L} \left(Q p _ {\text {d a t a}} \| Q p _ {\text {e b m}}\right) = \mathrm {E D} _ {q} \left(p _ {\text {d a t a}}, U\right) + c
$$
where  $c$  is a constant that contains the  $U$ -independent entropies of  $p_{\mathrm{data}}$  and  $Qp_{\mathrm{data}}$ .
B Aspects of Training Energy-Based Models with Energy Discrepancy
In this section, we discuss the  $w$ -stabilisation in depth. Furthermore, we give a proof for Theorem 3 and show energy discrepancy can be approximated for other diffusion processes via the Feynman-Kac formula.
B.1 Conceptual Understanding of the  $w$ -Stabilisation
The  $w$ -stabilisation is a useful stabilisation for all types of perturbation  $q$ . For example, we use the same stabilisation in the discrete setting (see Appendix D.3). To provide a better intuition for the stabilisation, we will investigate it more closely for the Gaussian perturbation.
The critical step for using energy discrepancy in practice is a stable approximation of the contrastive potential. For the Gaussian-based energy discrepancy, we can write the contrastive potential as  $U_{t}(\mathbf{y}) = -\log \mathbb{E}[\exp (-U(\mathbf{y} + \sqrt{t}\pmb{\xi}'))]$  with  $\pmb{\xi}' \sim \mathcal{N}(0,\mathbf{I})$  and  $\mathbf{y} \in \mathbb{R}^d$ . A naive approximation of the expectation with a Monte-Carlo estimator, however, is biased because of Jensen's inequality, i.e. for  $\pmb{\xi}', \pmb{\xi}'^j \sim \mathcal{N}(0,\mathbf{I})$  we have
$$
U _ {t} (\mathbf {y}) = - \log \mathbb {E} \left[ \exp \left(- U \left(\mathbf {y} + \sqrt {t} \boldsymbol {\xi} ^ {\prime}\right)\right) \right] <   - \mathbb {E} \left[ \log \frac {1}{M} \sum_ {j = 1} ^ {M} \exp \left(- U \left(\mathbf {y} + \sqrt {t} \boldsymbol {\xi} ^ {\prime j}\right)\right) \right].
$$
Our first observation is that the appearing bias can be quantified to leading order. For this, define  $v_{t}(\mathbf{y}) \coloneqq \mathbb{E}\left[\exp (-U(\mathbf{y} + \sqrt{t}\pmb{\xi}'))\right]$  and  $\widehat{v}_{t}(\mathbf{y}) \coloneqq \frac{1}{M}\sum_{j = 1}^{M}\exp (-U(\mathbf{y} + \sqrt{t}\pmb{\xi}^{\prime j}))$ . We use the Taylor-expansion of  $\log (1 + u)\approx u - 1 / 2u^2 +\mathrm{h.o.t.}$  which gives
$$
\begin{array}{l} \mathbb {E} \log \widehat {v} _ {t} (\mathbf {y}) - \log v _ {t} (\mathbf {y}) = \mathbb {E} \left[ \log \left(1 + \frac {\widehat {v} _ {t} (\mathbf {y}) - v _ {t} (\mathbf {y})}{v _ {t} (\mathbf {y})}\right) \right] \tag {11} \\ \approx - \frac {1}{2} \mathbb {E} \left[ \frac {\left(\widehat {v} _ {t} (\mathbf {y}) - v _ {t} (\mathbf {y})\right) ^ {2}}{v _ {t} (\mathbf {y}) ^ {2}} \right] + \mathrm {h . o . t .} \\ = - \frac {1}{2 M v _ {t} (\mathbf {y}) ^ {2}} \mathrm {V a r} \left[ \exp (- U (\mathbf {y} + \sqrt {t} \pmb {\xi} ^ {\prime})) \right] + \mathrm {h . o . t .} \\ \end{array}
$$
The linear term in the Taylor-expansion does not contribute because  $\mathbb{E}[\widehat{v}_t(\mathbf{y})] = v_t(\mathbf{y})$ . In the final equation we used that  $\mathrm{Var}(\widehat{v}_t(\mathbf{y})) = \mathrm{Var}[\exp (-U(\mathbf{y} + \sqrt{t}\pmb{\xi}'))] / M$  because all  $\pmb{\xi}^{\prime j}$  are independent. The Taylor expansion shows that the dominating contribution to the bias is the variance of the approximated convolution integral.
Our second observation is that this occurring bias can become infinite for malformed energy functions. For this reason, the optimiser may start to increase the bias instead of minimising our target loss. To illustrate how a high-variance estimator of the contrastive potential can be divergent, consider the energy function
$$
U (\mathbf {x}) = \left\{ \begin{array}{l l} 0 & \text {f o r} \mathbf {x} \leq 0 \\ b \mathbf {x} & \text {f o r} \mathbf {x} > 0 \end{array} \right..
$$
The energy function does not strictly adhere to our conditions that the energy based model should be normalisable. Our argument still holds when  $\exp(-U)$  is changed to be normalisable. In theory, the contrastive potential at 0 is upper bounded because
$$
U _ {t} (0) \leq - \lim  _ {b \rightarrow \infty} \log \mathbb {E} \left[ \exp (- U (\sqrt {t} \boldsymbol {\xi} ^ {\prime})) \right] = - \log \mathbb {P} (\boldsymbol {\xi} ^ {\prime} \leq 0) = - \log (1 / 2)
$$
because  $\exp(-U(\mathbf{x}))$  converges to an indicator function on  $\{\mathbf{x} \leq 0\}$  as  $b \to \infty$ . The Monte Carlo estimator of the contrastive potential, on the other hand, has upper bound
$$
\widehat {U} _ {t} (0) = - \log \frac {1}{M} \sum_ {j = 1} ^ {M} \exp (- U (\sqrt {t} \boldsymbol {\xi} ^ {\prime j})) \leq \min  [ U (\sqrt {t} \boldsymbol {\xi} ^ {\prime 1}), \dots , U (\sqrt {t} \boldsymbol {\xi} ^ {\prime M}) ] + \log (M)
$$
which can be seen by applying standard inequalities for the logsumexp function. Hence, as long as there exists a  $j$  such that  $\pmb{\xi}^{j}\leq 0$ , the estimated contrastive potential does not diverge. If, however,  $\pmb{\xi}^{\prime j} > 0$  for every  $j = 1,\dots ,M$ , then
$$
\widehat {U} _ {t} (0) \geq \min  [ U (\sqrt {t} \boldsymbol {\xi} ^ {\prime 1}), \dots , (\sqrt {t} \boldsymbol {\xi} ^ {\prime M}) ] = b \sqrt {t} \min  [ \boldsymbol {\xi} ^ {\prime 1}, \dots , \boldsymbol {\xi} ^ {\prime M} ] \xrightarrow {b \to \infty} \infty .
$$
Consequently, the approximate contrastive potential may attain diverging values at discontinuities in the energy function. Indeed, this phenomenon is observed for  $w = 0$  in Figure 2. Here, the learned energy becomes discontinuous at the edge of the support and the energy discrepancy loss diverges during training. In low dimensions, this problem can be alleviated by using variance reduction techniques such as antithetic variables or by using large enough values of  $M$  during the training. The stabilising effect of  $M$  is observed in our ablation studies in Figure 24. In high-dimensional settings, however, such variance reduction techniques are infeasible.
The idea of the  $w$ -stabilisation is that the value of the energy at non-perturbed data points  $U(\mathbf{x}_0)$  is guaranteed to stay controlled since it is minimised in the optimisation of ED. Hence, the diverging contrasting potential can be controlled by including  $U(\mathbf{x}_0)$  in the summation in the logsumexp operation which acts as a soft-min over all contrasting energy contributions. Indeed, this augmentation provides a deterministic upper bound to the approximated contrastive potential:
$$
\begin{array}{l} \widehat {U} _ {t, w} (\mathbf {x} _ {t}) = - \log \left(\frac {w}{M} \exp (- U (\mathbf {x} _ {0})) + \frac {1}{M} \sum_ {j = 1} ^ {M} \exp (- U (\mathbf {x} _ {t} + \sqrt {t} \boldsymbol {\xi} ^ {\prime j}))\right) \\ \leq \min  [ U (\mathbf {x} _ {t} + \sqrt {t} \pmb {\xi} ^ {\prime 1}), \dots , U (\mathbf {x} _ {t} + \sqrt {t} \pmb {\xi} ^ {\prime M}), U (\mathbf {x}) - \log (w) ] + \log (M) \\ \end{array}
$$
Additionally, the  $w$ -stabilisation introduces a negative bias to the approximated contrastive potential. Hence, if tuned correctly, it counteracts the bias introduced by the Jensen-gap of the logarithm.
To gain additional intuition on the effect of  $w$ , notice that by the same bounds as before,
$$
U \left(\mathbf {x} _ {0}\right) - \widehat {U} _ {t, w} \left(\mathbf {x} _ {t}\right) \leq \min  \left[ U \left(\mathbf {x} _ {0}\right) - U \left(\mathbf {x} _ {t} + \sqrt {t} \boldsymbol {\xi} ^ {\prime 1}\right), \dots , U \left(\mathbf {x} _ {0}\right) - U \left(\mathbf {x} _ {t} + \sqrt {t} \boldsymbol {\xi} ^ {\prime M}\right), \log (w) \right]
$$