Inflationary Flows: Calibrated Bayesian Inference with Diffusion-Based Models

Summary

This paper introduces inflationary flows, a diffusion model with deterministic denoising trajectories that allow mapping high-dimensional data to low-dimensional Normal distributions. The dimensionality reduction is controllable, and it is claimed that, as long as the score is correctly estimated, the method provides calibrated uncertainty estimates.

Strengths

### Originality * The method is new to me. ### Quality * I did not identify flaws in the mathematical derivations but did not check derivations in the appendices carefully. ### Clarity * The flow is good. * Most of the implementation details are provided. * The authors made an effort to be pedagogical and introduce all the necessary background. ### Significance * Having calibrated uncertainty estimates is of high significance in risk-sensitive tasks.

Weaknesses

### Originality * No concerns about originality. ### Quality * I have some concerns about the empirical evaluation of the method. My main concern is that there are no comparisons to existing methods while there exist plenty of generative models, such as other diffusion models, and normalizing flows,... Additionally, the experiments did not convince me that this method indeed provides calibrated uncertainty estimates, probably because I did not understand what Figure 4 shows (see clarity concerns). ### Clarity * To me, the central part of the submission, the calibrated uncertainty estimates, is not detailed enough. Only 7 lines at the beginning of section 5 mention how to obtain a calibrated uncertainty estimate, and after reading the manuscript, it is still not clear to me as to why this method provides calibrated uncertainty estimates while others don't. * I read Appendix B6 three times and still don't understand what is done, how Figure 4 is obtained, and what it is supposed to show. ### Significance * This paper presents a method that produces calibrated estimates as long as the score is well estimated. My intuition would be that estimating the score correctly is the root of the problem. other methods, such as normalizing flows, also provide calibrated estimates if the neural network is perfectly trained. The miscalibration of other methods, given perfect score estimates, has not been quantified in order to showcase the significance of their method. * It is not clear to me what this method brings in comparison to other diffusion models. The authors seem to claim that their method provides calibrated uncertainty estimates because it is based on a deterministic process, but as mentioned by the authors, equation (4) is also a deterministic process. What does this method bring that equation (4) doesn't have? Additionally, Normalizing flows are also deterministic transformations of data to a base (usually Normal) distribution and satisfy this property. It is also not clear to me why a deterministic process is required for calibrated sampling-based uncertainty estimates (probably linked to clarity issues in section 5).

Questions

* Do you have evidence that strong miscalibration can be observed when using other methods with perfectly learned scores that would strengthen the significance of your method? * Is there any comparison available between your method and other methods both on the quality of the generated objects and on the calibration of the uncertainty estimates? * Could you clarify exactly what properties this method has that no other existing methods have and how those properties lead to calibrated uncertainty estimates?

Rating

5

Confidence

3

Soundness

2

Presentation

2

Contribution

2

Limitations

Limitations are properly addressed.

Summary

The paper proposes inflationary flows (IFs) for sampling-based density estimation. IFs belong to the growing family of diffusion-based models (DBMs) that are trained to transform distributions via a continuous process. However, the current paper notes that typical DBMs actually increase the intrinsic dimensionality of the data in the process, which may be disadvantageous in certain cases. Thus the authors propose to constrain the time-varying covariance of the process in order to approximately preserve the intrinsic data dimensionality as indexed by the participation ratio metric. In addition, the constraint can be modified to encourage dimensionality reduction, with interesting implications for modeling nominally high-dimensional data. The properties and performance of IFs are evaluated across low-dimensional toy data sets and image benchmarks.

Strengths

- The paper is well-written and the flow of argument can be easily followed. I particularly appreciate Section 2. - The method is an interesting addition to the growing number of expressive sampling methods for density estimation, while allowing for preservation or reduction of the intrinsic data dimensionality. - The method can be useful way beyond the empirical settings considered in the paper. - The appendix is rich with useful details and additional results.

Weaknesses

