Annealing Flow Generative Model Towards Sampling High-Dimensional and Multi-Modal Distributions

## Abstract

Abstract Sampling from high-dimensional, multi-modal distributions remains a fundamental challenge across domains such as statistical Bayesian inference and physics-based machine learning. In this paper, we propose Annealing Flow (AF), a continuous normalizing flow-based approach designed to sample from high-dimensional and multi-modal distributions. The key idea is to learn a continuous normalizing flow-based transport map, guided by annealing, to transition samples from an easy-to-sample distribution to the target distribution, facilitating effective exploration of modes in high-dimensional spaces. Unlike many existing methods, AF training does not rely on samples from the target distribution. AF ensures effective and balanced mode exploration, achieves linear complexity in sample size and dimensions, and circumvents inefficient mixing times. We demonstrate the superior performance of AF compared to state-of-the-art methods through extensive experiments on various challenging distributions and real-world datasets, particularly in high-dimensional and multi-modal settings. We also highlight AF’s potential for sampling the least favorable distributions.

## 1 Introduction

Sampling from high-dimensional and multi-modal distributions is crucial for various fields, including physics-based machine learning like molecular dynamics (Miao et al., ; Salo-Ahen et al., ), quantum physics (Carlson et al., ; Lynn et al., ), and lattice field theory (Jay & Neil, ; Lozanovski et al., ). With modern datasets, it also plays a key role in Bayesian areas, including Bayesian modeling (Kandasamy et al., ; Balandat et al., ; Stephan et al., ) with applications in areas like computational biology (Stanton et al., ; Overstall et al., ), and Bayesian Neural Network sampling (Cobb & Jalaian, ; Izmailov et al., ).

Numerous MCMC methods have been developed over the past 50 years, including Metropolis-Hastings (MH) and its variants (Haario et al., ; Cornish et al., ; Griffin & Walker, ; Choi, ), Hamiltonian Monte Carlo (HMC) schemes (Girolami & Calderhead, ; Bou-Rabee & Sanz-Serna, ; Shahbaba et al., ; Li et al., ; Hoffman et al., ). HMC variants are still considered state-of-the-art methods. However, they require exponentially many steps in the dimension for mixing, even with just two modes (Hackett et al., ).

Annealing methods (Gelfand et al., ; Sorkin, ; Van Groenigen & Stein, ; Neal, ) are widely used to develop MCMC techniques like Parallel Tempering (PT) and its variants (Earl & Deem, ; Chandra et al., ; Syed et al., ). In annealing, sampling gradually shifts from an easy distribution to the target by lowering temperature. Annealed Importance Sampling (Neal, ) and its variants(Zhang et al., ; Karagiannis & Andrieu, ; Chehab et al., ) are developed for estimating normalizing constants with low variance using MCMC samples from intermediate distributions. Recent score-based annealing methods (Arbel et al., ; Doucet et al., ) optimize intermediate densities for lower-variance estimates, but still rely on MCMC for sampling. However, MCMC struggles with slow mixing, local mode trapping, mode imbalance, and correlated samples issues. These limitations are particularly pronounced in high-dimensional, multi-modal settings (Van Ravenzwaaij et al., ; Hackett et al., ).

Recently, particle-based optimization methods have emerged for sampling, including Stein Variational Gradient Descent (SVGD) (Liu & Wang, ), and stochastic approaches such as (Dai et al., ; Nitanda & Suzuki, ; Maddison et al., ; Liu, ; Pulido & van Leeuwen, ; Li et al., ; Detommaso et al., ). However, many of these methods rely on kernel computations, which scale polynomially with sample size and dimensions, and are sensitive to hyperparameters.

More recently, Neural network (NN)-based sampling algorithms (Wolniewicz et al., ; Bonati et al., ; Gu & Sun, ; Egorov et al., ; Li et al., ; Hackett et al., ) have been developed to leverage NN expressiveness for improving MCMC, but they still inherit some limitations like slow mixing and imbalanced mode exploration. Normalizing flow models have also been explored, with several works (Hackett et al., ; Gabrié et al., , ; Albergo & Vanden-Eijnden, ; Brofos et al., ) combining them with MCMC, though their performance depends heavily on the quality of MCMC samples, which can struggle in high dimensions. In contrast, Koehler & Vuong (); Chen et al. (, ); Wu et al. () introduced score-based learning for sampling that do not rely on neural networks. These papers offer strong theoretical guarantees for sampling from specific distributions, primarily Gaussian Mixture Models. However, these methods do not generalize well to arbitrary distributions, as the score functions are derived analytically.

Challenges persist with multi-modal distributions in high-dimensional spaces. This paper introduces Annealing Flow (AF), a new sampling scheme that learns a continuous normalizing flow map from an easy-to-sample distribution $\pi_{0}(x)$ to the target $q(x)$, guided by annealing principles. Unlike diffusion sampling (Bruna & Han, ; Shih et al., ; Zhou et al., ; Chung et al., ) which requires pre-learning from a dataset of unknown distribution, AF training does not require preliminary samples from the target $q(x)$. AF is not based on MCMC, thus avoiding issues like slow mixing, sample correlation, and mode imbalance. And unlike particle-based optimization methods, AF scales linearly with sample size and dimensions. Once trained, one simply samples from $\pi_{0}(x)$, and the learned transport map directly pushes these samples towards the target distribution.

## 2 Preliminaries

Neural ODE and Continuous Normalizing Flow.
A Neural ODE is a continuous model where the trajectory of data is modeled as the solution of an ordinary differential equation (ODE). Formally, in $\mathbb{R}^{d}$, given an input $x(t_{0})=x_{0}$ at time $t_{0}$, the transformation to the output $x(T)$ is governed by:

$$ $\frac{dx(t)}{dt}=v(x(t),t),$ (1) $$

where $v(x(t),t)$ represents the velocity field, which is of the same dimension as $x(t)$ and is parameterized by a neural network with input $x(t)$ and $t$.

A Continuous Normalizing Flow (CNF) is a class of normalizing flows where the transformation of a probability density from a base distribution $p(x)$ (at $t=0$) to a target distribution $q(x)$ (at $t=T$) is governed by a Neural ODE. The marginal density of $x(t)$, denoted as $\rho(x,t)$, evolves according to the continuity equation derived from the ODE in Eq. (1). This continuity equation is written as:

$$ $\partial_{t}\rho(x,t)+\nabla\cdot(\rho(x,t)v(x,t))=0,\quad\rho(x,0)=p(x)$ (2) $$

where the divergence $\nabla\cdot(\rho v)$ accounts for the change in density as the flow evolves over time.

Dynamic Optimal Transport.
The Benamou-Brenier equation (Benamou & Brenier, ) below provides the dynamic formulation of Optimal Transport $\mathcal{T}$.

$$ $\displaystyle\inf_{\rho,v}\mathcal{T}:=\int_{0}^{1}\mathbb{E}_{x(t)\sim\rho( \cdot,t)}\|v(x(t),t)\|^{2}dt$ (3) $\displaystyle\text{s.t.}\quad\partial_{t}\rho+\nabla\cdot(\rho v)=0,\quad\rho( \cdot,0)=p,\quad\rho(\cdot,1)=q,$ $$

The optimization problem seeks to find the optimal transport map that moves mass from the base density $p$ to the target density $q$, subject to the continuity equation (2) to ensure that $\rho(\cdot,t)$ evolves as a valid probability density over time. Additionally, the constraint $\rho(\cdot,1)=q$ ensures that the target density is reached by the end of the time horizon, which is scaled to $[0,1]$.

## 3 Annealing Flow Model

The annealing philosophy (Gelfand et al., ; Sorkin, ; Van Groenigen & Stein, ; Neal, ) refers to gradually transitioning an initial flattened distribution to the target distribution as the temperature decreases. Building on this idea, we introduce Annealing Flow (AF), a sampling algorithm that learns a continuous normalizing flow to gradually map an initial easy-to-sample density $\pi_{0}(x)$ to the target density $\tilde{q}(x)$ through a set of intermediate distributions.

We first define a sequence of unnormalized intermediate distributions $\tilde{f}_{k}(x)$, which interpolate between an easy-to-sample initial distribution $\pi_{0}(x)$ (e.g., a Gaussian) and the unnormalized target distribution $\tilde{q}(x)$. These intermediate distributions are formulated as:

$$ $\tilde{f}_{k}(x)=\pi_{0}(x)^{1-\beta_{k}}\tilde{q}(x)^{\beta_{k}},$ (4) $$

where $\beta_{k}$ is an increasing sequence with $\beta_{0}=0$ and $\beta_{K}=1$. This formulation ensures that $\tilde{f}_{0}(x)=\pi_{0}(x)$ and $\tilde{f}_{K}(x)=\tilde{q}(x)$. The sequence $0=\beta_{0}<\beta_{1}<\cdots<\beta_{K}=1$ controls the gradual transition between the two distributions.

The above construction aligns with the annealing philosophy. As $\beta_{k}$ increases, $\tilde{f}_{k}(x)$ gradually sharpens toward the target $\tilde{q}(x)$, starting from the initially flattened distribution around $\pi_{0}(x)$. These annealed densities serve as a bridge, providing a gradual flow path from the easy-to-sample distribution $\pi_{0}(x)$ to the target density $q(x)$. Figure 1 provides an intuitive illustration of this process, where $\pi_{0}(x)$ is a standard Gaussian, and $q(x)$ is a Gaussian mixture model with six modes.

Figure: (a) $\beta_{0}=0$
Refer to caption: extracted/5874394/d=2_GMM_c8_heatmap_beta0.png

### 3.1 Optimal transport map

We aim to learn a continuous optimal transport map between an easy-to-sample distribution $\pi_{0}(x)$ and the target distribution $q(x)$. Once trained, users simply sample $\{x^{(i)}(0)\}_{i=1}^{n}\sim\pi_{0}(x)$, and the transport map pushes them to $\{x^{(i)}(T=1)\}_{i=1}^{n}\sim q(x)$. The transport map $\mathcal{T}$ evolves $x(t)$ according to the ODE (1). Specifically,

$$ $\mathcal{T}(x_{t})=x(0)+\int_{0}^{t}v(x(s),s)ds,\quad t\in[0,1].$ (5) $$

We divide the time horizon $[0,1]$ of $\mathcal{T}$ into $K$ intervals $[t_{k-1},t_{k}]$ for $k=1,2,\dots,K$, where $t_{0}=0$ and $t_{K}=1$. Guided by the annealing flow path defined in (4), the continuous flow map $\mathcal{T}$ gradually transforms the distribution from $f_{0}(x)$ to $f_{1}(x)$ over $[0,t_{1}]$, and continues this process until $f_{K-1}(x)$ is transformed into $f_{K}(x)=q(x)$ over $[t_{K-1},t_{K}]$. Figure 1 shows this progression with two intermediate densities. For clarity, we denote $\mathcal{T}_{k}(x)$ as the segment of the continuous normalizing flow during $[t_{k-1},t_{k}]$, which pushes the distribution from $f_{k-1}(x)$ to $f_{k}(x)$.

### 3.2 Objective of annealing flow net

Annealing Flow aims to learn each transport map $\mathcal{T}_{k}$ based on dynamic optimal transport (3) over the time horizon $[t_{k-1},t_{k}]$, where the velocity field $v_{k}(x(t),t)$ is learned using a neural network. The terminal condition $\rho(\cdot,1)=q$ in (3) can be relaxed by introducing a KL divergence term (see, for instance, Ruthotto et al. ()). Consequently, minimizing the objective (3) of dynamic optimal transport $\mathcal{T}_{k}:f_{k-1}(x)\to f_{k}(x)$ is equivalent to solving the following problem:

$$ $\mathcal{T}_{k}=\arg\min_{\mathcal{T}}\left\{\text{KL}(\mathcal{T}_{\#}\tilde{ f}_{k-1}\|\tilde{f}_{k})+\gamma\int_{t_{k-1}}^{t_{k}}\mathbb{E}_{x(t)\sim\rho_ {k}(\cdot,t)}\|v_{k}(x(t),t)\|^{2}dt\right\},$ (6) $$

subject to $v_{k}(x(t),t)$ evolving according to (2). Here, $\gamma>0$ is a weight parameter, $v_{k}(x(t),t)$ denotes the velocity field during the $k$-th time interval $[t_{k-1},t_{k}]$, and $\text{KL}(\mathcal{T}_{\#}\tilde{f}_{k-1}\|\tilde{f}_{k})$ represents the KL divergence between the push-forward density $\mathcal{T}_{\#}\tilde{f}_{k-1}$ and the unnormalized intermediate target density $\tilde{f}_{k}$. Additionally, the constraint (2) ensures that $x(t)$ follows the ODE trajectory defined by (1) during $t\in[t_{k-1},t_{k}]$, which is given by:

$$ $x(t)=x(t_{k-1})+\int_{t_{k-1}}^{t}v_{k}(x(s),s)ds,\quad t\in[t_{k-1},t_{k}].$ (7) $$

We can rewrite the density $f_{k}(x)$ in terms of energy function: $f_{k}(x)=Z\cdot\tilde{f}_{k}(x)=Z\cdot e^{-\tilde{E}_{k}(x)}$,
where $\tilde{E}_{k}(x)=-\log\tilde{f}_{k}$ is the unnormalized energy function associated with $f_{k}(x)$. The following proposition shows that once we have obtained samples from $\tilde{f}_{k-1}(x)$, the KL divergence in (6) can be computed exactly based on $v_{k}(x(t),t)$ and $\tilde{E}_{k}(x)$. Therefore, learning an optimal transport map $\mathcal{T}_{\#}\tilde{f}_{k-1}$ reduces to learning the optimal $v_{k}(x(t),t)$. The proof is provided in A.1.

###### Proposition 1

Given the unnormalized density $\tilde{f}_{k-1}$, the KL-Divergence between $\mathcal{T}_{\#}\tilde{f}_{k-1}$ and $\tilde{f}_{k}$ is equivalent to:

$$ $\text{KL}(\mathcal{T}_{\#}\tilde{f}_{k-1}\|\tilde{f}_{k})=c+\mathbb{E}_{x(t_{k -1})\sim f_{k-1}}\left[\tilde{E}_{k}(x(t_{k}))-\int_{t_{k-1}}^{t_{k}}\nabla \cdot v_{k}(x(s),s)\,ds\right],$ (8) $$

up to a constant $c$ that is independent of $v_{k}(x(s),s)$.

Given $x(t_{k-1})$ from $f_{k-1}(x)$, the value of $x(t_{k})$ inside the energy function $\tilde{E}_{k}$ in (8) can be calculated as shown in equation (7). Additionally, according to the proposition below, the second term in the objective (6) can be relaxed to a dynamic $W_{2}$ regularization. The proof is provided in A.1.

###### Proposition 2

Let $x(t)$ be particle trajectories driven by a smooth velocity field $v_{k}(x(t),t)$ over the time interval $[t_{k-1},t_{k}]$, where $h_{k}=t_{k}-t_{k-1}$. Assume that $v_{k}(x,t)$ is Lipschitz continuous in both $x$ and $t$. By dividing $[t_{k-1},t_{k}]$ into $S$ equal mini-intervals with grid points $t_{k-1,s}$ (where $s=0,1,\ldots,S$ and $t_{k-1,0}=t_{k-1}$, $t_{k-1,S}=t_{k}$), the following approximation holds:

$$ $\int_{t_{k-1}}^{t_{k}}\mathbb{E}_{x(t)}\left[\|v_{k}(x(t),t)\|^{2}\right]dt= \frac{S}{h_{k}}\sum_{s=0}^{S-1}\mathbb{E}\left[\|x(t_{k-1,s+1})-x(t_{k-1,s})\| ^{2}\right]+O\left(\frac{h_{k}^{2}}{S}\right).$ (9) $$

As $h_{k}\to 0$ or $S\to\infty$, the error term $O\left(\frac{h_{k}^{2}}{S}\right)$ becomes negligible.

One can observe that the RHS of (9) can be interpreted as the discretized sum of the squared Wasserstein-2 distance. The dynamic $W_{2}$ regularization encourages smooth transitions from $f_{k-1}$ to $f_{k}$ with minimal transport cost, promoting efficient mode exploration.

Therefore, by incorporating Propositions 1 and 2 into objective (6), the final objective becomes:

$$ $\min_{v_{k}(\cdot,t)}\mathbb{E}_{x(t_{k-1})\sim f_{k-1}}\left[\tilde{E}_{k}(x( t_{k}))-\int_{t_{k-1}}^{t_{k}}\nabla\cdot v_{k}(x(s),s)ds+\alpha\sum_{s=0}^{S- 1}\|x(t_{k-1,s+1})-x(t_{k-1,s})\|^{2}\right].$ (10) $$

Here, $\alpha=\gamma\cdot S/h_{k}$ and $v_{k}(x(s),s)$ is learned by a neural network. We break the time interval $[t_{k-1},t_{k}]$ into $S$ mini-intervals, and $x(t_{k-1,s+1})$ is computed as in equation (7).

After learning, connecting the Annealing Flow nets together yields a smooth flow map $\mathcal{T}:\mathcal{T}_{1}\to\mathcal{T}_{2}\to\cdots\mathcal{T}_{K}$, which transforms samples from $\pi_{0}(x)$ to the target $q(x)$. Please see Section 4.2 for efficient sampling of Annealing Flow and its comparisons with other sampling methods.

###### Proposition 1

###### Proposition 2

### 3.3 Properties of learned velocity field

###### Proposition 3

By Taylor’s expansions, and assume that $h_{k}=t_{k}-t_{k-1}$, we have:

- 1.
$x(t_{k})-x(t_{k-1})=\int_{t_{k-1}}^{t_{k}}v_{k}(x(s),s)ds=h_{k}v_{k}(x(t_{k-1}
),t_{k-1})+O(h_{k}^{2})$
- 2.
$\tilde{E}_{k}(x(t_{k}))=\tilde{E}_{k}(x(t_{k-1}))+h_{k}\nabla\tilde{E}_{k}(x(t
_{k-1}))\cdot v_{k}(x(t_{k-1}),t_{k-1})+O(h_{k}^{2})$

Therefore, if $h_{k}\to 0$, then $O(h_{k}^{2})$ in Proposition 3 vanishes, and we no longer need to divide the time interval, i.e., $S=1$. By defining the score function as $s_{k}=\nabla\log f_{k}$, the objective function (10) can be approximated as:

$$ $\displaystyle\mathbb{E}_{x\sim f_{k-1}}[\tilde{E}_{k}(x(t_{k-1}))+h_{k}(-s_{k} (x)\cdot v_{k}(x,t_{k-1})-\nabla\cdot v_{k}(x,t_{k-1})+\frac{1}{2}\|v_{k}(x,t_ {k-1})\|^{2})+O(h_{k}^{2})]$ (11) $$

Since $\mathbb{E}_{x(t_{k-1})\sim f_{k-1}}[\tilde{E}_{k}(x(t_{k-1}))]$ is independent of $v_{k}(x,t)$, as $h_{k}\to 0$, the minimization of the leading term is equivalent to:

$$ $\min_{v_{k}=v_{k}(\cdot,0)}\mathbb{E}_{x\sim f_{k-1}}\left[-T_{f_{k}}v_{k}+ \frac{1}{2}\|v_{k}\|^{2}\right],\quad T_{f_{k}}v_{k}:=s_{k}\cdot v_{k}+\nabla \cdot v_{k}.$ (12) $$

Here, $T_{f_{k}}$ is referred to as the Stein operator. The term $T_{f_{k}}v_{k}$ reflects the fact that the derivative of the KL divergence with respect to the transport map yields the Stein operator (Liu & Wang, ). Therefore, if $\mathbb{E}_{x\sim f_{k-1}}[\|v_{k}\|^{2}]<\infty$, and both $s_{k-1}$ and $s_{k}$ belong to $L^{2}(f_{k})$, one can verify that the minimizer of (12) is:

$$ $v_{k}^{*}=s_{k}-s_{k-1}.$ (13) $$

Therefore, the infinitesimal optimal $v_{k}^{*}$ is equal to the difference between score function of the next density, $f_{k}$, and the current density, $f_{k-1}$. This suggests that when the two intermediate densities are sufficiently close, i.e., when the number of $\beta_{k}$ is large enough, the optimal velocity field equals the difference between the score functions. By adding more intermediate densities, one can construct a sufficiently smooth transport map $\mathcal{T}$ that exactly learns the mapping between each pair of densities.

Additionally, one can observe that when each $\tilde{f}_{k}(x)$ is set to the target $q(x)$, i.e., when all $\beta_{k}$ are set to 1, and the second term in the objective (6) is relaxed to static $W_{2}$ regularization, the objective of Annealing Flow becomes equivalent to Wasserstein gradient flow. This is detailed in A.2.

###### Proposition 3

## 4 Training and Sampling of Annealing Flow Net

In this section, we describe the training process for each block and provide an explanation of the sampling process in our algorithm.

### 4.1 Block-wise training

Training of the $k$-th flow map in Annealing Flow begins once the $(k-1)$-th block has completed training. Given the samples $\{x^{(i)}(t_{k-1})\}_{i=1}^{n}\sim f_{k-1}(x)$ produced after the $(k-1)$-th block, we can replace $\mathbb{E}_{x\sim f_{k-1}}$ with the empirical average. The divergence of the velocity field can be computed either by brute force or via the Hutchinson trace estimator (Hutchinson, ; Xu et al., ):

$$ $\nabla\cdot v_{k}(x,t)\approx\mathbb{E}_{\epsilon\sim N(0,I_{d})}\left[ \epsilon^{T}\frac{v_{k}(x+\sigma\epsilon,t)-v_{k}(x,t)}{\sigma}\right].$ (14) $$

This approximation becomes exact as $\sigma\to 0$. Further details are provided in A.3.2. Additionally, we apply the Runge-Kutta method for numerical integration, with details provided in A.3.3.

Our algorithm uses a block-wise training of the continuous normalizing flow map. Specifically, the training of Annealing Flow is summarized in Algorithm 1. The block-wise training approach of Annealing Flow significantly reduces memory and computational requirements, as only one neural network is trained at a time, independent of the other flow networks.

Figure: Algorithm 1 Block-wise Training of Annealing Flow Net

### 4.2 Efficient sampling and comparisons with other methods

Once the continuous normalizing flow map $\mathcal{T}$ is learned, the sampling process of the target $q(x)$ can be very efficient. Users can simply sample $\{x^{(i)}(t_{0}=0)\}_{i=1}^{n}$ from $\pi_{0}(x)$, and then directly calculate $\{x^{(i)}(t_{K}=1)\}_{i=1}^{n}\sim q(x)$ through Annealing Flow nets:

$$ $x^{(i)}(t_{k})=\mathcal{T}_{k}(x^{(i)}(t_{k-1}))=x^{(i)}(t_{k-1})+\int_{t_{k-1 }}^{t_{k}}v_{k}(x^{(i)}(s),s)ds,\quad k=1,2,\cdots,K.$ (15) $$

MCMC methods require long mixing times when sampling from complex distributions. In contrast, Annealing Flow (AF) pushes samples directly from $\pi_{0}(x)$ through the learned transport map, enabling faster sampling, especially for large sample sizes. MCMC also generates correlated samples, as each new sample depends on the previous one, reducing the effective sample size (ESS) and efficiency. AF avoids this by producing independent samples, improving overall sample quality.

Additionally, MCMC struggles with multimodal distributions, as chains often get trapped in local modes. While methods like Parallel Tempering (PT) may attempt to explore all modes, they do not ensure proportional time across them, leading to imbalanced sampling. In contrast, AF generated balanced samples across all modes in line with the target distribution, as illustrated in the below figure.

Figure: (a) AF
Refer to caption: extracted/5874394/d=1_exponential_c=10_Annealing.png

NN-based MCMC algorithms still struggle with issues like slow mixing and correlated samples. Particle-based optimization methods, such as SVGD and MIED, avoid the need for a burn-in period and produce less correlated samples, but their reliance on kernel computations leads to polynomial scaling with sample size and dimensions, and they are sensitive to kernel hyperparameters. In contrast, AF computes samples independently through (15), allowing the sampling process to scale linearly with both sample size and dimensions.

We comment that Annealing Flow indeed needs more expensive pre-training than MCMC, which, however, can be done offline and only needs to be done once and then deployed for sampling. Once trained, AF samplers are highly efficient, generating 10,000 samples in an average of 1.5 seconds in our experiments. In contrast, MCMC takes around 1 minute to sample 10,000, while particle-based methods take significantly longer—over 20 minutes. AF also performs well on multi-modal and high-dimensional densities, where other methods often struggle.
Detailed comparisons of algorithms, including the training and sampling times, are provided in A.5.

## 5 Importance Flow

Sampling from complex distributions is fundamental, which can benefit tasks like normalizing constant estimation, Bayesian analysis, and various machine learning problems. Here, we briefly discuss another aspect: using Annealing Flow to sample from the Least-Favorite-Distribution (LFD) and obtain a low-variance Importance Sampling (IS) estimator, referred to as Importance Flow.

### 5.1 Settings

Suppose we want to estimate $\mathbb{E}_{X\sim\pi_{0}(x)}\left[h(X)\right]$, which cannot be computed in closed form. A natural approach is to use Monte Carlo estimation by sampling $\{x_{i}\}_{i=1}^{n}$ from $\pi_{0}(x)$. However, if $x_{i}$ consistently falls in regions where $h(x)$ has extreme values, the estimator may exhibit high variance. For example, with $\pi_{0}(x)=N(0,I_{d})$ and $h(x)=1_{\|x\|\geq 6}$, almost no samples will satisfy $\|x\|\geq 6$, resulting in a zero estimate.

To address this situation, we can select an appropriate proposal distribution $q(x)$ and rewrite the expectation and MC estimator as:

$$ $\mathbb{E}_{x\sim\pi_{0}(x)}\left[h(x)\right]=\mathbb{E}_{x\sim q(x)}\left[ \frac{\pi_{0}(x)}{q(x)}\cdot h(x)\right]\approx\frac{1}{n}\sum_{i=1}^{n}\frac{ \pi_{0}(x_{i})}{q(x_{i})}\cdot h(x_{i}),\quad x_{i}\sim q(x).$ (16) $$

It is well-known that the theoretically optimal proposal for the importance sampler is: $q^{*}(x)\propto\pi_{0}(x)\cdot|h(x)|:=\tilde{q}^{*}(x)$. However, given the definition of $\tilde{q}^{*}(x)$, it is often difficult to sample from, especially when $\pi_{0}(x)$ or $h(x)$ is complex. Consequently, people typically choose a distribution that is similar in shape to the theoretically optimal proposal but easier to sample from.

Annealing Flow enables sampling from $q^{*}(x)$, allowing the construction of an IS estimator. A common choice can be the Normalized IS Estimator: $\hat{I}_{N}=\sum_{i=1}^{n}\frac{\pi_{0}(x_{i})}{\tilde{q}(x_{i})}h(x_{i})/\sum
_{i=1}^{n}\frac{\pi_{0}(x_{i})}{\tilde{q}(x_{i})}$. However, this estimator is often biased, as can be seen from Jensen’s Inequality.

### 5.2 Density ratio estimation

Using samples from $q^{*}(x)$ and those along the trajectory obtained via Annealing Flow, we can train a neural network for Density Ratio Estimation (DRE) of $\frac{\pi_{0}(x)}{q^{*}(x)}$. Inspired by works Rhodes et al. (); Choi et al. (); Xu et al. (), we can train a continuous-time neural network
$r(x)=r_{K}(x;\theta_{K})\circ r_{K-1}(x;\theta_{K-1})\circ\cdots\circ r_{1}(x;
\theta_{1})$, where samples $x_{i}\sim f_{K}=q^{*}(x)$ are inputs and the output is the density ratio $\frac{\pi_{0}(x_{i})}{q^{*}(x_{i})}$. Each $r_{k}(x;\theta_{k})$ is trained using the following loss:

$$ $\mathcal{L}_{k}(\theta_{k})=\mathbb{E}_{x(t_{k-1})\sim f_{k-1}}\left[\log(1+e^ {-r_{k}(x_{i}(t_{k-1}))})\right]+\mathbb{E}_{x(t_{k})\sim f_{k}}\left[\log(1+e ^{r_{k}(x_{i}(t_{k}))})\right].$ $$

After successful training, $r_{k}^{*}(x)=\log\frac{f_{k-1}(x)}{f_{k}(x)}$, and thus $r^{*}(x)=\sum_{k=1}^{K}r_{k}^{*}(x)=\log\frac{\pi_{0}(x)}{q^{*}(x)}$. Please refer to A.1 and A.3.5 for the proof and further details. To obtain the optimal importance sampling estimator, we can then directly use samples $\{x_{i}\}_{i=1}^{n}\sim q^{*}(x)$ from Annealing Flow and apply (16) together with the DRE: $\frac{1}{n}\sum_{i=1}^{n}\exp(r^{*}(x_{i}))\cdot h(x_{i})$. The estimator is unbiased and can achieve zero variance theoretically.

## 6 Numerical Experiments

In this section, we present numerical experiments comparing Annealing Flow (AF) with widely-used MCMC algorithms, including Hamiltonian Monte Carlo (HMC) and Parallel Tempering (PT), as well as other SOTA techniques, including particle-based methods: Stein Variational Gradient Descent (SVGD)  (Liu & Wang, ) and Mollified Interaction Energy Descent (MIED)  (Li et al., ), alongside NN-based MCMC approaches: AI-Sampler (AIS)  (Egorov et al., ). The experimental details can be found in A.3.3. Our code is publicly available on [https://github.com/StatFusion/Annealing-Flow-For-Sampling](https://github.com/StatFusion/Annealing-Flow-For-Sampling).

We test these algorithms on challenging distributions, including Exponential-Weighted Gaussian, Gaussian Mixture Models (GMM), funnel distributions, and Truncated Normal with extreme radii across varying dimensions. Maximum Mean Discrepancy (MMD) and Wasserstein Distance are used as evaluation metrics, but only reported for the GMM due to the need for true samples. For other experiments, we provide sample and density plots for easier comparison, as shown in A.4.

In addition, we compare our algorithm with others on Hierarchical Bayesian Logistic Regression across a range of datasets. We also report the preliminary results of the Importance Flow (discussed in Section 5) for estimating $\mathbb{E}_{x\sim N(0,I)}\left[1_{\|x\|\geq c}\right]$ with varying $c$ and dimensions.

Gaussian Mixture Models (GMM).
Figure 3 presents the sampling results of different methods on a 2D GMM, where the modes are distributed across circles with varying radii. We also experimented on a GMM with modes aligned on the vertices of a cube in higher dimensions, with the number of modes ranging from 8 to 64. Evaluation metrics and additional figures for these experiments are provided in Table 4 in A.4.

Figure: (a) True
Refer to caption: extracted/5874394/d=2_GMM_c10_heatmap.png

Truncated Normal Distribution.
Figure 4 shows the sampling results for $\tilde{q}(x)=1_{\|x\|\geq c}N(0,I_{d})$, based on 5000 samples for each method. SVGD, MIED, and AI-Sampler are designed for continuous densities. SVGD and MIED specifically require the gradient of the log-probability, given by $\nabla\log\left(1_{\|x\|\geq c}*N(0,I_{d})\right)$ in this experiment. Despite relaxing the indicator function to $\frac{1}{1+\exp(-k(\|x\|-c))}$ for large $k$, the algorithms failed to yield meaningful results (See Figure 8 in A.4 for the results of their algorithms). Therefore, we compare AF with MH, HMC, and PT. We also tested our algorithm on 10D space. Additional figures and tables are given in A.4.

Figure: (a) True
Refer to caption: extracted/5874394/d=2_indicator_c6_heatmap.png

Funnel Distribution. A well-known challenging distribution for sampling is the funnel distribution, defined as:

$$ $P(x_{1},x_{2},\ldots,x_{d})\propto\mathcal{N}(x_{1}\mid 0,\sigma^{2})\prod_{i= 2}^{d-1}\mathcal{N}(x_{i}\mid 0,\exp{(x_{1})}),$ $$

In this setup, $x_{i},i=2,\cdots,d$ has a variance that depends exponentially on $x_{1}$, forming a funnel-shaped distribution. Sampling is challenging due to this exponential dependence, causing extreme concentration for negative $x_{1}$ and wide dispersion for positive $x_{1}$, making exploration difficult, especially in high dimensions.

We tested our Annealing Flow together with other algorithms on $d=5$ case. Here, we present the sampling result projected onto a 3D space for a funnel distribution in a 5D space, with $\sigma^{2}=0.81$:

Figure: (a) AF
Refer to caption: extracted/5874394/d=5_funnel_c8_Annealing.png

Exp-Weighted Gaussian with an extreme number of modes in high-dimensional spaces.
We tested each algorithm on sampling from an extreme distribution:

$$ $p(x_{1},x_{2},\cdots,x_{10})\propto e^{10\sum_{i=1}^{10}|x_{i}|-\frac{1}{2}\|x \|^{2}},$ $$

which has $\mathbf{2^{10}=1024}$ modes arranged at the vertices of a 10-Dimension cube. The L2-distance between two horizontally or vertically adjacent modes is 10, while the diagonal modes are separated by up to $\mathbf{\sqrt{10\cdot 20^{2}}\approx 63.25}$. We also tested on the extreme distribution:

$$ $p(x_{1},x_{2},\cdots,x_{50})\propto e^{10\sum_{i=1}^{10}|x_{i}|+10\sum_{i=11}^ {50}x_{i}-\frac{1}{2}\|x\|^{2}},$ $$

which has $\mathbf{2^{10}=1024}$ modes arranged at the vertices of a 50-Dimension space.

Given the challenge of visualizing results in high-dimensional space, we present the projected results onto the first 3 dimensions in Figure 6. For comparisons in 10D space, please refer to Figure 12 in A.4. The performance of SVGD, MIED, and AIS is inferior to AF, as compared in Figures 13 and 14 in A.4.

Figure: (a) AF
Refer to caption: extracted/5874394/d=50_exponential_Annealing.png

**Table 1: The number of modes successfully explored by each algorithm across various dimensions.**
|  | $d=2$ | $d=5$ | $d=10$ | $d=\mathbf{50}$ |
| --- | --- | --- | --- | --- |
| True | 4 | 32 | 1024 | 1024 |
| AF | 4 | 32 | 1024 | 1024 |
| HMC | 3.1 | 24.3 | 213.5 | $<10$ |
| PT | 3.4 | 25.2 | 233.7 | $<10$ |

Table 1 presents the number of modes successfully explored by different algorithms across varying dimensions. Each algorithm was run 10 times, sampling 10,000 points per run, and the average number of modes explored by each algorithm was then calculated.

Bayesian logistic regression. We use the same Bayesian logistic regression setting as in Liu & Wang (), where a hierarchical structure is assigned to the model parameters. The weights $\beta$ follow a Gaussian prior $p_{0}(\beta|\alpha)=N(\beta;0,\alpha^{-1})$, and $\alpha$ follows a Gamma prior $p_{0}(\alpha)=\text{Gamma}(\alpha;1,0.01)$. Sampling is performed on the posterior $p(\beta,\alpha|D)$, where $D=\{x_{i},y_{i}\}_{i=1}^{n}$. The performance comparisons are shown in Table 2. Detailed settings are given in A.3.4.

**Table 2: Bayesian Logistic Regression: comparison of different algorithms across datasets. In the table $\cdot\pm\cdot/\cdot$ represents Accuracy(%)$\pm$std(%)/log-posterior**
| Dataset | Annealing | SVGD | MEID (ICLR2023) | AI-Sampler (ICML 2024) |
| --- | --- | --- | --- | --- |
| Diabetes ($d=8$) | $\mathbf{76.30\pm 2.12}/-0.496$ | $76.10\pm 2.5/-0.502$ | $75.80\pm 2.32/-0.503$ | $\mathbf{76.30\pm 2.18}/\mathbf{-0.493}$ |
| Breast Cancer ($d=10$) | $97.85\pm 1.12/-0.017$ | $\mathbf{98.83\pm 3.10}/\mathbf{-0.008}$ | $\mathbf{98.89\pm 2.12}/\mathbf{-0.008}$ | $97.83\pm 2.80/-0.019$ |
| Heart ($d=13$) | $\mathbf{88.46\pm 2.73}/\mathbf{-0.316}$ | $79.36\pm 3.78/-0.588$ | $86.70\pm 2.24/-0.321$ | $84.23\pm 2.54/-0.458$ |
| Australian ($d=14$) | $\mathbf{86.59\pm 1.20}/\mathbf{-0.361}$ | $84.56\pm 2.87/-0.365$ | $85.17\pm 1.34/-0.369$ | $84.62\pm 2.30/-0.375$ |
| Ijcnn1 ($d=22$) | $\mathbf{91.96\pm 0.05}/\mathbf{-0.195}$ | $89.44\pm 0.34/-0.209$ | $\mathbf{91.84\pm 0.15}/\mathbf{-0.198}$ | $88.32\pm 0.25/-0.334$ |
| Svmguide3 ($d=22$) | $80.04\pm 0.70/-0.472$ | $78.89\pm 1.20/-0.479$ | $\mathbf{80.12\pm 1.04}/\mathbf{-0.472}$ | $\mathbf{80.12\pm 0.98}/\mathbf{-0.468}$ |
| German ($d=24$) | $\mathbf{78.04\pm 1.70}/\mathbf{-0.473}$ | $76.43\pm 1.70/-0.483$ | $77.21\pm 1.80/-0.479$ | $76.89\pm 1.84/-0.484$ |

Importance flow. Here we also report the preliminary results of the importance flow (discussed in Section 5) for estimating $\mathbb{E}_{x\sim N(0,I)}\left[1_{\|x\|\geq c}\right]$ with varying $c$ and dimensions. This estimation uses samples from the experiment on the Truncated Normal Distribution, and thus the results for SVGD, MIED, and AIS cannot be reported. Please refer to A.3.5 for experimental settings. Additionally, we discussed a possible extension of the Importance Flow framework in A.6.

**Table 3: Comparison of Methods and Results for different radii (r) and dimensions (d).**
| Methods | Radius | $d=2$ | $d=3$ | $d=4$ | $d=5$ |
| --- | --- | --- | --- | --- | --- |
| True Probability | $r=4$ | 3.35E-04 | 1.13E-03 | 3.02E-03 | 6.84E-03 |
| $r=6$ | 1.52E-08 | 7.49E-08 | 2.89E-07 | 9.50E-07 |  |
| Importance Flow | $r=4$ | 4.04e-04(1.0e-04) | 1.30e-03(2.3e-04) | 3.36e-03(4.23e-04) | 7.86e-03(8.21e-04) |
| $r=6$ | 9.81e-08(4.02e-07) | 1.51e-07(1.23e-07) | 2.13e-07(8.71e-08) | 2.38e-07(3.48e-06) |  |
| DRE with HMC Samples | $r=4$ | 7.56e-04(4.99e-04) | 2.52e-03(6.33e-04) | 8.97e-03(9.05e-04) | 1.12e-02(1.55e-03) |
| $r=6$ | 4.35e-07(7.21e-07) | 9.01e-07(2.79e-06) | 1.82e-07(2.89e-06) | 2.31e-06(6.21e-06) |  |
| DRE with PT Samples | $r=4$ | 6.79e-04(3.58e-04) | 2.38e-03(5.40e-04) | 5.78e-03(7.98e-03) | 9.94e-03(1.13e-03) |
| $r=6$ | 5.37e-07(9.56e-07) | 8.78e-07(2.32e-06) | 9.23e-07(2.51e-06) | 1.98e-06(7.73e-06) |  |
| Naïve MC | $r=4$ | 2.75e-04(6.0e-04) | 1.18e-03(1.1e-03) | 2.71e-03(1.7e-03) | 7.44e-03(2.6e-03) |
| $r=6$ | 0 | 0 | 0 | 0 |  |

## 7 Discussions

In this paper, we have proposed the Annealing Flow (AF) framework, a novel and flexible approach for sampling from high-dimensional and multi-modal distributions. AF offers several advantages over existing methods, as thoroughly discussed in A.5. Additionally, we have also compared the training and sampling times in A.5. Extensive experiments demonstrate that AF performs well across a variety of challenging distributions and real-world datasets.

The Annealing Flow framework presented in this paper is highly flexible and accommodates various challenging distributions. The concept of ‘Annealing’ in sampling can be interpreted as gradually transitioning from an easy-to-sample distribution to the target distribution. Therefore, each intermediate distribution $f_{k}$ can be defined flexibly without adhering to (4), as long as the transitions between $f_{k-1}$ and $f_{k}$ are smooth and the sequence converges to the target $q(x)$. If the density modes are close enough, all $\tilde{f}_{k}(x)$ can simply be set to the target density $q(x)$, making the Annealing Flow objective equivalent to the Wasserstein gradient flow, as discussed in A.2. Additionally, we believe that by adding more intermediate distributions, one can obtain intermediate samples at various time points to construct a low-variance estimator for the normalizing constant. Finally, the importance flow discussed in Section 5 may be extended to a distribution-free model, allowing one to learn an importance flow from a dataset for sampling its Least-Favorable Distribution (LFD) with minimal variance, as further detailed in A.6.

## Appendix A Appendix

### A.1 Propositions and proofs

Proposition 1. Given the unnormalized density $\tilde{f}_{k-1}$, the KL-Divergence between $\mathcal{T}_{\#}\tilde{f}_{k-1}$ and $\tilde{f}_{k}$ is equivalent to:

$$ $\text{KL}(\mathcal{T}_{\#}\tilde{f}_{k-1}\|\tilde{f}_{k})=c+\mathbb{E}_{x\sim f _{k-1}}\left[\tilde{E}_{k}(x(t_{k}))-\int_{t_{k-1}}^{t_{k}}\nabla\cdot v_{k}(x (s),s)\,ds\right],$ $$

up to a constant $c$ that is independent of $v_{k}(x(s),s)$.

Proof:

The expression for KL-divergence is given by:

$$ $KL(T_{\#}\tilde{f}_{k-1}\parallel\tilde{f}_{k})=\mathbb{E}_{x\sim\tilde{f}_{k- 1}}\left[\log\frac{T_{\#}\tilde{f}_{k-1}(x)}{\tilde{f}_{k}(x)}\right]=\mathbb{ E}_{x\sim\tilde{f}_{k-1}}\left[\log T_{\#}\tilde{f}_{k-1}(x)-\log\tilde{f}_{k} (x)\right].$ $$

Now, recall that $\log\tilde{f}_{k}(x)$ corresponds to the energy term $\tilde{E}_{k}(x(t_{k}))$ up to an additive constant, so we substitute:

$$ $KL(T_{\#}\tilde{f}_{k-1}\parallel\tilde{f}_{k})=\mathbb{E}_{x\sim\tilde{f}_{k- 1}}\left[\log T_{\#}\tilde{f}_{k-1}(x)+\tilde{E}_{k}(x(t_{k}))+c\right],$ $$

where $c$ is a constant that is independent of $v_{k}(x(s),s)$. This constant $c$ is omitted for now since it doesn’t affect the gradient descent calculations.

Next, to compute $\log T_{\#}\tilde{f}_{k-1}(x)$, we use the fact that the dynamics of the pushforward density are governed by the velocity field $v_{k}(x(s),s)$:

$$ $\displaystyle\frac{d}{ds}\log\rho_{k}(x(s),s)$ $\displaystyle=\frac{\nabla\rho_{k}(x(s),s)\cdot\partial_{s}x(s)+\partial_{s} \rho_{k}(x(s),s)}{\rho_{k}(x(s),s)}$ $\displaystyle=\frac{\nabla\rho_{k}\cdot v_{k}-\nabla\cdot(\rho_{k}v_{k})}{\rho _{k}}\Big{|}_{(x(s),s)}\quad\text{(by (\ref{ODE equation}) and (\ref{Liouville }))}$ $\displaystyle=-\nabla\cdot v_{k}(x(s),s).$ $$

Integrating this equation over the interval $s\in[t_{k-1},t_{k}]$, we find:

$$ $\log T_{\#}\tilde{f}_{k-1}(x)=\log\rho_{k}(x(t_{k}))=\log\rho_{k}(x(t_{k-1}))- \int_{t_{k-1}}^{t_{k}}\nabla\cdot v_{k}(x(s),s)ds.$ $$

We now substitute this result back into the KL-divergence expression:

$$ $KL(T_{\#}\tilde{f}_{k-1}\parallel\tilde{f}_{k})=\mathbb{E}_{x\sim\tilde{f}_{k- 1}}\left[\log\rho_{k}(x(t_{k-1}))-\int_{t_{k-1}}^{t_{k}}\nabla\cdot v_{k}(x(s) ,s)ds+\tilde{E}_{k}(x(t_{k}))+c\right].$ $$

Notice that $\log\rho_{k}(x(t_{k-1}))$ is independent of $v_{k}(x(s),s)$ and acts as another constant term that can be absorbed into $c$. Therefore, the relevant terms for the KL-divergence are:

$$ $KL(T_{\#}\tilde{f}_{k-1}\parallel\tilde{f}_{k})=c+\mathbb{E}_{x\sim\tilde{f}_{ k-1}}\left[\tilde{E}_{k}(x(t_{k}))-\int_{t_{k-1}}^{t_{k}}\nabla\cdot v_{k}(x(s ),s)ds\right].$ $$

Proposition 2.
Let $x(t)$ be particle trajectories driven by a smooth velocity field $v_{k}(x(t),t)$ over the time interval $[t_{k-1},t_{k}]$, where $h_{k}=t_{k}-t_{k-1}$. Assume that $v_{k}(x,t)$ is Lipschitz continuous in both $x$ and $t$. By dividing $[t_{k-1},t_{k}]$ into $S$ equal mini-intervals with grid points $t_{k-1,s}$ (where $s=0,1,\ldots,S$ and $t_{k-1,0}=t_{k-1}$, $t_{k-1,S}=t_{k}$), the following approximation holds:

$$ $\int_{t_{k-1}}^{t_{k}}\mathbb{E}_{x(t)}\left[\|v_{k}(x(t),t)\|^{2}\right]dt= \frac{S}{h_{k}}\sum_{s=0}^{S-1}\mathbb{E}\left[\|x(t_{k-1,s+1})-x(t_{k-1,s})\| ^{2}\right]+O\left(\frac{h_{k}^{2}}{S}\right).$ $$

As $h_{k}\to 0$ or $S\to\infty$, the error term $O\left(\frac{h_{k}^{2}}{S}\right)$ becomes negligible.

Proof:

Consider particle trajectories $x(t)$ driven by a sufficiently smooth velocity field $v_{k}(x(t),t)$ over the time interval $[t_{k-1},t_{k}]$, where $h_{k}=t_{k}-t_{k-1}$. We divide this interval into $S$ equal mini-intervals of length $\delta t=\frac{h_{k}}{S}$, resulting in grid points $t_{k-1,s}=t_{k-1}+s\delta t$ for $s=0,1,\ldots,S$, where $\delta t=\frac{t_{k}-t_{k-1}}{S}$.

Within each mini-interval $[t_{k-1,s},t_{k-1,s+1}]$, we perform a Taylor expansion of $x(t)$ around $t_{k-1,s}$:

$$ $x(t_{k-1,s+1})=x(t_{k-1,s})+v_{k}(x(t_{k-1,s}),t_{k-1,s})\delta t+\frac{1}{2} \frac{dv_{k}}{dt}\delta t^{2}+O(\delta t^{3}),$ $$

where $\frac{dv_{k}}{dt}$ denotes the total derivative of $v_{k}$ with respect to time.

The squared displacement over the mini-interval $[t_{k-1,s},t_{k-1,s+1}]$ is given by:

$$ $\displaystyle\|x(t_{k-1,s+1})-x(t_{k-1,s})\|^{2}$ $\displaystyle=\left\|v_{k}(x(t_{k-1,s}),t_{k-1,s})\delta t+\frac{1}{2}\frac{dv _{k}}{dt}\delta t^{2}+O(\delta t^{3})\right\|^{2}$ $\displaystyle=\|v_{k}(x(t_{k-1,s}),t_{k-1,s})\|^{2}\delta t^{2}+O(\delta t^{3}),$ $$

as we assume that $v_{k}$ is $L$-Lipschitz continuous and it follows that $|\frac{dv_{k}}{dt}|\leq L$. The higher-order terms $O(\delta t^{3})$ become negligible as $\delta t\to 0$.

Summing the expected squared displacements over all mini-intervals, we obtain:

$$ $\sum_{s=0}^{S-1}\mathbb{E}\left[\|x(t_{k-1,s+1})-x(t_{k-1,s})\|^{2}\right]= \delta t^{2}\sum_{s=0}^{S-1}\mathbb{E}\left[\|v_{k}(x(t_{k-1,s}),t_{k-1,s})\|^ {2}\right]+O\left(S\cdot\delta t^{3}\right).$ $$

Now, we examine the L.H.S. of Proposition 2 by approximating the integral of the expected squared velocity using a Riemann sum:

$$ $\displaystyle\int_{t_{k-1}}^{t_{k}}\mathbb{E}_{x(t)}\left[\|v_{k}(x(t),t)\|^{2 }\right]dt$ $\displaystyle=\delta t\sum_{s=0}^{S-1}\mathbb{E}\left[\|v_{k}(x(t_{k-1,s}),t_{ k-1,s})\|^{2}\right]+O\left(S\cdot\delta t^{2}\right)$ $\displaystyle=\delta t\left[\frac{1}{\delta t^{2}}\sum_{s=0}^{S-1}\mathbb{E} \left[\|x(t_{k-1,s+1})-x(t_{k-1,s})\|^{2}\right]+O(S\cdot\delta t)\right]+O(S \cdot\delta t^{2})$ $\displaystyle=\frac{1}{\delta t}\sum_{s=0}^{S-1}\mathbb{E}\left[\|x(t_{k-1,s+1 })-x(t_{k-1,s})\|^{2}\right]+O\left(S\cdot\delta t^{2}\right),\quad\square$ $$

where the Riemann sum error term $O(S\cdot\delta t^{2})$ arises from a well-known result (for instance, see Chapter 1 of Axler ()), given the assumption that $v_{k}$ is $L-$Lipschitz continuous.

Proposition 3. By Taylor’s expansions, and assume that $h_{k}=t_{k}-t_{k-1}$, we have:
1. $x(t_{k})-x(t_{k-1})=\int_{t_{k-1}}^{t_{k}}v_{k}(x(s),s)ds=h_{k}v_{k}(x(t_{k-1}
),t_{k-1})+O(h_{k}^{2})$
2. $\tilde{E}_{k}(x(t_{k}))=\tilde{E}_{k}(X(t_{k-1}))+h_{k}\nabla\tilde{E}_{k}(X(t
_{k-1}))\cdot v_{k}(x(t_{k-1}),t_{k-1})+O(h_{k}^{2})$

Proof:

The first equation directly follows from a simple Taylor expansion around time $t_{k-1}$. Similarly, by performing Taylor expansion around $t_{k-1}$:

$$ $\displaystyle\tilde{E}_{k}(x(t_{k}))$ $\displaystyle=\tilde{E}_{k}(X(t_{k-1}))+(x(t_{k})-x(t_{k-1}))\cdot\nabla\tilde {E}_{k}(X(t_{k-1}))+O(h_{k}^{2})$ $\displaystyle=\tilde{E}_{k}(X(t_{k-1}))+(h_{k}v_{k}(x(t_{k-1}),t_{k-1})+O(h_{k }^{2}))\cdot\nabla\tilde{E}_{k}(X(t_{k-1}))+O(h_{k}^{2})$ $\displaystyle=\tilde{E}_{k}(X(t_{k-1}))+h_{k}\nabla\tilde{E}_{k}(X(t_{k-1})) \cdot v_{k}(x(t_{k-1}),t_{k-1})+O(h_{k}^{2})$ $$

Density Ratio Estimation (DRE) By optimizing the following loss function:

$$ $\mathcal{L}_{k}(\theta_{k})=\mathbb{E}_{x(t_{k-1})\sim f_{k-1}}\left[\log(1+e^ {-r_{k}(x_{i}(t_{k-1}))})\right]+\mathbb{E}_{x(t_{k})\sim f_{k}}\left[\log(1+e ^{r_{k}(x_{i}(t_{k}))})\right],$ $$

the model learns an optimal $r^{*}(x;\theta_{k})=\log\frac{f_{k-1}(x)}{f_{k}(x)}$.

Proof:

Express the loss function as integrals over $x$:

$$ $\mathcal{L}_{k}=\int f_{k-1}(x)\log\left(1+e^{-r_{k}(x)}\right)\,dx+\int f_{k} (x)\log\left(1+e^{r_{k}(x)}\right)\,dx\triangleq\int l_{k}(x)dx.$ $$

Compute the functional derivative of $l_{k}(x)$ with respect to $r_{k}(x)$:

$$ $\frac{\delta l_{k}}{\delta r_{k}(x)}=-f_{k-1}(x)\cdot\frac{e^{-r_{k}(x)}}{1+e^ {-r_{k}(x)}}+f_{k}(x)\cdot\frac{e^{r_{k}(x)}}{1+e^{r_{k}(x)}}.$ $$

Next, we can set the derivative $\frac{\delta l_{k}}{\delta r_{k}(x)}$ to zero to find the minimizer $r_{k}^{*}(x)$:

$$ $r_{k}^{*}(x)=\ln\left(\frac{f_{k-1}(x)}{f_{k}(x)}\right).$ $$

Therefore, by concatenating each $r_{k}^{*}(x)$, we obtain

$$ $r^{*}(x)=\sum_{k=1}^{K}r_{k}^{*}(x)=\log\frac{f_{K-1}(x)}{f_{K}(x)}\cdot\frac{ f_{K-2}(x)}{f_{K-1}(x)}\cdots\frac{f_{0}(x)}{f_{1}(x)}=\log\frac{f_{0}(x)}{f_{ K}(x)}=\log\frac{\pi_{0}(x)}{q^{*}(x)},$ $$

the log density ratio between $\pi_{0}(x)$ and $q^{*}(x)$.

### A.2 Equivalence to Wasserstein gradient flow when β = 1 𝛽 1 \beta=1 italic_β = 1 and static Wasserstein regularization is used

In this section, we demonstrate the equivalence of Annealing Flow to the Wasserstein Gradient Flow when all $\beta_{k},k=1,2,\dots,K$, are set to 1, and when using a static Wasserstein regularization, instead of the dynamic Wasserstein regularization derived in Proposition 9.

Langevin Dynamics and Fokker-Planck Equation: Langevin Dynamics is represented by the following SDE:

$$ $dX_{t}=-\nabla E(X_{t})\,dt+\sqrt{2}\,dW_{t},$ (17) $$

where $E$ is the energy function of the equilibrium density $f(x,T)=q(x)$. Standard generative model training typically focuses on the case of a normal equilibrium, i.e., $E(x)=\frac{x^{2}}{2}$ and $q(x)\propto e^{-E(x)}$. Let $X_{0}\sim p_{X}$ and denote the density of $X_{t}$ by $\rho(x,t)$. The Langevin Dynamics also corresponds to the Fokker-Planck Equation (FPE), which describes the evolution of $\rho(x,t)$ towards the equilibrium $\rho(x,T)=q(x)$, as follows:

$$ $\partial_{t}\rho=\nabla\cdot(\rho\nabla E+\nabla\rho),\quad\rho(x,0)=p_{X}(x).$ (18) $$

In our algorithm, we focus on sampling from any distribution using its energy function, requiring only the unnormalized density. Therefore, $E(X_{t})$ represents the potential of any target density $q(x)$. We initialize samples from an easy-to-sample distribution, $\rho(x,0)=\pi_{0}(x)$, such as $N(0,I_{d})$, and aim to learn the trajectory between $\pi_{0}(x)$ and the target $q(x)$. Therefore, sampling from $q(x)$ boils down to first drawing $x(0)$ from $\pi_{0}(x)$ and then moving $x(0)$ along the learned trajectory to finally obtain $x(T)\sim q(x)$.

JKO Scheme:
The Jordan-Kinderlehrer-Otto (JKO) scheme  (Jordan et al., ) is a time discretization scheme for gradient flows to minimize $\text{KL}(\rho\|q)$ under the Wasserstein-2 metric. Given a target density $q$ and a functional $\mathcal{F}(\rho)=\text{KL}(\rho\|q)$, the JKO scheme approximates the continuous gradient flow of $\rho(x,t)$ by solving a sequence of minimization problems. Assume there are $K$ steps with time stamps $0=t_{0},t_{1},\cdots,t_{K}=T$, at each time stamp $t_{k}$, the scheme updates $\rho_{k}$ at each time step by minimizing the functional

$$ $\rho_{k}=\arg\min_{\rho}\left(\mathcal{F}(\rho)+\frac{1}{2\tau}W_{2}^{2}(\rho, \rho_{k-1})\right),$ (19) $$

where $W_{2}(\rho,\rho_{k-1})$ denotes the squared 2-Wasserstein distance between the probability measures $\rho$ and $\rho_{k}$. It was proven in Jordan et al. () that as $h=t_{k}-t_{k-1}$ approaches 0, the solution $\rho(\cdot,kh)$ provided by the JKO scheme converges to the solution of (18), at each step $k$.

The later works Xu et al. () have further proved that solving for the transport density $\rho_{k}$ by (19) is equivalent to solving for the transport map $\mathcal{T}_{k}$ by:

$$ $\mathcal{T}_{k}=\arg\min_{\mathcal{T}:\mathbb{R}^{d}\to\mathbb{R}^{d}}\left(KL (\mathcal{T}_{\#}\rho_{k-1}\|q)+\frac{1}{2\tau}\mathbb{E}_{x\sim\rho_{k-1}}\|x -\mathcal{T}_{k}(x)\|^{2}\right)$ (20) $$

Therefore, we immediately see that the Wasserstein gradient flow based on the discretized JKO scheme is equivalent to (6) when we set each $\tilde{f}_{k}(x)$ as the target distribution $q(x)$, i.e., when all the $\beta_{k}$ are set to 1, and when the second term in the objective (6) is relaxed to a static $W_{2}$ regularization.

This suggests that when the modes of the densities are not too far apart, and it is difficult to find a proper sequence of $\beta_{k}$, one can simply set all $\tilde{f}_{k}(x)$ in our algorithm as the target density $q(x)$, to construct a discretized sequence of transport maps based on Wasserstein gradient descent.

### A.3 Experimental Details

#### A.3.1 Evaluation metrics

To assess the performance of our model, we utilized two key metrics: Maximum Mean Discrepancy (MMD) and Wasserstein Distance, both of which measure the divergence between the true samples and the samples generated by the algorithms.

Maximum Mean Discrepancy (MMD)

MMD is a non-parametric metric used to quantify the difference between two distributions based on samples. Given two sets of samples $X_{1}\in\mathbb{R}^{n_{1}\times d}$ and $X_{2}\in\mathbb{R}^{n_{2}\times d}$, MMD computes the kernel-based distances between these sets. Specifically, we employed a Gaussian kernel:

$$ $k(x,y)=\exp\{-\alpha\|x-y\|_{2}^{2}\},$ $$

parameterized by a bandwidth $\alpha$. The MMD is computed as follows:

$$ $\text{MMD}(X_{1},X_{2})=\frac{1}{n_{1}^{2}}\sum_{i,j}k(X_{1}^{i},X_{1}^{j})+ \frac{1}{n_{2}^{2}}\sum_{i,j}k(X_{2}^{i},X_{2}^{j})-\frac{2}{n_{1}n_{2}}\sum_{ i,j}k(X_{1}^{i},X_{2}^{j}),$ $$

where $k(\cdot,\cdot)$ represents the Gaussian kernel. In our experiments, we set $\alpha=1/\gamma^{2}$ and $\gamma=0.1\cdot\text{median_dist}$, where median_dist denotes the median of the pairwise distances between the two datasets.

Wasserstein Distance

In addition to MMD, we used the Wasserstein distance, which measures the cost of transporting mass between distributions. Given two point sets $X\in\mathbb{R}^{d}$ and $Y\in\mathbb{R}^{d}$, we compute the pairwise Euclidean distance between the points. The Wasserstein distance is then computed using the optimal transport plan via the linear sum assignment method (from scipy.optimize package):

$$ $W(X,Y)=\frac{1}{n}\sum_{i=1}^{n}\|X_{r(i)}-Y_{c(i)}\|_{2},$ $$

where $r(i)$ and $c(i)$ are the optimal row and column assignments determined through linear sum assignment.

In all experiments, we sample 10,000 points from each model and generate 10,000 true samples from the GMM to calculate and report both MMD and Wasserstein distance. Note that the smaller the two metrics mentioned above, the better the sampling performance.

#### A.3.2 Hutchinson trace estimator

The objective functions in (10) and (12) involve the calculation of $\nabla\cdot v_{k}(x,t)$, i.e., the divergence of the velocity field represented by a neural network. This may be computed by brute force using reverse-mode automatic differentiation, which is much slower and less stable in high dimensions.

We can express $\nabla\cdot v_{k}(x,t)=\mathbb{E}_{\epsilon\sim N(0,I_{d})}\left[\epsilon^{T}J
_{v}(x)\epsilon\right]$, where $J_{v}(x)$ is the Jacobian of $v_{k}(x,t)$ at $x$. Given a fixed $\epsilon$, we have $J_{v}(x)\epsilon=\lim_{\sigma\to 0}\frac{v_{k}(x+\sigma\epsilon)-v_{k}(x)}{\sigma}$, which is the directional derivative of $v_{k}$ along the direction $\epsilon$. Thus, for a sufficiently small $\sigma>0$, we can propose the following estimator  (Hutchinson, ; Xu et al., ):

$$ $\nabla\cdot v_{k}(x,t)\approx\mathbb{E}_{\epsilon\sim N(0,I_{d})}\left[ \epsilon^{T}\frac{v_{k}(x+\sigma\epsilon,t)-v_{k}(x,t)}{\sigma}\right].$ (21) $$

This approximation becomes exact as $\sigma\to 0$. In our experiments, we set $\sigma=\sigma_{0}/\sqrt{d}$ with $\sigma_{0}=0.02$.

#### A.3.3 Other Annealing Flow settings

Time stamps and numerical integration

By selecting $K$ values of $\beta$, we divide the original time scale $[0,1]$ of the Continuous Normalizing Flow (2) and (3) into $K$ intervals: $[t_{k-1},t_{k}]$ for $k=1,2,\dots,K$. Notice that the learning of each velocity field $v_{k}$ depends only on the samples from the $(k-1)$-th block, not on the specific time stamp. Therefore, we can re-scale each block’s time interval to $[0,1]$, knowing that using the time stamps $[(k-1)h,kh]$ yields the same results as using $[0,1]$ for the neural network $v_{k}(x,t)$. For example, the neural network will learn $v_{k}(x,0)=v_{k}(x,(k-1)h)$ and $v_{k}(x,1)=v_{k}(x,kh)$, regardless of the time stamps.

Recall that we relaxed the shortest transport map path into a dynamic $W_{2}$ regularization loss via Proposition 2. This requires calculating intermediate points $x(t_{k-1,s})$, where $s=0,1,\dots,S$. We set $S=3$, evenly spacing the points on $[t_{k-1},t_{k}]$, resulting in the path points $x(t_{k-1}),x(t_{k-1}+h_{k}/3),x(t_{k-1}+2h_{k}/3),x(t_{k})$. To compute each $x(t_{k-1,s})$, we integrate the velocity field $v_{k}$ between $t_{k-1}$ and $t_{k-1,s}$, using the Runge-Kutta method for numerical integration. Additionally, for each $x(t_{k-1,s})$, we calculate the velocity field at an intermediate time step between $t_{k-1,s-1}$ and $t_{k-1,s}$ to enable accurate numerical integration. Specifically, to calculate $x(t+h)$ based on $x(t)$ and an intermediate time stamp $t+\frac{h}{2}$:

$$ $x(t+h)=x(t)+\frac{h}{6}\left(k_{1}+2k_{2}+2k_{3}+k_{4}\right),$ $$

$$ $k_{1}=v(x(t),t),\quad k_{2}=v\left(x(t)+\frac{h}{2}k_{1},t+\frac{h}{2}\right)$ $$

$$ $k_{3}=v\left(x(t)+\frac{h}{2}k_{2},t+\frac{h}{2}\right),\quad k_{4}=v\left(x(t )+hk_{3},t+h\right)$ $$

Here, $h$ is the step size, and $v(x,t)$ represents the velocity field.

The choice of $\beta_{k}$

In the experiments on Gaussian Mixture Models (GMM) and Exp-Weighted Gaussians with various dimensions and radii, we set the number of intermediate $\beta_{k}$ values to 8, equally spaced such that $\beta_{0}=0$, $\beta_{1}=1/8$, $\beta_{2}=2/8$, $\dots$, $\beta_{8}=1$. We chose the easy-to-sample distribution $\pi_{0}(x)$ as $N(0,I_{d})$. Finally, we added 2 refinement blocks. The intermediate distributions are defined as:

$$ $\tilde{f}_{k}(x)=\pi_{0}(x)^{1-\beta_{k}}\tilde{q}(x)^{\beta_{k}}.$ $$

In the experiment on the Truncated Normal Distribution, we did not select $\beta_{k}$ in the same manner as for the GMM and Exp-Weighted Gaussian distributions. Instead, following the same Annealing philosophy, we construct a gradually transforming bridge from $\pi_{0}(x)$ to $\tilde{q}(x)=1_{|x|\geq c}N(0,I_{d})$ by setting each intermediate density as:

$$ $\tilde{f}_{k}(x)=1_{\|x\|\geq c/(k+1)}N(0,I_{d}).$ $$

This choice also demonstrates that our Annealing Flow is highly flexible and capable of handling a wide range of challenging distributions.

In the experiment on funnel distributions, we set all $\beta_{k}=1$. Therefore, as discussed in Appendix A.2, the algorithm becomes equivalent to a Wasserstein gradient descent problem. We also set the number of blocks to 8, consistent with the other experiments. This indicates that when the densities are largely concentrated in one region, one can simply set $\beta_{k}$ to 1 and use a few blocks to find the optimal transport path based on Wasserstein gradient descent.

The objective

During the experiments, we found that using the Taylor approximation (as described in Proposition 3, with a slight modification such that the expansion is around $x(t_{k})$, allowing the loss to include the velocity field term): $\tilde{E}_{k}(x_{k-1})-\tilde{E}_{k}(x_{k})=(-h_{k})\nabla E(x_{k})\cdot v_{k}$, and replacing the energy function $\tilde{E}_{k}(x_{k})$ generally led to better performance. In our experiments on the GMM, Funnel distribution, and Exp-weighted Gaussian, we consistently used this form. For the experiments on the Truncated Normal and Bayesian Logistic Regression, the original $\tilde{E}_{k}(x_{k})$ was used.

Neural networks and selection of other hyperparameters

The neural network structure in our experiments is consistently set with hidden layers of size 32-32-32. During implementation, we observed that when $d\leq 5$, even a neural network with a single hidden layer of size 32 can perform well for sampling. However, for consistency across all experiments, we uniformly set the structure to 32-32-32.

We sample 100,000 data points from $N(0,I_{d})$ for training, with a batch size of 1,000. The Adam optimizer is used with a learning rate of 0.0001, and the maximum number of iterations for each block $v_{k}$ is set to 1,000. An additional two blocks are added for refinement after $\beta_{K}=1$.

Different numbers of test samples are used for reporting the experimental results: 5,000 points are sampled and plotted for the experiment on Gaussian Mixture Models, 5,000 points for the experiment on Truncated Normal Distributions, 10,000 points for the experiment on Funnel Distributions, and 10,000 points for the experiment on Exp-Weighted Gaussian with 1,024 modes in 10D space.

#### A.3.4 Bayesian logistic regression

We use a hierarchical Bayesian structure for logistic regression across a range of datasets provided by [LIBSVM](https://www.csie.ntu.edu.tw/~cjlin/libsvmtools/datasets/binary.html). The detailed setting of the Bayesian Logistic Regression is as follows.

We adopt the same Bayesian logistic regression setting as described in Liu & Wang (), where a hierarchical structure is assigned to the model parameters. The weights $\beta$ follow a Gaussian prior, $p_{0}(\beta|\alpha)=N(\beta;0,\alpha^{-1})$, and $\alpha$ follows a Gamma prior, $p_{0}(\alpha)=\text{Gamma}(\alpha;1,0.01)$. The datasets used are binary, where $x_{i}$ has a varying number of features, and $y_{i}\in\{+1,-1\}$ across different datasets. Sampling is performed from the posterior distribution:

$$ $p(\beta,\alpha|D)\propto Gamma(\alpha;1,0.01)\cdot\prod_{d=1}^{D}N(\beta_{d};0 ,\alpha^{-1})\cdot\prod_{i=1}^{n}\frac{1}{1+\exp(-y_{i}\beta^{T}x_{i})},$ $$

We set $\beta_{k}=1$ and use 8 blocks to train the Annealing Flow.

During testing, we use all algorithms to sample 1,000 particles of $\beta$ and $\alpha$ jointly, and use $\{\beta^{(i)}\}_{i=1}^{1000}$ to construct 1,000 classifiers. The mean accuracy and standard deviation are then reported in Table 2. Additionally, the average log posterior in Table 2 is reported as:

$$ $\frac{1}{|D_{\text{test}}|}\sum_{x,y\in D_{\text{test}}}\log\frac{1}{|C|}\sum_ {\theta\in C}p(y|x,\theta).$ $$

#### A.3.5 Importance flow

We report the results of the importance sampler (discussed in 16) for estimating $\mathbb{E}_{x\sim N(0,I)}\left[1_{\|x\|\geq c}\right]$ with varying $c$ and dimensions, based on our Annealing Flow. To estimate $\mathbb{E}_{x\sim N(0,I)}\left[1_{\|x\|\geq c}\right]$ , we know that the theoretically optimal proposal distribution which can achieve 0 variance is $\tilde{q}^{*}(x)=1_{\|x\|\geq c}N(0,I)$. Then the estimator becomes:

$$ $\mathbb{E}_{X\sim\pi_{0}(x)}\left[h(X)\right]=\mathbb{E}_{X\sim q^{*}(x)}\left [\frac{\pi_{0}(x)}{q^{*}(x)}\cdot h(x)\right]\approx\frac{1}{n}\sum_{i=1}^{n} \frac{\pi_{0}(x_{i})}{q^{*}(x_{i})}\cdot h(x_{i}),\quad x_{i}\sim q^{*}(x),$ $$

where $\pi_{0}(x)=N(0,I_{d})$, $h(x)=1_{\|x\|\geq c}$ and $q^{*}(x)=Z\cdot\tilde{q}^{*}(x)$.

Therefore, the Importance Flow consists of two parts: First, using Annealing Flow to sample from $\tilde{q}^{*}(x)$; second, constructing a Density Ratio Estimation (DRE) neural network using samples from $\{x_{i}\}_{i=1}^{n}\sim\tilde{q}^{*}(x)$ and $\{y_{i}\}_{i=1}^{n}\sim N(0,I_{d})$, as discussed in Section 5.2. After successful training, the estimator becomes

$$ $\frac{1}{n}\sum_{i=1}^{n}DRE(x_{i})\cdot h(x_{i}).$ $$

The Naive MC results comes from directly using $\{y_{i}\}_{i=1}^{n}\sim N(0,I_{d})$ to construct estimator $\frac{1}{n}\sum_{i=1}^{n}1_{\|y_{i}\|\geq c}$. It can be observed that when $c\geq 6$, the Naive MC methods consistently output 0 as the result.

In our experiment, we use a single DRE neural network to construct the density ratio between $\pi_{0}(x)$ and $q^{*}(x)=Z\cdot 1_{\|x\|\geq c}N(0,I)$ directly. The neural network structure consists of hidden layers with sizes 64-64-64. The size of the training data is set to 100,000, and the batch size is set to 10,000. We use 30 to 70 epochs for different distributions, depending on the values of $c$ and dimension $d$. The Adam optimizer is used, with a learning rate of 0.0001. The test data size is set to 1,000, and all results are based on 200 estimation rounds, each using 500 samples.

#### A.3.6 Details of other algorithms

The Algorithm 2, 3, and 4 introduce the algorithmic framework of Metropolis-Hastings (MH), Hamiltonian Monte Carlo (HMC), and Parallel Tempering (PT) compared in our experiments.

Figure: Algorithm 2 Metropolis-Hastings Algorithm

Figure: Algorithm 3 Hamiltonian Monte Carlo (HMC)

Figure: Algorithm 4 Parallel Tempering Algorithm

In our experiments, we set the proposal density as $q(x^{\prime}|x)=\mathcal{N}(x;0,I_{d})$. We use 5 replicas in Parallel Tempering (PT), with a linear temperature progression ranging from $T_{1}=1.0$ to $T_{\text{max}}=2.0$, and an exchange interval of 100 iterations. For HMC, we set the number of leapfrog steps to 10, with a step size $(\epsilon)$ of 0.01, and the mass matrix $M$ is set as the identity matrix. Additionally, we use the default hyperparameters as specified in SVGD  (Liu & Wang, ), MIED  (Li et al., ), and AI-Sampler  (Egorov et al., ). In the actual implementation, we found that the time required for SVGD to converge increases significantly with the number of samples. Therefore, in most experiments, we sample 1000 data points at a time using SVGD, aggregate the samples, and then generate the final plot.

### A.4 More Results

We adopt the standard Annealing Flow framework discussed in this paper for experiments on Gaussian Mixture Models (GMM), Truncated Normal distributions, and Exp-Weighted Gaussian distributions. For experiments on funnel distributions, we set each $\tilde{f}_{k}(x)$ as the target $q(x)$, under which the Annealing Flow objective becomes equivalent to the Wasserstein Gradient Flow based on the JKO scheme, as discussed in A.2. Please refer to A.3.3 for $\beta_{k}$ selections.

Gaussian Mixture Models (GMM)

Figure: (a) True
Refer to caption: extracted/5874394/d=2_GMM_c8_heatmap.png

Evaluation Metrics: We report 1) the Maximum Mean Discrepancy (MMD) and 2) the Wasserstein Distance for the GMM experiments, as both metrics require access to true data samples. The results for these metrics are presented in Table 4. Please refer to A.3.1 for more details.

**Table 4: MMD and Wasserstein Distance results: $\cdot/\cdot$ represents MMD/Wasserstein. The first row corresponds to $d=\{\text{dimension}\}$ GMM-$\{\text{Number of Modes}\}$.**
|  | $d=2$ GMM-8 | $d=2$ GMM-12 | $d=3$ GMM-8 | $d=4$ GMM-16 | $d=5$ GMM-32 | $d=6$ GMM-64 |
| --- | --- | --- | --- | --- | --- | --- |
| AF | 2.32E-03/7.38E-01 | 3.01E-03/8.05E-01 | 5.82E-03/1.97E+00 | 1.25E-03/3.33E+00 | 1.57E-03/2.82E+00 | 4.31E-03/3.53E+00 |
| HMC | 7.33E-02/6.28E+00 | 9.06E-02/8.73E+00 | 9.92E-02/1.12E+01 | 9.76E-02/1.98E+01 | 2.14E-01/2.53E+01 | 2.15E-01/3.03E+01 |
| PT | 6.27E-02/5.71E+00 | 9.01E-02/7.91E+00 | 8.83E-02/1.07E+01 | 8.98E-02/1.53E+01 | 1.18E-01/1.83E+01 | 1.05E-01/2.13E+01 |
| SVGD | 9.35E-02/9.97E+00 | 1.85E-01/1.82E+01 | 9.81E-02/1.13E+01 | 9.63E-02/2.07E+01 | 1.98E-01/2.45E+01 | 1.32E-01/2.34E+01 |
| MIED | 2.34E-03/8.01E-01 | 6.28E-03/9.35E-01 | 8.01E-03/2.52E+00 | 3.88E-02/0.89E+01 | 9.88E-03/7.89E+00 | 2.03E-02/1.13E+01 |
| AIS | 2.33E-03/7.92E-01 | 4.02E-03/8.13E-01 | 7.55E-02/2.38E+00 | 5.26E-03/5.53E+00 | 6.37E-03/3.83E+00 | 1.87E-02/9.73E+00 |

Truncated Normal Distribution

Relaxations are applied to the Truncated Normal Distribution in all experiments except for MH, HMC, and PT. Specifically, we relax the indicator function $1_{\|x\|\geq c}$ to $\frac{1}{1+\exp(-k(\|x\|-c))}$. We set $k=20$ for all experiments. AIS is designed for continuous densities, and we similarly relax the densities in SVGD and MIED, following the approach used in AF. The resulting plots are as follows:

Figure: (a) True
Refer to caption: extracted/5874394/d=2_indicator_c6_heatmap.png

Each algorithm draws 5,000 samples. It can be observed that MCMC-based methods, including HMC and PT, produce many overlapping samples. This occurs because when a new proposal is rejected, the algorithms retain the previous sample, leading to highly correlated sample sets.

**Table 5: Proportion of Annealing Flow Samples Within c, Across Different Dimensions**
| Proportion Within c | $c=4$ | $c=6$ | $c=8$ |
| --- | --- | --- | --- |
| $D=2$ | 0.17% | 0.18% | 1.78% |
| $D=3$ | 0.20% | 0.23% | 3.23% |
| $D=4$ | 0.68% | 1.48% | 3.68% |
| $D=5$ | 1.46% | 3.37% | 4.12% |
| $D=\mathbf{10}$ | 2.13% | 4.68% | 7.13% |

Figure: (a) True
[Uncaptioned image]: extracted/5874394/d=2_indicator_c4_heatmap.png

For dimensions $d>2$, visualizing the results by comparing the sample positions using a red sphere surface becomes challenging. Therefore, we calculate the proportion of samples within radius $c$. A lower proportion indicates better sampling performance. Table 5 presents these results. We also calculate the proportion of the surface $\|x\|=c$ covered by the samples for AF, MH, HMC, and PT. In all experiments with the Truncated Normal distribution, AF covers more than 95% of the surface area. However, when $d\geq 3$ and $c\geq 6$, all other methods cover less than 70% of the surface area.

Funnel Distribution

In the main paper, we present the sampling methods for the funnel distribution with $d=5$, projected onto a 3D space. To assess the sample quality, here we present the corresponding results projected onto a 2D space, plotted alongside the density heat map.

Figure: (a) AF
Refer to caption: extracted/5874394/d=5_funnel_c8_Annealing.png

Figure: (a) True
Refer to caption: extracted/5874394/d=5_funnel_True_sigma_0.9_12.png

As seen from both figures, our AF method achieves the best sampling performance on the funnel distribution, while other methods, such as MIED and AIS, fail to capture the full spread of the funnel’s tail. Additionally, PT, SVGD, and AIS all fail to capture the sharp part of the funnel’s shape.

Exp-Weighted Gaussian

In the main paper, we present the sampling methods for the Exp-Weighted Gaussian distribution with 1024 modes in a 50D space, projected onto a 3D space. To better assess the sample quality, we now present the corresponding results projected onto 2D and 1D spaces, plotted alongside the heat map and the true density, respectively.

Figure: (a) AF
Refer to caption: extracted/5874394/d=10_exponential_Annealing.png

Figure: (a) True
Refer to caption: extracted/5874394/d=10_exponential_heatmap.png

Figure: (a) AF
Refer to caption: extracted/5874394/d=10_exponential_c10_Annealing_1.png

As seen in Figures 13 and 14, AF produces balanced samples, and its 1D projection closely matches the true density. While both SVGD and MIED captured around 800 to 900 modes, their samples across the modes are imbalanced, as observed in the figures.

### A.5 Comparisons

**Table 6: Comparison of Different Sampling Methods**
| Method | Key Characteristics | Advantages | Disadvantages |
| --- | --- | --- | --- |
| Annealing Flow (AF) | - Continuous Normalizing Flow-based approach.<br>- Leverages annealing principles for sampling challenging high-dimensional, multi-modal distributions.<br>- Uses transport maps to transform samples from a base distribution to the target distribution. | - Independent sampling.<br>- Balanced mode exploration.<br>- Handles multi-modal distributions effectively.<br>- Once trained, the sampling process is very fast<br>- Scales linearly with sample size and dimensionality. | - Requires pre-training, which can be computationally expensive. |
| MCMC | - Metropolis-Hastings, Parallel Tempering, Hamiltonian Monte Carlo (HMC) variants.<br>- Samples sequentially from the target distribution, with each sample depending on the previous one. | - Flexible, general-purpose.<br>- Doesn’t require pre-training. | - Slow mixing time.<br>- Struggles with multi-modal distributions.<br>- Sample correlation reduces effective sample size (ESS).<br>- Imbalanced mode exploration. |
| Particle-Based Optimization (SVGD, MIED) | - Relies on particle dynamics and kernel methods to sample from the target distribution. | - No burn-in period.<br>- Less sample correlation than MCMC.<br>- Encourages global search. | - Kernel computations scale polynomially with sample size and dimensionality.<br>- Sensitive to kernel hyperparameters. |
| NN-Assisted MCMC | - Uses neural networks to accelerate or guide MCMC methods.<br>- Combines the expressive power of neural networks with MCMC. | - Can speed up the explorations of MCMC methods.<br>- Leverages NN for improved sampling efficiency. | - Inherits some limitations of MCMC, such as slow mixing, correlated samples, and mode imbalance. |
| Score-based Diffusion | - Learns score functions to iteratively perturb samples towards the target distribution.<br>- Combines the expressive power of neural networks with MCMC. | - Strong theoretical guarantees for sampling specific distributions. | - Limited generalization to arbitrary distributions, as score functions are analytically derived.<br>- Challenging in complex, high-dimensional distributions |

Annealing Flow (AF) requires pre-training, typically taking 10-20 minutes for tasks with dimensions $<10$, and around 30 minutes for tasks around dimension 50. Once trained, AF samplers are very efficient: generating 10,000 samples in just 1.5 seconds. These pre-trained samplers can be reused at any time, offering significant speed advantages. In contrast, MCMC methods, such as Metropolis-Hastings or Hamiltonian Monte Carlo, require about 1 minute to sample 10,000 points, and their performance deteriorates in high-dimensional, multi-modal settings. Moreover, particle-based methods, like SVGD, struggle significantly when generating more than 3,000 samples, requiring about 20 minutes for that many samples. Therefore, we believe that users can take advantage of AF’s offline training, as it allows the samplers to be trained once and then efficiently reused for sampling whenever needed.

### A.6 Importance Flow

The importance flow discussed and experimented with in this paper requires a given form of $\pi_{0}(x)$, and thus, a given form of $\tilde{q}^{*}(x)=\pi_{0}(x)\cdot|h(x)|$ for estimating $\mathbb{E}_{X\sim\pi_{0}(x)}\left[h(X)\right]$. In our experimental settings, $\tilde{q}^{*}(x)=1_{\|x\|\geq c}N(0,I_{d})$ can be regarded as the Least-Favorite-Distribution (LFD). We conducted a parametric experiment for the case where $\tilde{q}^{*}(x)$ has the given analytical form.

However, we believe future research may extend this approach to a distribution-free model. That is, given a dataset without prior knowledge of its distribution, one could attempt to learn an importance flow for sampling from its Least-Favorite Distribution (LFD) while minimizing the variance. For example, in the case of sampling from the LFD and obtaining a low-variance IS estimator for $P_{x\sim\pi(x)}(\|x\|\geq c)$, one may use the following distribution-free loss for learning the flow:

$$ $\min_{\theta}\frac{1}{n}\sum_{i=1}^{n}\left[1\{\mathcal{T}(x_{i};\theta)\leq c \}\cdot\|\mathcal{T}(x_{i};\theta)-c\|^{2}\right]+\gamma\int_{0}^{1}\|v(x(t),t ;\theta)\|^{2},$ (22) $$

where the first term of the loss pushes the dataset $\{x_{i}\}_{i=1}^{n}$ towards the Least-Favorite tail region, while the second term ensures a smooth and cost-optimal transport map. Note that the above loss assumes no prior knowledge of the dataset distribution $\pi(x)$ or the target density $q(x)$.

Xu et al. () has also explored this to some extent by designing a distributionally robust optimization problem to learn a flow model that pushes samples toward the LFD $Q^{*}$, which is unknown and learned by the model through a risk function $\mathcal{R}(Q^{*},\phi)$. Such framework has significant applications in adversarial attacks, robust hypothesis testing, and differential privacy. Additionally, the recent paper by Ribera Borrell et al. () introduces a dynamic control loss for training a neural network to approximate the importance sampling control. We believe that by designing an optimal control loss in line with the approaches of these two papers, one can develop a distribution-free Importance Flow for sampling from the LFD of a dataset while minimizing the variance of the adversarial loss, which can generate a greater impact on the fields of adversarial attacks and differential privacy.

## Peer Reviews

confidence: 2
contribution: 3
initial_rating: 6
presentation: 2
questions: 1. The author should explain  the detail of training  including how the minimization problem is solved and how the parameters would affect the sampling performance. 
2. Comments or preliminary results  on how it can appplied to any real high dimensional applications (>>50), for example a real physical machine learning problem. 
3. Other comparisons results in terms of computational cost.
soundness: 3
strengths: The paper present a continuous Normalizing flow  based method for sampling from challenging high dimensional multi-modal distributions.  There are some theoretically insights on the proposed models using the  predefined intermediate and the numerical results show the superiority of the proposed sampling method under variety of criteria.
summary: The paper proposed a so-called Annealing flow (AF) sampling method for  high-dimensional multi-modal distributions.  The key idea is to learn a continuous normalizing flow (CNF)  based transport map, guided by annealing process, to transition samples from an easy-to-sample distribution to the target distribution. The feature of  AF training does not rely on samples from the target distribution. The results in dimensional up to 50  show that AF ensures effective and balanced mode exploration, especially for multi-model distribution by comparing  AF to other SOTA methods though various challenging distributions, particularly in high-dimensional and multi-modal settings, including  AF’s potential for sampling the least favorable distributions.
weaknesses: 1. The details of training process is missing. Especially, the cost of the training process of CNF flow is not mentioned,  and also how many intervals are needed in order to obtain a stable sampling process. It seems also crucial especially how the overall sampling results depend on the flow interval. 

2. It is unclear how the method can be generalized to real high dimensional problems arising from real applications and distribution free problems where only samples are available. 

3. The sample efficiency (both sampling speed and quality) is not commented and compared with different methods.

---

confidence: 4
contribution: 1
initial_rating: 3
presentation: 3
questions: 1. How sensitive is the weight hyperparameter $\gamma$ that controls the balance between the two terms in the dynamic optimal transport loss? Does the value of $\gamma$ vary a lot across different experiments? How the values of $\gamma$ were chosen in the experiments?

2. How does the hyperparameters $K$ (the number of intermediate distributions) and $S$ (the number of discretization grids within each transition between two intermediate distributions) scale with the dimensionality and multi-modaility of the target distribution?

3. [Typo] What is "MEID" in table 2? Is it "MIED"?
soundness: 2
strengths: 1. The paper is generally well-written. The description of the proposed method is clear and easy to follow.
2. The proposed sampler is able to explore modes in multi-modal target distributions, which is a desired property in many applications.
3. Once the model is trained, the sampling cost is low since it does not require running MCMC chains and sampling scales linearly with the sample size and dimensionality.
4. The proposed method is evaluated on various experimental settings and has comparable or better results than some of the tradition MCMC, particle-based and NN-based samplers, showing its good applicability.
summary: This paper proposes to fit continuous normalizing flows (CNFs) to multi-modal unnormalized target distributions using annealing and optimal transport techniques. The proposed method first defines a sequence of annealed distributions interpolating the target distribution and a simple base distribution. Leveraging the dynamic optimal transport objective, the authors propose to train a sequence of CNFs to model the transition between each pair of neighboring annealed distributions. This method encourages the model to explore modes without relying on any target samples or running MCMC chains. Various experiments show that the proposed method has comparable or better performance than some of the traditional MCMC, particle-based and NN-based methods.
weaknesses: 1. The manuscript contains a GitHub link in Line 369-370, which reveals the authors' identity.

2. The proposed method lacks technical novelty and is a bit incremental, since it directly applies an optimal transport loss with KL relaxation to train a sequence of CNFs to model the transitions between neighboring annealed distributions, which is a trivial combination of existing techniques. 

3. The scalability and reliability of the proposed training algorithm is questionable, since it requires training K seperate CNFs sequentially, each modeling the transition between one of the K neighboring pairs of annealed distributions. Furthermore, within each transition, the interval is further discretized into S grid points in order to estimate the integral.
- If K is small, the transition will be difficult and the model can fail to capture some modes. If S is small, the discretization error in numerical integration will be large and the gradient for updating the model will be inaccurate. If K and S are large, the proposed training algorithm will be very expensive as it requires backpropagating through all these steps and hierarchies.
- The approximation error in each transition will accumulate as the sampler gets closer to the target distribution due to the bootstrap-style training method.

4. It is also unclear how the proposed method should be positioned among other flows-based or NN-based samplers with similar techniques, since many important recent related works are not discussed in the paper, including 
- flow-based sampler trained with annealing: [1, 2, 3]
- score/diffusion-based sampler: [4, 5, 6, 7, 9]
- optimal control-based sampler: [8]

5. Several closely related baselines are missing in the experiment section:
- At the moment, most baselines are MCMC/particle-based samplers.
- Since the proposed method is a flow-based model trained with the annealing technique, it should be empirically compared to at least the following models which use similar flow-annealing techniques: AFT [1], CRAFT [2], FAB [3]. 
- In addition, since the proposed method is broadly an NN-based sampler, ideally it should also be compared with some of other types of SOTA NN-based samplers, such as iDEM [6], DDS [7], PIS [8], PDIS [9].

6. For the toy 1D experiment in Figure 2, it would be more convincing if 
- the two modes in the target distribution are unbalanced (i.e., one has a larger density than the other), as suggested by [5] to test whether the proposed method exhibits a similar blindeness issue as in some of the score-based samplers.
- it also includes results of other flow and annealing-based baselines in this comparison to get a sense of how these similar methods compare to each other, since they are more related to the proposed method in this paper.

7. [Optional] The proposed method is only tested on synthetic problems. Perhaps for future work, it would be nice if the authors could consider some real-world problems, such as sampling from Boltzmann distributions (i.e., Boltzmann generator, see e.g., [3, 6]) which share exactly the same problem setting as in this paper.

**Reference**:

[1] Arber et al., “Annealed flow transport Monte Carlo”, ICML 2021

[2] Matthews et al., “Continual repeated annealed flow transport Monte Carlo" ICML 2022.

[3] Midgley et al., “Flow annealed importance sampling bootstrap”, ICLR 2023

[4] Huang et al., “Reverse diffusion Monte Carlo”, ICLR 2024

[5] Chen et al., “Diffusive Gibbs sampling”, ICML 2024

[6] Akhound-Sadegh et al., “Iterated denoising energy matching for sampling from Boltzmann densities”, ICML 2024

[7] Vargas et al., “Denoising diffusion samplers”, ICLR 2023

[8] Zhang et al., “Path integral sampler: a stochastic control approach for sampling”, ICLR 2022

[9] Phillips et al., “Particle denoising diffusion sampler”, ICML 2024

---

confidence: 5
contribution: 2
initial_rating: 3
presentation: 2
questions: - Is the condition of optimality of the transport crucial to the success of the method? 

- I am surprised by PT’s result in Figure 2. I would expect that this exact sampler designed to sample from multimodal distribution works perfectly in this 1d example. Can the author justify underwhich circumstances the imbalance result was obtained?

- Table 3, what are the expected ground truth values?

- The notations between $\tilde f_k$ (the target annealing path of distributions)  and $f_k$ (the intermediary push forwards of the CNF) is maybe inconsistent in some places. For instance in (8) and (10), should the expectation be over $\tilde f_{k-1}$ instead of $f_{k-1}$ ?
soundness: 3
strengths: - The proposed methods addresses the important point of mode coverage in the task of sampling multimodal distributions in high dimension, which is demonstrated in numerical experiments.
summary: The present paper proposes a sampler based on a continuous normalizing flow. The idea is to decompose the transport into sub-transport tasks following an annealing scheme between a simple base distribution and the target. Velocity fields are learned independently for each such transport task so as to minimize a Kullback-Leibler divergence with a transport cost penalization. Once the networks are trained, the sampling is operated by simply sampling from the continuous NF. 

The paper also points that the method can be employed to learn a flow as a proposal distribution in a rare event  importance sampling (IS) scenario. It is proposed that the IS weights are estimate by training a neural network to estimate the density ratio between the flow and the target. 

Numerical experiments report a good mode exploration and qualitative good coverage for the proposed method compared to diverse samplers not relying on deep learning.
weaknesses: - The paper lacks a related works section. Namely, the connections to  NN-assisted MCMCs, which are in my opinion the most related methods, are not discussed and the performances of these approaches are not reported. In my experience, these methods have better performance than what the paper seems to indicate. In particular, they typically do correct for the mode imbalance that a traditional MCMC based on local updates could not do. 

- In this regard, the novelty of the paper is lesser than what it appears. Annealed Flow Transport  (Arbel et al 2021 in the paper) proposes a very related strategy. 

- Limitations are not properly discussed: 
	- The proposed sampling methods lacks theoretical guarantees of the accuracy of the sampling conversely to NN-assisted (Markov Chain) Monte Carlo samplers, such as Neural IS [3], Flow MC (Gabrie et al 2021, 2022 in the paper) or Annealed Flow Transport (Arbel et al 2021 in the paper). 
	- The cost of training only one velocity field at the time is not discussed, nor how to choose the number of these intermediary steps.
	- For the importance flow method as well, the paper claims “The estimator is unbiased and can achieve zero variance theoretically” but this is only true if the density estimation ratio is predicted exactly by the neural network. 



Minor: 
- About “score based sampling that do not rely on neural networks” (line 66), the paper should also cite RDMC [1] and SLIPS [2].



[1] Xunpeng Huang and Hanze Dong and Yifan HAO and Yian Ma and Tong Zhang. Reverse Diffusion Monte Carlo. The Twelfth International Conference on Learning Representations, 2024, https://openreview.net/forum?id=kIPEyMSdFV

[2] Grenioux, Louis, Maxence Noble, Marylou Gabrié, and Alain Oliviero Durmus. “Stochastic Localization via Iterative Posterior Sampling.” In Proceedings of the 41st International Conference on Machine Learning, 16337–76. PMLR, 2024. https://proceedings.mlr.press/v235/grenioux24a.html.

[3] Müller, Thomas, Brian Mcwilliams, Fabrice Rousselle, Markus Gross, and Jan Novák. “Neural Importance Sampling.” ACM Transactions on Graphics 38, no. 5 (October 31, 2019): 1–19. https://doi.org/10.1145/3341156.

---

confidence: 4
contribution: 3
initial_rating: 5
presentation: 2
questions: - in the introduction in connection with normalizing flows I am missing stochastic normalizing flows, see
\\
H. Wu, J. Köhler, and F. Noe, Stochastic normalizing flows, in Advances in Neural Information
Processing Systems 33, 5933–5944, 2020.
\\ 
P. Hagemann, J. Hertrich, G. Steidl,
   Stochastic normalizing flows for inverse problems: a Markov Chains viewpoint
SIAM Journal on Uncertainty Quantification 10 (3), 1162-1190,2022
	
- in (3): don't call the \textbf{value} of the integral $\mathcal T$; later this is your OT transport map.
- line 126: hint to Fig. 1 is too early here since $t_1,t_2$  is not explained so far
- line 152: ,,pushes density from $f_{k-1}(x)$ to $f_{k}(x)$, since these are function values, please skip the $x$;
this appears also at other places (but is a minor remark)
- line 159: I would not say that (3) is equivalent to (6)
- in (6) and in the following: what is the pushforward of $\tilde{f}_{k-1}$ by $\mathcal{T}$  for an unnormalized density? Definition of KL for 
unnormalized densities ?; see Prop 1 
The authors should write the definitions down, then they would see that this notation does not work.
- In general the authors switch between $f_k$ and $\rho_k$

- Prop 1 (Appendix A): Despite that the tilde notation does not work  and that $T$ is indeed $\mathcal T$ the authors have to correct
  line 770: first equality is wrong; you need here and in the following $\mathbb E_{x \sim \mathcal T _\#  \tilde f_{k-1} }$ (unfortunately the system does not translates this latex formula correctly, but the authors hopefully see what I mean);
 - from 774 to 776 is wrong and becomes correct with my correction above.
- First equality in 794 is wrong and in the final equality it should be $\log \rho_{k-1} (x(t_{k-1})) - \int ...$

- Prop. 2 is folklore; in formula (9) in the expectation value $x(t) \sim ??$ is missing

- Prop 3: Do not start ,,By Taylor expansion''. This is already the proof.
Write the correct assumptions on x and $\tilde E$ for this Taylor expansion.
Indeed 1. and 2. are folklore and there is nothing to prove. 
However, the authors  wrongly replaced $x$ by $X$ in part 2 of the appendix.
- from formula (19) to (20) there is nothing to prove; in general the heuristic App. B appears superfluous to me,
but maybe it could be of interest for people not directly working in the field.
- The results of the experiments seem to depend heavily on the symmetry of the target density, Please show a modified experiment of Fig. 6 where you shift the target density in the first dimension e.g. by 5 to the left (and keep the variance of the latent distribution).
I guess that the reconstruction misses modes.
- the comparison with HMC is a little unfair. Please redo the HMC experiment with a different chain for each sample starting in the same latent distribution as your model. I would expect that for your symmetric target distribution this woks fine.
soundness: 2
strengths: The basic idea is straightforward, but interesting and the numerical results are quite convincing.
summary: A continuous normalizing flow-based transport map guided by annealing is designed to sample from high-dimensional
and multi-modal distributions with unknown normalization factor. The approach supports effective exploration
of modes in high-dimensional spaces. The paper contains several comparison with state-of-the-art methods.
The method can be extended to a distribution-free model that allows 
to learn an importance flow from a dataset for sampling its least-favorable distribution with
minimal variance.
weaknesses: The mathematical notation is in many parts cursorily or wrong as the proof of Proposition 1.
The numerical experiments which are restricted to symmetric multimodal targets must be improved. 
See questions and experimental suggestions  below.

---

confidence: 3
contribution: 2
initial_rating: 1
presentation: 2
questions: - Could you compare your work against [1] and [2] which have solved the same problem with very similar tools ?
- Could you rewrite sections 2 and 3 but defining the normalizing constants from the beginning ? I think it will make the paper much clearer. Moreover, it would avoid writing KL divergences between unnormalised distributions as done extensively in section 3 and appendix 1.
- I suspects mistakes have been made in the proof of proposition 1. Could you explain the following points ?
1. In common litterature, the Kullback-Leiber divergence is defined between normalized distributions as $D(p || q) = \mathbb{E}_{X \sim p}[\log p(X) / q(X)]$. However, L760 shows a completely different definition (different order and unormalized densities);
2. According to the beginning of section 3.2, the quantity $\rho_k$ is the density of the flow following $v_k$ which means that $\log \rho_k(x(t_{k-1}))$ is not independent of $v_k(x(s),s)$  which would contradict L800.
- The block-wise training procedure (Sec. 4.1) uses AF samples from step $k-1$ to train for step $k$. However, those samples are approximate and have no reasons to be calibrated according to $f_{k-1}$. Can we trust those samples to compute expectations with respect the $f_{k-1}$ ? (Note that the same remark goes for the Importance Flow procedure)
- At L260, you claim that your training procedure only requires a single neural network. Doesn't your sampling procedure require to keep all the intermediates $\{v_k\}_{k = 1}^K$ ?
- Could you perform the experiment on Fig. 1 with differently weighted modes ? I feel like recovering the weights of the modes is the true challenge.
- Could you provide numerical comparisons with recent VI approaches [3,4,5] and with similar approaches [1,2] ? Those methods are important competitors.
- Could you compute the metrics given in [6] which seem computable in your case ? Those metrics really focus on multimodal distributions.
- Could you provide metrics's mean and standard deviation (as well as the number of samples used) for each experiment ?
- Could you display true samples from the Funnel distribution alongside yours ? It would clarify the pros and cons of each algorithm.


[1] Tian, Y., Panda, N., & Lin, Y. (2024). Liouville Flow Importance Sampler. In Proceedings of the 41st International Conference on Machine Learning (pp. 48186–48210). PMLR.

[2] Fan, M., Zhou, R., Tian, C., & Qian, X. (2024). Path-Guided Particle-based Sampling. In Proceedings of the 41st International Conference on Machine Learning (pp. 12916–12934). PMLR.


[3] Francisco Vargas, Will Sussman Grathwohl, & Arnaud Doucet (2023). Denoising Diffusion Samplers. In The Eleventh International Conference on Learning Representations .

[4] Qinsheng Zhang, & Yongxin Chen (2022). Path Integral Sampler: A Stochastic Control Approach For Sampling. In International Conference on Learning Representations.

[5] Richter, L., & Berner, J. (2024). Improved sampling via learned diffusions. In International Conference on Learning Representations.

[6] Blessing, D., Jia, X., Esslinger, J., Vargas, F., & Neumann, G. (2024). Beyond ELBOs: A Large-Scale Evaluation of Variational Methods for Sampling. In Proceedings of the 41st International Conference on Machine Learning (pp. 4205–4229). PMLR.
soundness: 1
strengths: - The importance flow technique is original
summary: This paper addresses the challenge of sampling from unnormalized probability distributions by introducing a novel approach called Annealed Flow (AF). The authors define a sequence of annealed distributions and propose learning continuous normalizing flows between each distribution in the sequence. Furthermore, the paper presents an additional method, termed Importance Flow (IF), which allows for computing expectations by first sampling the optimal importance distribution using AF and then learning the associated importance ratio. Experimental results validate the effectiveness of AF and IF across a diverse set of target distributions.
weaknesses: - The paper is poorly written (and follows unusual notations for instance with the normalizing constants)
- The paper suffers from theoretical issues/typos (KL between unnormalised distributions, proof of proposition 1, ...)
- The authors don't mention or compare to direct competitors [1, 2] who solve the same problem with the same tools
- Poor experimental part with very few comments