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

We appreciate your detailed comments and suggestions and will revise our draft accordingly. We would also like to clarify some of your questions related to our draft:

• Could you compare your work against [1] and [2] which have solved the same problem with very similar tools?

We have read the two papers carefully, and noted that their methods are based on gradient flows, instead of Annealing-guided flow with Wasserstein regularization. Through experiments, we observed that without $W_{2}$-regularization, the paths become less organized and more dispersed, which means the flow often misses the optimal path; and if the density modes are widely separated or in high-dimensional spaces, without Annealing-guidance, the gradient flow often fails. These two key features help ensure success in high-dimensional multi-modal cases (an example is in Figure 6, a density with 1024 modes in 50D). We will develop theories on the guiding effects of $W_{2}$ and Annealing. Thank you for identifying these two papers; we will include comparison experiments with theirs.

• Could you rewrite sections 2 and 3 but defining the normalizing constants from the beginning? I think it will make the paper much clearer. Moreover, it would avoid writing KL divergences between unnormalised distributions as done extensively in section 3 and appendix 1.

Thank you for your suggestions. We will revise accordingly.

• In common litterature, the Kullback-Leiber divergence is defined between normalized distributions as $D(p \parallel q) = \mathbb{E}_{X \sim p} \left[ \log \frac{p(X)}{q(X)} \right]$. However, L760 shows a completely different definition.
• According to the beginning of section 3.2, the quantity $\rho_k$ is the density of the flow following $v_k$, which means that $\log \rho_k(x(t_{k-1}))$ is not independent of $v_k(x(s), s)$, which would contradict L800.

We're sorry for the confusion; there is a typing error in the proof. We should have written the first step as $KL(T f_{k-1} \parallel f_k) = E_{X\sim \rho(t_{k})}\left[ \log T f_{k-1}(X) / f_{k}(X) \right]$, where $\rho(t)$ represents the density evolution governed by equation (2).

Here, $\rho(t_{k}) = T f_{k-1}$, and from the constraints in (3), we see that the density evolution is subject to the equations (1) and (2). Then subject to (1) and (2), we can write the objective $E_{X\sim \rho(t_{k})}\left[ \log T f_{k-1}(X) / f_{k}(X) \right] = E_{X\sim \rho(t_{k-1})}\left[ \log T f_{k-1}(X(t_{k})) / f_{k}(X(t_{k})) \right]$, where $\log T f_{k-1}(X(t_{k})) =\log \rho(X(t_{k})) = \log \rho(X(t_{k-1})) - \int_{t_{k-1}}^{t_k} \nabla \cdot v_k(X(s), s) ds$.

For the second question: $v_{k}$ flows density from $\rho(X(t_{k-1}))$ to $\rho(X(t_{k}))$, where $X(t_{k}) = X(t_{k-1}) +\int_{t_{k-1}}^{t_{k}} v_k(X(s), s) ds$. Therefore, $E_{X\sim \rho(t_{k-1})} [ \log \rho(X(t_{k-1})) ]$ is a constant that is independent of $v_{k}$. (It depends on $v_{1},...,v_{k-1}$.) We are sorry again for the typing typo that leads to the misunderstanding.

• At L260, you claim that your training procedure only requires a single neural network. Doesn't your sampling procedure require to keep all the intermediates

Our algorithm adopts a block-wise training procedure instead of end-to-end training. Our algorithm keeps the intermediate $v_{k}$. However, during training, only one NN (one $v_{k}$) is trained at each time, thus greatly enhancing training efficiency.

• Could you perform the experiment on Fig. 1 with differently weighted modes ? I feel like recovering the weights of the modes is the true challenge.

Thank you for the suggestion. After submission, we implemented modes with different weights, and the results are also convincing. We will include them in our future submission.

• Could you compute the metrics given in [6] which seem computable in your case ? Those metrics really focus on multimodal distributions.
• Could you provide metrics's mean and standard deviation (as well as the number of samples used) for each experiment ?
• Could you display true samples from the Funnel distribution alongside yours ? It would clarify the pros and cons of each algorithm.

Thank you! We will.

We wish to conclude by highlighting that we are the first to use Annealing-guided CNFs with $W_{2}$ regularization, which are the key to efficient explorations of all modes in high-dimensional spaces. Our experiments focus more on extreme distributions (high-dimensional densities with a large number of widely separated modes), which are distinct from many other CNFs studies. In our recent draft, we proved that the infinitesimal optimal velocity field represents the score difference between two consecutive densities. We will continue refining the theories, and thank you for these useful suggestions.

Thank you very much for your detailed review and kind suggestions on improvements! We acknowledge that the submitted draft was prepared in haste and contained some typographical errors and unclear notations, particularly in the proof of Proposition 1. We have carefully revised Proposition 1 and will upload a revised version shortly. Additionally, we would like to address some of your questions:

• in the introduction in connection with normalizing flows I am missing stochastic normalizing flows

Thank you for bringing such works to our attention. We will do literature review carefully and include the representative works of stochastic normalizing flows appropriately.

