Continuously Tempered Diffusion Samplers

TL;DR

This paper introduces continuously tempered diffusion samplers that enhance exploration during training by using a temperature-variant distribution, leading to improved sampling performance for neural samplers of unnormalized distributions.

Contribution

It proposes a novel sampler that employs a temperature-variant distribution to improve exploration and training of neural samplers, addressing limitations of previous annealing-based methods.

Findings

Enhanced exploration leads to better sampler performance.

Empirical validation shows improved sampling quality.

Code availability facilitates reproducibility.

Abstract

Annealing-based neural samplers seek to amortize sampling from unnormalized distributions by training neural networks to transport a family of densities interpolating from source to target. A crucial design choice in the training phase of such samplers is the proposal distribution by which locations are generated at which to evaluate the loss. Previous work has obtained such a proposal distribution by combining a partially learned transport with annealed Langevin dynamics. However, isolated modes and other pathological properties of the annealing path imply that such proposals achieve insufficient exploration and thereby lower performance post training. To remedy this, we propose continuously tempered diffusion samplers, which leverage exploration techniques developed in the context of molecular dynamics to improve proposal distributions. Specifically, a family of distributions across different temperatures is introduced to lower energy barriers at higher temperatures and drive exploration at the lower temperature of interest. We empirically validate improved sampler performance driven by extended exploration. Code is available at https://github.com/eje24/ctds.

Figures

Equations

$$\pi (x)=\frac{1}{Z}\hat{\pi }(x),$$

$$\{{\hat{\pi }}_t{\}}_{t\in \mathcal{T}},{\hat{\pi }}_0\propto {\pi }_{\text{source}},\text{and}{\hat{\pi }}_1\propto {\pi }_{\text{target}},$$

$$d{X}_t$$

$$d{P}_t$$

$$d{X}_t=\left[{\mu }_t^{\theta }(X_t)+{\epsilon }_t\nabla \log p_t(X_t)\right]\mathrm{d}t+\sqrt{2{\epsilon }_t}\mathrm{d}{W}_t,$$

$${\partial }_t{\pi }_t(x_t)=-\nabla \cdot \left[{\mu }_t^{\theta }(x_t){\pi }_t(x_t)\right],\forall (x_t,t)\in {\mathbb{R}}^d\times \mathcal{T},$$

$$U_t^{\text{linear}}(x)\triangleq (1-t)U_0(x)+t{U}_1(x).$$

$$U_t^{\text{linear}}(x)\triangleq (1-t)U_0(x)+t{U}_1(x)+t(1-t)U_t^{\theta }(x),$$

$${\mathcal{L}}_{\text{PINN}}({\mu }_t^{\theta },F_t^{\theta };{\tilde{\pi }}_t)\triangleq {\int }_{\mathcal{T}\times {\mathbb{R}}^d}|{\partial }_t{F}_t^{\theta }-{\partial }_t{U}_t(x)+{\nabla }_x\cdot {\mu }_t^{\theta }(x)-({\nabla }_x{U}_t(x){)}^T{\mu }_t^{\theta }(x){|}^2{\tilde{\pi }}_t(x)\mathrm{d}x\mathrm{d}t$$

$${\hat{\pi }}_t^{\beta }(x)\triangleq e^{-U_t^{\beta }(x)},U_t^{\beta }(x)\propto \beta U_t^1(x)\forall (t,b)\in \mathcal{T}\times \mathcal{B},$$

$$U_t^{\beta }(x)$$

$${\partial }_t{\pi }_t^{{\xi }^{\star }}(x)={\nabla }_x\cdot [{\mu }_t^{\theta }(x,{\xi }^{\star }){\pi }_t^{{\xi }^{\star }}(x)],$$

$${\mathcal{L}}_{\mathrm{PINN}}^{\text{MT}}(F^{\theta },{\mu }^{\theta };\tilde{\pi })\triangleq {\int }_{\mathcal{T}\times \Xi \times {\mathbb{R}}^d}|{\partial }_t{F}_t^{\theta }-{\partial }_t{U}_t^{\xi }+{\nabla }_x\cdot {\mu }_t^{\theta }-({\nabla }_x{U}_t^{\xi }{)}^T{\mu }_t^{\theta }{|}^2{\tilde{\pi }}_t^{\xi }(x,\xi )\mathrm{d}x\mathrm{d}\xi \mathrm{d}t,$$

$${\mathcal{L}}_{\mathrm{PINN}}(F^{\star },{\mu }^{\star };\tilde{\pi })=0,$$

$${\hat{\pi }}_t(x,\xi )\triangleq e^{-U_t^{\xi }(x)+{\psi }_t(\xi )}=e^{-{\tilde{U}}_t(x,\xi )},\forall (t,\xi )\in \mathcal{T}\times \Xi$$

$${\hat{\pi }}_t^{\theta }(x,\xi )=e^{-U_t^{\xi }(x)+F_t^{\theta }(\xi )+{\psi }_t^{\prime }(\xi )}=e^{-{\tilde{U}}_t^{\theta }(x,\xi )},$$

$${\mathcal{H}}_t^{\theta }(q_t,p_t)\triangleq {\tilde{U}}_t^{\theta }(q_t)+K(q_t,p_t),K(q_t,p_t)\triangleq \frac{\beta ({\xi }_t)}{2M_x}\Vert p_t^x{\Vert }^2+\frac{1}{2M_{\xi }}\Vert p_t^{\xi }{\Vert }^2.$$

$$\mathrm{d}{X}_t$$

$$\mathrm{d}{\xi }_t$$

$$\mathrm{d}{P}_t^x$$

$$\mathrm{d}{P}_t^{\xi }$$

$$(X_0,P_0^x,{\xi }_0,P_0^{\xi })$$

$${\pi }_t^{†}(x_t,{\xi }_t,p_t^x,p_t^{\xi })={\pi }_t^{\theta }(x_t,{\xi }_t)\mathcal{N}(p_t^x;0,\frac{M_x}{\beta ({\xi }_t)}{I}_d)\mathcal{N}(p_t^{\xi };0,M_{\xi })\propto e^{-{\mathcal{H}}_t^{\theta }(q_t,p_t)}.$$

$$A_T(Q)\triangleq {\int }_0^T{\nabla }_x\cdot {\mu }_t^{\theta }(x_t,{\xi }_t)-{\partial }_t\log {\tilde{U}}_t^{\theta }(x_t,{\xi }_t)-{\mu }_t^{\theta }(x_t,{\xi }_t{)}^T{\nabla }_x{\tilde{U}}_t^{\theta }(x_t,{\xi }_t)\mathrm{d}t,$$

$${\mathbb{E}}_{(x,\xi )\sim {\pi }_t(x,\xi ))}\left[h(x,\xi )\right]=\frac{{\mathbb{E}}_{Q\sim {\mathbb{P}}_T}\left[h(Q)\exp \left(A_T(Q)\right)\right]}{{\mathbb{E}}_{Q\sim {\mathbb{P}}_T}\left[\exp \left(A_T(Q)\right)\right]}.$$

$${\mathbb{E}}_{(x,\xi )\sim {\pi }_t(x,\xi ))}\left[\exp (A_T(Q))\right]=\exp (F_0-F_T),$$

$${\mathcal{L}}_{\mathrm{PINN}}^{\text{MT}}(F,\mu ;{\pi }^{\theta })$$

$$\mathrm{d}{X}_t={\mu }_t^{\theta }(X_t)\mathrm{d}t,X_0\sim {\pi }_0\triangleq \mathcal{N}(0,5.0I_2).$$

$${\partial }_t{\pi }_t(x_t)=-\nabla \cdot \left[{\mu }_t^{\theta }(x_t){\pi }_t(x_t)\right],\forall (x_t,t)\in {\mathbb{R}}^d\times \mathcal{T},$$

$${\partial }_t\log {\pi }_t(x_t)=-\nabla {\mu }_t^{\theta }(x_t)-{\mu }_t^{\theta }(x_t{)}^T\nabla \log {\pi }_t(x_t).$$

