Annealing Flow Generative Models Towards Sampling High-Dimensional and Multi-Modal Distributions
Authors: Dongze Wu
Yao Xie
Venue: ICML 2025 poster (2025)
Keywords: Statistical Sampling
Optimal Transport
Continuous Normalizing Flows
Primary area: probabilistic_methods->monte_carlo_and_sampling_methods

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 method built on Continuous Normalizing Flows (CNFs) for sampling from high-dimensional and multi-modal distributions. AF is trained with a dynamic Optimal Transport (OT) objective incorporating Wasserstein regularization, and guided by annealing procedures, facilitating effective exploration of modes in high-dimensional spaces. Compared to recent NF methods, AF significantly improves training efficiency and stability, with minimal reliance on MC assistance. 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.

Annealing Flow Generative Models Towards Sampling High-Dimensional and Multi-Modal Distributions
Dongze Wu!
Yao Xie
Abstract
2021; Izmailov et al., 2021; Wu et al., 2024).
Sampling from high-dimensional, multi-modal distributions remains fundamental challenge across domains such as statistical Bayesian in- ference and saysics-basedd machine learning. In this paper, we propose Annealing Flow (AF), a method built on Continuous Normalizing Flow (CNF) for sampling from high- Eigh-dimmesiooal and multi-modal distributions. AF is trained with dy- namic Optimal Transport (OT) objective incorpo- rating Wasserstein regularization, and guided by annealing procedures, facilitating effective explo- ration of modes in high- -igeesimmnsi spaces. Com- pared to recent NF methods, AF greatly improves training efficiency and stability, with minimal re- liance on MC assistance. We demonstrate the superior performance of AF compared state-of- the-art methods through experiments various challenging distributions and real-world datasets, particularly in high-dimensional and multi-modal settings. We also highlight AF's potential for sam- pling the least favorable distributions.
MCMC and Neural Network Variants: Numerous MCMC methods have been developed over the past 50 years, includ- ing M (MH) and its variants (Haario et al., 2001; Cornish et al., 2019; Griffin & Walker, 2013; Choi, 2020), Hamiltonian Monte Carlo (HMC) schemes (Girolami & Calderhead, 2011; Bou-Rabee & Sanz-Serna, 2017: Shah- baba et al., 2014; Li et al., 2015; Hoffmman et al., 2021; 2014). HMC variants are still considered state-of-the-an methods. However, they require exponentially many steps in the di- mension for mixing, even with just two modes (Hackett et al., 2021). More recently, Neural network -aN)sassitted sampling algorithms (Wolniewiez et al., 2024; Bonati et al., 2019; Gu & Sun, 2020; Egorov et al., 2024; Li et al., 2021; Hackett et al., 2021) have been developed to leverage NN expressiveness for improving MCMC, but they still inherit some limitations like slow mixing and imbalanced mode exploration, particularly in high-dimensional spaces.
Annealing Variants: Annealing methods (Gelfand et al., 1990; Sorkin, 1991; Van Groenigen Stein, 1998; Neal, 2001) are widely used develop MCMC techniques like Parallel Tempering (PT) and its variants (Earl & Deem, 2005; Chandra et al., 2019; Syed et al., 2022). In anneal- ing, sampling gradually shifts from an easy distribution to the target by lowering the temperature. Annealed Impor- tance Sampling (Neal, 2001) and its variants(Zhang et al., 2021; Karagiannis & Andrieu, 2013; Chehab et al., 2024) are developed for estimating normalizing constants with low variance using MCMC samples from intermediate distribu- tions. Recent Normalizing Flow and score-based annealing methods (Arbel et al., 2021: Doucet et al., 2022) 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 par- ticularly pronounced in high-dimensional, multi-modal set- tings (Van Ravenzwaaij et al., 2018: Hackett et al., 2021).
1. Introduction
Sampling from figh-dimensional and multi-modal distribu- tions is erucial for various fields, including physics-based machine learning such molecular dynamics (Miao al., 2015; Salo-Ahen et al., 2020), quantum physics (Carlson et al., 2015; Lynn et al.. 2019), and lattice field theory (Jay & Neil, 2021; Lozanovski et al., 2020). With modern datasets, also plays a key role in Bayesian areas, including Bayesian modeling (Kandasamy al., 2018; Balandat et al., 2020: Stephan et al., 2017) with applications in areas like compu- tational biology (Stanton et al., 2022; Overstall et al., 2020), and Bayesian Neural Network sampling (Cobb & Jalaian,
'H. Milton Stewart School of Industrial and Systems Engineer- ing (ISyE), Georgia Institute Technology, USA. Correspondence to: Yao Xie Eyao.xie@isye waaaeeeiseggatcc.cuu...... gatech.edu>.
Particle Optimization Methods: Recently, particle- basiele-bbsed optimization methods have emerged for sampling, includ- ing Stein Variational Gradient Descent (SVGD) (Liu & Wang, 2016), and stochastic approaches such as Dai et al. (2016); Nitanda & Suzuki (2017); Maddison et al. (2018);
Proceedings of the 4nnd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

Annealing Flow Generative Models Towards Sampling Hligh-Dimensionll and Multi-Modal Distributions
Liu (2017); Pulido & van Leeuwen (2019); Li et al. (2023); Detommaso al. (2018); Maurais & Marzouk (2024). How- ever, many of these methods rely kernel computations, which scale polynomially with sample size, and are sensitive to hyperparameters.
where v(x(1),)) represents the velocity field, which is of the same dimension as xx() and is parameterized by neural
network with input =z() and t. The time horizon is rescaled to [0, 1] without loss of generality.
Continuous Normalizing Flow (CNF) is a class of normaliz- ing flows where the transformation of a probability density from a base distribution To (2) (at 0) to a target distri- bution g(z) (at t 1) is governed by Neural ODE. The marginal density of z=(t), denoted as p(zz, t),, evolves accord- ing to the continuity equation derived from the ODE in Eq. (I). This continuity equation is written as:
Normalizing Flows: Recently, discrete Normalizing Flows (NFs) (Rezende & Mohamed, 2015) and Stochastic NFs (Wu et al., 2020; Hagemann et al., 2022) have been actively explored for sampling tasks. Discrete NFs sometimes suffer from mode collapse, and methods relying on them (Arbel et al., 2021: Matthews et al.. 2022: Gabrié et al., 2021: 2022; Albergo & Vanden-Eijnden, 2023; Brofos et al., 2022; Cabezas et al., 2024; Qiu & Wang. 2024) attempt to mitigate this issue using MCMC corrections or design specialized loss functions. More recently, Fan et al.; Tian et al. (2024) introduced path-guided NFs, which utilize training losses conceptually similar to score matching. However, these methods may require substantial number of discretized time steps, with score estimation at each step. Besides, the quality of these estimations substantially influences the overall training performance.
Oep(a, 6) + V (o(x, ()) = 0, p(x, TT(),
(2)
where V. denotes the divergence operator.
Dynamic Optimal Transport (OT): The Benamou-Brenier equation (Benamou & Brenier, 2000) below provides a dy- namic formulation of Optimal Transport T, which transports probability mass from Ty to q via a velocity field v(x(t), t):
inf I/'d a
Challenges persist with multi-modal distributions in high- dimensional spaces. This paper introduces Annealing Flow (AF), Continuous Normalizing Flow trained with dynamic Optimal Transport (OT) objective incorporating Wasserstein regularization, and guided by annealing proce- dures. Our key contributions are as follows:
s.t. 0pp+V (pu) = p(- no.
pl. 1) = q.
(3)
The optimization problem seeks to find the optimal T that moves mass from the base density To to the target density g. subject to the continuity equation (2) to ensure that p(: 6) evolves as valid probability density over time. Addition- ally, the constraint p(,, 1) = 9 ensures that the target density is reached by the end of the time horizon. The time horizon is rescaled to [0, 1] without loss of generality.
We present a dynamic OT objective that greatly reduces the number of intermediate annealing steps and im- proves training stability, compared to recent annealing- like NFs (Tian et al., 2024; Fan et al.), which require score estimation and matching during training.
3. Annealing Flow Model
The proposed annealing-assisted procedure is crucial distributions with
The annealing philosophy (Gelfand et al., 1990; Sorkin, 1991; Van Groenigen & Stein, 1998; Neal, 2001) 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 sam- pling algorithm that learns a continuous normalizing flow to gradually map an initial easy-to-sample density xo (z) to the target density (Ix) through set of intermediate distri- butions.
for success on widely separated modes, outperforming state-of-ther art ba-ed methods in challenging experiments with notably fewer annealing steps.
Theoretically, Theorem 3.3 establishes that the in- finitesimal optimal velocity field equals the difference in the score between consecutive annealing densities, desirable property unique to our dynamic OT objective.
We define the target distribution as g(x) = Zg(z), where ((x) is the unnormalized form given explicitly, and Z > 0 is an unknown normalizing constant. To interpolate between an say-toe initial distribution To(xr) (e.g., a standard Gaussian) and the target q(z), we construct sequence of intermediate distributions {=)], with corresponding unnormalized forms (f))..). These are defined as:
2. Preliminaries
Neural ODE and Continuous Normalizing Flow: A Neural ODE is continuous model where the trajectory of data is modeled as the solution of an ordinary differential equation (ODE). Formally, in Idd, given an input z0)) = zo at time t t=00 0, the transformation to the output (1) is governed by:
fx((x) = =0 ((r)
dx(t) dt v(z(t),t), t), t [0. 1],
(4)
(1)
fz(z) (x).
2

---

Anncaling Flow Generative Models Towards Sampling Hlig--Dimensionll and Multi-Modal Distributions
Here, 1/2x- San (x)dd is an unknown normalizing constant, and By is increasing sequence with (0 = 0 and 3x = L. Therefore, fo(x)) xo (x) and [x(x)) = g(x). Im- portantly, our algorithm depends only on the unnormalized fa (x) and is independent of Zg.
we define vs (x((), as the velocity field of v(x(),,) dur- ing the time horizon E (kx-1,4x)-
Annealing Flow aims to learn each transport map Te based on dynamic OT objective (3) over the time horizon ( where the velocity field u (xx(t), 1) is learned us- ing a neural network. The terminal condition p(:, = q in (3) can be imposed approximately using the Kull- back-Leibler (KL) divergence (see. for instance, Ruthotto al. (2020)). Consequently, minimizing the objective (3) for dynamic optimal transport Ta ( - Sa (2) over the time horizon [šx-1, fx) can be reduced to solving the following problem:
annealing trajectory
fs
f6
flww
f
Th = arg min T
(L
diffusion
+ Ex(0~ (o)))
le-1
(6)
h
E (k)),, subject to p(z(t), t) and (x(((t), t) evolving according to (2). Here, (L fx-i represents the KL divergence between the push-forward density Tạ fx-1 (by transport map T) and the target density fu. and 7 > is a regularization parameter that controls the smoothness and efficiency of the transport path. Additionally, the constraint (2) ensures that (t) follows the ODE trajectory defined by (1) during t E [tx-1, tx),, which is given by:
Figure 1. Illustrations and comparisons of annealing trajectories: Annealing Flow (based optimal transport maps) versus other methods based on diffusion and score matching.
3.1. Algorithm
The objective of Annealing Flow (AF) is to learn a continu- ous optimal transport flow that transitions between interme- diate distributions fx-1 and fü for k = ......... K, ultimately mapping an initial distribution no the target distribu- tion qzx), illustrated in Figure 1. A key observation is that during AF training, the normalizing constant does not affect the training objective. Once trained, users simply sample (0) ~ (T)), and the transport map pushes them to (1) ~ (xx). The transport map T evolves the density according (2), which in turn drives the evolution of the sample (t) following the ODE in (I):
/"' J (x(s))) EE [1k-1,
s(1) x((2-1)
The following proposition shows that, once samples x(t=-1) from fA-1 have been obtained through the previously learned Tı, T the KL divergence in (6) can be equivalently expressed using va (=(()), and - log (x(),, up to an additive constant. Therefore, learning the optimal transport map T. reduces to learning the optimal uy ())
The proof provided in Appendix A.1.
<<(( = T(x(0) = x(0) + // u(a(s), LE [0,
Proposition 3.1. Given the samples from SK-1, we have:
[- llog (a(()()))
(5)
KL(T St- (1)) = c Ex(ta-
We divide the time horizon [0, 1] into K intervals: 14-1, (x) for k = 1,..., K- - 1, and [t txl with to = 0 and tK = 1. We thus define x(4)) as the forwad sam- ple at time step 44 Following the annealing flow path in (4), the continuous flow map T gradually transforms the density from fo(x) to fi (x) over [0, t1). and continues this process until fK - (2:) is transformed into fx (x) g(x) over 11.
pl . (x())))
(8)
up to a constant e that is independent of vu (x(x(()) s).
Given =(tx-1) from fx-1 the value of r((x)) inside - llog f-) can be calculated as shown in equation (7). Ad- ditionally, according to the lemma below, the second term in the objective (6) can be approximated as a discretized sum. The proof provided in Appendix A.1.
Figure 2 shows this progression with two intermediate dis- tributions. For clarity, we denote Tr as the segment of the continuous normalizing flow during [t ts), which pushes the density from (x) to (x) Consequently,
Lemma 3.2. Let x(t) be particle trajectories driven by a smooth velocity field us (x(x(t), t) over the time interval
3

---

Annealing Flow Generative Models Towards Sampling High- nigh--Dimensionl and Multi-Modal Distributions
To [0,4) (1)
Ti
T2
(2)
[12,
(a) Bg = 0
(b) 81 1/3
(c) B2 = 2/3
(d) B3
Figure Illustration of Annealing Flow Map, with a set of annealing densities from mo(a) (0, 12) (z)) a GMM with modes.
(fk-1,tk), (k), where hy: = 1- Assume that vg (I. 1) is Lipschitz continuous in both x and By dividing [xx-15 (x) into equal mini-intervatt with grid points 1,-1,s- where ........ and 1k-1.0 tx-1, 1.5 tg. we have:
Theorem 3.3. (Optimal Velocity Field as Score Difference)
Suppose hi - O. Let fe-1 and fü be continuously dif- ferentiable on Rd. Assume that vx(x) exists for all x E Rd, and V ux(x), Sk-1 and sy belong (fx-1) Assume that the components of Uk are independent and lim [a]-20 Jn-1 (1r) (1))2)) 0. Under these conditions, the minimizer of (10) is:
Be(()~((-1)) [11=)(
4--1
S h4
8-1
E - ( O(h/s))
lim 2i 8k 8k-1- hy -0
(11)
(9)
Therefore, the infinitesimal optimal vy equal to the differ- ence between the score function of the next density, fú, and the current density, ft-1- This suggests that by adding more intermediate densities, one can learn a sufficiently smooth transport map T that exactly learns the mapping between each pair of densities.
One can observe that the RHS of (9) is discretized sum of stepwise Wasserstein-2 distances, reflecting the dynamic nature of the transport. Dynamic W2 regularization en- courages smooth transitions from fx-1 to f& with minimal transport cost, ensuring stable sampling performance.
Next, by incorporating Proposition 3.1 and Lemma 3.2 into (6), the final objective becomes:
Additionally, one can observe that when each fx is set to the target q(x), i.e., when all 3g are set 1, and the second term in the objective (6) is relaxed static Wa regularization, the objective of Annealing Flow becomes equivalent to Wasserstein gradient flow. Details, along with its convergence theorems, are provided in Appendix B.
log fa(x(te))) f"h
min 1 oz(-.))
V (x(s)))
4-1
8-1
+a Σ |= a((t=11.)11? ==0
(10)
4. Algorithm Implementation
Remark We empirically observed that replacing the first term - log fu (z(tx)) in the objective (10) with the first- order Taylor approximation -/4 V log fa (x((t)) us led to improved empirical performance, as detailed Ap- pendix C.I.
Algorithm Block-wiis Training of Annealing Flow Net Require: Unnormalized target density q(æ); an easy-to- sample To(x);; (8,,8,,". .......3. .3x-13; total number of blocks K.
In the final objective, learned by neural network. We break the time interval [f=-1, 4x) into S mini-intervals, and x(12-1,44) is com- puted as in equation (7).
7S/hu, and un (z(s), is
Set B0 0 and 3x =1
2: for k = to K do
3: Set fa( = To (-x
4: Sample (0)) from mn(x)
3.2. Properties of learned velocity field
5: Compute the pushed-forward samples (fx-1)
from the trained (I: 1) blocks via (13): (Optional Langevin adjustments)
As hu = 14: - - - 0, the final objective in (10) converges to form involving the Stein operator. This is shown in A.2 as part of the prerequisite proof for Theorem 3.3.
for several steps, update each ( (ta-1) via
(ta + (te-a) 3 V log fx-1 vn 9: end for
8:
< x0) as 1 (x) dr). Define 8 = V log fx as the score function of the density ft- We can then establish the following property, with proofs provided in A.2:
Define (f²((t-))= = {v: fau (()
the 2 space over (Rd, fx-
10: Optimize (g( (e(()) () by minimizing the final objective.
11: end for

---

Annealing Flow Generative Models Towards Sampling Highh-Dmennsonal and Multi-Modal Distributions
4.1. Block-wise training
dynamic W₂ regularization. This unique loss facilitates op- timal paths with notably fewer time steps. Detailed results, algorithmic analyses, and ablation studies are thoroughly presented in Section 6 and Appendix D.
Training of the k-th flow map Tk begins once the (k 1)-th block has completed training. Given the samples (ts-)))=
(x) after the (k - 1)-th block, we can replace f with the empirical average. The diver- gence of the velocity field in (10) can be computed either by brute force or via the Hutchinson trace estimator (Hutchin- son, 1989; Xu et al., 2024b):
5. Importance Flow
Sampling from complex distributions fundamental, which can benefit tasks like Bayesian analysis and statistical physics. Here, briefly discuss another aspect: us- ing Annealing Flow to sample from the Least-Favorable Distribution (LFD) and obtain a low-variance Importance Sampling (IS) estimator, referred to as Importance Flow.
Vo0x(z,t) 2o EeNp, [c. 1x(2 + tc,tt vx(a,t) (12)
0
where we select p(c) = (O, Ta). The Hutchinson estima- tor significantly improves computational efficiency while maintaining good performance. More details are provided in C.2. Furthermore, we apply the Runge-Kutt method for numerical integration, with details provided in C.3.
5.1. Settings
Suppose we want to estimate Ex-=o(() [h(X)], which can- not be computed in closed form. A natural approach is to use Monte Carlo estimation by sampling {mi]== from no(x). However, if xr, consistently falls in regions where h(x) has extreme values, the estimator may exhibit high variance.
Our method adopts block-wise training for the continuous normalizing flow, summarized in Algorithm 1. After training each block, optional Langevin adjustments (Lines 6-9) can be applied to enhance performance. The block- wise training approach of Annealing Flow significantly re- duces memory and computational requirements, as only one neural network is trained a time.
For example, with (z (0, Ia) and h(z) ||-1|/- almost no samples will satisfy ||>||| 2 6, resulting in zero estimate.
To address this situation, we can select an appropriate pro- posal distribution (xx) and rewrite the expectation and MC estimator as:
4.2. Efficient sampling and methods comparisons
Once T' is learned, the sampling process of the target g(x) becomes very efficient. Users simply sample (z(1 (0)))=
To(x) M((z)
E Ex)To( [k(x:)] E Bxog(x)
from (I)), and then directly calculate (ax(*) (1) ((xr)
((x))
through Annealing Flow nets (k: = 1,..,,,,........., KK):
(14)
1
To To( h(xz), ~ q(x).
(t2) = Ta (tk-1)
((24)
(tx-1) + 1' (x (s) syds. Jra
(13)
It is well-known that the theoretically optimal proposal for the importance sampler is: 9" (a) x +o(=))](<)) == q (z). However, q" (z) is often difficult to sample from, especially when (z) or h(x)) is complex. Consequently, people typ- ically choose a distribution that is similar in shape to the theoretically optimal proposal but easier to sample from.
MCMC struggles with poor performance on multi-modal distributions, long mixing times, and correlated samples, es- pecially in high dimensions, In contrast, NFs push samples from (x)) through a learned transport map, enabling faster and independent sampling. Although Annealing Flow re- quires more expensive pre-training, this can be done offline and only once. Once trained, AF samplers are highly effi- cient, generating 10,000 samples in 2. seconds on average in our 50D experiments.
Amnealing Flow enables sampling from q" (x). allowing the construction of Importance Sampling (IS) estimator. However, q" (æ) is only known up the normalizing con- stant 7, where (xr) 27' (I:) and Z = ( [?()) is our target. Therefore, assuming no knowledge of Z, a common choice can be the Normalized IS Estimator:
Compared to recent NFs, our annealing procedure enables more direct and efficient exploration of transport trajecto- ries. The lower part of Figure/ illustrates a single transport trajectory between two consecutive annealing densities. Un- like many recent annealing-like like NFs (Tian et al., 2024; Fan et al.), which depend on score estimation and matching for training, our Annealing Flow is based on a dynamic OT ob- jective, free from score matching, and uniquely supported by
mo(2s) ÎN Σ Emrte) h(x:)/ 1 (x)))
Σ To mo(#;)
(15)
(x:)
However, this estimator is often biased, as can be seen from Jensen's 's Inequality.
5

---

Annealing Flow Generative Models Towards Sampling High--Dmeensionll and Multi-Modal Distributions
5.2. Density ratio estimation
use the same number of training steps and intermediate den- sities for CRAFT, LFIS, PGPS, and our AF, with identical network architectures and training iterations. All AF results are reported without the optional Langevin adjustments, and PGPS is also run without them for fairness. Additional re- sults- intlusinc LFIS with 256 steps, and CRAFT/PGPS with 128 steps-are are provided in Appendix D.
Using samples from (ax) and those along the trajectory obtained via Annealing Flow,
train neural network for Density Ratio Estimation (DRE) of To(2)// (z). In- spired by works Rhodes et al. (2020); Choi et al. (2022); Xu et al. (2025), we can train a continuous neural network
r(z) =
(x; (z; 01), (16)
6.1. Experiments and results
where samples z;, ~ fx = q (z) are imputs and the output is the density ratio no ozx))// //" (zrz). Each r& (2; 0s) is trained using the following loss:
Gaussian Mixture Models (GMMs): Figure (main script) and Figures 4 and 5 (Appendix) show the visualized sam- pling results of various methods on GMMs, Additionally, we tested unequally weighted GMMs, where two modes have double the weight of the others. Numerical evaluation metrics for experiments with different numbers of modes and dimensions are presented in Tables 1, 4, and 5. Addi- tional figures and tables are provided in Appendix D.
Lx (0x) = Eg((e- [lo((1 + +
(17)
+ +Ian(ta)ff log(1 + e=((ai((())))]
The theoretically optimal rx) = log(fx-1(z)/ and thus r*(x) >m k- (ā()) (1g (Ap-
Truncated normal distribution: Figure 7 the Appendix shows the sampling results for the truncated normal dis- tribution ((x)) | 22 Ia)- Notably, even af- ter relaxing ||=||>> to 1/(1 + exp(-A([|=]I
pendices A.3 and C.5 contain proof and further details.) To obtain the optimal importance sampling estimator, one can then directly use samples { ~ (x) from Annealing Flow and apply (14) together with the DRE: n Ei=1 exp(r" (z+)) h(z)). The estimator is unbiased and can achieve zero variance theoretically.
c))) in V log [1 Ta)). methods including CRAFT, LFIS, PGPS, SVGD, MIED, and Al-Sampler fail to produce mean- ingful results. Please refer to Figure 6 in Appendix D for the results of these algorithms.
6. Numerical Experiments
Funnel distribution: We tested each algorithm for the funnel distribution
In this section, we present numerical experiments compar- ing Annealing Flow (AF) with the following methods: (1) Annealing-based MCMC: Parallel Tempering (PT); (2) NN- assisted MCMC: Al-Sampler (AIS) (Egorov et al., 2024); (3) Particle Optimization methods: Stein Variational Gra- dient Descent (SVGD) (Liu & Wang, 2016) and Mollified Interaction Energy Descent (MIED) (Li et al., 2023); and (4) Normalizing Flow approaches: Continual Repeated An- nealed Flow Transport Monte Carlo (CRAFT) (Matthews et al., 2022), Liouville Flow Importance Sampler (LFIS) (Tian al., 2024), and Path-Guided Particle-based Sam- pling (PGPS) (Fan et al.). Experimental details are provided in Appendix C.3.
q 22, ... xa) N 0, 0,
I exp (21)).
(18)
Figure 9 in Appendix D shows the sampling results.
Wxigheeghted Gaussian (ExpGauss) with an extreme num- ber of modes in high--dmenssonal spaces:
We tested each algorithm on sampling from an extreme distribution:
p(11, 22, . ,,,0) x e 10 E |2.1- 1
(19)
The number of time steps for Annealing Flow (AF) is set per distribution as follows: (1) GMMs: 12 steps-10 10 inter- mediate densities and refinement blocks using (zx) as the target: (2) Truncated Normal: steps with 7 intermediate densities; (3) Funnel: 8 steps with each intermediate density as the target; (4) Exigheigtted Gaussian: 20 s1eps-15 in- termediate densities and 5 refinement blocks; (5) Bayesian Logistic Regression: 6 steps with each intermediate density as the target.
which has 200 1024 modes arranged at the vertices of a 10-D cube. The Euclidean distance between two horizon- tally or vertically adjacent modes is 20, while the diagonal modes are separated by up to v10 20 Rs 63.25. We also tested on the extreme distribution:
10 10 1لg +10
5 -
p(z1, T2, I30)
(20)
which has 200 = 1024 modes arranged at the vertices of 50-dimensional space, with imbalanced variances across modes along different dimensions.
For fair comparison, all figures and tables (except Table 3)
Our code is publicly available https://githht com/ Stust.FFu on/ Anneal ing-F1
Given the challenge of visualizing results in high- dimensional space, we first present the number of modes
ow- Flo Sampl ing.
6

---

Annealing Flow Generative Models Towards Sampling High- High-DDmeensionl and Multi-Modal Distributions
RR
(a) True
(b) AF
(c) CRAFT
(d) LFIS
(c) PGPS
(f) PT
(g) SVGD
(h) MIED
(i) AIS
Figure 3. Sampling from a 10-mode Gaussian Mixture Model (GMM) arranged in a circle of radius r = 12. Annealing Flow (AF) always
uses 12 time steps. Top row: AF, CRAFT, LFIS, PGPS are trained using 12 time steps; PGPS is trained without I Langevin correction. Bottom row: AF CRAFT, LFIS, and PGPS are trained using 12, 128, 256, and 128 time steps, respectively; PGPS is trained with Langevin correction. See Figures 4 and 5 for additional visualizations.
Table 1. Mode- Weight Mean Squared Error across distributions. The number of time steps for CRAFT, LFIS, and PGPS is set equal to that of AF. Three additional types metrics are presented in Tables 4. 5, and
Distributions
AF
CRAFT
LHSS
PGPS
SVGD
MEED
Al-Sampler
GMM-6-8
(7.72 m))
10
[1:21.
(7.4%.
(1.62a1) 10"
(144.)))
10
4n)
10
GMM-8-10
(24+
x 10
(277
x 10
(6.84.gg
x 10
(9.8t x 10
(1.7h x 10-2
(1.5t.r)
10
(2.3+.32) x0
d
GMM-10-12
5.72..))
(2.62.0)
10
(2.33.4) 10
(2:22 x
(1.7kx)) x 10
(5.61
10-3
(102.as) 10
WGMM-110112
9.5- x 10
(9.7 sa) x 10-
(3.8..g)
x 10
(4.7=as) x100
(2.4h x 10--2
(6:2+..)
x 10
(3.6t..) 10
GMM-6-8
132.
(1.52.aa)
(5.41.
(1.2karr 100
(2.11 10
(1.11.
10-3
(8.62. 10-3
d
GMM-10-12
GMM-8 8-10
(12.o)
2.2 x 10
(8.81, 10-3
(1.1k.) x 10
(9:24.a)
(3.3d
10
10
(5.22.m) x
(4.82 x100
(3.4k x 10-2
(3.221 10-2
(8.12
(4.8
10-3
10
(2.7k.m) 10
(1.5t.u) K1-5
wGMM-10-12
7 (.33..))) 10
(9.74.p) x 10
(7.6k.))
x 10
(5.9t..s) 10
(6.4k 10
(6.01)
10
(7.94.22 10
ExpGauss- 1024
(8.84.a) x 100
(1.5+ 10
(1.2=ans)
(8.5=
(4.6+ 10
(8.6+.a)
Exptians--102 1024
(9.8 10
(9:4+:a0) x 10
(2:2 10
(3.5m x 10"
(3.3raa)) x 10
(2.1+.z) 10
(1.1+,0) 10
ExpGaussUV.. 1V-2-1024
(1.8+ x
(7.84
10-6
(7.14 10
(4.8= 10
(7.6+ x 10-6
(9.24).) 10-6
(1.74
d 50
PxpCaasssVV100-100 101-11024
(2.4+.aa) x 10-7
(1.9t.a))
10-5
(6.9t.r)
10-
(0.9+:00) x 10-5
(1.1+as) 10-5
(9.3r) 10-5
(1.4ha)) 10-1
ExpGamssUV-15...
(2.0+ 10
(2.24
x 10-5
(8:94 ad)
(1.9 an) 10m5
(7.0+ 10-5
(3.74 10-5
(1.6t.) 10-1
ExpGanss 1V-20-004
(3.6+ m0) x 10
(4.7+.0) 10-5
(9.7 10-
(4.8 so) 10-
(8.7+ 10-5
(7.tt
x 10-
(1.8ta)) 10
the details of these metrics. In the tables, GMM refers to Gaussian Mixture Models, while wGMM denotes un- equally weighted GMM. The notation (W))MMM(nnubbb of modes)-(radinn of the circle) represents (W)GMM with the specified number of modes arranged on circle of the given radius. ExpGauss- 1024 refers to exponentially weighted Gaussian experiments involving 1,024 widely sep- (number of UV dimensions }- 1024 refers to ExpGauss with 1024 modes that have Unequal Variances (UV) across the specified number of dimensions. Alongside Table 7, we provide more results for ExpGauss. with details on the selection of of across dimensions.
Table 2. Number of modes explored by different methods in a 50-dimensional Exponentially Weighted Gaussian with 1024 far- separated modes. The number of time steps for AP, CRAFT, LFIS, and PGPS is to 20.
True
AF
CRAFT
LFIS
PGPS
PT
SVGD
MIED
AIS
2=2
4
2.6
3.8
3/4
3.9
3.8
3.8
d=5
5
32
32
22.3
27.4
31.2
25.2
28.5
28.0
28.3
d=
10
1024
1024
515.4
387.0
826.0
253.7
957.3
923.4
501.2
arated modes. ExpGaussUV-nnn UV-umbbe
50
1024
473.2
298.2
813.6
10
916.4
8906
125.6
successfully explored by different algorithms across varying dimensions in Table 2. Each algorithm was run 10 times, with 20,000 points sampled per run, and the average number of modes explored was calculated. Additionally, Table presents the Mode Weights Mean Squared Error along with results of other distributions.
Bayesian logistic regression: We use the same Bayesian lo- gistic regression setting as in Liu & Wang (2016), where a hi- erarchical structure is assigned to the model parameters. The weights 3 follow a Gaussian prior po (3]a) = N (3; 0, a ) and a follow a Gamma prior Pu Pa(a) = Gamma(a: 1, 0.01). Sampling is performed the posterior p(3, ox D), where D = (x,, 34 J The performance comparisons shown in Table 3. Detailed settings are given in C.4.
Evaluation metrics: We report the following metrics for each applicable experiment: (1) Mode-Weight Mean Squared Error, (2) Maximum Mean Discrepancy (MMD), (3) Wasserstein Distance, and (4) Sample-Variance Variance Mean Squared Error across dimensions. The results for these metrics are presented in Table 1 (main script) and in Ta- bles 4, 5 and in the Appendix. Appendix C.7 explains
Importance flow: Table 8 in the Appendix reports the prelim- inary results of the importance flow (discussed in Section 5) for estimating Eg (0.1) [[ with varying radii and
7

---

Annealing Flow Generative Models Towards Sampling Hligh-Dimensionll and Multi-Modal Distributions
Table Bayesian logistic regression: comparison of different algorithms across datasets. In the table icc represents Accuracy (%) + std (%), and feoa represents the averaged Posttrrri. The number of time steps is set to 6, 128, 256, and 128 for AF, CRAFT, LFIS, and PGPS, respectively.
Datasee \ Methods
AF
CRAFT
LIFIS
PGPS
SVGD
MIED
Al-Sampler
1pe0
fyer
Acc
1pe
fpo
Accusu
Acc_sa
Accusa
fyee
Diabetes = 5)
30.3
-0.50
76.2
-0.51
70.2.a.
-0.50
g0.3g2a
-0.50
70.1
-0.50
75.6
-0.50
70.3
-0.49
Cancer (d 10)
97.9,a.1
-0.02
98 9.1.2
-0.01
96.4
-0.05
98 9,1.1
-0.01
97.8
-0.02
98. 98.9:10
-0.01
07.8228
-0.02
Heart (d 13)
88.5
-0.32
88 88.4gan
-0.22
86.6ga
-0.43
87.2
-0.40
78.4
-0.50
86.7
-0.32
8422255
-0.46
Australiza (d - 14)
86.611.2
0.36
85.0g1.
0.30
84.1
0.36
85.4
0.30
84.612s
0.37
86.2
0.37
$46123
0.38
Ijcnal
82.0.0.1
-0.20
68.8ca.
-0.23
89.8
-0.31
91
89.4
-0.21
91.8102
-0.20
88.3.u.3
-0.33
Svmguide3 (d = 22)
80.0g1.0
-0.47
78.0ga9
-0.50
60 60.3
-0.48
80.0e10
-0.19
78.9212
-0.48
80.6s19
-0.47
80.1210
-0.47
German (f 24)
78.021.
77.5
-0.48
76.7
-0.49
77.7
-0.48
76.4g1
-0.48
77.2218
-0.48
76.921
-0.48
Splice (d = (1)
80.9
-0.41
81.1,22
-0.19
52.0g1.
-0.48
82.8
-0.17
52.4,94
-0.47
83.1
-0.40
80.1,19
-0.50
dimensions. Please refer to C.5 for detailed experimental settings. Additionally, we discussed a possible extension of the Importance Flow framework in D.4.
achieving much higher efficiency, and enabling successful sampling in high high-dimensional and multi-modal settings.
Comparisons with MCMC: Compared to MCMC methods, AF offers three key advantages. First, AF enables efficient sampling after a one-time offline training phase, whereas MCMC requires no training but suffers from long mixing times and high runtime per effective sample. Second, AF performs reliably in high-dimensional, multimodal settings where MCMC often fails to explore all modes, as evidenced by its poor coverage in the 50D ExpGauss experiment (Table 2). Third, AF generates balanced and independent samples by construction, while MCMC frequently produces imbal- anced mode visitation due stochastic mode-hopping.
6.2. Results discussions and comparisons
Significance of Annealing Procedures: We comment that anncaling procedures play crucial role in the success of our AF in high-dimensional settings with widely separated modes. Appendix 1.3.1 and D.3.2 present ablation studies for cases with no or very few annealing steps. This unique feature allows AF to succeed on challenging distributions, unlike most NFs used for sampling-
Little dependence on Monte Carlo sampling: During im- plementations, we observe that CRAFT heavily depends the choice of the MC kernel, one of its major functioning components. LFIS depends on accurate score estimation for training, necessitating many more time steps. PGPS relies heavily on Langevin adjustments after each step, with- out which its performance becomes poor. In contrast, AF achieves success with minimal assistance.
7. Conclusions
In this paper, we have proposed the Annealing Flow (AF) algorithm, a novel approach for sampling from high- dimensional and multi-modal distributions. With the unique Annealing
objective
eeeeediig-guided guided dynamic optimal transport objective, AF offers multiple advantages over existing methods, including superior performance on extreme distributions, greatly im- proved training efficiency compared to other NF methods, minimal reliance on MC assistance, and enhanced training stability compared to other NF methods.
Computational and storage efficiency: In Tables 9 and 10 of Appendix D. we report the training and sampling times for AF, CRAFT, LFIS, and PGPS. Notably, AF requires signifi- cantly fewer intermediate time steps than other normalizing flow methods while achieving superior performance in our experiments. We also provide ablation studies in D.3.2 with even fewer time steps. In addition, in our experiments, AF with single 32-unit hidden layer and sigmoid activation suf- fices for most tasks, except Exp-weighted Gaussian, which requires only 32-32 hidden layers. This also yields notable storage savings-for instance, MBB total for 20 networks trained in the 50D ExpGauss experiment.
We establish in Theorem 3.3 that the infinitesimal optimal velocity field corresponds to the score difference between consecutive annealing densities, a property unique to our AF, This desirable feature sets our AF's objective apart from recent amnealing-like methods (Tian et al., 2024; Fan et al.). In Appendix B, we demonstrate the equivalence of the AF objective with the Wasserstein gradient flow and its associated convergence theorems. A range of experiments demonstrate that AF performs well across a variety of chal- lenging distributions and real-world datasets. Finally, the importance flow discussed in Section 5 may be extended to distribution-free model, allowing one to learn an importance flow from a dataset for sampling its Least-Favorabi Distri- bution (LFD) with minimal variance, as further discussed in D.4.
Training stability: Our unique dynamic optimal transport (OT) objective with W2 regularization is key ensuring much more stable performance, even with far fewer interme- diate time steps. In Appendix 1.3.1 and D.3.2, we provide ablation studies highlighting the significance of our OT objective. Unlike the score-matching training, the unique dynamic OT objective of AF is crucial for ensuring stability,
8

---

Annealing Flow Generative Models Towards Sampling High-Dimensional and Multi-Modal Distributions
Acknowledgment
Carlson, J., Gandolfi, S., Pederiva, F., Pieper, S. C., Schi- avilla, R., Schmidt, K. E., and Wiringa, R. B. Quantum monte carlo methods for nuclear physics. Reviews of modern physics, 87(3):1067-1118, 2015.
This work is partially supported by an NSF DMS-2134037, CMMI-21 12533, and the Coca-Cola Foundation. This work also partially supported by the Stewart Fellowship.
Chandra, R., Jain, K., Deo, R. V., and Cripps. S. Langevin- gradient parallel tempering for bayesian neural learning. Neurocomputing, 359:315-326, 2019.
Impact Statement
Our work contributes to statistical sampling and machine learning, offering both theoretical and practical advances. Sampling is a fundamental and ubiquitous problem with broad applications statistical physics, Bayesian modeling, and machine learning. While our method has broad poten- tial societal implications, none of which we feel must be specifically highlighted here.
Chehab, O., Hyvarinen, A., and Risteski, A. Provable ben- efits of annealing for estimating normalizing constants: Importance sampling, noise-contrastive estimation, and beyond. Advances Neural Information Processing Sys- tems, 36, 2024.
Cheng, X., Lu, J., Tan, Y., and Xie, Y. Convergence flow- based generative models via proximal gradient descent in wasserstein space. IEEE Transactions Information Theory, 2024.
References
Albergo, M. S. and Vanden-Eijnden, E. Learning to sam- ple better, 2023. URL https: /tops::/arxy: //aaiiii...///.a org/ abs/ 2310. 11 232. Les Houches 2022 Summer School on Statistical Physics and Machine Learning.
Choi, K., Meng, C., Song, Y., and Ermon, S. Density ratio estimation via infinitesimal classification. In Interna- tional Conference on Artificial Intelligence and Statistics, pp. 2552-22522222 2573. PMLR, 2022.
Arbel, M., Matthews, A., and Doucet, A. Annealed flow transport monte carlo. In International Conference Machine Learning, pp. 318-330. PMLR, 2021.
Choi, M. C. Metropolis hastings reversiblizations of non- reversible markov chains. Stochastic Processes and their Applications, 130(2):1041-1073, 2020.
Axler, S. Measure, integration real analysis. Springer Nature, 2020.
Cobb, A. D. and Jalaian, B. Scaling hamiltonian monte carlo inference for bayesian neural networks with symmetric splitting. In Uncertainty in Artificial Intelligence, pp. 675 PMLR, 2021.
Balandat, M., Karrer, B., Jiang. D., Daulton, S., Letham, B., Wilson, A. G., and Bakshy, E. Botorch: A framework for efficient monte-carle bayesian optimization. Advances in neural information processing systems, 33:21524-21538, 2020.
Cornish, R., Vanetti, P., Bouchard-Côté, A., Deligianni- dis, G., and Doucet, A. Scalable metropolis-hasting for exact bayesian inference with large datasets. In Interna- tional Conference Machine Learning, pp. 1351 PMLR, 2019.
Benamou, J.-D. and Brenier, Y. computational fluid me- chanics solution the monge-kantorovich mass transfer problem. Numerische Mathematik, 84(3):375-393, 2000.
Bonati, I.., Zhang, Y.YY, and Parrinello, M. Neural networks-based variationally enhanced sampling. Pro- ceedings of the National Academy of Sciences, 116(36); 17641-17647, 2019.
Dai, B., He, N., Dai, H., and Song, L. Provable bayesian inference via particle mirror descent. In Artificial Intelli- gence and Statistics, pp. 985-994, PMLR, 2016.
Bou-Rabee, N. and Sanz-Sema,) J. M. Randomized hamilto- nian monte carlo. 2017.
Detommaso, G., Cui, T., Marzouk, Y., Spantini, A., and Scheichl, R. A stein variational newton method. Advances in Neural Information Processing Systems, 31, 2018.
Brofos, J., Gabrié, M., Brubaker, M. A., and Lederman, R. R. Adaptation of the independent metropolis- hastinns sampler with normalizing flow proposals. In International Conference Artificial Intelligence and Statistics, pp- 5949-5986. PMLR, 2022.
Doucet, A., Grathwohl, W. S.S Matthews, A. G. d. G., and Strathmann, H. Annealed importance sampling meets score matching. In ICLR Workshop on Deep Generative Models for Highly Structured Data, 2022.
Cabezas, A., Sharrock, L., and Nemeth, C. Markovian flow matching: Accelerating meme with continuous normaliz- ing flows. arXiv preprint arXiv:2405. /4392, 2024.
Earl, D. J. and Deem, M. W. Parallel tempering: Theory, applications, and new perspectives. Physical Chemistry Chemical Physics, 7(23):391-3991 2005.

---

Anncaling Flow Generative Models Towards Sampling High-Dimennional and Multi-Modal Distributions
Egorov, E., Valperga, R., and Gavves, E. Ai-sampler: Adver- sarial learning of markov kernels with involutive maps. In Proceedings of the International Conference Machine Learning (ICML), 2024.
Hutchinson, M. F. A stochastic estimator of the trace of the influence matrix for laplacian smoothing splines. Com- munications in Statistics-Simulattion and Computation, 18(3):1599-076, 1989.
Fan, M., Zhou, R., Tian, C., and Qian, X. Path-guided particle-basee sampling. In Forty-first International Con- ference on Machine Learning.
Izmailov, P., Vikram, S., Hoffman, M. D., and Wilson, A. G. G. What are bayesian neural network posteriors really like? In International conference machine learning, PMLR, 2021.
pp.
Gabrié, M., Rotskoff, G. M., and Vanden-Fiinden, E. Ef- ficient bayesian sampling using normalizing flows to as- sist markov chain monte carlo methods. arXiv preprint arXiv:2107.08001, 2021.
Jay, W. and Neil, E. T. Bayesian model averaging for analysis of lattice field theory results. Physical Review D, 103(1 10:(11):114) 114502, 2021.
Jordan, R., Kinderlehrer, D., and Otto, F. The variational formulation of the fokker plankkk equation. SIAM journal on mathematical analysis, 29(1):1-17, 1998.
Gabrié, M., Rotskoff, G. M., and Vanden-Eijnden E. Adap- tive monte carlo augmented with normalizing flows. Pro- ceedings of the National Academy of Sciences, 119(;0): c2109420119, 2022.
Kandasamy, K., Krishnamurthy, A., Schneider, J., and Póczos, B. Parallelised bayesian optimisation via thomp- son sampling. In International conference on artificial intelligence and statistics, pp. 133-142. PMLR, 2018.
Gelfand, S. B., Mitter, S. K., et al. On sampling methods and annealing algorithms. 1990.
Girolami, M. and Calderhead, B. Riemann manifold langevin and hamiltonian monte carlo methods. Jour- nal of the Royal Statistical Society Series Statistical Methodology, 73(2):123-2144 2011.
Karagiannis, G. and Andrieu, C. Annealed importance sampling reversible jump meme algorithms. Journal of Computational and Graphical Statistics, 22(3):623-66 2488, 2013.
Griffin, J. E. and Walker, G.
On adaptive metropolis- hastings methods. Statistics and Computing, 13:123-134, 2013.
Li, L.., Liu, Q., Korba, A., Yurochkin, M., and Solomon, J. Sampling with mollified interaction energy descent. In Proceedings of the International Conference Learning Representations (ICLR), 2023.
Gu, M. and Sun, Neural langevin dynamical sampling. IEEE Access, 8:31595-31605, 2020.
Li, T.-M., Lehtinen, J., Ramamoorthi, R., Jakob, W., and Durand, Anisotropic gaussian mutations for metropo- lis light transport through
Haario, H., Saksman, E., and Tamminen, J. An adaptive metropolis algorithm. 2001.
dynamics. ACM Transactions Graphics (TOG), 34(6):1-13, 2015.
Hackett, D. C., Hsieh, C.-C., Albergo, M. S., Boyda, D., Chen, J.-W., Chen, K.-F., Cranmer, K., Kanwar, G., and Shanahan, P. E. Flow-based sampling for multi- modal distributions in lattice field theory. arXiv preprint arXiv:2107.00734, 2021.
Li, 7., Chen, Y., and Sommer, T. A neural network meme sampler that maximizes proposal entropy. Entropy, 23(3): 269, 2021.
Liu, Q. Stein variational gradient descent as gradient flow. Advances in neural information processing systems, 30, 2017.
Hagemann, P., Hertrich, J., and Steidl, G. Stochastic nor- malizing flows for inverse problems: markov chains viewpoint. SIAM/ASA Journal on Uncertainty Quantifi- cation, 10(3):1 2022.
Liu, Q. and Wang, D. Stein variational gradient descent: general purpose bayesian inference algorithm. Advances in neural information processing systems, 29, 2016.
Hoffman, M., Radul, A., and Sountsov, P. An adaptive- mome scheme for setting trajectory lengths in hamiltonian monte carlo. International Conference on Artificial Intelligence and Statistics, pp. 3907-3915. PMLR, 2021.
Lozanovski, B., Downing, D., Tran, P., Shidid, D., Qian, M., Choong, P., Brandt, M., and Leary, M. A monte carlo simulation-basee approach realistic modelling of additively manufactured lattice structures. Additive Manafacturing, 32:101092, 2020.
Hoffman, M. D., Gelman, A., et al. The no-u turu sampler: adaptively setting path lengths in hamiltonian monte carlo. Mach. Learn. Res., (5(1):1593316)) 2014.
Lynn, J. E., Tews, I., Gandolfi, S., and Lovato, A. Quantum monte carlo methods in nuclear physics: recent advances.
10

---

Annealing Flow Generative Models Towards Sampling High--Diennsinnll and Multi-Modal Distributions
Annual Review of Nuclear and Particle Science, 69(1): 279-305, 2019.
Ruthotto, L., Osher, S. J., Li, W., Nurbekyan, L.,, and Fung, S. W. A machine learning framework for solving high- dimensional mean field game and mean field control prob- lems. Proceedings of the National Academy of Sciences, 1 117(17):9183 83-9193, 2020.
Maddison, C. J., Paulin, D., Teh, Y. W. O'Donoghue, B., and Doucet, A. Hamiltonian descent methods. arXiv preprint arXiv:/809. wrxv:/809.00542,, 050400 2018.
Salo-Ahen, O. M., Alanko, L., Bhadane, R., Bonvin, A. M., Honorato, R. V., Hossain, S., Juffer, A. H., Kabedev, A., kahtela-Kakkonen M., Larsen, A. S., al. Molecular dy- namics simulations in drug discovery and pharmaceutical development. Processes, 9(1):71, 2020.
Matthews, A., Arbel, M., Rezende, D. J., and Doucet, A. Continual repeated annealed flow transport monte carlo. In International Conference on Machine Learning, pp- 15196-15219. PMLR, 2022.
Shahbaba, B., Lan, S., Johnson, W. O., and Neal, R. M. Split hamiltonian monte carlo. Statistics and Computing, 24:339-349, 2014.
Maurais, A. and Marzouk, Y. Sampling in unit time with kernel fisher-rao flow. arXiv preprint arXiv:2401.33892 2024.
Sorkin, G. B. Efficient simulated annealing on fractal energy landscapes. Algorithmica, 6:367-415 3667-188,7 1991.
Miao, Y., Feher, A., and McCammon, J. A. Gaussian accelerated molecular dynamics: unconstrained enhanced sampling and free energy calculation. Journal of chemicul theory and computation, 118::3588433555 2015.
Stanton, S., Maddox, W., Gruver, N., Maffettone, P., De- laney, E.,, Greenside, P., and Wilson, A. G. Accelerating bayesian optimization for biological sequence design with denoising autoencoders. International Conference on Machine Learning, pp. 20459-24400 204799. PMLR, 2022.
Neal, M. Annealed importance sampling. Statistics and computing, 11:125-139, 2001.
Nitanda, A. and Suzuki, T. Stochastic particle gradi- ent descent for infinite ensembles. arXiv preprint 1r1iv:1712.05455 2017.
Stephan, M., Hoffman, M. D., Blei, D. M., al. Stochas- tic gradient descent as approximate bayesian inference. Journal of Machine Learning Research, 18(134):1-35, 2017.
Overstall, A. M., Woods, D. C., and Parker, B. M. Bayesian optimal design for ordinary differential equation models with application biological science. Journal of the American Statistical Association, 2020.
Syed, S., Bouchard-Cöté, A., Deligiannidis, G., and Doucet, Non-reversible parallel tempering: a scalable highly parallel meme scheme. Journal of the Royal Statistical Society Series B: Statistical Methodology, 84(2):321-350 2022.
Pulido, M. and van Leeuwen, J. Sequential monte carlo with kernel embedded mappings: The mapping particle filter. Journal of Computational Physics, 396:400-415 2019.
Tian, Y., Panda, N., and Lin, Y. T. Liouville flow importance sampler. arXiv preprint 1rxxi:2405.06672 2024.
Van Groenigen, J. and Stein, A. Constrained optimization of spatial sampling using continuous simulated annealing. Technical report, Wiley Online Library, 1998.
Qiu, and Wang, X. Efficient multimodal sampling via tempered distribution flow. Journal the American Sta- tistical Association, 1
2024.
Van Ravenzwaaij. D., Cassey, P., and Brown, S. D. A sim- ple introduction to markov chain monte-carlo sampling. Psychonomic bulletin & review, 25(1): 1511)1443,154, 2018.
Rezende, D. and Mohamed, S. Variational inference with normalizing flows. In International conference on ma- chine learning, pp. 1530-1538. PMLR, 2015.
Wolniewicz, L. M., Sadowski, P., and Corti, C. Neural surrogate hme: Accelerated hamiltonian monte carlo with neural network surrogate likelihood. arXiv preprint arXiv:2407.204. 2024.
Rhodes, B., Xu, K., and Gutmann, M. U. Telescoping density-ratio estimation. Advances in neural information processing systems, 33:4905 4916, 2020.
Ribera Borrell, E., Quer, J., Richter, I.., and Schütte, C. Improving control based importance sampling strate- gies for metastable diffusions via adapted metadynamics. SIAM Journal on Scientific Computing, 46(2):8298-S33 S3233, 2024.
Wu, D., Gao, J., and Yin, F. Bayesian-booster metaloc: Ef- ficient training and guaranteed generalization for indoor localization. In ICASSP 2024-2024 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP), pp. 6070-6074. IEFF, 2024.

---

Anncaling Flow Generative Models Towards Sampling Hligh-Dimensional and Multi-Modal Distributions
Wu, H., Köhler, J., and Noé, Stochastic normalizing flows. Advances in Neural Information Processing Systems, 33: 5933-5944, 2020.
Xu, C., Cheng, X., and Xie, Y. Local flow matching genera- tive models. arXiv preprint arXiv:2410.02544, 2024a.
Xu, C., Cheng, X., and Xie, Y. Normalizing flow neural networks by jko scheme. Advances in Neural Information Processing Systems, 36, 2024b.
Xu, C., Lee, J., Cheng, X., and Xie, Y. Flow-based distribu- tionally robust optimization. IEEE Journal on Selected Areas in Information Theory, 2024c.
Xu, C., Cheng, X., and Xie, Y. Computing high-dimensional optimal transport by flow neural networks. In The 28th International Conference on Artificial Intelligence and Statistics, 2025. URL https: / http://peeeeeee.eetttteee net/ forum? -.EWNNNsSV......
Zhang, G., Hsu, K., J., Finn, C., and Grosse, R. B. Dif- ferentiable annealed importance sampling and the perils of gradient noise. Advances in Neural Information Pro- cessing Systems, 34: Ne.19988-19410 2021.
12

---

Annealing Flow Generative Models Towards Sampling High-Dimennional and Multi-Modal Distributions
A. Proofs
A.1. Proofs in Section 3.1
Proposition 3.1. IKL-Divergence Decomposition) Given the samples from fk-1, the KL-Divergence between T*fk- and fx equivalent to:
(L((#fi-illfa) = + Ee- - log fa((tt))) - -fth V (x
sk-2
up to constant that is independent of ub (x((),
Proof:
Let p(x,() denote the density evolution under the transport map T. defined in the continuity equation (2). By the constraints in the transport map objective (3), have T* f (x) p(xr, 4x). The expression for KI.-divergence is given by:
Tx fu (x))
K || Sk) = Eemp(a:))) log
Ex-p(x,)) [log Tw Sk (x) - log Sa (x)]
(21)
fx(
Recall that in equation (4), we express f& (x) = (z)) Substituting this into the equation yields:
KL(T fk) = Exmple,(es [log T+ fa-1 (x) - log Ja [x() log Zk
(22)
Exxp(t,tte log T# J (zx(4)) - log fa (zx(t=)) log Z
(23)
where the second equality holds because the flow maps samples x(th-1) to xttx) in way that preserves the probability measure by the continuity equation (2). That is,
Ex-p(stt2
Ex-p(e,te-n) - -etee)))-tt)))
(24)
where the particles a=ft evolve according to the Neural ODE (1).
Next, to compute log T fa (x((+)),, we use the fact that the dynamics of the pushforward density are governed by the velocity field is (z))),,, s):
d
log p(gp((s)),s) = Vp(x(s), 0,z(s) 0sp(x(s),
(25)
ds
s))()),s)
Vp.uz Uk (pux)
(by (1) and (2))
(26)
(m(a)..)
Vp.u. Uk -VVp Uk PV Uk
(27)
(m(s)).)
- (x()))), 8).
(28)
Integrating this equation over the interval s E (itt-,,t)) tx), we find:
frh V vn (x(s))),
log Tw ( log p()()))) (x) log
(29)
We now substitute this result back into the Kl.-divergence expreeesgncceeppessiiinn.
KL fa) = Ex-oladee 1) (ogp(x(ta-)))taa1 - 1'te
V ux(c()), s)ds log fa((((t))) - log Zg. (30)
1--1
Note that Bep-p(((t--1). 1)) Dogg ((x((1(1));) 4-1) is independent of (x((x(s), 8) t and thus acts as constant term, along with log Ze. which we denote as c:
c E Bappfu(ts 1) logg((((44-))) 4x-14) log 7
(31)
13

---

Anncaling Flow Generative Models Towards Sampling High-Dimenniinall and Multi-Modal Distributions
Besides, after successfully training the previous velocity fields, we have p(x((e-1); (x-1(1). Therefore, the relevant terms for the KL--divergence are:
fL fa) = Sa-1 - log fa (x(tx)) - f
FR V V. (u(r(s),))ds
(32)
18-1
Lemma 3.2. (Wasserstein Distance Discretization) Let x(t) be particle trajectories driven by a smooth velocity field Uk (x=(t).t) 1) over the time interval [fk-1, where h& ix tk-1- Assume that vg (x. t) Lipschitz continuous in both
and I. By dividing 1 1x) into S equal mini-intervals with grid points 1, where 1,..., S and 1 1,0 lk- 1
1k-1,5 1,8 1g, have:
(((()) (1)) S hx
E [1]=(ta-1,011) =(tx-1,.)]1 +0 (Ht/S))
1
Proof:
Consider particle trajectories xtt) driven by a sufficiently smooth velocity field es (x((t,,, t) over the time interval (f
where hi 1s 1--1- We divide this interval into equal mini-intervals of length 8t 5 resulting in grid points
f4-1, = 11-1-+ + sōt for s = 0, 1, where öt f--1--1
Within each mini-interval [Lk-1,s1 11-1, we perform Taylor expansion of s(1) around 1--1,0
1 dua
x((R-1,+11) m15x-1,s) + uk (x(x(((1111)) 41 +
51 O(81))
(33)
2 dt
where des denotes the total derivative of U3 with respect to time. dt
The squared displacement over the mini-interval [tx 11s: 1.s: :. tu 1, is given by:
jōt + dup 8t2 + 0(833) 2 dt
||x(te-1.041) x 1f)= (z((t)
(34)
[1uu(
1t O
(35)
as we assume that 0g L-Lipschitz continuous and it follows that negligible 5t ->
des de
L. The higher-order terms O(6t²) become
Summing the expected squared displacements all mini-intervals, obtain:
Σ E [10((1-1,041) - x((+=11,) 512
Σ E (1
(36)
Now, we examine the L.H.S. of Lemma 3.2 by approximating the integral of the expected squared velocity using Riemann sum:
Exte [llex(x(f) dt St
E
1.0 + 0 (8-
==0
5t
E [11(((=1144)) - O(S 51)
(37)
8t2
E [||0
x((6-1,,))11" O 612),
where the Riemann sum error term OSS 8t arises from well-known result (for instance, see Chapter Axler (2020)),
given the assumption that vg is T.- Lipschit continuous.
14

---

Annealing Flow Generative Models Towards Sampling High-Dimennional and Multi-Modal Distributions
A.2. Proofs in Section 3.2
Before formally proving Theorem 3.3, we first reformulate the objective in (10) form involving the Stein operator, as stated in Lemma 3.3.
Lemma 3.3. (Objective Reformulation) Denote hy = - - and let V log fx( denote the score function of Jk- As lig- - and with y = 2' the objective in (10) becomes equivalent to the following:
As uk + 110x113
min Ex (-.0)
Ase us = 3y Uk + V uk
Proof:
From the Neural ODE (I) and using Taylor's expansion, obtain:
1 (x((s)) hhoa(o((4-11)) + 0
<(lx) x1(x-1)
(38)
A4-1
Next, by performing Taylor expansion of - log fa (z((t)) around 1x-1
log fa((((t))
log fa (ottaa-o-(---)))00ooooooooaa-aa-
log fo (x h4 V log (z((t= ( mz + 0(h2)
(39)
log (x h4 sx(x(tx-1)) (10(0)))) O
where we define the score function 8% v Vlogja = Vlog fr. Besides, we also have that:
(x((), h& (((x((x-)),441)) +
(40)
As h4 + 0, we no longer need to divide the time interval, i.e., S = The objective function (10) can be then approximated
as:
- log fa(((t) - f" v V. s)
Ex-
|lx(t)) wtttttt-000t 1)1?
2hv
= -onsse (-log log (x - h sx (x(((x211) u (1r( (1- 1), 0(12))
(41)
(hx - (=(( 11, + + + 2/18
1 (2 (x(4)))))()))))
log (a(x) + he (-sa(a)) (x 4-1) V uz 4 4-1 + 2 |, (a=1))l)) + 0(h2)
= Egnsy
Since Bn(t)-)) log F (x((k-1))] is independent of vx(x, t), as h O, the minimization of the leading term equivalent to:
min
Ex- AAp.
As. uk + 2 |l A : 2x V
(42)
(s.))
Here, Ag, Uh denotes the Stein operator applied the vector field eg under the target density fk-
15

---

Annealing Flow Generative Models Towards Sampling Hligh-Dimensional and Multi-Modal Distributions
Theorem 3.3 (Optimal Velocity Field as Score Difference) Suppose hy - O. Let fk-1 and fr. be continuously differentiable on R. Assume that V un (x) exists for all E IR and V Uu (z), 5k-1 and sg belong to L²(fe-1)- Assume that the components of Up are independent and |1 fa (2-1(2))llxx (z) ll2 = 0. Under these conditions, the minimizer of (10) is:
lim uk = Sk 86-1- h4-0
Proof:
Under the assumptions that hy - 0 and = 2/2 we begin by considering the equivalent minimization objective derived in Lemma 3.3:
min J(rk) == min Ex-fo-n AAp, uk + |112112] As, Uk = 2o ok + V lk-
Expanding the objective functional, we have:
1 Ex- fx-1 [SR 5k- UK- V uu + 2 ||0011
fn- fn-1(x)) / Rd
(-su(t)) (-su(x)-ex(x)) vx()) V vx(xx)+ + _[lva(e))P)) dx. (43)
Define B, = fx || I r). and let 0B, denote the boundary of B,, which is the sphere radius r. Under the assumption that lim - fa- ( = 0, we have the following:
1/4 3d V. (fa-10x) = = lim
/ V
(fa-11k) dzl
(44)
r-0
B,
lim
fa (x)(()) m(x)dS(x))
(45)
lim
fa-1 12 (e-1l1//ll/m(t>)ldd5d
(46)
lim J
fo U0 (2d)(((()
(47)
(48)
Therefore, Ja V (fr-1 Ux) dr = 0. Next, we further expand the divergence theorem:
0 gd V. (o)))())
fo-1 (x:) V (x(xddr + ux(x) (ir--(
(49)
(6-1(1) vx(x)do + vx(x). 84- (x)) (k-1(x) dx
Substitute the result back into the objective functional, we have:
Ex- - S3 - Uy Uk + [|uu|13] = [ fx-)(x)) (-s(()) (05(x))-x(() -x(x)) - V vx(2) + [11x(2)11)) de
(50)
[ ( (SK-1(0) - -x(()) vx(x) + ||x(2)))1) da.
(51)
The integrand does not involve (z) 1, ..... d and higher-order derivatives. Assuming the components Uk.j. 1.-.. d of Up are independent, we can take the functional derivative component-wise and set them to zero:
8J = ft-1 (eu + (st-1 - -s+) = O,
(8k-1-5t4)
(52)
dCs
Since fx-1 > 0 for all x, this implies:
2/ 8/ 8k- 1-
(53)
16

---

Annealing Flow Generative Models Towards Sampling High-Dimennional and Multi-Modal Distributions
A.3. Proofs in Section 5.2
Density Ratio Estimation (DRE) By optimizing the following loss function:
Lx (Ox) = Ex(te-)>~se- [log(1 + eatattt ])1}] + Extesmn [log(1 + enulertte)))
the model learns an optimal r" (2; 02) - log (ur)(a))
Jx(x)
Proof:
Express the loss function integrals over
5==1 fx- (x) log (1 (1+ (1+e)))))) dx / fx(())lgg (1 (1te)a(t)) dx.
L*
(54)
Compute the functional derivative of LA with respect to ri: 8Cu(rk) fx-1(())
(a)
eta(a)
1 + fx(x))
(55)
6rk
1 etx(#)
Next, we can set the derivative (lx/)r((2) to zero to find the minimizer ri (x)):
(k=i(xx)
(x)) = In
(56)
fx(x))
Therefore, by concatenating each (z). we obtain
K
(x) = log [K-1(x) [K-2(1)
fo(x) fi(x)
log So(x) logg mr
(x)
(57)
fx(x) fK-1(x)
fK(x) q" (x)
the log density ratio between (o()) and (x).
B. Equivalence to Wasserstein gradient flow when
In this section, we demonstrate the equivalence between the dynamic optimal transport (OT) objective of AF and the .....,,) K) are set to and a static Wasserstein regularization is used in place of the dynamic Wasserstein regularization introduced in
Wasserstein Gradient Flow, under the condition that all 3g (k = 1, 2..,...
Langevin Dynamics and Cokker-Plance Equation: Langevin Dynamics is represented by the following SDE.
dX, VEE + 22dW,,
(58)
where E(x) = log (x) is the energy function of the equilibrium density f(x, T) = g(x). Let Xo and denote the density of X, by p(x, t). The Langevin Dynamics corresponds to the Pokker-Plande Equation (FPE), which describes the evolution of p(x, 1) towards the equilibrium p(x, T) =q(x). as follows:
d.p v (pVE Vp). p(z, px)x).
(59)
JKO Scheme and Wasserstein Gradient Flow: The Kinden--inndrfleeee-Otttto (JKO) scheme (Jordan al., 1998) time discretization scheme for gradient flows minimize KL-(plgg under the Wasserstein-2 metric. Given target density q and functional F(p, KL(pllq), the JKO scheme approximates the continuous gradient flow of p(z, () by solving sequence of minimization problems. Assume there are K steps with time stamps 0 to, 1,, ......,, stamp 4. the scheme updates at each time step by minimizing the functional
..," .tK = T, at each time
1 W (p (P-PP)
Pk arg min F(p.g) +
(60)
2T
17

---

Anncaling Flow Generative Models Towards Sampling High-Dimennsinal and Multi-Modal Distributions
where (2 ((,PPp-1 Pk-1) denotes the squared 2- 2-Wasserttein distance between the probability measures p and Pk- It was proven in Jordan et al. (1998) that as h i 1 approaches 0, the solution k(:, provided by the JKO scheme converges the solution of (59), at each step k.
It is straightforward to see that solving for the transport density Pk: using (60) is equivalent to solving for the transport map Ty via:
Tk = arg min 7:Re RRd
(L + Bx-pn 1| Txx 2T
1
(61)
Therefore, is immediately evident that the Wasserstein gradient flow based on the discretized JKO scheme is equivalent to (6) when each fi (I) is set as the target distribution ((x), i.c., when all 6g are set to I, and the second term in the objective (6) is relaxed to a static W2 regularization instead of a dynamic W/2 regularization.
There are some well-established properties regarding the convergence densities under the Wasserstein Gradient Flow objective when the target distribution ((x) is log-concave. Define P2 {P: S ()))) < oc) and P (P P2
P < d
Assumption B.1. For all n, the learned velocity field é, guarantees that the mappings Ta is non-degenerate. Additionally, for the time interval [tn-1; fn]. the integrated squared deviation of the velocity field satisfies the inequality
S ||v - 61|2 p S e, e (0.
Assumption B.2. F(p. q): : + (-o0, 6o) where Sas ||x2||dp(x)) < oc, is lower semi-continuous; Dom F)) C P3; F(p.g) X-convex a.g.g. P2 = {P: Sea 1| [x)))(((>) < <o0..
When the energy function E of q e is strongly convex, and F(p, q) is chosen as in (61) in our method, Assumption 2 holds true. Under the Assumptions 1-2, the Wasserstein gradient flow converges toward the target distribution at polynomial rate in terms of the Wasserstein W2 distance. Specifically, the distance between the iterates and the target distribution decays geometrically with each iteration, up to a fixed error term (Cheng et al., 2024).
Furthermore, when additional regularity and integrability conditions are inposed-ensuring boundedness, smoothness, and suitable tail behavior- the convergence behavior improves stronger sense. Under these conditions, the gradient flow is shown to converge exponentially fast in the x² -iveergee measure, meaning that the discrepancy between the iterates and the target decays at exponential rate (Xu et al., 2024a).
Therefore, when all 3g are set to I and a static Wasserstein regularization is used in place of the dynamic Wasserstein regularization introduced in (9), our AF retains the well-established convergence properties for log-concave in terms of both the Wy-distance and the -ivieergee measure. Furthermore, for non-log-concave densities q, we have established in Theorem 3.3 that the difference in the optimal velocity field between two consecutive annealing densities equals the score difference, under the unique dynamic optimal transport (OT) objective of AF.
C. Experimental Details
C.1. Alternative Loss
We empirically observed that replacing the first term - log (x((t)) in the objective (10) with the approximation - h V log (x(xx)) Uk leads to improved performance. This substitution is motivated by the first-order Taylor expansion of 11-log log Ja (x((x=1)) around v((2), given by:
log fa(c(te-1)) = log fa(((tt)) Vo fx(ta((((())) (x(tz-1) x(tx)) +
(62)
= log fa(a(tt)) + h4 V log ((gg + (lh2).
Therefore, the alternative objective (10) is given by:
min Ex(tt-1)<-fo> C1 - h4 v log (x((t)) 2q f/ V (x(ss) 1os+oo a (xt).)
S
Σ
| - 2(tx-1,))1
(63)
18

---

Anncaling Flow Generative Models Towards Sampling High-Dimennional and Multi-Modal Distributions
where Gj -lo log (x((((/))))) - (hh)), noting that log (=((((=))) is independent of vy given samples sr(4-1) SK-1-
We found that this alternative loss led to slightly better performance across several distributions. The alternative objective (63) was used consistently for experiments on the Gaussian Mixture Model, Funnel distribution, and Exponentially Weighted Gaussian. For the Truncated Normal and Bayesian Logistic Regression tasks, we retained the original objective (10) formulation.
C.2. Hutchinson trace estimator
The objective function in (10) involve the calculation of V (,(2,, 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 much slower and less stable in high dimensions.
[e Jz())c). where (x)x) is the Jacobian of uy (2, at æ. Given a fixed e, we have the (e (m)) which is the directional derivative of up along the direction e. Therefore, for a sufficiently small a 0, we can propose the following estimator (Hutchinson, 1989; Xu et al., 2024b):
We can express (x)) lim
(2,
2 (0 + oe,t) - (x(x
V vg(2, x Bewp(e)
(64)
where p(c) is distribution in Re satisfying E[c] = 0 and Cov(e) I (e.g., standard Gaussian). This approximation becomes exact + 0. In our experiments, we set 0.02/Vdd
C.3. Other Annealing Flow settings
Time steps and numerical integration
By selecting K values of B, we divide the original time scale [0, 1] of the Continuous Normalizing Flow (2) and (3) into K intervals: [f fx) for k = 1, ...,,.. K. Notice that the learning of each velocity field ek depends only on the samples from the (I: 1) block, not on the specific time stamp. Therefore, we can re-scale each block's time interval to [o, 1], knowing that using the time stamps [(A 1)h, kh] yields the same results as using [0, 1] for the neural network (x, t). For example, the neural network will learn 2% (z,, 0) = 2% (2r, (k- 1)h) and 4% (z. 1) = dy. (1, kh), regardless of the time stamps.
Recall that we relaxed the shortest transport map path into a dynamic W₂ regularization loss via Lemma 3.2. This requires calculating intermediate points x(tx-1..), where = 0, 1,..,,,.... S. We set $ = 3, evenly spacing the points on [fx-1, fx), (ttt=11 + 2h./3), x (txx). To compute each =.f4-11.); we integrate resulting in the path points x(fx-1), (ft-1111111 h2/3). 33). xx
the velocity field uk between f and 1.s1 using the Runge- Kutta method for numerical integration. Specifically, to calculate x(l + h) based on x() and an intermediate time stamp / + h.
2.
+(t+h) x(t) + (ka + 2k2 + 2ka ka);
kz, + h
ln = v()(t),t),
l = xx())
2 2
1 kg =v 1 (a(1) (1() + Th k2, L +
kG = v (x(()) + hks, ++ 1))
Here, is the step size, and t(r,t) represents the velocity field.
The choice of Bk
experiments on Gaussian Mixture Models (GMMs), we set the number intermediate 3e values to 8, equally spaced such that f0 0, 31 1/8, 82 = 2/8, , Agg = 1. We chose the casy-to-sample distribution (T) as (0, Ta). Finally, we added 2 refinement blocks. The intermediate distributions are defined as:
fx(x) (=(====++)(())
In the experiment on the Truncated Normal Distribution, we did not select B& in the same manner as for the GMM and Exp- Sep-Weighted Gaussian distributions. Instead, following the same Annealing philosophy, construct a gradually transforming bridge from Tu( to ((r) NO, Ia) by setting each intermediate density
(x) = 111124/(44)) Ta).
19

---

Anncaling Flow Generative Models Towards Sampling Hligh-Dimensional and Multi-Modal Distributions
The number of intermediate 3 values set to 8.
experiment on funnel distributions, we set all 3 1, with the number of time steps set to discussed in Appendix B, the algorithm becomes equivalent to a Wasserstein gradient descent problem.
In the experiment on 50D Sei-Weeghted Gaussian, 20 time steps are used, with 15 intermediate densities and 5 refinement blocks.
The choice a
In the experiments Gaussian Mixture Models (GMMs), funnel distributions, truncated normal, and Bayesian Logistic Regressions, is uniformly set [3 3
"...... the experiments on Exp-Weighted Gaussian, a set to 1,
13 20
10
1,1,
3 31
Neural networks and selection other hyperparameters
The neural network structure in our experiments is consistently set with hidden layers of size 32-32. During implementation, we observed that when S 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.
We sample 100,000 data points from (0, Ia) for training, with batch size of 1,000. The Adam optimizer is used with learning rate of 0.0001, and the maximum number of iterations for each block en set 1,000.
Different numbers of test samples are used 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 See-Weeghted Gaussian with 1,024 modes in 10D space.
C.4. Bayesian logistic regression
We use a hierarchical Bayesian structure for logistic regression across a range datasets provided by LIBSVM. The detailed setting the Bayesian Logistic Regression is as follows.
We adopt the same Bayesian logistic regression setting as described in Liu & Wang (2016), where hierarchical structure assigned to the model parameters. The weights 3 follow Gaussian prior, Pa (3 = (8; 0, a a= and follows Gamma prior, po(a) Gamma(a; 1, 0.01). The datasets used binary. where has varying number of features, and y, E {+1, -1} across different datasets. Sampling is performed from the posterior distribution:
a II
a(3,a(D x Gamma(a; 1, 0.01) II 0,
+ xx)
We Be and use blocks train the Annealing Flow.
During testing, we use all algorithms to sample 1,000 particles of 8 and a jointly, and use (18(0) 1000 to construct 1,000 classifiers. The mean accuracy and standard deviation are then reported in Table 3. Additionally, the average log posterior in Table 3 is reported as:
Σ
log
Σ p(ulx,0)
Duest
C
E. Dina
DEC
C.5. Importance flow
We report the results of the importance sampler (discussed in Section 5) for estimating E2-N((.1) [А]=]e]] with varying and dimensions, based on Annealing Flow. To estimate Ry~N(o..) [1 [=] we know that the theoretically optimal proposal distribution which can achieve 0 variance 4" (x) I1 1). Then the estimator becomes:
To( (x)
1
Σ To Ta(x))
Ex=egla) [x(X)] = Ex-
h(x)
M(xs), i=1 ge (2:)
as s q
n
where <x) No0, Ta), h(>) = 1112 and (x) 7 7 (x))
Therefore, the Importance Flow consists of parts: First, using Annealing Flow to sample from g" (æ): second, constructing Density Ratio Estimation (DRE) neural network using samples from {xr) ~ q" (2) and {ur] ~ [0, La). as discussed Section 5.2. The estimator becomes:
DRE(2.) h(x:).
The Naive MC results comes from directly using ( (O, Ia) construct estimator 1 E. 1b.. When the Naive MC methods consistently output as the result.
20

---

Anncaling Flow Generative Models Towards Sampling High-Dimennional and Multi-Modal Distributions
experiment, we use single DRE neural network construct the density ratio between TT())) and q (x) Z 1 2e (0, 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 e and dimension d. The Adam optimizer is used, with 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.
C.6. Details of other algorithms
The Algorithms 2 and 3 introduce the algorithmic framework of Metropolis-HHane Hasstnng (MHl) and Parallel Tempering (PT) compared in our experiments.
Algorithm 2 Metsopolis-Hastings Algorithm
I: Initialize
2: fort = 1==l to N do
3: Propose x² ((=1(=))))
4: Compute acceplance ratio min (1, aae" "ellleeeee (x")
5: Sample Uniform(O, 1)
6: i a then
8: else
end F1-1
for
12:
return {2x)/
Algorithm Parallel Tempering Algorithm
1: Initialize replicas { R2.. Foum replicas) with Gaussian noise
2: Initialize temperatures [Ti, T2...., Ne....,Tmmmpliiss Timmsppiczs)
3: for i = to iterations do 4: for j = 1 num replicas do
5: Propose 23 ~ (zz] {Using Metropolis-Hasttnnss Hastings for cach replica)
mlx2)
Compute acceptance ratio a, REP57
7: Sample Uniform(O, 1)
if < aj then = 10: end if
8:
Store zy sumples for replica j end for
13:
mod exchange interval then
14:
for j = to num num replicas do
15:
Compute energies F,- = log(s(xx)) + +c), Es41 -lg((((15))) +
16:
Compute = } T
(Ez-1 Ey)
Sample Uniform(0.) 1)
18:
u exp(A) then
Swap x, ++ 1141 20: end
19:
21: end for
22:
end
23:
end for
24: return samples from all replicas
In experiments, we set the proposal density q(a'lx)= =N(x;; 0, Iz). We use replicas in Parallel Tempering (PT), with linear temperature progression ranging from
1.0 Taax 2.0, and an exchange interval 100 iterations. For HMC, we set the number of leapfrog steps to 10, with a step size (c) of 0.01, and the mass matrix is set as the identity matrix. Additionally, we use the default hyperparamcters as specified in SVGD (Liu & Wang. 2016), MIED (Li et al., 2023), and Al-Sampler (Egorov et al., 2024). In the actual implementation, found that the time required for SVGD to converge increases significantly with the number of samples. Therefore, in most experiments, sample 1000 data points time using SVGD, aggregate the samples, and then generate the final plot. The settings for CRAFT, LFIS, PGPS are discussed alongside table figure in Appendix
21

---

Annealing Flow Generative Models Towards Sampling Hligh-Dimensionll and Multi-Modal Distributions
C.7. Evaluation metries
To assess the performance of our model, we utilized two key metries: Maximum Mcan 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 non-parametric metric used quantify the difference hetween distrihutions based samples. Given of samples X1 E TR"1 xd and X₂ E R=axx MMD computes the kernel-basce distances between these sets. Specifically, we employed Gaussian
kernel:
k(x, v) exp(-ai|æ vll),,
parameterized by a bandwidth The MMD is computed as follows:
MMD(X1,
Σ M(x ni
x
1(x2, 71
Σ X(x,, 1,3
22 Fd
where k(,,)) represents the Gaussian kernel. In our experiments, the median of the pairwise distances between the two datasets.
1/7² and 7 0.1 median_didttt where median_disg denotes
Wasserstein Distance
In addition to MMD, we used the Wasserstein distance, which measures the cost of transporting mass between distributions. Given two point X Rd and Y R", we compute the pairwise Euclidean distance between the points. The Wasserstein distance then computed using the optimal transport plan via the linear sum assignment method (from scipy.optimize package):
W(X, Y) 1
Σ
Σ X - Ya
where r(i) and c(i) are the optimal row and column assignments determined through linear sum assignment.
Sample-Variance Variance Mean Squared Error across dimensions
This metric is reported for Seeihee Gaussian distribution in our experiments. Recall that the distribution of a 50D ExpGauss is:
Ind of
10 Ii St - 11111
p(z1, 22, (50) x
To calculate the Mean Squared Error of sample variances across dimensions, first take the absolute values of the first 10 dimensions, while leaving the values of the remaining dimensions unchanged. The metric is calculated as:
MSE
1 D
0o) Σ
oz)
aid 22)
In all experiments, 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.
D. Additional Experiment Results
We adopt the standard Annealing Flow framework discussed in this paper for experiments Gaussian Mixture Models (GMM), Truncated Normal distributions, Exp-Weighted Gaussian distributions, and Bayesian Logistic Regressions. For experiments funnel distributions, we set each (z) the target g(z), under which the Annealing Flow objective becomes equivalent to the Wasserstein Gradient Flow based on the JKO scheme, as discussed in B. Please refer C.3 for St selections.
D.I. More Results
We begin by presenting the numerical tables for Gaussian Mixture Models (GMMs). In the tables, GMM refers to Gaussian Mixture Models, while wGMM denotes unequally weighted Gaussian Mixture Models. The notation (W)GMM- (OMM-Mnumber (number of modes) (radii- of the circle) represents (w)GMM with the specified number of modes arranged circle of the given radius. Note that the Mode Mode-Weight Mean Squared Error (MSE) also reported in Table (Main Script).
MMD and Wasserstein Distance
22

---

Annealing Flow Generative Models Towards Sampling High-Dimennsinall and Multi-Modal Distributions
Table MMD results of different methods the GMM distributions
Disebbuions
CRAFT
LAIS
PGPS
SCGD
MIED
Al Al-Sper
GMM-6-8
2,38 10
2,50 10
1.15 x 10-2
7.12 10
6.27
10-
9.35
10
9.52 10
2.34 x 10
10-
131
GMM-10-12
3.01 1000"0000
9.06 10
8.97 10
701 10
9.11
10
1 10
6.28 100
4.02 111
wCIMM-10-122
4.96 10
9.96 10
1.14 10
895 10
7.02 10
2.72
8.31 10
3.19 10
GMM-58
10
10
d=
GMM 5:10
6MM-10-12
1.25 10
1.67 1.67x
9.76 10
2,14 10
7.20 x 10
4.52 x 10-2
7.76 10"2
831 10
8.98 10
1,18 x 10
9.63 10
158 10
9.88 x 10
3.88 10
5.26 x 111
6.37 x 10
10
Table 5, Wasserstein distance results of different methods across the GMM distributions
Distrbttions
CRAFT
PGPS
SCGD
MIED
Al-Sampler
GMM-663
1:38 x 10
9.28 10
8.21 x 10
6,35 10
6.71 11
9357
30
8.01
x0
7.92
10
GMM-8-10
T.22 x
7.67 x 10
8.90 x 10
5.82 1010
5,98 10
1.34 x 10
8.15 x 10
8.35
10
0MM-10-12
10+3
10-6
10-1
10
WGMM-110112
9.94x10-1
9.78
10
1,02 x 10
7,88
10""
6.48 x 10
2:5 1072
1.06 10
8.06
101
GMM-6-8
1.97 10
1,12 10
1.1 x 10
7.76 T0
1,07 10
1.38 10
2,52 10
2.58 x 10
GMM-8-10
10
10+1
10+3
10
10
GMM 10-12
2.82
10
2.53
4.89 x 10
8.75 10"h
183 x 10
2:45
1073
7.89
3.89
18
WGMMM-10-12
3.53
1010
3.03
x 10
5.21 x 30
8.64
10
2,13 x 10
2.34
10
1,13
10
9.78 10
M5
ou1
(a) True
(b) AF
(c)
CRAFT
(d) LFIS
(c) PGPS
(l) PT
(g)
SVGD
(h) MIED
ii AIS
Figure Sampling methods for Gaussian Mixture Models (GMM) with 8, and 10 modes arranged on circles with radii 8, 10, 12. The number of time steps for CRAFT, LFIS, and PGPS is set to 12, the same for AF.
Mx2
605
900
(a) True
(b) AF
(c) CRAFT
(d) LFIS
(e) PGPS
(f) PT
(g) SVGD
(h) MIED
ii AIS
Figure 5. Sampling methods for Gaussian Mixture Models (GMM) with 6, 8, and 10 modes arranged on circles with radii = 8, 10, 12. The number of time steps for AF is 12. The number of time steps for CRAFT, LFIS, and PGPS 128, 256, and 128, respectively.
23

---

Anncaling Flow Generative Models Towards Sampling High-Dimennional and Multi-Modal Distributions
Gaussian Mixture Models (GMMs)
We tested each algorithm on GMMs with dimensions ranging from to S. In the 2D GMM, the modes are arranged in circles with radii v 8, 10, 12. For dimensions higher than 2, the coordinates of the additional dimensions set to r/2.
Figure 5 shows the results when the number of time steps for CRAFT, LFIS, and PGPS is set 10, the same as for AF. Additionally, Figure presents results where the number of time steps for CRAFT, LFIS, and PGPS is 128, 256, and 128, respectively, while the time step for AF remains at 10.
Truncated Normal Distribution
Relaxations applicd to the Truncated Normal Distribution in all experiments except for MH, HMC, and PT. Specifically. we relax the indicator function ||=]2e 2e T4exp(-Kii -ej) We set & 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:
-Rob
(a) True
(b) AF
(c) LFIS
(d) PGPS
(e) PT
(I) SVGD
(g) MIED
Figure 6. Sampling Methods for Truncated Normal Distributions with Radius 6, together with the failure cases of SVGD and MIED.
Each algorithm draws 5,000 samples. It can be observed that MCMC- HCMCC-based methods, including HMC and PT, produce many overlapping samples. This occurs because when new proposal is rejected, the algorithms retain the previous sample, leading to highly correlated sample sets.
-02g
Hogh Ma
Final Dempeto
(a) True
(b) AF
(c) MH
(d) HMC
(e) PT
Figure 7. Sampling methods for truncated normal distributions with radii
4, 6 in 2D space for the first three rows. The last row
presents sampling results SD with a radius 8, projected onto 3D space.
Funnel Distribution
In the main paper, present the sampling methods for the funnel distribution with projected onto 3D space. To the sample quality, here we present corresponding results projected onto space, plotted alongside density heat map.
24

---

Annealing Flow Generative Models Towards Sampling High-Dimentn Siah--Dimnsional and Multi-Modal Distributions
(a) AF
(b) CRAFT
(c) LFIS
(d) PGPS
(c) PT
(f) SVGD
g MIED
(h) AIS
(a) True
(b) AF
(c) CRAFT
(d) LFIS
(e) PGPS
(1) PT
(g) SVGD
(h) MIED
ii) AIS
Figure 9. Sampling Methods for Funnel Distribution with o² 0.81 in Dimension d 5, projected onto 3 Space.
As seen from both figures, AF method achieves the best sampling performance on the funnel distribution, while other methods, such MIED and AIS, fail capture the full spread the funnel's ' tail. Additionally, PT, SVGD, and ATS all fail capture the sharp part the funnel's shape.
Exp- Weighted Gaussian Distribution
Table in the main manuscript presents results for CRAFT, LFIS, and PGPS, using the number of intermediate time steps AF, Below, in Table 6, we report the number of modes explored 50D Exp- Weighted Gaussian distribution, which unequal variances across 10 dimensions, when CRAFT, LFIS, and PGPS are trained with 128, 256, and 128 time steps, respectively.
Table Number of modes explored in the Exp Weighted Gaussian distribution by different methods, with CRAFT, LFIS, and PGPS trained with 128, 256, and 128 time steps, respectively. AF trained with 20 time steps.
True
CRAFT
LFIS
PGPS
PT
SVGD
MIED
d=2
4
4
3.4
3.9
3.8
d
32
32
30.3
31.6
32
25.2
28.5
28.0
28.3
d=
10
1024
1024
984.0
993.2
1002.8
233.7
957.3
923.4
301.2
d= 50
1024
1024
886.5
923.4
994.0
10
916.4
890.6
125.6
We conducted more experiments on asymmetric 50D ExpGaussian distributions with unequal variances across dimensions. addition mode weights' MSE results shown in Table the main script, we provide an additional metric: mean squared error sample variances across dimensions.
Recall that the distribution of 50D ExpGaussian is:
110
+10
500 11n
3 11=11
p( .55 x
2
To calculate the Mean Squared Error of sample variances across dimensions, first take the absolute values of the first 10 dimensions. while leaving values the remaining dimensions unchanged. The metric calculated as:
MSE
1 D 0 -
0
Σ (rdd x;)²
Table 7 presents the sample-variance MSE across ExpGauss experiments. The fourth and fifth rows correspond to 50D ExpGauss where the variances of and 10 randomly selected dimensions correspond to 50D ExpGauss where variances 15 randomly selected dimensions are uniformly chosen between of = 0.5 and of 2, while the rest are set to I.
of 0.5, with the remaining dimensions to 1. The last two
Importance Flow
25

---

Annealing Flow Generative Models Towards Sampling Hligh-Dimensional and Multi-Modal Distributions
Table Sample-Variance Variance Mean Squared Error across dimensions for ExpGauss (1024 modes). The number of time steps for CRAFT, LFIS, and PGPS is set to 20, the same as for AF.
Distrabations
AF
CRAFT
LFIS
PGPS
SVGD
MIED
Al Al-SSmpler
d 10
ExpCinmst- 1004
112
23)
x 10
[9.8 + 49) 10
[9.8 +
10
(2.5 37] 1
5.2 +
50)
10
[1.7 + 28)
10
(3.9 + 76]
x 10
ExpCinass 1024
1.5
[3.7 71) 10
(3.5 1
(3.0
10
3.1
49)
10-2
[2.5 1 68)
10-
[8.2 2.00
10
[1.6+22 26)
[7.0
(8.6
10
(6.8 -3]
10-2
[6.0
56)
10-2
[3.2+ -7)
10-2
(9.9
10-2
d' 50
ExpGoassU..10 V-10-1025
(2.0 28
x 10-
[8.8 80) 10-2
(9.2 96]
10
(8.9 73)
10
88.1
90
10-2
(4.8 + 82)
10
(1. + 26]
10
Evp
121 31)
1111
(1.41
(10
19.9
10-2
[7.8
10
(1.8
10
ExpGonseU (V-20-/UVV2000002
V-20-1024
22. 30)
x 10
[2.1 23) 10
(2. +
44)
10
(19
39)
x 10
133
1 10
[9.8 +
86)
10
(2.3 +
.2)
x 10
Table 8 reports the preliminary results of the importance flow (discussed in Section 5) for estimating Ez-N(o.r) [1]=[2e] with varying radii e 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. Additionally, we discussed possible extension of the Importance Flow framework in D.4.
Table 8. Comparison of estimation results for (2,NN [e] across different radii and dimensions d. Values are reported as mean + standard deviation.
Methods
Radius
d= 2
d
4
d-
True Probability
e 4
(1.52) x 10-
(3.36) x 10-
(1.13) x 10-3
(7.49) 10-
(3.12) 10
(2.89) 10-7
(9.50) x 10-7
(6.84) x 10
Importance Flow
==4
(4.04 1.00) x 10
[9.81 4.00) 10-
[1:30 0.20) 10
11:111129
10-7
(2.13+0.88)
(3.36 x 10-1
10-7
(7.8610.82) x 10-3
(2.38 3.50) 10--
[7.56 4.0) x 10
(2.0200.63) x 10
[8:77 +091) x 10
[112 1.60) x 10
DRE with HMC Samples
(1.35 (:231)
(9.01
-0
(1.821290) 10-7
[2.1 2:21))))11))) 10
DRE with PT Samples
e 4
[6.79 (.558) x 10-1
(2.38 0.54) x 10-3
(5.78 0.78)7998111 x 10-3
(9.94 1.13) x 10
(5.37 + 9.56) x 10
(8.78 2.02) 10-7
(9:23 4 2.50)
[1.98 + 1.75)
[2.75 +600 x 10"
(1.18 1.10 10
(2:71 +.70) 10
(7.04 + 2.00) 10
MC
0
D.2. Training Efficiency
Table 9 presents the training and sampling times for AF, CRAFT, LFIS, and PGPS in experiments 50D Exp- Weighted Gaussian distribution, conducted on V100 GPU. The training setup includes 100,000 samples, 1,000 training iterations per time step. and a batch size of 1,000 for AF, CRAFT, and PGPS. For LFIS, batch size of 5,000 is used to ensure good performance. The sampling time measured for generating 10,000 samples. AF achieves optimal sampling performance with only 20 time steps, compared to other methods requiring up to 256 time steps. Specifically, CRAFT and PGPS were trained with 128 time steps, while LFIS used 256 time steps, ensuring these methods achieved the results shown in Figure 5 and Table 6.
Table 9. Total Training and Sampling Time Comparison for the 50D Exp- Weighted Gaussian
Methods
Training time (mins)
Sampling time (8)
AF
14.5
2.1 0.5
CRAFT
51.2 1 1.8
4.9
LFIS
86.4 + 3.5
6.4 0.6
PGPS
39.7 2.1
5.2 5.2.100 0.4
Table 10 reports the training and sampling times per time step (block) for each method. Notably, as AF requires numerical integration over the velocity field, its training time per time step is slightly higher compared to other methods.
Table 10. Training and Sampling Time Comparison Per Time Step for the 50D Weighte Gaussian
Methods
Training time (mins)
Sampling time (s)
AF
0.70 0.10
0.10 0.02
CRAFT
0.45 0.09
0.06 0.01
LFIS
0.37 + 0.07
0.04 + 0.01
PGPS
0.44 1: 0.10
0.04 : 0.01
D.3. Ablation Studies
As reported in the main manuseript, we use intermediate densities and refinement blacks for the GMM experiments, and 15 intermediate densities with 5 refinement blocks for the 50D Exp- Weighted Gaussian experiments. For GMMs, the regularization constant a is set to
26

---

Annealing Flow Generative Models Towards Sampling Hlgg--immensional and Multi-Modal Distributions
3: ...... For Exp- Weighted Gaussian, set to
29 3
16
1, 1,,,,,,
our main experiments, we note CRAFT, LFIS, and PGPS require 128, 256, and 128 time steps, respectively, to achieve comparable performance. In this section, we conduct additional ablation studies to investigate the role of annealing densities and Wasserstein regularization in ensuring the smoothness and success of Annealing Flow, particularly when fewer time sleps are used. We also compare the performance of all NF methods under further reduced time steps.
D.3.1. SIGNIFICANCE OF ANNEALING DENSITIES AND WASSERSTEIN REGULARIZATION
Here, conduct experiments GMMs without intermediate densities (i.e., all (x((r) = ((x)) and using 5 blocks. In Figure 10, the second column shows AF with the regularization constant a set L3 31 3 2.. the third column with a set [3 31 a 4| and the fourth column with set to
Figure Ablation Studies for Gaussian Mixture Models (GMM) with 6, 8, and modes arranged on circles 8, 10, 12. AF is trained with no intermediate densities and 5 training blocks. The regularization constant a used in the experiments for the three columns is described the first paragraph D.3.1.
The figure illustrates an extreme training scenario with no intermediate annealing densities, where target modes are far separated from the initial density To = N (0, Z). By comparing this to the successful case, where AF is trained with 8 intermediate densities and 2 refinement blocks (Figure 5), one can immediately see that annealing procedures are essential for successfully handling far-separated modes.
Furthermore, in Figure 10, when intermediate densities used, increasing the Wasserstein regularization constant- parttant--articularly the initial blocks-leade to improved results. This experimentally highlights the importance Wasserstein regularization our AF objective for ensuring stable performance and significantly reducing the number of intermediate time steps.
Table 11. Number of modes explored in the 50D Exponentially-Wcighhe Gaussian by various methods with different numbers of annealing densities K, while keeping the Wasserstein normalizing constant fixed.
AF
CRAFT
LFIS
PGPS
K
18.4
4.8
3.2
6.0
K 86.7
18.0
14.2
22.8
K
284.3
128.6
108.0
148.0
K 808.0
256.8
186.4
424.5
K =
996.2
382.0
208.4
578.8
K 10
1024
406.2
234.0
689.0
For the challenging 50D Exponentially Weighted Gaussian with 1024 widely separated modes, where the two farthest modes are 63.25
La distance apart, annealing procedures are mandatory to ensure success. Table 11 shows the number of modes explored in the 50D Exp- Weighted Gaussian by AF, CRAFT, LFIS, and PGPS as the number of annealing steps K increases. Together with Table 2, it is evident that AF consistently requires the fewest annealing steps achieve success in highly challenging scenarios, owing the W22 regularization our unique dynamic OT loss.
27

---

Annealing Flow Generative Models Towards Sampling High-Dimennional and Multi-Modal Distributions
D.3.2. PERFORMANCE OF ALGORITEMS WITH EVEN FEWER ANNEALING STEPS
In Figures and 4, the number of time steps for AF, CRAFT, LFIS, and PGPS is set to 10. Hcre, in Figure 11, we present ablation studies comparing the performance of these four methods when the number of time steps is reduced to 5.
175
Pine
ONT
U8
(a) True
(b) AF
(c) CRAFT
(d) LFIS
(c) PGPS
Figure 11. Ablation Studies for Gaussian Mixture Models (GMM) with 6, 8, and modes arranged on circles with radii = 8, 10, 12. All methods are trained with 4 intermediate densities and refinement block.
It can he observed that even with half the numher of time steps, AF maintains competitive performance GMMs with 6 and 8 modes and significantly outperforms other methods on GMMs with 10 modes.
Table 12. Mode- Sede-Weight Mcan Squared Error across distributions. The number of time steps for AF, CRAFT, LFIS, and PGPS is reduced 15 from 20.
Distribations
AF
CRAFT
LFTS
PGPS
d 10
ExpGauss 1024
(9.5 0.50) x 10
(48.1.0.33) x 10
(7.3 0.74) x 10
(7.81.0.70) x 10
ExpGauss. 1024
01.2+ 10-7
07.2+0.00) 10-5
(93++
(9.8+0070)) 10-7
d ExpGaustUV-2 2--V--1-000 1024
(3.3 0.07)
(1.710070 10
(2.71
10
(3.5.1
ExpGaustU/ 10 1024
(3.8 1 0.11) 10
(5.0 0.82)
10-5
(9.8 4 0.33) 10
(7.6 10
We conducted similar ablation studies the 50D -xe-hWeghted Gaussian with 1024 widely separated modes, reducing the number of time steps from 20 to 15. Table presents mode-weight MSE in Exp-weighted Gaussian across different dimensions and unequal variances, when the number of time steps is reduced 15 from 20 for all NF methods. AF still successfully captures all 1024 modes, with slightly higher Mode- Mede-Weight MSEs compared to the values reported in Table 1. In contrast, other methods, including CRAFT, LFIS, and PGPS, perform much worse than AF.
D.4. Possible Extensions of Importance Flow
The importance flow discussed and experimented with in this paper requires given form of mu(z). and thus, given form of (xx) Tu(z) [h(z)] for estimating Ex-ms(x) [N(X)]. In our experimental settings, q" (x) [|=][eeN (0, Ia) be regarded as the Least- Covorabee-Distrbutte (LFD). We conducted a parametric experiment for the case where q (z) has the given analytical form.
However, we believe future research may extend this approach distribution-fre 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-Pavorable Distribution (LFD) while minimizing the variance. For example, in the case of sampling from the LFD and obtaining low-variance IS estimator for Pawe(x)) ([|=]] 2 e). one may the following distribution-fin free loss for learning the flow:
min Σ (I(((a:)) ([T(, cl1] - +7 l [b(=())40)))
(65)
where the first term the loss pushes dataset {a:) towards the Least-Pavorable tail region, while the second term ensures a smooth
28

---

Anncaling Flow Generative Models Towards Sampling Highh-iimensionll and Multi-Modal Distributions
and cost-optimal transport map. Note that the above loss assumes no prior knowledge of the dataset distribution (xx) or the target density (xx).
Xu et al. (2024c) 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 R(Q", .)) Such framework has significant applications in adversarial attacks, robust hypothesis testing, and differential privacy. Additionally, the recent paper by Ribera Borrell al. (2024) introduces dynamic control loss training 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 generate greater impact the fields of adversarial attacks and differential privacy.
29


Reviews:
[{"review_id": "JKW8I5GQNG", "reviewer": "Reviewer_aCqJ", "summary": "The authors devise a new technique for sampling from high-dimensional multi-modal distributions. The assumed setting is that we are given an unnormalized analytical form of the density we wish to sample from. The proposed technique, dubbed Annealing Flow, is based on a continuous normalizing flow guided by an annealing procedure and is trained with a dynamic optimal transport objective. This has several benefits such as enabling stable training without a score-matching objective and cutting down the number of annealing steps needed to achieve state-of-the-art performance on the tested datasets.\n\n## Update after rebuttal:\nI find the authors answers satisfying and appreciate the extra effort put in to improve this work clarity and benchmarking. Therefore, I'm happy to change my score to \"Accept\" to increase their chances for acceptance.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you sincerely for your time and thoughtful review! Below, we respond to each of your questions in detail. We also provide a summary of our key contributions and the representativeness of our comparison experiments at the end of this response.\n\n> Nonetheless, the writing could be made clearer and more concise for the benefit of the average reader, and some parts could be relayed to the appendix and replaced with more intuitive explanations.\n\nThanks for the helpful comments! We agree that the algorithmic development and theoretical claims can be made more concise and accessible. Following your suggestions, we have made the following revisions: \n\n1. *Reorganized Fig. 2*:  We moved Fig. 2 earlier in the paper to accompany the introduction of $f_k$ and $\\tilde{f}_k$, making the exposition more illustrative. \n\n2. *Clarified definitions*: We corrected the definition of $Z_k$ associated with $f_k$, as pointed out by **Reviewer zGCn**; We added Wasserstein distance $W_2$ definition in Section 2 to help general readers follow the paper more easily.\n\n3. *Removed redundancy*: We eliminated the redundant definition of $\\tilde{E}_k$ and replaced it with $-\\log \\tilde{f}_k$; We removed the redundant definition of $\\tilde{Z}_k$;\n\n4. *Streamlined the objective*: We made the final objective (Eq. 10) more explicit and moved Proposition 3.3 (Objective Reformulation) into Proposition 3.4 to clarify that the reformulation is used solely for the derivation of Proposition 3.4, avoiding unnecessary repetition.\n\n\n> How do you calculate $\\tilde{E}_k (x(t_k))$?\n\nWe originally defined $\\tilde{E}_k (x) = -\\log \\tilde{f}_k$ under (7), but upon revision, found this definition unnecessary and have replaced it with the direct expression $-\\log \\tilde{f}_k$ throughout the paper. Given the unnormalized target density $\\tilde{q}(x)$ and the definition of $\\tilde{f}_k$ in Eq. (4), the form $\\tilde{E}_k (x(t_k))=-\\log \\tilde{f}_k$ is known. \n\nIn objective (10), the expectation $\\mathbb{E}$ is taken over sample $x(t_{k-1})$ from the density $f_{k-1}$, while the evaluation of $\\tilde{E}_k (x(t_k))=-\\log \\tilde{f}_k (x(t_k))$ requires the sample at the later time step $x(t_k)$. Therefore, a numerical integration is needed to compute $x(t_k)$, as shown in Eq. (13):\n\n$$x(t_k) = x(t_{k-1}) + \\int_{t_{k-1}}^{t_k} v_k (x,s) ds,$$\nwhere $v_k$ is the learnable velocity field optimized in the objective (10).\n\nWe referenced the computation of $x(t_{k-1})$ in Line 5 of Alg. 1. We have now clarified this step further in the revised paper.\n\n> You mentioned your method requires more expensive pertaining, is this quantified somewhere?\n\nYes, we reported training and sampling times for AF and other Normalizing Flows (NFs) in Tables 10 and 11. While AF requires one-time pre-training, it can be done offline and reused for efficient sampling—producing 10,000 samples in 50D space in just 2.1 seconds (given in Tab. 10). In contrast, MCMC methods require significantly longer mixing times and yield poor performance.\n\nWe have also included AF vs. MCMC comparisons in our response to **Reviewer Eiy4**. We would greatly appreciate it if you could take a look. We hope these have address your concerns!\n\nBesides, We have now further implemented **NUTS** (mentioned by Reviewer Eiy4) and **TemFlow** (mentioned by Reviewer mLHv), and would sincerely appreciate it if you could take a look at the anonymous links provided in our responses to them!\n\n> Below, we summarize our **key contributions** and further highlight the **depth of our experimental evaluations**:\n\n1. *Algorithmic novelty*: \n\nOur annealing-based dynamic Optimal Transport (OT) objective is novel in the current literature, combining the strengths of both annealing and dynamic OT to yield several key advantages outlined below.\n\n2. *Greatly improved training efficiency and sampling performance*: \n\nWith our dynamic OT loss, AF achieves superior performance using only 10 annealing steps—compared to 256 in LFIS, 128 in CRAFT and PGPS, and 100 in TemFlow. As shown in Tables 10 and 11, this results in significantly faster training.\n\n3. *Theoretical advantages*: \n\nIn Proposition 3.4, we establish that the optimal velocity field satisfies $\\lim_{h \\to 0} v_{k}^* = s_k - s_{k-1}$, i.e., the score difference between successive annealing densities. This is a desirable property uniquely enabled by our annealing-based dynamic OT objective.\n\n4. *Experimental evaluation*: \n\nWe benchmark AF against leading methods across three sampling mainstreams: MCMC (PT, AI-Sampler, and the new NUTS mentioned by Reviewer Eiy4), gradient-based methods (SVGD, MIED), and NF-based methods (CRAFT, LFIS, PGPS, and the new TempFlow noted by Reviewer mLHv). Notably, we also take a big step towards sampling very challenging cases, the 50D ExpGauss distribution with 1024 far-separated modes (separated by distances ranging from 20 to 63.25).\n\n\nWe sincerely appreciate your valuable time and feedback! We look forward to any further comments you may have!", "claims_and_evidence": "Yes, the experiments do seem to prove empirically that methods/claims are valid in comparison to the baselines. Nonetheless, as I'm unfamiliar with the tasks in this field, I'm not certain how representative these unnormalized densities/sampling tasks are.", "methods_and_evaluation_criteria": "I'm not very familiar with this line of work, but after skimming through some of the references in the paper, I noticed that the authors performed extensive comparisons with quite a few representative recent baselines.", "theoretical_claims": "I did go over some of the proofs in Appendix A and as far as I can tell they seem sound. Also, the theoretical results regarding the optimal velocity field with infinitesimal annealing steps make sense and feel overall intuitive.", "experimental_designs_or_analyses": "I haven't thoroughly checked the details as I'm unfamiliar with the examined densities, but the overall experiments section seems sound.", "relation_to_broader_scientific_literature": "The paper aims to provide an easier-to-train and overall more efficient sampler for multi-modal high-dimensional distributions. The paper seems theoretically solid as far as I can tell, however, I'm not sure how novel and impactful are the contributions regarding the dynamic optimal transport objective, and if these had been proposed elsewhere in prior works (other than the referenced ones).", "essential_references_not_discussed": "I'm not very familiar with this line of work and do not follow it closely. However, it does seem the authors did their due diligence and cited the existing literature properly.", "other_strengths_and_weaknesses": "In terms of mathematical rigor, the paper is written with a very high standard. Nonetheless, the writing could be made clearer and more concise for the benefit of the average reader, and some parts could be relayed to the appendix and replaced with more intuitive explanations. For example, how do you calculate $\\tilde{E}_{k}(x(t_k))$ is still unclear to me in the current version.", "other_comments_or_suggestions": "General suggestion: I think overall the method goes through so many objectives until it arrives at the final one, and this part could be made significantly more concise and clearer for the reader. For instance, a small paragraph explaining the task at the beginning of the method section would be very helpful.\nCaught preceding shorthand: In line 26 in the abstract the MC shorthand is used before being defined as Monte Carlo.", "questions_for_authors": "1) How do you calculate $\\tilde{E}_{k}(x(t_k))$?\n2) You mentioned your method requires more expensive pertaining, is this quantified somewhere?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "YU7WO2EF61", "reviewer": "Reviewer_zGCn", "summary": "The paper proposes a new method to learn a vector field v(x,t) such that the neural ODE dx = v(x,t) dt approximates the optimal transport between two given distributions p and q. This is laudable, because scalable and accurate optimal transport solvers in high dimensions are still an important research topic. The paper derives a tractable loss that leads to an efficient learning algorithm. The new method appears to be superior to existing algorithms on several benchmarks.\n\nUnfortunately, the presentation contains severe errors and shortcomings that might not be fixable in a rebuttal phase, see below.\n\nEDIT: In light of the rebuttal discussion, I raised my score.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your thoughtful review and careful attention to our mathematical developments! We acknowledge that some symbols lacked rigor and appreciate you pointing them out! Please see below for our response to your concerns on both math rigor and experimental comparisons.\n\n> The paper defines $f_k(x) = \\pi_{0}(x)^{1-\\beta_{k}} q(x)^{\\beta_k}$ in Eq. (4). [...] However, f_k(x) are not normalized distributions. Consequently, the KL divergence in Eq. (6) is undefined, as it requires normalized distributions.\n\nThanks for the careful review! You are absolutely right—our original definition of $f_{k}(x)$ in (4) was unnormalized. We have now slightly revised $f_k(x)$ in (4) by:\n $$\\tilde{f}_k(x)=\\pi_0 (x)^{1-\\beta_k} \\tilde{q} (x)^{\\beta_k},$$\n\n$$f_k (x)= Z_k \\tilde{f}_k (x),$$ \nwhere $\\tilde{q} (x)$ is the unnormalized target density and $Z_k = \\frac{1}{\\int \\tilde{f}_k (x)dx}$ is an unknown constant. \n\n- **The only change that needs to be made is in how we define $Z_k$.** Our algorithm solely relies on $\\tilde{f}_k$, which is defined the same as original. The constant $Z_k$ of the normalized $f_k$ appears only in the constant term $c$ in Eq. (8) of Prop. 3.1, given by:\n\n$$c = E_{x\\sim \\rho(x(t_{k-1}),t_{k-1})}[\\log \\rho(x(t_{k-1}),t_{k-1})]- \\log Z_k,$$\nwhich is a constant independent of our algorithm.\n\n- With the revised definition of $f_k$, the KL divergence in Eq. (6) is well-defined, and all subsequent propositions and algorithmic steps remain correct.\n\nThanks for the detailed observation!\n\n> Following $f_k (x)$, the $E_k $, $\\tilde{E}_k$, and $\\tilde{Z}_k $ are unclear [...]. Contrary to Eq. (4), the f_k in the proof are defined by the probability path (2). And it requires the vector field $v(x,t)$ already known.\n\nBased on $f_k (x)$ above, the energy $E_k (x)$ is still defined as $E_k (x) = -\\log f_k (x)$, and we've now made this clear in the draft. We have also removed the redundant definitions of $\\tilde{E}_k$ and $\\tilde{Z}_k $, now consistently used $Z_k$ in Prop. 3.1, and replaced all occurrences of $\\tilde{E}_k$ in the draft with:\n$\\tilde{E}_k \\mapsto -\\log \\tilde{f}_k (x).$\n\nThe density $\\rho(\\cdot, t)$ evolves according to Eq. (2). Under the constraints of the dynamic OT (3), we have: $\\rho(\\cdot, t_{k-1}) = f_{k-1},\\ \\rho(\\cdot, t_k) = f_k$. We adopt an efficient block-wise training (as opposed to end-to-end), in which each $v_k (x,t)$ is trained sequentially after $v_1,\\cdots, v_{k-1}$ have been learned, as given in Alg. 1.\n\n> Your concerns about Methods And Evaluation Criteria:\n\nWe would like to clarify: **Fig. 3** shows GMM results for CRAFT, LFIS, and PGPS using the **same number of annealing steps (8–10)** as our AF, ensuring a fair comparison on **computational efficiency**. In contrast, **Fig. 5** presents these same methods with their **official settings**—using significantly more steps (128 for CRAFT and PGPS, 256 for LFIS) with PGPS using Langevin adjustments. They perform well in Fig. 5, but at a higher computational cost, as detailed in Tables 10 and 11. \n\nOur experiments focus exclusively on *well-separated modes*, where methods without annealing - PT, SVGD, MIED, and AIS—naturally struggle. In contrast, *prior works commonly evaluate on closely aligned GMM modes*: e.g., radius 1 in LFIS (their paper's Fig. 2), radius 4 in both AI-Sampler (their Fig. 3) and the TemFlow (their Fig. 3) referenced by Reviewer mLHv.\n\nBy comparison, in our GMM settings (Fig. 3 and 5), modes are separated by radius 12. And in the extremely challenging 50D ExpGauss, by distances ranging from 20 to 63.25.\n\nWe also did ablation studies on **AF without annealing** (Fig. 10 and Tab. 12), where our performance drops notably. This shows the necessity of Annealing.\n\nTo fairly evaluate AF’s *efficiency and performance*, results reported in **Tab. 3** (main text), **Tab. 7**, and **Fig. 5** (appendix) for other NFs uses their official (and much higher) annealing steps to reflect their full performance.\n\nThe purpose of **Fig. 3** is only to *compare performance under matched annealing steps to highlight AF's efficiency*. We've now added clearer references to Fig. 5 in the main text to make this distinction.\n\n> Response to other concerns:\n\n1. \"Optimal Flow Matching\" is designed for image generation and requires target samples for training, not target densities as in statistical sampling. We've further implemented **TemFlow** mentioned by **Reviewer mLHv** and **NUTS** mentioned by **Reviewer Eiy4**. Please see the link in the responses!\n2. $t_{k-1,s}$ is defined for $W_2$ discretization purpose in Prop. 3.2. We've made it clearer and more concise.\n3. Our code implemented both torchdiff and Hutchinson for the divergence. torchdiff requires more computation with little improvement.\n\n> Finally:\n\nWe have summarized our key contributions at the end of our response to Reviewer aCqJ, and we’d be grateful if you could take a look!\n\nThank you sincerely for your detailed review—it has indeed helped make our paper stronger!", "claims_and_evidence": "The paper defines the \"annealing flow\" as the sequence of functions f_k(x) = p(x)^{1-beta_k} q(x)^{beta_k} with two predefined distributions p and q and a sequence of inverse temperatures beta_0 = 0 < beta_1 < ... < beta_K = 1 (equation (4)). However, the claim that the f_k(x) are distributions is false, because the RHS is not normalized for 0 < beta_k < 1 (this is easily seen by setting p = normal(0,1) and q = normal(1,1)). Consequently, the KL divergence in equation (6) is undefined, as it requires normalized distributions. \n\nThis mistake may not have severe consequences, because the authors subsequently allow q to be unnormalized, which requires the introduction of an explicit normalization constant. Since this constant is independent of the model, it has no influence on the gradient of the loss in equation (8) -- an incorrect normalization constant would thus be inconsequential. However, this is not quite clear, because the definition of tilde(E)_k after equation (7) makes no sense: The constant tilde(Z)_k and the energy function E_k(x) are both undefined at this point. The proof of proposition 3.1 in the appendix does not resolve the riddle: Contrary to equation (4), the f_k in the proof are defined by the probability path induced by the continuity equation (2). The thus defined f_k are normalized, but require that the vector field v(x,t) is already known. In addition, the proof uses a normalization constant Z_k instead of tilde(Z)_k, which is much more plausible.\n\nThese contradictions must be fixed.", "methods_and_evaluation_criteria": "The new method is compared to seven alternatives, using challenging synthetic data distributions, and wins most of the time. However, it seems that only one of these experiments uses an existing benchmark protocol, so results are hard to compare with the literature. \n\nIn the qualitative experiment on a 2D Gaussian mixture (figure 3), all methods except the proposed one perform terribly. This is highly implausible, as similar experiments with satisfactory results are frequently reported in the literature. It looks as if the hyperparameters of the competition have not been tuned properly -- this would be unacceptable. To be fair, it could also be a consequence of forcing all methods to use the same network architecture for comparability, but even if this were the case, the choice of network architecture would be suspicious. Of course, this also sheds a bad light on the other experimental comparisons.", "theoretical_claims": "See above.", "experimental_designs_or_analyses": "See above.", "relation_to_broader_scientific_literature": "The paper reviews a considerable body of existing literature. However, it is surprising that \"Optimal flow matching\" (Kornilov et al. arXiv:2403.13117) is not mentioned and not compared against.\n\nThe claim in the introduction \"Discrete normalizing flows often suffer from mode collapse\" is highly questionable. This only happens when either the hyperparameters are not chosen properly, or the network is trained with reverse KL. Properly chosen NF architectures trained with the forward KL (the standard approach) do not exhibit mode collapse.", "essential_references_not_discussed": "See above.", "other_strengths_and_weaknesses": "In additional to the errors discussed above, I'd like to raise the following issues:\n* The authors hierarchically discretize time first into segments [t_{k-1}, t_k] and then into subsegments t_{k-1,s}. However, it is never discussed why this setup is necessary or beneficial or superior to alternatives, and what the underlying trade-offs are.\n\n* In order to drop the constant c from equation (8), the objective in equation (10) must be an argmin. Moreover, it would be nice to illustrate how the divergence term and the energy term counterbalance each other such that this objective has a fixed point at the optimum (instead of converging to a trivial solution).\n\n* In section 4.1, the authors implement the Hutchinson estimator in terms of discrete directional derivatives. However, the underlying expression epsilon^T J_v epsilon (see appendix C.1) can be evaluated exactly using Jacobian-vector primitives from the autodiff library. The proposed implementation introduces an unnecessary additional approximation error.\n\n* Section 5 on importance flows seems out of scope, given that there are only preliminary results which are only reported in the appendix. Leave this topic for future work!", "other_comments_or_suggestions": "See above.", "questions_for_authors": "See above.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "QKCXuVXWsW", "reviewer": "Reviewer_Eiy4", "summary": "This paper proposed a new flow-based sampler from continuous target density functions via combining several ideas.\nMore specifically, they introduce a method that they call Annealing Flow (AF) by using continuous normalizing flows trained with dynamic optimal transport objective function. They use Wasserstein regularization and their method is guided by annealing the target pdf.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you very much for your helpful review and thoughtful comments! We address your concerns point by point below:\n\n> It is hard for me to convince myself that the proposed density ratio estimation (Section 5.2) will perform better than the formula in line 289.\n\nThank you for your detailed observation on the extended Importance Flow framework. The traditional normalized Importance Sampling (IS) estimator in line 289 is generally biased. For example, in our IS experiment on estimating $\\mathbb{E}_{X \\sim \\pi_0}[h(X)]$ with $h(x) = \\text{1}(||x||>c)$, $\\pi_0 (x) = N(x;0,I)$, the resulting optimal proposal $\\tilde{q}^* (x) =  \\text{1}(||x||>c) \\cdot N(x;0,I)$. We can immediately see the self-normalized IS estimator is:\n\n$$\\hat{I}_{N}= \\sum \\frac{\\pi_0 (x_i)}{\\tilde{q}^* (x_i)} h (x_i) / \\sum \\frac{\\pi_0 (x_i)}{\\tilde{q}^* (x_i)} = \\sum 1 / \\sum \\frac{1}{\\text{1}(||x||>c)} =  \\sum 1 / \\sum 1 = 1, \\text{where} \\ x_i \\sim q^* (x).$$\n\nThe estimator always equals to 1 in this case. However, the true probability we want to estimate is $P_{X \\sim N(0,I)} (||X||>c)$.\n\nHowever, our Density Ratio Estimation (DRE) network introduced in Section 5.2 can theoretically recover the exact ratio $r^* (x) = \\log \\frac{\\pi_0 (x)}{q^* (x)}$, where $q^* (x)$ is the normalized target density. This allows us to directly construct an importance sampling (IS) estimator:\n$$\\hat{I} = \\sum \\frac{\\pi_0 (x_i)}{q^* (x_i)} h (x_i)=\\sum \\exp (r^* (x_i))\\cdot h(x_i).$$\nAs shown in the preliminary results in Table 9, the Importance Flow yields a rare-event probability estimate that closely matches the true value, rather than defaulting to 1 as in $\\hat{I}_{N}$.\n\n\n\n> My main criticism is not testing the performance of their method on high-dimensional models (e.g. d > 100) as well as not comparing it with MCMC. It would be good to provide experimental results that give an idea of cases where AF outperforms MCMC (say NUTS) and cases where the opposite is true.\n\nThank you very much for your thoughtful feedback! We included comparisons with two MCMC-based methods in our experiments: **Parallel Tempering (PT)**, a classical tempered MCMC, and **AI-Sampler**, a neural-network-assisted MCMC. A functional comparison with MCMC is also briefly discussed in Section 4.2. That said, we agree that further comparison is valuable and have added more detailed analysis now in our draft. To summarize:\n\n- *Sampling Efficiency*: \n\nUnlike MCMC, which does not require pre-training but often suffers from long mixing times, AF requires one-time offline training, after which sampling becomes a deterministic, efficient numerical computations. We reported **training and sampling time** of AF in Tab. 10 and Tab. 11.\n- *Performance in High Dimensions*: \n\nIn our 50D ExpGauss experiment with 1024 far-separated modes, PT captures <10 modes and AIS captures around 100 (Table 3), illustrating MCMC’s limitations in exploring highly multimodal, high-dimensional spaces. AF significantly outperforms in this setting.\n- *Sample Balance*: \n\nEven when MCMC reaches multiple modes, it often fails to maintain proper sample weighting due to the stochastic nature of mode-hopping.\n\nFurther, we have now implemented **NUTS**, and below is a link of figures and tables:\n\n**Anonymous Link:** https://drive.google.com/file/d/1vyciBojsTKuRfMfw0GZ6LYfMXx49g7Fe/view?usp=sharing\n\n*It took us great efforts so we sincerely appreciate your time in reviewing it!* \n\nAs shown in link~Tab. 1 and 2, MCMC methods—NUTS, PT, and AIS—fail immediately on the 50D target with 1024 well-separated modes. This is due to the exponential growth of the effective search space (roughly $O(\\exp (d))$) for a MC, making exploration of all far-separated modes infeasible in high dimensions.\n\n> For your main criticism on not testing the performance on high-dimensional models (e.g. d > 100): \n\nNotably, the 50D ExpGauss (nearest two separated by 20, farthest by 63.25) we study is substantially more challenging than prior benchmarks. \n\nFor context, among very recent methods: **TemFlow** (referred by Reviewer mLHv) experiments on a 2D GMM with 8 close modes, a Copula distribution with 4 close modes, and samples from up to a 32D latent space (normally uni-modal). **Fisher-rao** (referred by Reviewer mLHv) explores up to 20D funnel and 2D closely aligned modes. **LFIS** focuses on a 2D GMM with 9 close modes and a 10D funnel. **AI-Sampler** tests on a 2D circular GMM (radius 5) and Bayesian logistic regression up to 21 dimensions.\n\nIn summary, we compared the three major sampling paradigms—MCMC (NUTS, AI-Sampler, PT), particle-based methods (SVGD, MIED), and normalizing flows (CRAFT, LFIS, PGPS, TempFlow)—on significantly more challenging densities than prior works. \n\nWe have also summarized our key contributions at the end of our response to **Reviewer aCqJ**, and we’d be grateful if you could take a look!\n\nThank you sincerely for your thoughtful feedback—it helped us clarify both our MCMC comparisons and overall contributions!", "claims_and_evidence": "To some extent. \nMy criticism is as follows: \nWhile in the title and throughout the paper they claim their method is suitable for high-dimensional targets, their experimental models are quite toyish and low-dimensional. Clearly, reporting results for high-dimensional models would support their claim better (or would show its limitations that would still be valuable).\nAlso, (apart from one visual experiment in the supplementary) they do not compare their method against MCMC-based methods. For instance, it would be interesting to see the mode switches and MSE versus the model dimension for the AF and NUTS sampler.  Also, it would be good to report the AF's training time and compare MCMC versus AF (and other flow-based methods) if the total time budget is limited and fixed.", "methods_and_evaluation_criteria": "Yes but see my  comments on \"Claims And Evidence\"", "theoretical_claims": "I did not closely check the correctness of the proofs. Also, it is hard for me to convince myself that the proposed density ratio estimation (Section 5.2) will perform better than the formula in line 289. It would be valuable if they would compare the performance of these two approaches.", "experimental_designs_or_analyses": "Yes but see my  comments on \"Claims And Evidence\"", "relation_to_broader_scientific_literature": "Combining several state-of-the-art algorithms in an smart way and proposing a more robust flow-based method that (apparently) can be used in high-dimensional and multi-model models.", "essential_references_not_discussed": "Not that I am aware of", "other_strengths_and_weaknesses": "As I mentioned earlier, my main criticism is not testing the performance of their method on high-dimensional models (e.g. d > 100) as well as not comparing it with MCMC. \nClearly, every algorithm has its limitations, but it would be good to provide experimental results that give an idea of cases where AF outperforms MCMC (say NUTS) and cases where the opposite is true.", "other_comments_or_suggestions": "N/A", "questions_for_authors": "N/A", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {"ethical_review_concerns": "N/A"}}, {"review_id": "HRnHGtDTKn", "reviewer": "Reviewer_mLHv", "summary": "This paper proposes Annealing Flow, a method based on continuous normalizing flow to sample from high dimensional multi-modal distribution. The authors provide a efficient training and sampling algorithm, which can be also applied to Monte-Carlo estimation. Various experiments are conducted to verify the efficiency of AF compared to previous sampling methods.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely thank the reviewer for your valuable comments and suggestions! Please see below for our response to your concerns.\n\n> Summary of Review: The paper lacks novelty and misses key related works. [1] proposes an annealed NF using L2 distance. [2] employs kernel methods instead of neural networks to reduce training costs. The authors should provide more thorough comparisons.\n\nThank you for bringing these works to our attention! The novelty of our work lies in the 1. computational efficiency, 2. stability, and 3. unique theoretical property enabled by our annealing dynamic OT objective, along with 4. strong sampling performance—most notably on the challenging 50D ExpGauss with 1024 far-separated modes, well beyond existing sampling densities baselines.\n\nWhile we tried best to survey a broad range of NFs/CNFs, and compared three very recent NF methods (plus four beyond NF/CNF scope), we acknowledge that some relevant works have been missed.\n\nWe've now implemented **[1] (\"TemFlow\")** and evaluated it on GMMs (with radius 8, 10, 12) and the extreme 50D ExpGauss distribution. Below is a link to the figures and tables:\n\n**Anonymous link**: https://drive.google.com/file/d/1M_o4XF5440n0G_k_PaHioncx-ro1mHhd/view?usp=sharing. \n\nWe sincerely appreciate your time in reviewing the link—*it took considerable effort and we would be grateful for your feedback!*\n\nSummary on [1]: \n\n- *Architecture sensitivity*: \n\nTemFlow is an NF method whose performance depends on its spline layer design (e.g., number/type of bins, bounds), which must be adapted for each target density. In contrast, AF consistently uses a fixed 32-32 MLP without any network-specific tuning. As shown in link~Fig. 3, suboptimal spline design can significantly degrade TempFlow’s performance.\n\n- *Annealing steps*: \n\nTemFlow’s official implementation uses 100 annealing on GMMs (radius=4), while our AF achieves comparable performance using 10 steps on a harder GMM (radius=12). Fig. 1 compares TempFlow when matched to our step count.\n\n- *Efficiency*: \n\nTemFlow trains each annealing block quickly (0.39 min vs. our 0.70 min), but needs more steps—its total training time on 50D ExpGauss reaches 49.6 min versus 14.5 min for AF.\n\n\nAdditionally, [2] is conceptually similar to SVGD and MIED, and we compared the latter two in our paper. [2] is a kernel-based, gradient-driven method and not directly related to NF/CNF. For their comparison with SVGD, please see the first paragraph of their Page 5. We’ll provide further context as we investigate further.\n\nWe also evaluated **NUTS** mentioned by Reviewer Eiy4. We'd be grateful if you could also review that response!\n\n> The experiments are not very convincing due to lack of more high-dimensional challenging distributions [...].\n\nThanks for the thoughtful question! Our primary contribution lies in advancing statistical sampling methods, but not targeting specific real-world datasets. To our knowledge, prior work has not attempted sampling from the highly challenging distributions such as 50D ExpGauss distribution:\n$$p(x_1,x_2,\\cdots,x_{50}) \\propto e^{10\\sum_{i=1}^{10}\\frac{|x_{i}|}{\\sigma_{i}^{2}}+10\\sum_{i=11}^{50}\\frac{x_{i}}{\\sigma_{i}^2}-\\frac{1}{2}\\|x\\|^{2}},$$\nwhich has 1024 far-separated modes (nearest two separated by 20, farthest by 63.25). \n\nFor context, prior works primarily evaluate on lower-dimensional or less multimodal distributions: **TempFlow** [1] experiments on a 2D GMM with 8 close modes, a Copula distribution with 4 modes, and samples from up to a 32D latent space (usually unimodal). **FisherRao** [2] explores up to 20D funnel and 2D close modes. **LFIS** (ICML 2024, compared in our work) focuses on a 2D GMM with 9 close modes and a 10D funnel. **AI-Sampler** (ICML 2024, compared in our work) tests on a 2D circular GMM (radius 5) and Bayesian regression up to 21 dimensions.\n\nStatistical sampling is fundamental in areas like statistical physics, rare event analysis, and Bayesian modeling. Beyond testing on challenging distributions, we introduce the novel *Importance Flow framework* for rare-event probability estimation (Table 9); and through *Hierarchical Bayesian Modeling* on real datasets ranging from 8 to 61 dimensions (Table 3).\n\n> Other Comments:\n\nWe have replaced the Eq. (12) in Prop. 3.4 with the following:\n\n$$\\lim_{h_k \\to 0} v_{k}^*=s_k-s_{k-1}.$$\n\nThank you for the detailed observation! \n\n> **Summary of our Contributions** (summarized in more detail at the end of our response to **Reviewer aCqJ**):\n\n1. Our dynamic OT objective significantly reduces the number of intermediate annealing steps, compared to most recent NF/CNFs works.\n2. Prop. 3.4 establishes $\\lim_{h_k \\to 0} v_{k}^*=s_k-s_{k-1}$, a desired property unique to our AF.\n3. Our experiments on highly challenging densities show AF's superior performance towards multi-modal and high-dimensional sampling - on tasks that go beyond prior works.\n\nThank you once again for your valuable time and thoughtful feedback—we truly appreciate it!", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "Yes", "theoretical_claims": "Yes. Some claims are not mathematically rigorous, see weakness part", "experimental_designs_or_analyses": "Yes, in Section 6", "relation_to_broader_scientific_literature": "It is generally related to machine learning community", "essential_references_not_discussed": "[1] Qiu, Yixuan, and Xiao Wang. \"Efficient multimodal sampling via tempered distribution flow.\" Journal of the American Statistical Association 119.546 (2024): 1446-1460.\n\n[2] Maurais, Aimee, and Youssef Marzouk. \"Sampling in unit time with kernel fisher-rao flow.\" arXiv preprint arXiv:2401.03892 (2024).\n\nsee weakness part for details", "other_strengths_and_weaknesses": "## Strengths\n\nThe paper is written clearly and all the mathematic derivations are easy to follow. Apart from developing algorithms, the authors also present some theoretical insights of annealing flow, in the limit of infinitesimal time. The proposed method is useful to handle multi-modal distribution.\n\n## Weaknesses\n### 1. \nThe idea is not very novel overall, and the paper lacks discussions and comparisons with some important previous works. For example, [1] considers similar annealing flow, utilizing L2 distance to train the map between $f_{k-1}$ and $f_k$ and claiming that KL divergence used in this paper will suffer from mode collapse. [2] uses kernel trick to estimate the map instead of neural network, which may save the training cost to some extent. The authors should include more comparisons, either from theoretical side or empirical side.\n\n### 2.\nThe numerical experiments are not very convincing due to lack of more high-dimensional challenging distributions. The highest dimension in the experiments is merely 50, which is still too low in modern applications of machine learning. I suggest the authors include more high-dimensional experiments such as Bayesian neural networks. \n\n### 3.\nThe theoretical claims in Proposition 3.3 and 3.4 are not mathematically rigorous. For example, eq (12) holds only in the sense of limit when $h\\to 0$ and thus it's not mathematically sound to directly write eq (12). A better way to express this is, e.g., $\\lim_{h\\to 0} v_k^*-(s_k-s_{k-1}=0$.\n\n[1] Qiu, Yixuan, and Xiao Wang. \"Efficient multimodal sampling via tempered distribution flow.\" Journal of the American Statistical Association 119.546 (2024): 1446-1460.\n\n[2] Maurais, Aimee, and Youssef Marzouk. \"Sampling in unit time with kernel fisher-rao flow.\" arXiv preprint arXiv:2401.03892 (2024).", "other_comments_or_suggestions": "see weakness part", "questions_for_authors": "Are the neural networks trained in block-wise manner also stored sperately? If the annealing timesteps are high and the size of neural network is large, the storage cost is not negligible.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: This paper proposes Annealing Flow, a normalizing flow method for sampling from high-dimensional and multi-modal distributions. It builds upon continuous normalizing flows (CNFs) by introducing an annealing-guided dynamic optimal transport (OT) objective. It improves training efficiency and stability compared to previous NF methods. Empirical results demonstrate strong performance across a diverse set of baselines, particularly on high-dimensional benchmarks. 

Despite some concerns regarding clarity and novelty, most reviewers recognized the method’s technical depth and consistent empirical improvements.
Decision: Accept (poster)