- Generic claims that NFs cannot operate on compressed representations and invalidated by recent work on relaxing architectural constraints [1-3]. These architectures need to be at least considered in the related work and possibly make it into the benchmarks as contenders. - The abstract and introduction begin as if the paper is situated in the context of sampling methods for Bayesian estimation (i.e., conditional estimation), but doesn’t feature any examples relevant to the field (i.e., only unconditional density estimation) and only references generic generative models. Moreover, with this framing, it is written as if the field of simulation-based inference (SBI, [4]), or amortized Bayesian inference (ABI) in particular, does not exist. However, SBI has long been concerned with calibrated Bayesian inference (especially for non-linear, rather low-dimensional problems) and several recent papers have used continuous probability models, such as diffusion models [5] or consistency models [6], to name just a few. - Section 3 can profit from a bit more consistency overall, for example some symbols change meaning, such as $\boldsymbol{\Sigma}_0$ and $\boldsymbol{\Sigma}$, P4L141 seems to imply that $\boldsymbol{\Sigma}(t)$ is diagonal for any $t$, but in my understanding this will not hold for $t \approx 0$, PR is defined as a function over matrices (7) and vectors (8), etc. - An empirical example where the dimensionality-preserving property of IFs actually brings about a palpable advantage over a standard DBM would further strengthen the paper.

Questions

- Does the model admit a procedure for density estimation in addition to sampling? Is there a danger of this density ending up degenerate due to errors in compression?

Rating

7

Confidence

3

Soundness

3

Presentation

3

Contribution

3

Limitations

The authors openly discuss limitations of their method and future work.

Summary

This paper presents a novel approach to Bayesian inference using diffusion-based models (DBMs). Traditional Bayesian inference methods struggle with high-dimensional integrals, and existing approximation methods either scale poorly (sampling-based) or offer limited theoretical guarantees (variational methods). This work leverages DBMs to develop "inflationary flows": these flows utilize a deterministic map from high-dimensional data to a lower-dimensional Gaussian distribution, achieved through ODE integration. The proposed method ensures local neighborhood preservation and controllable numerical error, enabling accurate uncertainty quantification. The approach also includes novel noise schedules and demonstrates high generative performance, even under significant data compression.

Strengths

* The article is very well written and structured. Moreover, the authors provide good explanations that aid in developing intuition around the model, s.a. the diagram presented in Figure 1. * The idea behind the inflationary flows is a very interesting one from the point of view of the properties of the method. It does seem to provide new key features that are missing in the literature and that may be helpful in many other problems. * The theoretical contributions made by the authors are strong, with enough support behind the main points. The discussion makes it so that the inflationary flows model seems a natural consequence of previous works, filling previous gaps in the literature. * The method produces calibrated, identifiable Bayesian inference by leveraging a deterministic map between the original data and a Gaussian distribution. This ensures accurate propagation of uncertainties from the data to the latent space, leading to reliable predictive distributions.

Weaknesses

* While it is true that the KL divergence used in VI induces a mode-seeking behavior, this can be easily fixed alongside the properties of the posterior distribution obtained via VI. This is explored in articles s.a. [1], although it has become a widely-known fact for some time [2] and its properties are usually exploited in the context of VI-based methods. $\alpha$-divergences improve the performance of VI-based methods with very simple changes in the formulation. I think this should be discussed in the context of the introduction and the related work of the article. * The accuracy of the method heavily depends on the quality of the estimated score function. Errors in this estimation can propagate through the model and impact the final results. * The training process for diffusion-based models requires significant computational resources, particularly due to the need for training over larger noise ranges. This can limit the practical applicability of the method to very large datasets. * I would appreciate it if the authors provided further metrics to fully ensure the quality of the predictive distributions. These can range from the usual log-likelihood estimates to the CRPS or other proper scoring rules [3], alonside the ECE compared to other baseline models (e.g. some based on VI, regular DBMs and MCMC-based approaches). This would help in comparatively assessing the quality of the predictions obtained in a broader sense. ## Minor: * Check again for typos, e.g. "asympototically" (first paragraph of the introduction). I only found this one, but it is worth checking again. * On the first equation, add the definition of $\mathbf{W}$ for completeness (as is done later for the time-reversed $\bar{\mathbf{W}}$). * Although I do like the representation, I would add some short descriptions on the arrows of Figure 1 connecting the three paradigms (diffusion model $\rightarrow$ Fokker-Plank $\leftarrow$ pfODE), making it more explicit. _(please see the 'Limitations' section for the references)_