$${\partial }_t{F}_t-{\partial }_t{U}_t(x)+\nabla {\mu }_t^{\theta }(x_t)-({\nabla }_x{U}_t(x){)}^T{\mu }_t^{\theta }(x)=0$$

$${\mathcal{L}}_{\text{PINN}}({\mu }_t^{\theta };{\tilde{\pi }}_t)\triangleq {\int }_{\mathcal{T}\times {\mathbb{R}}^d}|{\partial }_t{F}_t-{\partial }_t{U}_t(x)+{\nabla }_x\cdot {\mu }_t^{\theta }(x)-({\nabla }_x{U}_t(x){)}^T{\mu }_t^{\theta }(x){|}^2{\tilde{\pi }}_t(x)\mathrm{d}x\mathrm{d}t$$

Full text

1 Introduction

A challenging task in Bayesian statistics and natural sciences is to sample from distributions of the form

[TABLE]

where $\hat{\pi}(x)$ denotes a known unnormalized density and $Z=\int_{\mathbb{R}^{d}}\hat{\pi}(x)\,\mathop{}\!\mathrm{d}x$ denotes the unknown partition function. Traditional MCMC-based sampling approaches, including Metropolis-Hastings and Hamiltonian Monte-Carlo (HMC), have been observed to suffer from pseudo-ergodic behavior for high-dimensional or otherwise complex choices of the density $\hat{\pi}$. Techniques such as annealed importance sampling (AIS) (Neal, 1998), and sequential Monte-Carlo-based approaches (Del Moral et al., 2006) work instead over an annealed sequence of densities $\{\hat{\pi}_{t}\}$ which interpolates a tractable source $\hat{\pi}_{0}$ and the desired target $\hat{\pi}_{1}=\hat{\pi}$. Recently, neural samplers have been proposed which leverage a wide array of learned transport techniques, parameterized with neural networks, to optimize the annealed auxiliary distribution (Thin et al., 2021; Wu et al., 2020; Geffner & Domke, 2021; Zhang et al., 2021; Doucet et al., 2022; Geffner & Domke, 2023; Vargas & Nüsken, 2023; Richter & Berner, 2024; Albergo & Vanden-Eijnden, 2024).

In this work, we consider neural samplers which learn a control or transport to generate the evolution of a prespecified annealing path $\hat{\pi}_{t}$ via a so-called physics-informed neural network (PINN) loss, which penalizes the degree to which the learned vector field fails to satisfy the continuity equation (Raissi et al., 2019; Sun et al., ; Máté & Fleuret, 2023; Albergo & Vanden-Eijnden, 2024). The off-policy nature of the PINN loss means that a crucial design choice in PINN-based sampler training is the proposal distribution $\tilde{\pi}_{t}$ by which samples are generated at which to evaluate the PINN loss. To this end, recent work has proposed to construct $\tilde{\pi}_{t}$ as the marginals of annealed Langevin dynamics steered with the partially-learned control. By augmenting the learned control in this manner, samples are driven towards high density regions of the density path, allowing for boostrapped training in which exploration gradually improves with the quality of the learned control. However, for many choices of density path, including the commonly used linear interpolation (see eq. 6), isolated modes, high energy barriers, and other problematic features of the energy landscape obstruct this exploration process, resulting in poor sampler performance Máté & Fleuret (2023).

To alleviate this problem, we turn toward tempering techniques, which are widely used in e.g., molecular dynamics to overcome high energy barriers between separate metastable states. Their success is based primarily on the insight that such barriers are lowered when the temperature is raised (Marinari & Parisi, 1992; Swendsen & Wang, 1986; Hansmann, 1997; Sugita & Okamoto, 1999; Lenner & Mathias, 2016). Specifically, we draw inspiration from recently proposed continuous tempering approaches, in which Langevin dynamics is run over an augmented state space given by both position and a continuous temperature variable Gobbo & Leimkuhler (2015); Graham & Storkey (2017); Luo et al. (2019). We propose to use such a continuously tempered dynamics as the basis for the proposal distribution, in doing so exploiting lower energy barriers at higher temperatures to drive exploration and thus better training. Specifically, we make the following contributions:

We extend the notion of a time-indexed density path to a time-and-temperature-indexed density continuum. We propose a multi-temperature PINN objective over such a continuum which allows for sampler training to be amortized across temperatures. 2. 2.

To learn dynamics across the time-temperature continuum, we propose the continuously tempered diffusion sampler (CTDS), a scheme that allows for efficient exploration during sampler training by leveraging continuous tempering techniques. 3. 3.

We validate the performance of CTDS-trained samplers against those trained with existing proposals and demonstrate improved sampler performance.

2 Background

Notation

We will denote normalized and unnormalized probability densities by $\pi$ and $\hat{\pi}$ respectively. Further, we will call $F=-\log Z$ the free energy and $U=-\log\hat{\pi}$ the energy, so that $\pi(x)=\frac{1}{Z}\hat{\pi}(x)=e^{-U(x)+F}$. Additionally, we will denote by $\mathcal{T}$ the unit-time interval $[0,1]$, and by $\mathcal{P}(\mathcal{X})$ the space of twice-differentiable, fully-supported, normalized densities on the space $\mathcal{X}$. When we work over phase space coordinates, we will denote the momentum as $p$.

2.1 Controlled Langevin Annealing

To overcome pseudo-ergodic behavior in traditional MCMC-based sampling approaches (Metropolis et al., 1953; Neal, 2012), techniques such as annealed importance sampling (AIS) Neal (1998), and later sequential Monte Carlo (SMC) based approaches Del Moral et al. (2006), work by simulating MCMC over an annealed density path

[TABLE]

interpolating between $\pi_{\text{source}}$ and $\pi_{\text{target}}$. Recent work has since emphasized combining the sampling process with a learned transformation using e.g., normalizing flows (Wu et al., 2020; Arbel et al., 2021; Matthews et al., 2023), and, in the continuous time setting, a learned vector field (Vargas & Nüsken, 2023; Tian et al., 2024; Fan et al., 2024; Albergo & Vanden-Eijnden, 2024). Following the second approach, and given a prescribed density path as in eq. 2, we may consider a controlled (underdamped) Langevin dynamics

[TABLE]

as in (Zhong et al., 2024; Albergo & Vanden-Eijnden, 2024), as well as the overdamped limit given by

[TABLE]

as in Albergo & Vanden-Eijnden (2024); Vargas & Nüsken (2023), where in both cases $\mu_{t}^{\theta}$ is the learned, time-dependent control.111The underdamped parameterization is slightly different than analogous formulations presented in Geffner & Domke (2023); Vargas & Nüsken (2023). Here, $W_{t}$ denotes a standard Brownian motion on $\mathbb{R}^{d}$, ($d=d_{x}$ being the dimensionality of the data), and $\gamma_{t}$, $\varepsilon_{t}$, and $M$ denote scaling and damping coefficients, and particle weight, respectively. The advantage of the controlled Langevin schemes in eq. 3 and eq. 4 is that the Langevin dynamics can be seen as correcting for an imperfectly learned control, and thus allowing for a bootstrapped training procedure in which the marginals of the forward process converge on the desired density path as the control improves (Albergo & Vanden-Eijnden, 2024). In particular, perfect sampling is recovered as when $\gamma_{t}\to\infty$ in eq. 3 and $\varepsilon_{t}\to\infty$ in eq. 4, or when the the control is learned perfectly so that it satisfies the continuity equation

[TABLE]

and thus facilitates an instantaneous transition along the density path for any choice of $\gamma_{t}$ or $\epsilon_{t}$. Several popular families of training objective have emerged to learn such a control, including minimizing a path KL divergence (Vargas & Nüsken, 2023; Zhang & Chen, 2022; Berner et al., 2022; Vargas et al., 2023), log-variance objectives (Richter et al., 2020; Nüsken & Richter, 2023; Vargas & Nüsken, 2023; Richter & Berner, 2024), and physics-informed neural network (PINN) losses (Raissi et al., 2019; Wang et al., 2023; Sun et al., ). The last of these is the focus of this work and will be discussed more below.