• All your kind suggestions related to the proof of Propositions 1,2, and 3:

We have addressed all the points and sincerely thank you for your detailed review and valuable comments.

• The results of the experiments seem to depend heavily on the symmetry of the target density, Please show a modified experiment of Fig. 6 where you shift the target density in the first dimension e.g. by 5 to the left (and keep the variance of the latent distribution). I guess that the reconstruction misses modes.

Thank you for your comment! After submission, we conducted additional experiments on densities with differently weighted modes and will include these in our future submissions. Additionally, we will incorporate experimental comparisons with other normalizing flow methods as suggested by other reviewers, and try to include more metrics.

• the comparison with HMC is a little unfair. Please redo the HMC experiment with a different chain for each sample starting in the same latent distribution as your model. I would expect that for your symmetric target distribution this works fine.

Our current HMC experiment initializes the chain from the same latent distribution, N(0,I), as used in our model. We followed the pseudo-code outlined in Algorithm 3 in Appendix, running a single continuous Markov chain where each sample $x_{t}$ depends on the location of the previous sample $x_{t-1}$. We are concerned that starting a different chain for each sample from N(0,I), would break the Markov property and detailed balance of Markov chain. Besides, we observed that in high-dimensional spaces, despite our efforts in adjusting the momentums and increasing leapfrog steps, HMC often fails to take $x_{t}$ to a new mode starting from $x_{t-1}$, as there is too much space to explore in 50D and the chain can hardly explore on the correct direction.

We will carefully review the implementations of different algorithms and include additional experiments and comparisons in our future submissions. Once again, we sincerely thank you for your detailed review and thoughtful suggestions. We will upload a revised version of draft here very soon and hope it addresses at least some of your concerns. We will also continue refining the draft in preparation for future submissions.

Thank you for these comments! We would like to reply to some of your questions and concerns:

• The details of training process is missing. Especially, the cost of the training process of CNF flow is not mentioned, and also how many intervals are needed in order to obtain a stable sampling process. It seems also crucial especially how the overall sampling results depend on the flow interval.
• The sample efficiency (both sampling speed and quality) is not commented and compared with different methods.

Thank you for pointing these out. The costs of training and sampling are discussed in Appendix D.1, and the choice of intermediate annealing densities is discussed in Appendix C.3. We will include more details and discussions on these aspects in our future submissions.

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

The algorithm presented in this paper is designed for parametric statistical sampling in high-dimensional, multi-modal settings, a fundamental and important task in statistical fields such as Bayesian analysis and statistical physics. We have included experiments on hierarchical Bayesian logistic regression across a range of datasets.

The task discussed in this paper focuses on statistical sampling from parametric distributions. We acknowledge that exploring data-driven scenarios—i.e., cases where only samples are available and the parametric form of the data densities is unknown—is indeed an interesting direction. In this work, we introduced a preliminary framework called Importance Flow for sampling rare regions of distributions and constructing low-variance estimators for rare probabilities. Additionally, we discussed the potential for extending this framework into a data-driven method, which we plan to focus on in future work.

Thank you once again for your advice! We will revise the manuscript carefully based on your suggestions. In future submissions, we will include more detailed comparisons of training and sampling costs, as well as additional comparison experiments with other algorithms, as suggested by other reviewers.

We appreciate your comments and suggestions. We would like to clarify some weaknesses and questions you are concerned about:

• The paper lacks a related works section. Namely, the connections to NN-assisted MCMCs, which are in my opinion the most related methods, are not discussed and the performances of these approaches are not reported. In my experience, these methods have better performance than what the paper seems to indicate. In particular, they typically do correct for the mode imbalance that a traditional MCMC based on local updates could not do.

In the second-to-last paragraph, we discussed recent NN-assisted MCMC methods, including those that use neural networks as transition kernels and MCMCs based on normalizing flows. Notably, the "AI-Sampler" compared in our experiments is an NN-assisted MCMC method, which leverages NN to construct the transition kernels. We will include additional experimental comparisons with other NN-assisted methods, including "Neural Importance Sampling" you mentioned and the papers mentioned by Reviewer 7EWK and HDGZ.

A unique advantage of our algorithm is that, unlike most NN-assisted sampling methods, it does not involve MCMC as part of the training process. This approach offers several benefits, including avoiding the long mixing times of MCMC (which can increase exponentially with dimensionality) and producing independent samples. NN-methods that lack Annealing-guidance and $W_{2}$ regularization often suffer from mode-seeking behaviors and often fail with densities that have widely-separated modes (like in the settings of our experiments).

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

Thank you for pointing out this work. We noted that it focuses on discrete normalizing flows (NFs), which inherently exhibit mode-seeking behavior, rather than continuous normalizing flows (CNFs). To address such behavior, their paper uses SMC and importance weights to adjust the loss. Additionally, the experiments presented are primarily focused on scenarios with a low number of modes. Moreover, the mode-seeking tendencies of discrete NFs can become more pronounced as the number of target modes increases, and the SMC procedure may fail in high-dimensional density with separated modes, as in our experiment with 1024 widely separated modes in 50D space.