Questions

* Could you please discuss the potential impact of errors in the score function estimation on the accuracy and stability of the pfODE integration? How do you address the sensitivity of your method to these errors, and what strategies can be employed to mitigate their effects? This would be particularly interesting for complex datasets. * When is it the PR measure for dimensionality expected to perform best? Are there any cases where this measure can lead to a bad performance due to the reliance only on second-order statistics?

Rating

7

Confidence

4

Soundness

4

Presentation

4

Contribution

3

Limitations

* The article relies on the participation ratio (PR) as a measure of dimensionality, which might not capture more complex data structures effectively. * Accuracy of the method depends heavily on the quality of the estimated score function, which can introduce errors if not precisely learned. * The need for training diffusion-based models over larger ranges of noise can lead to trade-offs in compression performance and score estimation accuracy. * The experimental phase can be bolstered with the addition of a few extra comparisons with other VI-based models, as well as regular DBMs and MCMC-based models in order to provide a fuller comparative picture of the performance of the proposal. ## References: [1] Rodríguez Santana, S., & Hernández-Lobato, D. (2022). Adversarial $\alpha$-divergence minimization for Bayesian approximate inference. Neurocomputing, 471, 260-274. [2] Minka, T. (2005). Divergence measures and message passing (p. 17). Technical report, Microsoft Research. [3] Gneiting, T., & Raftery, A. E. (2007). Strictly proper scoring rules, prediction, and estimation. Journal of the American statistical Association, 102(477), 359-378.

Summary

This paper proposes a new version of a probability flow-based diffusion model that allows for the effective reduction in dimensionality of the data after it has been diffused. Furthermore, the proposed approach allows for proper uncertainty measures of the data by preserving local neighborhoods from latent space to data space---in both dimension preserving and dimension reducing setups. Extensive empirical results are generated demonstrating the effects of the proposed scheme.

Strengths

This is a unique and interesting problem to try and solve with diffusion models. The approach is well justified and presented clearly. The experimental setups are well motivated and the corresponding results are largely convincing.

Weaknesses

For the experiments with high-dimensional data, I found the trend of having a higher percentage of dimension preservation resulting in worse performance to be concerning and the lack of investigation on this to be unsatisfactory. Only values of 2%, 10%, and 20% were chosen with the explanation of the results being "retaining more dimensions in our PR-Reducing schedules leads to larger scale gaps between our preserved and compressed dimensions (i.e., larger $e^{ρ(g_{*}−g_i)T}$), thus increasing the required noise range over which networks must optimally estimate scores;" however, I do not believe that this tells the full story. The dimension-reducing flow approaches the original dimension-preserving flow as the number of $K$ reduced dimensions approaches the original data dimensionality $N$. As such, one would expect the performance to at some point start improving again. Should it not, and experience a discontinuity in performance from $K=N-1$ to $K=N$ then that would be interesting to note. Additionally, clearly it stands to reason that we cannot keep reducing the dimensionality further and further without eventually experiencing a decrease in performance, i.e., $K=1$ is most likely going to perform poorly for images. Finding out where this inflection point would be interesting to note as that is naturally going to be the optimal choice. With the lack of comparison to other methods, which is understandable in and of itself given the focus of the paper, it is even more important then to fully explore the failure modes and general behavior of the proposed method.

Questions

No additional questions, please address the main weakness I brought up. Please let me know if I misunderstood something as well. I am more than happy to be convinced otherwise on my points.

Rating

6

Confidence

4

Soundness

4

Presentation

4

Contribution

4

Limitations

The authors adequately discussed the limitations.