The success of controlled annealing approaches depends heavily on the choice of density path, for which the most common choice is linear path

[TABLE]

While nearly universally available, the linear path has been shown to suffer from certain regularity issues including the so-called teleportation of mass phenomenon described in (Máté & Fleuret, 2023), in certain cases causing any control satisfying eq. 5 to exhibit singularities making it hard to integrate numerically, let alone learn. One approach toward alleviating these issues is proposed in Máté & Fleuret (2023), in which an additional learned term is added to the linear path from eq. 6 to obtain

[TABLE]

which we shall refer to as the learned path.

2.2 PINN-Based Objectives

In the sampler setting, PINN-based objectives seek to learn the control $\mu_{t}^{\theta}$ by minimizing the expected pointwise residual error of the continuity equation eq. 5. While such an objective, as originally formulated, requires access to the generally intractable free energy $F_{t}$, the free energy may be jointly learned along with the control, yielding the modified PINN objective

[TABLE]

for which the true free energy $F_{t}$ can be shown to be the unique minimizer (Máté & Fleuret, 2023; Albergo & Vanden-Eijnden, 2024). More details on the PINN loss, including a derivation, can be found in appendix A. Besides the choice of density path $\pi_{t}$, the fundamental object underpinning eq. 8 is the proposal distribution $\tilde{\pi}_{t}$, by which samples are drawn at which to evaluate the objective. A good choice of proposal distribution $\tilde{\pi}_{t}$ is given by the density path itself (i.e., when $\tilde{\pi}_{t}=\pi_{t}$), but such a choice is nearly always intractable. As an alternative, recent work has proposed to construct such a proposal distribution using a controlled annealed Langevin scheme (Albergo & Vanden-Eijnden, 2024).

Key Idea 2.1** (Proposal via controlled Langevin annealing).**

Construct the PINN proposal $\tilde{\pi}$ as the marginals of a controlled Langevin annealing process, as is given in eq. 3 and eq. 4. Intuitively, as the quality of the learned control increases, so will the mode-coverage of the proposal. Optionally, one may additionally reweigh samples toward the ground-truth marginal $\pi$ using a Jarzynski-like equality (Albergo & Vanden-Eijnden, 2024).

2.3 Tempering Techniques for Enhanced Exploration

One family of techniques to improving exploration are tempering-based methods, based on the fundamental insight that energy barriers are lowered and mixing times are improved when one raises the temperature of the target to obtain $\pi_{t}^{\beta}(x)\propto\pi_{t}(x)^{\beta}$, for inverse temperature $\beta\triangleq\frac{1}{T}<1$. To exploit this insight, approaches such as parallel tempering Swendsen & Wang (1986); Hansmann (1997); Sugita & Okamoto (1999) and simulated tempering Marinari & Parisi (1992) have been proposed. These are MCMC-based sampling schemes over a temperature ladder along with mechanisms by which samples may jump between temperatures, thereby allowing faster-mixing dynamics at higher temperatures to drive convergence at a lower temperature of interest.

More recently, continuous-tempering schemes have been proposed involving a Hamiltonian dynamics over a temperature-augmented extended state thereby allowing for a smoothly-varying temperature variable Gobbo & Leimkuhler (2015); Graham & Storkey (2017); Luo et al. (2019), and demonstrating improvements over simulated and parallel tempering Graham & Storkey (2017). Such continuous tempering schemes have found utility in MD simulations Lenner & Mathias (2016), Bayesian posterior inference Luo et al. (2019), and in improving the convergence of piece-wise deterministic Markov processes Sutton et al. (2022). In this work, we extend the application of continuous tempering to the annealed neural sampler setting.

3 Continuously Tempered Diffusion Samplers

In this section, we present our main contribution: combining the controlled Langevin annealing schemes of section 2.1 with the continuous tempering approaches outlined in section 2.3, obtaining a new, tempered class of annealing schemes over a continuously-varying temperature variable, which we refer to as continuously tempered diffusion samplers (CTDS).

3.1 Tempered Density Continuums

Recall that a density path is a $t$-indexed annealing path $\hat{\pi}_{t}(x)\triangleq e^{-U_{t}(x)}$ which agrees with a prescribed source density $\hat{\pi}_{0}$ and target density $\hat{\pi}_{1}$. It is then natural that we extend the notion of density path to the tempered setting by affixing an additional indexing variable - the inverse-temperature variable $\beta$ - to obtain the annealing continuum

[TABLE]

where we have defined $\mathcal{B}\triangleq[\beta_{\text{min}},1]$ as the range of inverse-temperatures $\beta$ of interest. We refer to the resulting continuously-indexed collection of densities $\{\hat{\pi}_{t}^{\beta}(x)\}_{t,\beta\in\mathcal{T}\times\mathcal{B}}$ as a density continuum. Given a target density $U^{\text{target}}$ of interest, two such density continuums are the linear continuum and the learned continuum defined respectively for a prescribed target density $U_{1}^{1}(x)=U^{\text{target}}(x)$ by $\pi_{0}^{\beta}(x)=\mathcal{N}(0,\frac{1}{\beta}I_{d})$ (thereby defining $U_{0}^{\beta}$) and

[TABLE]

for all $(t,b)\in\mathcal{T}\times\mathcal{B}$. In the next two sections, we will construct an annealed Langevin dynamics over such a tempered density continuum. Doing so directly over $\mathcal{B}$ presents two challenges: First, it is difficult to handle the boundaries of $\mathcal{B}$ (at $\beta_{\text{min}}$ and $1.0$) without introducing some sort of boundary condition on the Langevin dynamics, or with the introduction of some confining potential, neither of which are particularly desirable. Second, samples from any continuous marginal density over $\beta$ would almost surely not be equal to either $\beta=\beta_{\text{min}}$ (thereby discouraging mixing behavior at high temperatures) or $\beta=1.0$ (thereby not providing training signal at the original temperature of interest). We therefore follow the lead of Gobbo & Leimkuhler (2015) and reparameterize the inverse-temperature coordinate by pulling back to the space $\Xi\triangleq\mathbb{R}$ via a carefully defined continuous mapping $\beta(\xi):\Xi\to\mathcal{B}$ in a manner which resolves both issues; see section B.1 for details.

3.2 PINN Objectives Across Temperature

For fixed $\xi=\xi^{\star}$ (the case of $\beta(\xi^{\star})=1$ being of particular interest), we recover the standard annealed sampling problem, in which we would like to learn a control $\mu_{t}^{\theta}(x,\xi^{\star})$ which satisfies the continuity equation

[TABLE]

corresponding to $\{\hat{\pi}_{t}^{\xi^{\star}}(x)\}_{t\in\mathcal{T}}$ (and in turn, to the path $\{\hat{\pi}_{t}^{\beta(\xi^{\star})}(x)\}_{t\in\mathcal{T}}$). To obtain such a control, let us consider a multi-temperature variant of the PINN loss over the space $\mathbb{R}^{d}\times\mathcal{T}\times\Xi$ from eq. 8, given by

[TABLE]

for free energy estimate $F:\mathcal{T}\times\Xi\to\mathbb{R}$, control $\mu:\mathcal{T}\times\mathbb{R}^{d}\times\Xi\to\mathbb{R}^{d}$, and proposal distribution $\tilde{\pi}_{t}(x,\xi)\in\mathcal{P}(\mathbb{R}^{d})$. Now, defining the free energy $F_{t}(\xi)$ as $F_{t}(\xi)\triangleq\int_{\mathbb{R}^{d}}e^{-U_{t}^{\xi}(x)}\,\mathop{}\!\mathrm{d}x$, we propose the following result which characterizing the minimizing $F$ of eq. 13.

Theorem 3.1**.**