Besides, through experiments, we observed that without $W_{2}$-regularization (as in their NFs), the transport paths become less organized and more dispersed, which means the flow often misses the optimal path. WWe are currently working on developing theoretical insights into the convergence properties of Annealing guidance and $W_{2}$ regularization.

Thank you for your comment. We will make the comparisons clearer in our future submission.

• The proposed sampling methods lacks theoretical guarantees of the accuracy of the sampling conversely to NN-assisted (Markov Chain) Monte Carlo samplers, such as Neural IS [3], Flow MC (Gabrie et al 2021, 2022 in the paper) or Annealed Flow Transport (Arbel et al 2021 in the paper).

Our algorithm does not rely on an MCMC process, so fewer theoretical properties, such as the law of large numbers, CLT, and variance properties, can be proved. Nonetheless, after submission, we provided an official proof of the infinitesimal optimal velocity field, demonstrating it as the score difference between two consecutive annealing densities. We will further expand on the convergence properties of annealing guidance and $W_{2}$ regularization, which are key to our algorithm’s ability to explore the entire space effectively. Thanks for bringing us attention to develop more theories!

• The cost of training only one velocity field at the time is not discussed, nor how to choose the number of these intermediary steps.

Thank you for your detailed review. We will discuss these in the revised draft.

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

For PT to balance samples between the modes, it must allow transitions between different modes frequently enough at different temperatures. When the two modes are far-separated, particles can have difficulty moving between them even at a high-temperature.

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

This is shown in the first line of the table.

• For the importance flow method as well, the paper claims “The estimator is unbiased and can achieve zero variance theoretically” but this is only true if the density estimation ratio is predicted exactly by the neural network.

Thank you again for all these comments! We will revise our draft accordingly.

Thank you for your detailed review and these kind suggestions. We want to reply to some of your concerns:

• The manuscript contains a GitHub link in Line 369-370

We're sorry for the oversight, but we believe that the current GitHub link does not directly reveal authors' information (The link and the Annealing Flow repository does not contain any name or institution of authors). We will be very cautious in future submissions and will not include any link.

• The proposed method lacks technical novelty and is a bit incremental, since it directly applies an optimal transport loss with KL relaxation to train a sequence of CNFs to model the transitions between neighboring annealed distributions, which is a trivial combination of existing techniques.
• The scalability and reliability of the proposed training algorithm is questionable, since it requires training K separate CNFs sequentially, each modeling the transition between one of the K neighboring pairs of annealed distributions. Furthermore, within each transition, the interval is further discretized into S grid points in order to estimate the integral.

Thank you for these comments. We wish to comment that Annealing principle is the key to ensure smooth transitions and efficient explorations of widely-separated density modes, especially in high-dimensional spaces, as demonstrated in our numerical experiments, which distinct our algorithm from others that also rely on Normalizing Flows. Without annealing procedures, the existing NF methods may easily fail when it comes to densities with widely separated modes.

Besides, during experiments, we observed that without $W_{2}$-regularization (as in the most recent NF works), the transport paths become less organized and more dispersed, which means the flow often misses the optimal path. We are currently working on developing theories on the convergence properties of Annealing guidance and regularization.

Indeed, more blocks (i.e., annealing densities) are required as the number of modes increases and becomes more separated. However, as mentioned in the last paragraph of Section 4.2, the training of AF can be performed offline, requiring only a single training process. Once trained, it can be deployed for sampling anytime as needed. In contrast, methods like MCMC and particle-based approaches often fail regardless of the training duration and must be reimplemented each time sampling is required. Additionally, recent normalizing flow (NF) works without annealing or $W_2$ regularization frequently encounter mode collapse in high-dimensional settings.

Across all of our experiments, the discretization step $S$ is consistently set to 3, i.e., two intermediate $x(t_s)$ are used for numerical integration.

• It is also unclear how the proposed method should be positioned among other flows-based or NN-based samplers with similar techniques, since many important recent related works are not discussed in the paper
• Several closely related baselines are missing in the experiment section

Thank you for your comments. We will include more experimental comparisons with other algorithms. We noted that these Annealing works focus on discrete normalizing flows (NFs), which inherently exhibit mode-seeking behavior, rather than continuous normalizing flows (CNFs). To address such behavior, they use MC procedures and importance weights to correct the loss. Additionally, the experiments presented are primarily focused on scenarios with a low number of modes. Moreover, the mode-seeking tendencies of discrete NFs can become more pronounced as the number of target modes increases, and the SMC procedure may fail in high-dimensional density with separated modes, as in our experiment with 1024 widely separated modes in 50D space.

Thank you for these suggestions. We will include clearer comparisons in future submission.

• For the toy 1D experiment in Figure 2, it would be more convincing if ...

Since the submission, we have conducted additional experiments on densities with differently weighted modes, which we will include in future submissions. We have also updated Figure 2 in our latest draft, which will be uploaded shortly after. However, we need more time to incorporate the results of other flow methods.

Thank you once again for your detailed review and suggestions!