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

We sincerely thank the reviewer for your valuable comments and suggestions! Please see below for our response to your concerns.

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.

Thank 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.

While 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.

We'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:

Anonymous link: https://drive.google.com/file/d/1M_o4XF5440n0G_k_PaHioncx-ro1mHhd/view?usp=sharing.

We sincerely appreciate your time in reviewing the link—it took considerable effort and we would be grateful for your feedback!

Summary on [1]:

• Architecture sensitivity:

TemFlow 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.

• Annealing steps:

TemFlow’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.

• Efficiency:

TemFlow 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.

Additionally, [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.

We also evaluated NUTS mentioned by Reviewer Eiy4. We'd be grateful if you could also review that response!

The experiments are not very convincing due to lack of more high-dimensional challenging distributions [...].

Thanks 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: $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}},$ which has 1024 far-separated modes (nearest two separated by 20, farthest by 63.25).

For 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.

Statistical 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).

Other Comments:

We have replaced the Eq. (12) in Prop. 3.4 with the following:

$\lim_{h_k \to 0} v_{k}^*=s_k-s_{k-1}.$

Thank you for the detailed observation!

Summary of our Contributions (summarized in more detail at the end of our response to Reviewer aCqJ):

1. Our dynamic OT objective significantly reduces the number of intermediate annealing steps, compared to most recent NF/CNFs works.
2. Prop. 3.4 establishes $\lim_{h_k \to 0} v_{k}^*=s_k-s_{k-1}$, a desired property unique to our AF.
3. 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.

Thank you once again for your valuable time and thoughtful feedback—we truly appreciate it!

Thank you very much for your helpful review and thoughtful comments! We address your concerns point by point below:

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.

Thank 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:

$\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).$

The 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)$.

However, 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: $\hat{I} = \sum \frac{\pi_0 (x_i)}{q^* (x_i)} h (x_i)=\sum \exp (r^* (x_i))\cdot h(x_i).$ As 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}$.

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.

Thank 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:

• Sampling Efficiency:

Unlike 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.

• Performance in High Dimensions:

In 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.

• Sample Balance:

Even when MCMC reaches multiple modes, it often fails to maintain proper sample weighting due to the stochastic nature of mode-hopping.

Further, we have now implemented NUTS, and below is a link of figures and tables:

Anonymous Link: https://drive.google.com/file/d/1vyciBojsTKuRfMfw0GZ6LYfMXx49g7Fe/view?usp=sharing

It took us great efforts so we sincerely appreciate your time in reviewing it!

As 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.

For your main criticism on not testing the performance on high-dimensional models (e.g. d > 100):

Notably, the 50D ExpGauss (nearest two separated by 20, farthest by 63.25) we study is substantially more challenging than prior benchmarks.

For 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.

In 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.

We 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!

Thank you sincerely for your thoughtful feedback—it helped us clarify both our MCMC comparisons and overall contributions!

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.

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.

Thanks 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: $\tilde{f}_k(x)=\pi_0 (x)^{1-\beta_k} \tilde{q} (x)^{\beta_k},$

$f_k (x)= Z_k \tilde{f}_k (x),$ where $\tilde{q} (x)$ is the unnormalized target density and $Z_k = \frac{1}{\int \tilde{f}_k (x)dx}$ is an unknown constant.

• 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:

$c = E_{x\sim \rho(x(t_{k-1}),t_{k-1})}[\log \rho(x(t_{k-1}),t_{k-1})]- \log Z_k,$ which is a constant independent of our algorithm.

• 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.

Thanks for the detailed observation!

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.

Based 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: $\tilde{E}_k \mapsto -\log \tilde{f}_k (x).$

The 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.

Your concerns about Methods And Evaluation Criteria:

We 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.

Our 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.

By 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.

We also did ablation studies on AF without annealing (Fig. 10 and Tab. 12), where our performance drops notably. This shows the necessity of Annealing.

To 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.

The 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.

Response to other concerns:

1. "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!
2. $t_{k-1,s}$ is defined for $W_2$ discretization purpose in Prop. 3.2. We've made it clearer and more concise.
3. Our code implemented both torchdiff and Hutchinson for the divergence. torchdiff requires more computation with little improvement.

Finally:

We 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!

Thank you sincerely for your detailed review—it has indeed helped make our paper stronger!

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.

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.

Thanks 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:

1. 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.

2. 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.

3. 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$;

4. 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.

How do you calculate $\tilde{E}_k (x(t_k))$?

We 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.

In 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):

$x(t_k) = x(t_{k-1}) + \int_{t_{k-1}}^{t_k} v_k (x,s) ds,$ where $v_k$ is the learnable velocity field optimized in the objective (10).

We 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.

You mentioned your method requires more expensive pertaining, is this quantified somewhere?

Yes, 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.

We 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!

Besides, 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!

Below, we summarize our key contributions and further highlight the depth of our experimental evaluations:

1. Algorithmic novelty:

Our 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.

2. Greatly improved training efficiency and sampling performance:

With 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.

3. Theoretical advantages:

In 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.

4. Experimental evaluation:

We 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).

We sincerely appreciate your valuable time and feedback! We look forward to any further comments you may have!