Let $\tilde{\pi}_{t}(x,\xi)\in\mathcal{P}(\mathbb{R}^{d}\times\Xi)$ and suppose there exists a pair $(F^{\star},\mu^{\star})$ which satisfies

$\mathcal{L}_{\operatorname{PINN}}(F^{\star},\mu^{\star};\tilde{\pi})=0,$

as well as the boundary condition $F_{0}^{\star}(\xi)=F_{0}(\xi)$ for all $\xi\in\Xi$. In this case, $F^{\star}$ is unique and is given by the free energy $F_{t}(\xi)$.

The result from theorem 3.1 extends an existing result in the single-temperature setting (Máté & Fleuret (2023), Lemma 1) to the multi-temperature setting, and a short proof can be found in section B.3. Furthermore, theorem 3.1 establishes that we may learn the control $\mu_{t}^{\theta}(x,\xi)$ by simultaneously learning the free energy $F_{t}^{\theta}(\xi)$. To construct such a proposal $\tilde{\pi}_{t}(x,\xi)$, we look to finally exploit the tempered nature of the density continuum $\{\hat{\pi}_{t}^{\xi}(x)\}$, as we discuss next.

3.3 Continuously Tempered Diffusion Samplers

In this section, we will at last formulate our main contribution. First, we will realize the density continuum as a density path by introducing a fictitious, time-dependent density over $\xi$. Second, we will construct the continuously tempered diffusion sampler, a controlled annnealed Langevin dynamics over such a density path. By introducing some time-dependent, normalized potential $\psi_{t}(\xi):\mathcal{T}\times\Xi\to\mathbb{R}$, we may realize the continuum $\{\hat{\pi}_{t}^{\xi}(x)\}$ as a joint density

[TABLE]

over the augmented state $(x,\xi)$, where we have defined $\tilde{U}_{t}(x,\xi)\triangleq U_{t}(x,\xi)-\psi_{t}(\xi)$. Then it is easily seen that the marginal $\pi_{t}(\xi)$ is given by $\pi_{t}(\xi)\propto e^{\psi_{t}(\xi)-F_{t}(\xi)}$. As a consequence, we may rewrite $\psi_{t}^{\theta}(\xi)=\psi_{t}^{\prime}(\xi)+F_{t}^{\theta}(\xi)$, where $F_{t}^{\theta}(\xi)$ is a learned approximation of the free energy as in eq. 8, and so that as $F_{t}^{\theta}(\xi)\to F_{t}(\xi)$, $\pi_{t}(\xi)$ becomes distributed proportionally to $e^{\psi_{t}^{\prime}(\xi)}$. In doing so, we obtain

[TABLE]

so as to “cancel-out” the bias of the unknown free energy $F_{t}(\xi)$ as the quality of our learned approximation $F_{t}^{\theta}(\xi)$ increases. In section B.2, we characterize more explicitly the construction of $\psi_{t}^{\prime}(\xi)$ in practice, and in particularly the inclusion of a confining potential term. Denoting by $q_{t}=(x_{t},\xi_{t})$ and $p_{t}=(p_{t}^{x},p_{t}^{\xi})$ we may now define the non-separable Hamiltonian

[TABLE]

With the inclusion of the learned control $\mu_{t}^{\theta}$, we obtain the controlled annealed underdamped Langevin dynamics given by

[TABLE]

where the extended joint density $\pi_{t}^{\dagger}$ is given by

[TABLE]

In eq. 17, we have used $(\gamma_{t}^{x},\gamma_{t}^{\xi})$, $(\varepsilon_{t}^{x},\varepsilon_{t}^{\xi})$, and $(M_{x},M_{\xi})$ to denote the respective scaling, damping, and mass coefficients for each of $x$ and $\xi$, and have denoted by $W_{t}^{x}$ and $W_{t}^{\xi}$ standard Brownian motions on $\mathbb{R}^{d}$ and $\Xi$ respectively. We refer to this scheme as a continuously tempered diffusion sampler (CTDS). We may then construct the proposal $\tilde{\pi}_{t}^{\xi}$ as the marginals of the forward process given by eq. 17.

3.4 A Continuously Tempered Controlled Jarzynski Result

For imperfectly learned control $\mu_{t}^{\theta}(x,\beta(\xi))$ and finite scaling coefficients $\gamma_{t}^{x}$ and $\gamma_{t}^{\xi}$, the proposal distribution $\tilde{\pi}_{t}^{\beta}$ induced by the marginals of the CTDS forward process in eq. 17 will lag behind the true density path $\pi_{t}(x,\xi)$. To correct for this discrepancy, we derive a controlled Jarzynski equality, as in (Vargas & Nüsken, 2023; Albergo & Vanden-Eijnden, 2024), so that we may reweigh samples from eq. 17 to the correct joint density $\pi_{t}^{\theta}(x,\xi)$ described in eq. 15. Specifically, for fixed $T\in\mathcal{T}$, we may consider the work functional as in (Vaikuntanathan & Jarzynski, 2008; Zhong et al., 2024; Albergo & Vanden-Eijnden, 2024), given by

[TABLE]

where we have denoted $q_{t}=(x_{t},\xi_{t})$ and $Q=\{q_{t}\}_{t\in[0,T]}$. Defining $\mathbb{P}_{T}$ as the path measure of the forward process eq. 17 on the interval $[0,T]\subseteq\mathcal{T}$, and letting $F_{t}$ denote the free energy of the density path $\hat{\pi}_{t}(x,\xi)$, we have the following controlled Jarzynski identity.

Theorem 3.2** (Continuously Tempered Controlled Jarzynski Equality).**

For $T\in(0,1]$, let $h:\mathbb{R}^{d}\times\Xi\to\mathbb{R}$ denote some observable. Then

$\mathbb{E}_{(x,\xi)\sim\pi_{t}(x,\xi))}\left[h(x,\xi)\right]=\frac{\mathbb{E}_{Q\sim\mathbb{P}_{T}}\left[h(Q)\exp\left(A_{T}(Q)\right)\right]}{\mathbb{E}_{Q\sim\mathbb{P}_{T}}\left[\exp\left(A_{T}(Q)\right)\right]}.$

(20)

In particular,

$\mathbb{E}_{(x,\xi)\sim\pi_{t}(x,\xi))}\left[\exp(A_{T}(Q))\right]=\exp(F_{0}-F_{T}),$

(21)

where we have defined $F_{t}\triangleq\int_{\Xi\times\mathbb{R}^{d}}e^{-\tilde{U}_{t}^{\theta}(x,z)}\,\mathop{}\!\mathrm{d}x\mathop{}\!\mathrm{d}z$.

We note that theorem 3.2 closely resembles recently established controlled Jarzynski results from (Vargas & Nüsken, 2023; Albergo & Vanden-Eijnden, 2024), that it is in fact a corollary of a more general controlled variant of the Crooks fluctuation theorem as established in (Vaikuntanathan & Jarzynski, 2008; Vargas & Nüsken, 2023; Zhong et al., 2024), and that this resemblance is due entirely to the CTDS construction as an annealed Langevin dynamics. We present a self-contained proof of theorem 3.2 by way of an analogous Crooks identity in appendix C. We conclude by demonstrating how theorem 3.2 can be used to reweigh samples from the forward process eq. 17 to match the joint density $\pi_{t}^{\theta}(x,\xi)$ defined in eq. 15. This also allows us to reweight the loss and up-weigh “important” samples during training. Observe that by theorem 3.2, it follows that

[TABLE]

Details can be found in section C.4. We conclude from eq. 22 that we may reweight samples from the forward process eq. 17 so as obtain the proposal $\pi^{\theta}_{t}(x,\xi)$, which as we recall from eq. 15, converges to $\pi_{t}^{\xi}(x)\pi_{t}(\xi)=\pi_{t}^{\xi}(x)e^{-\psi^{\prime}(\xi)}$ as $F_{t}^{\theta}(\xi)\to F_{t}(\xi)$. At the start of training, or when $F_{t}^{\theta}(\xi)\approx F_{t}(\xi)$ otherwise fails to hold, it is observed that samples tend to “pool” at either higher or lower temperatures (depending on the target). This effect diminishes as the quality of the free energy approximation increases; see fig. 2. It should also be noted that since theorem 3.2 reweights toward the joint distribution $\pi_{t}^{\theta}(x,\xi)$, a poor estimate of the free energy will bias the reweighting toward particular values of $\xi$.

4 Experiments

4.1 40-Mode Gaussian Mixture

To investigate the benefit of the CTDS proposal, we compare the performance of samplers trained with various choices of proposal on a modified version challenging 40-mode Gaussian mixture from (Midgley et al., 2023).222We take the standard deviation of the isotropic mixture components to be $\sigma=0.25$. We emphasize upfront, and to the strongest degree possible, that great care must be taken in comparing the results we report to existing work. Different works consider different density paths and various degrees of complexity for this problem. For example, learning a sampler for such a target can be made artificially easier by e.g., increasing the scale of the source distribution $\pi_{0}$. In this experiment, we choose to take our source distribution $\pi_{0}=\mathcal{N}(0,5.0I_{2})$, thereby ensuring that modes discovered during training are done so on account of the proposal, rather than an advantageously chosen initialization. As our density path, we take the learned path given in eq. 7 (and for CTDS, the multi-temperature analog given by eq. 11).

Training.

As a baseline, we consider the proposal obtained by simulating the control by itself with no Langevin dynamics, as in (Máté & Fleuret, 2023), and which we label as the reference proposal. We additionally consider proposals obtained from the controlled overdamped (eq. 4, with $\varepsilon_{t}=50.0$) and controlled underdamped/inertial (eq. 3, with $(\gamma_{t},\varepsilon_{t}=50.0,2.0)$) Langevin annealing dynamics as in (Albergo & Vanden-Eijnden, 2024), which we train with and without use of the Jarzynski equality. Finally, we consider our CTDS proposal (eq. 17, with $(\gamma_{t}^{x},\gamma_{t}^{\xi},\varepsilon_{t}^{x},\varepsilon_{t}^{\xi})=(50.0,2.0,5.0,2.0)$) both with and without reweighting with the Jarzynski equality (see section 3.4). We train using a replay buffer, re-sampling once per epoch for 1250 epochs for a total of 125000 training iterations at which the PINN loss is evaluated at 6250 randomly sampled elements of the buffer. The control, learnable density path, and free energy are all parameterized as simple MLPs. Integration is performed with an Euler-Maruyama solver with step size $\Delta t=0.002$. We additionally utilized Gaussian Fourier features as in Tancik et al. (2020) to encode spatial, temporal, and temperature variables, and which we found to improve training stability and sampler performance. More details can be found in section D.1

Evaluation.

Post-training, we sample by simulating the dynamics

[TABLE]

In the CTDS case, we take $\mu_{t}^{\theta}(X_{t})=\mu_{t}^{\theta}(X_{t},\beta(\xi)=1.0)$. We report mean and std. values for both 2-Wasserstein ($\mathcal{W}_{2}$), evidence lower bound (ELBO), and evidence upper bound (EUBO), for each using $N=2500$ samples for each of $10$ trials, as can be found in fig. 3. We observe that the CTDS outperform their baseline, overdamped and underdamped counterparts, achieving both lower $\mathcal{W}_{2}$ values and ELBOs, and that the Jarzynski-reweighted CTDS achieves the lowest EUBO. Curiously, while reweighted NETS based proposals seem to uniformly outperform non-reweighted counterparts, the same trend does not hold for CTDS, and in particular the non-reweighted CTDS proposal achieves both lower $\mathcal{W}_{2}$ and higher ELBO. We hypothesize that this is due in part to the effects discussed in section 3.4 (and visualized in fig. 2), by which an initially poor free energy approximation biases both samples and importance weights are initially biased towards high temperatures. A visualization comparing the baseline, reweighted overdamped NETS, and reweighted CTDS is provided in fig. 1.

5 Conclusion

In this work, we have proposed continuously tempered diffusion samplers (CTDS), a novel family of controlled Langevin annealing processes based on continuous tempering techniques from e.g., molecular dynamics. When combined with PINN-based sampler training, CTDS-based proposals outperform existing proposals based on overdamped and underdamped Langevin annealing on the challenging 40-mode Gaussian mixture from Midgley et al. (2023). We leave it to future work to extend CTDS to higher dimensional or otherwise complex targets and other scientific applications.

6 Acknowledgements

We thank Jiajun He, Yuanqi Du, and Michael Albergo for helpful discussions regarding PINN-based sampler training over the course of this project. We additionally acknowledge support from the Machine Learning for Pharmaceutical Discovery and Synthesis (MLPDS) consortium, the DTRA Discovery of Medical Countermeasures Against New and Emerging (DOMANE) threats program, and the NSF Expeditions grant (award 1918839) Understanding the World Through Code

Appendix A A Brief Overview of PINN Objectives

We briefly derive the PINN objective from eq. 8. Recall the the continuity equation

[TABLE]

from eq. 5. Dividing both sides by $\pi_{t}(x_{t})$ yields

[TABLE]

Finally, plugging. in $\log\pi_{t}(x_{t})=-U_{t}(x)+F_{t}$ and rearranging yields

[TABLE]

from which we obtain the PINN objective

[TABLE]

for fully-supported proposal distribution $\tilde{\pi}_{t}$ (Raissi et al., 2019; Sun et al., ; Wang et al., 2023). The free energy $F_{t}$ however is generally unknown, and thus must be estimated (Tian et al., 2024; Fan et al., 2024), or jointly-learned with the control Máté & Fleuret (2023); Albergo & Vanden-Eijnden (2024). In line with this second approach, subsequent work has established a modified PINN objective by which the free energy may be simultaneously learned along with the control, viz.,

[TABLE]

for which the true free energy $F_{t}$ can be shown to be the unique minimizer Máté & Fleuret (2023); Albergo & Vanden-Eijnden (2024). To better facilitate the learning of the free energy, past work has proposed a curriculum training procedure in which the domain of integration along the time dimension is slowly annealing from zero to one, so as to initially learn the free energy on some smaller interval before all of $\mathcal{T}$.

Appendix B Extended Details on Continuously Tempered Diffusion Samplers

In this section, we elaborate an additional design choices for continuously tempered diffusion samplers, including the choice of reparameterization, the use of confining and biasing potentials for temperature. Additionally, we provide a proof of theorem 3.1.

B.1 Temperature Reparameterization Details

Let us recall the definition of a density continuum over $(t,\beta)\in\mathcal{T}\times\mathcal{B}$ frorm eq. 9, given by

[TABLE]

where we have defined $\mathcal{B}\triangleq[\beta_{\text{min}},1]$ as the range of inverse-temperatures $\beta$ of interest. Our goal of constructing an annealed Langevin dynamics over such a continuum necessitates a joint density over both $x$ and $\beta$, and in turn, an associated marginal distribution over $\beta$. Working directly over $\beta$ is thus undesirable for two reasons: First, assuming that such a marginal over $\beta$ is continuous, samples from this marginal would almost surely not be equal to either $\beta_{\text{min}}$ or $1.0$, the two temperatures of particular relevance. This is problematic because it would mean that any proposal constructed with e.g., Langevin dynamics over such a joint density would provide little training signal at the desired temperature of $1.0$, nor would it spend much exploring at $t=\beta_{\text{min}}$. Secondly, we would like to constrain the temperature range to only the interval $[\beta_{\text{min}},1]$, and enforcing these boundary conditions directly presents complications. We therefore follow the lead of Gobbo & Leimkuhler (2015) by reparameterizing via

[TABLE]

for all $\xi\in\Xi\triangleq\mathbb{R}$. A visualization of eq. 30 is given in fig. 4.

B.2 Biasing and Confining Potentials

Recall the joint density $\hat{\pi}_{t}(x,\xi)$ given in eq. 15 as

[TABLE]

In practice, and to encourage the $\xi$-component to stay close to the critical interval $[-\Delta^{\prime},\Delta^{\prime}]$ (see eq. 30), we may introduce choose $\psi_{t}^{\prime}(\xi)$ so as to include a confining potential $\psi^{\text{conf}}(\xi)$. In this paper, we take

[TABLE]

where $\tilde{\Delta}$ (chosen to be slightly larger than $\Delta^{\prime}$ from eq. 30) and $\eta>0$ (denoting the sharpness of the confining potential) are hyperparameters. Additionally, one may consider alleviating a poorly initialized $F_{t}^{\theta}$ using an additional biasing force term $\psi_{t}^{\text{bias}}$ updated continuously throughout training to balance out the measured marginal distribution over $\xi$, so as to obtain the general form

[TABLE]

In practice, we find that utilizing a biasing potential $\psi_{t}^{\text{bias}}(\xi)$ to be cumbersome and to provide little benefit over simpler solutions such as taking e.g., a larger value of $\sigma_{\xi}^{2}$ in eq. 18.

B.3 Characterizing Minimizers of the Multi-Temperature PINN Loss

Proof of theorem 3.1.

We apply the basic argument of (Máté & Fleuret, 2023) in a $\xi$-pointwise fashion. Let us suppose that $(\mu^{\star},F^{\star})$ satisfies $\mathcal{L}_{\operatorname{PINN}}(F^{\star},\mu^{\star};\tilde{\pi})=0$. Then, for all $(x,\xi,t)\in\mathbb{R}^{d}\times\Xi\times\mathcal{T}$ (except perhaps on a set of measure zero), we must have

[TABLE]

Now, observe that the true free energy $F_{t}(\xi)$ satisfies

[TABLE]

However, by eq. 34, we have that

[TABLE]

Plugging eq. 36 into eq. 35 yields

[TABLE]

where the second equality follows from the divergence theorem and the third from the definition of $Z_{t}(\xi)$. We conclude that $\partial_{t}F_{t}^{\star}(\xi)=\partial_{t}F_{t}(\xi)$ on $\mathcal{T}$, and together with the boundary condition $F_{0}^{\star}(\xi)=\partial_{0}F_{t}(\xi)$, this is enough to prove the desired result. ∎

Appendix C Controlled Crooks and Jarzynski Results for Continuously Tempered Dynamics

In this section, we establish a controlled Crooks fluctuation theorem for a controlled underdamped Langevin dynamics given by a non-separable Hamiltonian, and as a corollary obtain an associated controlled Jarzynski equality. Our results partially generalize existing work for separable Hamiltonians (Vaikuntanathan & Jarzynski, 2008; Vargas & Nüsken, 2023; Zhong et al., 2024; Albergo & Vanden-Eijnden, 2024).

C.1 Background on Stochastic Integration

In the remainder of this section, we’ll make heavy use of forward and backward Itô integrals, as well as the Stratonovich integral. We largely follow the notation of (Vargas & Nüsken, 2023), and refer thereto for a more elaborate and technical discussion. For now, we summarize the three important Stochastic integrals - forward and backward Itô, and Stratonovich - in fig. 5, and provide some useful identities which relate this integrals, and which we shall make heavy use of. In particular, let us consider the prototypical forward SDE

[TABLE]

as shorthand for the forward Itô integral $\overrightarrow{X}_{t}=\overrightarrow{X}_{0}+\int_{0}^{t}\mu_{s}(\overrightarrow{X}_{s})\mathop{}\!\mathrm{d}s+\int_{0}^{t}\sigma_{s}\mathop{}\!\mathrm{d}W_{s}$, and

[TABLE]

as shorthand for the backward Itô integral $\overleftarrow{X}_{t}=\overleftarrow{X}_{T}-\int_{t}^{T}\mu_{s}(\overleftarrow{X}_{s})\mathop{}\!\mathrm{d}s-\int_{t}^{T}\sigma_{s}\overleftarrow{\mathop{}\!\mathrm{d}}W_{s}$. We conclude with two useful identities which we make heavy use of, and which we provide intuitive, discretized intuition for. We again direct the reader to e.g., (Vargas & Nüsken, 2023; Kunita, 2019) for a more technical treatment. First, as is directly observed from the discretizations shown in fig. 5, we may relate the forward and backward Itô integrals to a corresponding Stratonovich integral via

[TABLE]

Second, and less intuitively obvious, is the fact that for $X=\overrightarrow{X}$ from eq. 37,

[TABLE]

Writing

[TABLE]

plugging this into eq. 40, and allowing the contribution of all $o(\Delta t)$ terms to vanish, we obtain

[TABLE]

C.2 Radon-Nikodym Derivatives of Forward and Backward Langevin Dynamics

Let us start by considering some time-dependent non-separable Hamiltonian of the form

[TABLE]

Then, we may define the controlled underdamped Langevin dynamics as the forward process $\{\overrightarrow{X}_{t}\}_{t\in[0,1]}$ given by

[TABLE]

where $X_{0}\sim\pi_{0}\propto e^{-\mathcal{H}_{0}}$ and where $\{W_{t}\}$ is a standard Brownian motion on $\mathbb{R}^{d}$. Here, we have defined the scaling and damping coefficients

[TABLE]

Then, for $T\in(0,1]$ may similarly consider the backward process $\{\overleftarrow{X}_{t}\}_{t\in[0,1]}$.

[TABLE]

where $\overleftarrow{X}_{T}\sim\pi_{T}\propto e^{-\mathcal{H}_{T}}$ and where $\overleftarrow{\mathop{}\!\mathrm{d}}$ denotes the backward Itô differential (see section C.1). Let $\mathbb{P}_{T}$ and $\mathbb{Q}_{T}$ denote the path measures of $\{\overrightarrow{X}_{t}\}_{t\in[0,T]}$ and $\{\overleftarrow{X}_{t}\}_{t\in[0,T]}$. In what follows, we will denote by $Q_{t}$ and $P_{t}$ the $q$ and $p$ components respectively of the process $X_{t}$, and when necessary, by $\overleftarrow{Q}_{t}$ and $\overleftarrow{P}_{t}$ the respective $q$ and $p$ components of the backward process $\overleftarrow{X}_{t}$. In this section we will establish the following result characterizing the Radon-Nikodym derivative (RND) of the forward and backward processes in eq. 45 and eq. 47.

Theorem C.1** (RND of Forward and Backward Processes).**

The path measures $\mathbb{P}_{T}$ and $\mathbb{Q}_{T}$ are absolutely continuous with respect to one another, and for a sufficiently well-behaved proces $\{X_{t}\}=\{(Q_{t},P_{t}\}$, the Radon-Nikodym derivative $\frac{\mathop{}\!\mathrm{d}\mathbb{Q}_{T}}{\mathop{}\!\mathrm{d}\mathbb{P}_{T}}(X)$ is given by

[TABLE]

One challenging aspect of eq. 45 (and, in turn, of eq. 47) is the fact that we are only injecting noise into the momentum component, so that the driving Brownian motion $W_{t}$ has dimension only half that of the phase space on which these two processes are defined. Accordingly, we recall the following Girsanov result which accommodates such diffusion processes involving a lower dimensional driving Brownian motion.

Theorem C.2** (Multivariate Extension of Theorem 7.19, Liptser & Shiryaev (2013)).**

Let $T\in[0,1]$ and consider the two diffusion processes $\{Y\}_{t\in[0,T]},\{Z\}_{t\in[0,T]}$, on $\mathbb{R}^{n}$ given by

[TABLE]

with $Y_{0}=Z_{0}$ fixed, so that $\mu,\phi,\Sigma$ satisfy mild regularity conditions, where $\{W_{t}\}$ is a standard Brownian motion on $\mathbb{R}^{k}$, and where $\Sigma_{t}\in\mathbb{R}^{n\times k}$. Additionally, let $\mathbb{Q}_{T}^{Y}$ and $\mathbb{Q}_{T}^{Z}$ denote the respective path measures of $Y$ and $Z$. Then, if

[TABLE]

both $\mathbb{Q}_{T}^{Y}$ and $\mathbb{Q}_{T}^{Z}$ almost surely, then $\mathbb{Q}_{T}^{Y}\sim\mathbb{Q}_{T}^{Z}$ are absolutely continuous with respect to one another and their Radon-Nikodym derivative is given by

[TABLE]

where $A^{+}$ denotes the Moore-Penrose pseudoinvere of a matrix $A$.

We may now proceed to the proof of Theorem C.1.

Proof of theorem C.1.

First, let us recall the forward process from eq. 45, given as

[TABLE]

and a corresponding backward process given by

[TABLE]

whose path measures on $[0,T]$ we have denoted by $\mathbb{P}_{T}$ and $\mathbb{Q}_{T}$ respectively. At a high level, we will employ the same trick as was utilized concurrently in both Vargas & Nüsken (2023) and Richter & Berner (2024) of introducing forward and backward references processes with identical path measures, so as to obtain an expression, via cancellation, of the desired RND between the original forward and backward processes. However, the approach must be adapted to handle the singular diffusion coefficient $\Sigma_{t}$. Explicitly, for some choice of drift $\nu^{\prime}:\mathcal{T}\times\mathbb{R}^{2d}\to\mathbb{R}^{d}$, let us first consider a forward reference process $\{\overrightarrow{X}_{t}^{\nu^{\prime},\operatorname{ref}}\}$ of the form

[TABLE]

where $\Lambda$ denotes the Lebesgue measure on phase space $\mathbb{R}^{2d}$. Let us now choose $\nu$ so that $X_{t}^{\nu^{\prime},\operatorname{ref}}\sim\Lambda$ is distributed like the Lebesgue measure for all $t\in[0,1]$. To do so, it is sufficient to take $\nu^{\prime}$ so that the drift in eq. 54 is divergence free, or equivalently, that

[TABLE]

This is easily achieved by setting

[TABLE]

where $\partial_{q}\mu_{t}(q_{t})$ denotes the Jacobian of $\mu_{t}$, and where the diagonal operator $\operatorname{diag}$ acts on vectors and (square) matrices as

[TABLE]

respectively embedding a vector as a diagonal, or extracting from a matrix its diagonal. Plugging $\nu^{\prime}=\nu$ from eq. 56 into eq. 54, we obtain

[TABLE]

Denoting by $\mathbb{P}^{\operatorname{ref}}$ the path measure of the process $\overrightarrow{X}^{\operatorname{ref}}$, we may verify that the regularity conditions of theorem C.2 hold, so that we may then apply the result to $Y=\overrightarrow{X}$ and $Z=\overrightarrow{X}^{\operatorname{ref}}$ with

[TABLE]

It follows that $\mathbb{P}_{T}\sim\mathbb{P}_{T}^{\text{ref}}$ are absolutely continuous with respect to one another. Observe now that $\Sigma_{t}\Sigma_{t}^{T}=\begin{bmatrix}0&0\\ 0&2\Gamma_{t}E_{t}\end{bmatrix}$, so that its pseudo-inverse is given by

[TABLE]

Writing $A_{t}=A_{t}(q_{t},p_{t})\triangleq\Gamma_{t}\nabla_{q}\mathcal{H}_{t}(q_{t},p_{t})+\Gamma_{t}E_{t}\nabla_{p}\mathcal{H}_{t}(q_{t},p_{t})$, eq. 51 then immediately yields

[TABLE]

Let us now turn our attention towards the backward process from eq. 47, which we recall below as

[TABLE]

from which we may construct the reference process

[TABLE]

whose path measure we shall denote by $\mathbb{Q}_{T}^{\text{ref}}$ and so that $\overleftarrow{X}_{t}^{\text{ref}}\sim\Lambda$ is distributed like the Lebesgue measure for all $t\in[0,T]$. By introducing $s\triangleq T-t$, we may instead consider the forward (in $s$) process $\{Y_{s}\}\triangleq\{\overleftarrow{X}_{1-s}\}$ and $\{Y^{\operatorname{ref}}_{s}\}\triangleq\{\overleftarrow{X}_{1-s}^{\operatorname{ref}}\}$ given by

[TABLE]

and the associated reverse-time reference process

[TABLE]

where we have defined e.g., $\tilde{\mu}_{s}=\mu_{T-s}$, and where $Y_{s}^{\text{ref}}\sim\Lambda$ is inherited from eq. 61. Note that both equations now involve the forward Itô differential $\mathop{}\!\mathrm{d}$ rather than the backward differential $\overleftarrow{\mathop{}\!\mathrm{d}}$. Denoting by $\tilde{\mathbb{Q}}_{T}$ and $\tilde{\mathbb{Q}}_{T}^{\operatorname{ref}}$ the path measures of the time-reversed processes $\{Y_{s}\}$ and $\{Y^{\operatorname{ref}}_{s}\}$, we may again verify the necessary regularity conditions and apply theorem C.2 with $Y=\{Y_{s}\}$ and $Z=\{Y_{s}^{\text{ref}}\}$ so that we have

[TABLE]

yielding

[TABLE]

where we have introduced $\tilde{C}_{s}(q_{s},p_{s})=C_{T-s}(q_{s},p_{s})$ with $C_{t}(q_{t},p_{t})\triangleq A_{t}(q_{t},p_{t})-2\Gamma_{t}\nabla_{q}\mathcal{H}_{t}(q_{t},p_{t})$. Denoting by $\mathbb{Q}^{\operatorname{ref}}$ the path measure of $\{\overleftarrow{X}^{\operatorname{ref}}_{t}\}$ , it is easily shown that

[TABLE]

where $Y_{s}=\overleftarrow{X}_{T-s}$. It therefore follows from eq. 66, the relation $\tilde{f}_{s}\mathop{}\!\mathrm{d}\tilde{P}_{s}=-f_{t}\overleftarrow{\mathop{}\!\mathrm{d}}P_{t}$, and eq. 65 that

[TABLE]

To finish, we may exploit the fact that since $\mathbb{Q}_{T}^{\operatorname{ref}}=\mathbb{P}_{T}^{\operatorname{ref}}$ are the same to find then

[TABLE]

To see that the last term is zero, it suffices to note that $\overleftarrow{X}_{t}^{\text{ref}}\sim\Lambda$ is given by the Lebesgue measure for all $t\in[0,T]$, and since the Lebesgue measure has vanishing score, the forward reference process eq. 58 is exactly the time-reversal of the backward reference process eq. 61, and thus $\log\frac{\mathop{}\!\mathrm{d}\mathbb{P}_{T}^{\text{ref}}}{\mathop{}\!\mathrm{d}\mathbb{Q}_{T}^{\text{ref}}}(X)=0$. We may therefore plug in eq. 59 and eq. 67 into eq. 68 to obtain

[TABLE]

We may now simplify the portions of the integrand correspond to each of the red and blue terms.

Red Terms: The red terms combine to

[TABLE]

Writing $D_{t}\triangleq\Gamma_{t}\nabla_{q}\mathcal{H}_{t}(q_{t},p_{t})$, so that $A_{t}=C_{t}+2D_{t}$, the first term simplifies as

[TABLE]

where in the first and second equalities we have utilized the definitions of $A_{t}$, $C_{t}$, and $D_{t}$, in the second we have used eq. 43, and in the third eq. 39. The second term simplifies as

[TABLE]

where we have utilized eq. 43 and eq. 55. Plugging eq. 73 and eq. 74 into eq. 72 yields

[TABLE]

Blue Terms: The blue terms combine to

[TABLE]

which readily simplifies to

[TABLE]

Combining eq. $\bigstar$ with eq. $\bigstar$, we obtain

[TABLE]

The desired result then follows from the observation that $\log\left(\frac{\mathop{}\!\mathrm{d}\pi_{T}}{\mathop{}\!\mathrm{d}\Lambda}\right)(X_{T})-\log\left(\frac{\mathop{}\!\mathrm{d}\pi_{0}}{\mathop{}\!\mathrm{d}\Lambda}\right)(X_{0})=\log\pi_{T}(X_{T})-\log\pi_{0}(X_{0})$. ∎

C.3 Controlled Crooks and Jarzynski Equalities for Non-Separable Langevin Dynamics

We now state and prove a controlled generalization of the Crooks fluctuation theorem (see (Crooks, 1999)) relating the dynamics from eq. 45 and eq. 47 for the non-separable Hamiltonian from eq. 44.

Theorem C.3** (Controlled Crooks Fluctuation Theorem for Non-Separable Langevin Dynamics).**

For any $T\in(0,1]$, and with $\mathbb{P}_{T}$ and $\mathbb{Q}_{T}$ defined as as the path measures of eq. 45 and eq. 47 on the interval $[0,T]$, then for any $X\sim\mathbb{P}_{T}$, we have

[TABLE]

Proof.

By Theorem C.1,

[TABLE]

Observe now that via the Stratonovich formulation of Itô’s lemma applied to $\log\hat{\pi}_{t}=\log\pi_{t}+F_{t}$,

[TABLE]

Since $X\sim\mathbb{P}_{T}$ is given by eq. 45, we may plug in

[TABLE]

to obtain

[TABLE]

Equation 83 thus implies that

[TABLE]

Plugging Equation 87 back into Equation 80 we have

[TABLE]

Simplifying then yields

[TABLE]

from which the desired result immediately follows. ∎

We finish by noting that the Radon-Nikodym derivative of Theorem C.3 allows us to importance sample from $\mathbb{Q}_{T}$ using $\mathbb{P}_{T}$, as is formalized in the following controlled Jarzynski equality.

Corollary C.4** (Controlled Jarzynski Equality).**

For $T\in(0,1]$, let $h\in\mathcal{C}^{1}([0,T],\mathbb{R})$ denote some observable. Then

[TABLE]

where the work functional $A_{T}(X)$ is given by

[TABLE]

In particular,

[TABLE]

A similar result is established in (Albergo & Vanden-Eijnden, 2024), for an overdamped Langevin dynamics, and in the special case that $h(X)=h(X_{T})$ is a function of the terminal point.

Proof.

Observe that by theorem C.3,

[TABLE]

Thus,

[TABLE]

as desired. ∎

C.4 Specialization to the Continuously Tempered Setting

Proof of theorem 3.2.

Recall the CTDS dynamics from eq. 17, which may be rewritten in the form of eq. 45 as

[TABLE]

where $q_{t}=(X_{t},\xi_{t})$, $p_{t}=(P_{t}^{x},P_{t}^{\xi})$, $U_{t}(q_{t})=\tilde{U}_{t}^{\theta}(X_{t},\xi_{t})$, and $K(q_{t},p_{t})=\frac{\beta(\xi_{t})}{2M_{x}}\lVert P_{t}^{x}\rVert^{2}+\frac{1}{2M_{\xi}}\lVert P_{t}^{\xi}\rVert^{2}$, and

[TABLE]

Thus, theorem 3.2 follows immediately from eq. 91 in the special case that $h(X)=h(X_{T})$ is a function of the terminal point. ∎

Now recall from eq. 15 that

[TABLE]

and define by $\pi_{t}^{\theta}(x,\xi)$ the normalized version of $\hat{\pi}_{t}^{\theta}(x,\xi)$. Additionally recall that $\mathbb{P}_{t}$ and $\mathbb{Q}_{t}$ denote the forward and backward path measures on the interval $[0,t]$, for $t\in[0,1]$. In particular, if $X\sim\mathbb{Q}_{t}$, then $X_{t}\sim\pi_{t}^{\theta}$. Then, theorem 3.2 may be applied to reweight the multi-temperature PINN objective from eq. 13 via

[TABLE]

yielding eq. 22. To arrive at the expression for $A_{t}(Q)$ (from eq. 19) from the general work functional in eq. 90, we have used the fact the $\xi$-component of the control $\mu_{t}^{\theta}(q_{t})$ is zero to obtain $\nabla_{q}\cdot\mu_{t}^{\theta}=\nabla_{x}\cdot\mu_{t}^{\theta}$, and this together with the fact that the kinetic energy depends only on the temperature and momenta to obtain that $(\mu_{t}^{\theta})^{T}\nabla_{q}\mathcal{H}_{t}^{\theta}=(\mu_{t}^{\theta})^{T}\nabla_{x}\tilde{U}_{t}^{\theta}$.

Appendix D Experimental Details

D.1 40-Mode Gaussian Mixture

Additional Training Details.

At train time, we consider the four proposal types as outlined in fig. 6. For each, we parameterize the learned control $\mu_{t}^{\theta}(x)$ ($\mu_{t}^{\theta}(x,\beta(\xi))$ for CTDS), the free energy $F_{t}^{\theta}$ ($F_{t}^{\theta}(\beta(\xi))$ for CTDS), and the learned potential $U_{t}^{\theta}$ as feed-forward neural networks with width $256$ and depth three, and using the SiLU non-linearity (Elfwing et al., 2017). We train using a replay buffer, re-sampling once per epoch for 1250 epochs for a total of 125000 training iterations at which the PINN loss is evaluated at 6250 randomly sampled elements of the buffer. We utilize the Adam optimizer (Kingma & Ba, 2017) with learning rate $1\times 10^{-3}$, reducing by a factor of $\gamma=0.97$ every $1000$ iterations, and after an initial burn-in period of $15000$ training iterations. Additionally, we follow the lead of (Wang et al., 2023; Albergo & Vanden-Eijnden, 2024) in utilizing curriculum-based training whereby at $T\in\{0.1,0.2,0.3,0.4,0.5,0.6,0.7,0.8,0.9\}$ (in that order) we spend $\{1000,1000,1000,1000,2000,2000,2000,3000,3000,3000\}$ iterations respectively, and the remaining iterations at $T=1.0$. For CTDS, we reparameterize using eq. 30 with $\beta_{\min}=0.2$, $\Delta=0.25$, and $\Delta^{\prime}=1.9$, and additionally set $\psi^{\prime}_{t}(\xi)=\psi^{\text{conf}}(\xi)$ as in eq. 32 with $\eta=10.0$ and $\tilde{\Delta}=2.0$. Finally, we utilize Gaussian Fourier features as in (Tancik et al., 2020) to encode position, time, and temperature with $100$, $20$, and $20$ features respectively, drawn from isotropic Gaussians with standard deviations $0.1$, $5$, and $1$, respectively.

Sampling Details.

In practice, we discretize using an Euler solver with with $\Delta t=0.004$ (250 timesteps). Letting $\phi^{\theta}:\mathbb{R}^{d}\to\mathbb{R}^{d}$, $\phi^{\theta}:X_{0}\mapsto X_{1}$ denote the corresponding unit-time flow of eq. 23, we sample from the pushforward $\tilde{\pi}^{\theta}\triangleq[\phi^{\theta}]_{\sharp}\mathcal{N}(0,5.0I_{2})$. The density $\tilde{\pi}^{\theta}$ may then be computed via the continuous change of variables formula as

[TABLE]

Evidence Lower Bound.

Letting $\tilde{\pi}^{\theta}$ denote our sample density as given in eq. 96, and $\hat{\pi}_{1}$ the unnormalized target, we define the evidence lower bound (ELBO) as

[TABLE]

where $Z=\int\hat{\pi}_{1}(x)\,\mathop{}\!\mathrm{d}x$ denotes the partition function of $\hat{\pi}_{1}$.

Evidence Upper Bound.

Following Blessing et al. (2024), we define the evidence upper bound (EUBO) as

[TABLE]

where $Z$ is defined as before.