Inflationary Flows: Calibrated Bayesian Inference with Diffusion-Based Models

Authors: daniela  de albuquerque
john  pearson
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-14 16:12:34

Abstract: Beyond estimating parameters of interest from data, one of the key goals of statistical inference is to properly quantify uncertainty in these estimates. In Bayesian inference, this uncertainty is provided by the posterior distribution, the computation of which typically involves an intractable high-dimensional integral. Among available approximation methods, sampling-based approaches come with strong theoretical guarantees but scale poorly to large problems, while variational approaches scale well but offer few theoretical guarantees. In particular, variational methods are known to produce overconfident estimates of posterior uncertainty and are typically non-identifiable, with many latent variable configurations generating equivalent predictions. Here, we address these challenges by showing how diffusion-based models (DBMs), which have recently produced state-of-the-art performance in generative modeling tasks, can be repurposed for performing calibrated, identifiable Bayesian inference. By exploiting a previously established connection between the stochastic and probability flow ordinary differential equations (pfODEs) underlying DBMs, we derive a class of models, \emph{inflationary flows,} that uniquely and deterministically map high-dimensional data to a lower-dimensional Gaussian distribution via ODE integration. This map is both invertible and neighborhood-preserving, with controllable numerical error, with the result that uncertainties in the data are correctly propagated to the latent space. We demonstrate how such maps can be learned via standard DBM training using a novel noise schedule and are effective at both preserving and reducing intrinsic data dimensionality. The result is a class of highly expressive generative models, uniquely defined on a low-dimensional latent space, that afford principled Bayesian inference.

Main paper: 
 introduction: 
 In many fields of science, the aim of statistical inference is not only to estimate model parameters of interest from data but to quantify the uncertainty in these estimates. In Bayesian inference, for data x generated from latent parameters z via a model p(x|z), this information is encapsulated in the posterior distribution p(z|x), computation of which requires evaluation of the often intractable normalizing integral p(x) = p(x, z) dz. Where accurate uncertainty estimation is required, the gold standard remains sampling-based Markov Chain Monte Carlo (MCMC) methods, which are guaranteed (asymptotically) to produce exact samples from the posterior distribution [1] . However, MCMC methods can be computationally costly and do not readily scale either to large or highdimensional data sets.Alternatively, methods based on variational inference (VI) attempt to approximate posterior distributions by optimization, minimizing some measure of divergence between the true posterior and a parameterized set of distributions q ϕ (z|x) [2] . For example, methods like the variational autoencoder (VAE) [3, 4] minimize the Kullback-Leibler (KL) divergence between true and approximate posteriors, producing bidirectional mappings between data and latent spaces. In vanilla VAEs, posterior uncertainty estimates are typically overconfident due to minimization of the reverse (mode-seeking) KL divergence [5, 6] . While some lines of work have sought to mitigate this posterior mismatch problem by utilizing different divergences 7-10, VAEs still tend to produce blurry data reconstructions and non-unique latent spaces without additional assumptions [11] [12] [13] .By contrast, normalizing flow (NF) models [14, 15] work by applying a series of bijective transformations to a simple base distribution (usually uniform or Gaussian) to deterministically convert samples to a desired target distribution. While NFs have been successfully used for posterior approximation [16] [17] [18] [19] [20] and produce higher-quality samples, the requirement that the Jacobian of each transformation be simple to compute often requires a high number of transformations and, traditionally, these transformations do not alter the the dimensionality of their inputs, resulting in latent spaces with thousands of dimensions. More recent lines of work on injective flow models 21-25 address this limitation by allowing practitioners to use flows to learn lower dimensional manifolds from data, but most compression-capable flow models still fail to reach high generative performance on key benchmark image datasets (cf. [23] ).More recently, diffusion-based models (DBMs) [26] [27] [28] [29] [30] [31] [32] [33] have been shown to achieve state-of-the-art results in several generative tasks, including image, sound, and text-to-image generation. These models work by stipulating a fixed forward noising process (e.g., a forward stochastic differential equation (SDE)), wherein Gaussian noise is incrementally added to samples of the target data distribution until all information in the original data is degraded. To generate samples from the target distribution, one then needs to simulate the reverse de-noising process (reverse SDE [34] ) which requires knowledge of the score of the intermediate "noised" transitional densities. Estimation of this score function across multiple noise levels is the key component of DBM model training, typically using a de-noising score matching objective [35, 28, 30] . Yet, despite their excellent performance as generative models, DBMs, unlike VAEs or flows, do not readily lend themselves to inference. In particular, because DBMs use a diffusion process to transform the data distribution, they fail to preserve local structure in the data (Figure 1 ), and uncertainty under this mapping is high at its endpoint because of continuous noise injection and resultant mixing. Moreover, because the final distribution-Gaussian white noise of the same dimension-must have higher entropy than the original data, there is no data compression.Finally, emerging work on flow matching models [36] [37] [38] [39] [40] [41] [42] has achieved impressive generative performance on several benchmark image datasets. Such models utilize simple conditional distribution families to learn a vector field capable of transporting points between two pre-specified densities. These are closely related to the probability flow ODE (pfODE) view of DBMs, and, in fact, have been shown to be equivalent to such models for specific choices of "interpolant" functions and conditional distributions. Despite their exceptional generative performance and deterministic nature, existing flow matching approaches do not allow for compression and, therefore, do not allow practitioners to infer a lower dimensional latent space from data. Thus, despite tremendous improvements in sample quality, modern generative models do not lend themselves to one of the key modeling goals in scientific applications: calibrated Bayesian inference. Note that while many works focus on predictive calibration, how well the inferred marginal p(x) matches real data [43] [44] [45] [46] [47] , our focus here is on posterior calibration, how well q(z|x) matches the true posterior p(z|x). We address this challenge by demonstrating how a novel DBM variant that we call inflationary flows can, in fact, produce calibrated Bayesian inference in this sense. Specifically, our contributions are: First, focusing on the case of unconditional generative models, we show how a previously established link between the SDE defining diffusion models and the probability flow ODE (pfODE) that gives rise to the same Fokker-Planck equation [30] can be used to define a unique, deterministic map between the original data and an asymptotically Gaussian distribution. This map is bidirectional, preserves local neighborhoods, and has controllable numerical error, making it suitable for rigorous uncertainty quantification. Second, we define two classes of flows that correspond to novel noise injection schedules in the forward SDE of the diffusion model. The first of these preserves a measure of dimensionality, the participation ratio (PR) [48] , based on1 Introduction @tpt(x) = X i @i[fi(x, t)pt(x)] + 1 2 X ij @i@j " X kGik(x, t)Gjk(x, t)pt(x) The forward (noising) SDE defining the DBM (left) gives rise to a sequence of marginal probability densities whose temporal evolution is described by a Fokker-Planck equation (FPE, middle). But this correspondence is not unique: the probability flow ODE (pfODE, right) gives rise to the same FPE. That is, while both the SDE and the pfODE possess the same marginals, the former is noisy and mixing while the latter is deterministic and neighborhood-preserving. Both models require knowledge of the score function ∇ x log p t (x), which can learned by training either model.# 1second-order data statistics, preventing an effective increase in data dimensionality with added noise, while the second flow reduces PR, providing data compression. We demonstrate experimentally that inflationary flows indeed preserve local neighborhood structure, allowing for sampling-based uncertainty estimation, and that these models continue to provide high-quality generation under compression, even from latent spaces reduced to as little as 0.03% of the nominal data dimensionality. As a result, inflationary flows offer excellent generative performance while affording data compression and accurate uncertainty estimation for scientific applications. 
three views of diffusion-based models: 
 As with standard DBMs, we assume a data distribution p 0 (x) = p data (x) at time t = 0, transformed via a forward noising process defined by the stochastic differential equation [e.g., 26, 28] :EQUATIONwith most DBMs assuming linear drift (f = f (t)x) and isotropic noise (G = σ(t)1) that monotonically increases over time [49] . As a result, for T 0 σ(T )dt ≫ σ data , p T (x) becomes essentially indistinguishable from an isotropic Gaussian (Figure 1 , left). DBMs work by learning an approximation to the reverse SDE [34, [28] [29] [30] 50] ,EQUATION)where W is time-reversed Brownian motion. In practice, this requires approximating the score function ∇ x log p t (x) by incrementally adding noise according to the schedule σ(t) of the forward process and then requiring that denoising by (2) match the original sample. The fully trained model then generates samples from the target distribution by starting with x T ∼ N (0, σ 2 (T )1) and integrating (2) in reversed time.As previously shown, this diffusive process gives rise to a series of marginal distributions p t (x) satisfying a Fokker-Planck equation (Figure 1 , middle) [30, 49] ,EQUATIONwhere ∂ i ≡ ∂ ∂xi . In the "variance preserving" noise schedule of [30] , (3) has as its stationary solution an isotropic Gaussian distribution. This "distributional" perspective views the forward process as a means of transforming the data into an easy-to-sample form (as with normalizing flows) and the reverse process as a means of data generation.However, in addition to the SDE and FPE perspectives, Song et al. [30] also showed that (3) is satisfied by the marginals of a different process with no noise term, the so-called probability flow ODE (pfODE):EQUATIONUnlike (1) , this process is deterministic, and data points evolve smoothly (Figure 1 , right), resulting in a flow that preserves local neighborhoods. Moreover, the pfODE is uniquely defined by f (x, t), G(x, t), and the score function. This connection between the marginals satisfying the SDEs of diffusion processes and deterministic flows described by an equivalent ODE has also been recently explored in the context of flow matching models [36] [37] [38] [39] [40] [41] [42] , a connection on which we elaborate in Section 7.In the following sections, we show how this pfODE, constructed using a score function estimated by training the corresponding DBM, can be used to map points from p data (x) to a compressed latent space in a manner that affords accurate uncertainty quantification. 
inflationary flows: 
 As argued above, the probability flow ODE offers a means of deterministically transforming an arbitrary data distribution into a simpler form via a score function learnable through DBM training.Here, we introduce a specialized class of pfODEs, inflationary flows, that follow from an intuitive picture of local dynamics and asymptotically give rise to stationary Gaussian solutions of (3).We begin by considering a sequence of marginal transformations in which points in the original data distribution are convolved with Gaussians of increasingly larger covariance C(t):EQUATIONIt is straightforward to show (Appendix A.1) that this class of time-varying densities satisfies (3) when f = 0 and GG ⊤ = Ċ. This can be viewed as a process of deterministically "inflating" each point in the data set, or equivalently as smoothing the underlying data distribution on ever coarser scales, similar to denoising approaches to DBMs [51, 52] . Eventually, if the smoothing kernel grows much larger than Σ 0 , the covariance in the original data, total covariance Σ(t) ≡ Σ 0 + C(t) → C(t), p t (x) ≈ N (0, C(t)), and all information has been removed from the original distribution. However, because it is numerically inconvenient for the variance of the asymptotic distribution p ∞ (x) to grow much larger than that of the data, we follow previous work in adding a time-dependent coordinate rescaling x(t) = A(t) • x(t) [30, 49] , which results in an asymptotic solution p ∞ (x) = N (0, AΣA ⊤ ) of the corresponding Fokker-Planck equation when Σ = Ċ and ȦΣA ⊤ + AΣ Ȧ⊤ = 0 (Appendix A.2). Together, these assumptions give rise to the pfODE (Appendix A.3):EQUATIONwhere the score function is evaluated at x = A -1 • x. Notably, ( 6) is equivalent to the general pfODE form given in [49] in the case both C(t) and A(t) are isotropic (Appendix A.4), with C(t) playing the role of injected noise and A(t) the role of the scale schedule. In the following sections, we will show how to choose both of these in ways that either preserve or reduce intrinsic data dimensionality.In standard DBMs, the final form of the distribution p T (x) approximates an isotropic Gaussian distribution, typically with unit variance. As a result, these models increase the effective dimensionality of the data, which may begin as a low-dimensional manifold embedded within R d . Thus, even maintaining intrinsic data dimensionality requires both a definition of dimensionality and a choice of flow that preserves this dimension. In this work, we consider a particularly simple measure of dimensionality, the participation ratio (PR), first introduced by Gao et al. [48] :Swirl Inflation GenerationPR(Σ) = tr(Σ) 2 tr(Σ 2 ) = ( i σ 2 i ) 2 i σ 4 i (7)where Σ is the covariance of the data with eigenvalues {σ 2 i }. PR is invariant to linear transforms of the data, depends only on second-order statistics, is 1 when Σ is rank-1, and is equal to the nominal dimensionality d when Σ ∝ 1 d×d . In Appendix C.1 we report this value for several benchmark image datasets, confirming that in all cases, PR is substantially lower than the nominal data dimensionality.To construct flows that preserve this measure of dimension, following (5), we write total variance as Σ(t) = diag(σ 2 (t)) = C(t) + Σ 0 , where Σ 0 is the original data covariance and C(t) is our time-dependent smoothing kernel. Moreover, we will choose C(t) to be diagonal in the eigenbasis of Σ 0 and work in that basis, in which case Σ(t) = diag(σ 2 (t)) and we have (Appendix A.6):EQUATIONThe simplest solution to this constraint is a proportional inflation, d dt (σ 2 ) = ρσ 2 , along with a rescaling along each principal axis:EQUATION)As with other flow models based on physical processes like diffusion [26] or electrostatics [53, 54] , our use of the term inflationary flows for these choices is inspired by cosmology, where a similar process of rapid expansion exponentially suppresses local fluctuations in background radiation density [55] . However, as a result of our coordinate rescaling, the effective covariance Σ = AΣA ⊤ = diag(A 2 0j ) remains constant (so dσ 2 = 0 trivially), and the additional conditions of Appendix A.2 are satisfied, such that N (0, Σ) is a stationary solution of the relevant rescaled Fokker-Planck equation. As Figure 2 shows, these choices result in a version of (6) that smoothly maps nonlinear manifolds to Gaussians and can be integrated in reverse to generate samples of the original data.In the previous section, we saw that isotropic inflation preserves intrinsic data dimensionality as measured by PR. Here, we generalize and consider anisotropic inflation at different rates along each of the eigenvectors of Σ: d dt (σ 2 ) = ρg ⊙ σ 2 . In addition, we denote g * ≡ max(g), so that the fastest inflation rate is ρg * . Then, if we take g i = g * for i ∈ {i 1 , i 2 , . . . i K } and g i < g * for the other dimensions,GenerationEQUATIONwhich is the dimension that would be achieved by simply truncating the original covariance matrix in a manner set by our choice of g. Here, unlike in (9), we do not aim for rescaling to compensate for expansion along each dimension, since that would undo the effect of differential inflation rates. Instead, we choose a single global rescaling factor α(t) ∝ A 0 exp(-ρg * t/2), leading to a Gaussian asymptotic solution with the original data covariance in dimensions i ∈ {i 1 , i 2 , . . . i K }.Two additional features of this class of flows are worth noting: First, the final scale ratio of preserved to shrunken dimensions for finite integration times T is governed by the quantity e ρ(g * -gi)T in (10) . For good compression, we want this number to be very large, but as we show in Appendix A.4, this corresponds to a maximum injected noise of order e ρ(g * -gi)T /2 in the equivalent DBM. That is, the compression one can achieve with inflationary flows is constrained by the range of noise levels over which the score function can be accurately estimated, and this is quite limited in typical models. Second, despite the appearance given by (10) , the corresponding flow is not simply a linear projection to the top K principal components: though higher PCs are selectively removed by dimension-reducing flows via exponential shrinkage, individual particles are repelled by local density as captured by the score function (6) , and this term couples different dimensions even when C and A are diagonal. Thus, the final positions of particles in the retained dimensions depend on their initial positions in the full space, producing a nonlinear map (Figure 3 ).Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof? Answer: [Yes] Justification: We provide most important set of assumptions and equations needed to understand the work (in main text) and provide full assumptions, proofs, and theoretical detail in Appendices A, B. Guidelines:• The answer NA means that the paper does not include theoretical results.• All the theorems, formulas, and proofs in the paper should be numbered and crossreferenced. • All assumptions should be clearly stated or referenced in the statement of any theorems.• The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition. • Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material. • Theorems and Lemmas that the proof relies upon should be properly referenced. 
score function approximation from dbms: 
 Having chosen inflation and rescaling schedules, the last component needed for the pfODE ( 6) is the score function s(x, t) ≡ ∇ x log p t (x). Our strategy will be to exploit the correspondence described above between diffusion models (1) and pfODEs (4) that give rise to the same marginals (3) . That is, we will learn an approximation to s(x, t) by fitting the DBM corresponding to our desired pfODE, since both make use of the same score function.Briefly, in line with previous work on DBMs [49] , we train neural networks to estimate a de-noised version, D(x, C(t)), of a noise-corrupted data sample x given noise level C(t) (cf. Appendix A.4 for the correspondence between C(t) and noise). That is, we model D θ (x, C(t)) using a neural network and train it by minimizing a standard L 2 de-noising error:EQUATIONDe-noised outputs can then be used to compute the desired score term using ∇ x log p(x, C(t)) = C -1 (t) • (D(x; C(t))x) [30, 49] . Moreover, as in [49] , we also adopt a series of preconditioning Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)? Answer: [Yes] Justification: We provide detailed information about all of our experiments (including additional experiments, not included in main text) in Appendices B, C. Additionally, we provide entire code needed to reproduce results of paper in this repository [63] . All datasets utilized are publicly available and we provide details on how to download and pre-process these data in our repository and in the appendices. Guidelines:• The answer NA means that the paper does not include experiments.• If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not. • If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable. • Depending on the contribution, reproducibility can be accomplished in various ways.For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed. Guidelines:• The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.• If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics. • The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction). 10. Broader Impacts Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed? Answer: [Yes] Justification: We include discussion of potential societal impacts of the work presented herein as part of section 7.Guidelines:• The answer NA means that there is no societal impact of the work performed.• If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact. • Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations. • The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster. • The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology. • If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML). 
calibrated uncertainty estimates from inflationary flows: 
 Several previous lines of work [43] [44] [45] [46] [47] have focused on assessing how well model-predicted marginals p(x) match real data (i.e., the predictive calibration case). Though we do compare our models' predictive calibration performance against existing injective flow models (Table 3 ), here we are primarily focused on quantifying error in unconditional posterior inference. That is, we are interested in quantifying the mismatch between inferred posteriors q(z|x) and true posteriors p(z|x), especially in contexts where the true generative model is unknown and must be learned from data. This is by far the most common scenario in modern generative models like VAEs, flows, and GANs.As with other implicit models, our inflationary flows provide a deterministic link between complex data and simplified distributions with tractable sampling properties. This mapping requires integrating the pfODE (6) for a given choice of C(t) and A(t) and an estimate of the score function of the original data. As a result, sampling-based estimates of uncertainty are trivial to compute: given a prior π(x) over the data (e.g., a Gaussian ball centered on a particular example x 0 ), this can be transformed into an uncertainty on the dimension-reduced space by sampling {x j } ∼ π(x) and integrating (6) forward to generate samples from p(x T |x 0 )π(x 0 ) dx 0 . As with MCMC, these samples can be used to construct either estimates of the posterior or credible intervals. Moreover, because the pfODE is unique given C, A, and the score, the model is identifiable when conditioned on these choices.The only potential source of error, apart from Monte Carlo error, in the above procedure arises from the fact that the score function used in (6) is only an estimate of the true score. To assess whether integrating noisy estimates of the score could produce errant posterior samples, we conducted the experiment showcased in Figure 4A (Appendix B.7). Briefly, we constructed a Gaussian Mixture Model (GMM) prior with three pre-specified components (Appendix B.7) from which we drew samples of z, integrating backwards in time using our trained pfODE networks to construct corresponding observed data points x. We then utilized Hamiltonian Monte Carlo (HMC) 1, 56-58 to obtain posterior samples for the GMM component weights. As shown in Figure 4B , C, the resulting posterior correctly covers the original ground-truth values, suggesting that numerical errors in score estimates, at least in this simplified scenario, do not appreciably accumulate. This is likely because, empirically, score estimates do not appear to be strongly auto-correlated in time (Appendix C.3), suggesting that ŝ(x, t) is well approximated as a scaled colored noise process and the corresponding pfODE as an SDE. In such a case, standard theorems for SDE integration show that while errors due to noise do accumulate, these can be mitigated by a careful choice of integrator and ultimately controlled by reducing step size [59, 60] . In addition, we verified this empirically in both lowdimensional examples (Figure 4 , Appendices B.7, C.2.1) and with round-trip integration of the pfODE in high-dimensional datasets (Tables 1, 2 , Appendix B.5.1). 
experiments: 
 For the PR-Reducing flows, the final scale ratio between preserved vs. shrunken dimensions for finite integration times is dependent on the quantity e ρ(g * -gi)T . Therefore, for fixed end integration time T and rate ρ, this scaling is dictated by g *g i , which we call the "inflation gap" (IG), Appendix B.2. As this inflation gap increases, compressed dimensions are shrunken to a greater extent, and the denoising networks are required to amortize score estimation over wider noise scales, a harder learning problem. Therefore, for our proposed model, compression should be understood both in terms of the number of dimensions being preserved and the size of this inflation gap.To assess how these two factors affect model performance, we performed two sets of experiments on two benchmark image datasets (CIFAR-10 [61] and AFHQv2 [62] ; Appendix B.4.2; code: [63]; project website: [64]). In the first set of experiments, we fixed T , ρ, and the inflation gap (IG = 1.02) while varying only the number of preserved dimensions d between d = 1 (compression to ≈ 0.03%) and d = 3072 (no compression) for both datasets. For the second set of experiments, we worked with the AFHQv2 dataset and fixed T , ρ, and d = 2, while varying the inflation gap (IG = 1.10, 1.25, 1.35, 1.50). In Tables 1 and 2 we showcase Frechet Inception Distance (FID) scores [65] (mean ±2σ over 3 independently generated sets of images, each with 50,000 samples) and round-trip integration mean squared errors (mean MSE ±2σ over 3 randomly sampled sets of images, each with 10,000 samples) for each (d, IG) combination explored (Appendices B.5.1, B.5.2, B.6). Figures 5, 6 , and 7 showcase 24 randomly generated images (top rows) along with round-trip integration results for 8 randomly sampled images (bottom rows), across select (d, IG) combinations.Finally, we also compared our inflationary flows (IFs) model generative performance on CIFAR-10 against three existing injective flow model baselines (Appendix B.5.2) -M-Flows [21] , Rectangular Flows (RFs) [22] , and Canonical Manifold Flows (CMF) [23] -for different numbers of preserved dimensions (d = 30, 40, 62). Table 3 showcases best FID scores (out of 3 independently generated sets of images, each with 10,000 samples) for each such experiment. For these comparison experiments, we fixed IG=1.02 when training our networks for the different d values.As a general trend, increasing the number of preserved dimensions at a constant inflation gap led to improvements in generative quality (lower FID scores) and reduced MSE (Table 1 ). However, some schedules we assessed are not entirely consistent with this trend. We hypothesize this is at least partially due to variance arising from different network initializations for each schedule, as well as differences between the two datasets explored here. As expected, increasing inflation gap while maintaining the number of preserved dimensions leads to worsened generative performance (higher FID scores, Table 2 ). Finally, in terms of predictive calibration, our model provides substantial gains when compared to existing injective flow model baselines (Table 3 ).  
discussion: 
 Here, we have proposed a new type of implicit probabilistic model based on the probability flow ODE (pfODE) in which it is possible to perform calibrated, identifiable Bayesian inference on a reduceddimension latent space via sampling and integration. To do so, we have leveraged a correspondence between pfODEs and diffusion-based models by means of their associated Fokker-Planck equations, and we have demonstrated that such models continue to produce high-quality generated samples even when latent spaces are as little as 0.03% of the nominal data dimension. More importantly, the uniqueness and controllable error of the generative process make these models an attractive approach in cases where accurate uncertainty estimates are required.Limitations: One limitation of our model is its reliance on the participation ratio (7) as a measure of dimensionality. Because PR relies only on second-order statistics and our proposals ( 9) are formulated in the data eigenbasis, our method tends to favor the top principal components of the data when reducing dimension. However, as noted above, this is not simply a truncation to the lowest principal components, since dimensions still mix via coupling to the score function in (6) . Nonetheless, solutions to the condition (8) that preserve (or reduce) more complex dimensionality measures might lead to even stronger compressions for curved manifolds (Appendix C.2.2), and more sophisticated choices for noise and rescaling schedules in (6) might lead to compressions that do not simply remove information along fixed axes, more similar to [66] . That is, we believe much more interesting classes of flows are possible. A second limitation is that mentioned in Section 3.2 and in our experiments: our schedule requires training DBMs over much larger ranges of noise than are typically used, and this results in noticeable tradeoffs in compression performance as the inflation gap and number of preserved dimensions are varied. Related work: This work draws on several related lines of research, including work on using DBMs as likelihood estimation machines [50, 67, 31] , relations with normalizing flows and hierarchical VAEs [67, 33, 68] , injective flow models [21] [22] [23] [24] [25] , and generative flow networks [69] . By contrast, our focus is on the use of DBMs to learn score functions estimates for implicit probabilistic models, with the ultimate goal of performing accurate posterior inference. In this way, it is also closely related to work on denoising models [51, 52, 66, 70] that cast that process in terms of statistical inference and to models that use DBMs for de-blurring and in-painting [71, 72] . However, this work is distinct from several models that use reversal of deterministic transforms to train generative models [73] [74] [75] [76] . Whereas those models work by removing information from each sample x, our proposal relies critically on adjusting the local density of samples with respect to one another, moving the marginal distribution toward a Gaussian.Our work is also similar to methods that use DBMs to construct samplers for unnormalized distributions [77] [78] [79] [80] [81] . Whereas we begin with samples from the target distribution and aim to learn latent representations, those studies start with a pre-specified form for the target distribution and aim to generate samples. Other groups have also leveraged sequential Monte Carlo (SMC) techniques to construct new types of denoising diffusion samplers for, e.g., conditional generation [82] [83] [84] . While our goals are distinct, we believe that the highly simplified Gaussian distribution of our latent spaces may potentially render joint and conditional generation more tractable in future models. Finally, while many prior studies have considered compressed representations for diffusion models [85] [86] [87] [88] , typically in an encoder-decoder framework, the focus there has been on generative quality, not inference. Along these lines, the most closely related to our work here is [89] , which considered diffusion along linear subspaces as a means of improving sample quality in DBMs, though there again, the focus was on improving generation and computational efficiency, not statistical inference.Yet another line of work closely related to ours is the emerging literature on flow matching [36] [37] [38] 90] models, which utilize a simple, time-differentiable, "interpolant" function to specify conditional families of distributions that continuously map between specified initial and final densities. That is, the interpolant functions define flows that map samples from a base distribution ρ 0 (x) to samples from a target distribution ρ 1 (x). Typically, these approaches rely on a simple quadratic objective that attempts to match the conditional flow field, which can be computed in closed form without needing to integrate the corresponding ODE. As shown in Appendix A.5, the pfODEs obtained using our proposed scaling and noising schedules are equivalent to the ODEs obtained by using the "Gaussian paths formulation" from [36] when the latter are generalized to full covariance matrices. As a result, our models are amenable to training using flow-matching techniques, suggesting that faster training and inference schemes may be possible through leveraging connections between flow matching and optimal transport [40, 42, 41, 38] Broader impacts: Works like this one that focus on improving generative models risk contributing to an increasingly dangerous set of tools capable of creating misleading, exploitative, or plagiarized content. While this work does not seek to improve the quality of data generation, it does propose a set of models that feature more informative latent representations of data, which could potentially be leveraged to those ends. However, this latent data organization may also help to mitigate certain types of content generation by selectively removing, prohibiting, or flagging regions of the compressed space corresponding to harmful or dangerous content. We believe this is a promising line of research that, if developed further, might help address privacy and security concerns raised by generative models.A.1 Derivation of the inflationary Fokker-Planck EquationWe start with derivatives of the smoothing kernel κ(x, t) ≡ N (x; µ, C(t)):EQUATIONEQUATIONand combine this with (5) to calculate terms in (3):EQUATIONEQUATIONEQUATION)1 2 ij ∂ i ∂ j k G ik G jk p = 1 2 p 0 * ij ∂ i ∂ j k G ik G jk κ (18) -2∂ j k G ik G jk (C -1 (x -µ)) i κ + k G ik G jk [C -1 (x -µ)] i [C -1 (x -µ)] j -(C -1 ) ij κ .Assuming f = 0 and ∂ i G jk (x, t) = 0 then gives the condition- 1 2 tr(C -1 Ċ) + 1 2 tr C -1 (x -µ)(x -µ) ⊤ C -1 Ċ = - 1 2 tr(C -1 GG ⊤ ) + 1 2 tr C -1 (x -µ)(x -µ) ⊤ C -1 GG ⊤ (19)which is satisfied when GG ⊤ (x, t) = Ċ(t).Starting from the unscaled Fokker-Planck Equation corresponding to the process of Appendix A.1EQUATIONwe introduce new coordinates x = A(t) • x, t = t, leading to the change of derivativesEQUATION= (A ĊA ⊤ ) kl ∂k ∂l (26) and the Fokker-Planck EquationEQUATIONwhere pt (x) = p t (x) is simply written in rescaled coordinates. However, this is not a properly normalized probability distribution in the rescaled coordinates, so we define q(x, t) ≡ J -1 ( t)p t(x), which in turn satisfiesEQUATIONNow consider the time-dependent Gaussian densityEQUATIONwith rescaling factor J( t) = |A ⊤ A(t)|. We then calculate the pieces of (28) as follows:∇q = -(AΣA ⊤ ) -1 (x -Aµ)q ∂i ∂j q = (AΣA ⊤ ) -1 (x -Aµ) i (AΣA ⊤ ) -1 (x -Aµ) j q -[(AΣA ⊤ ) -1 ] ij q ∂t log J = ∂t log |AA ⊤ | = tr( ∂t log AA ⊤ ) = tr (AA ⊤ ) -1 ( ∂t A)A ⊤ + A( ∂t A ⊤ ) ∂t q = - 1 2 tr((AΣA ⊤ ) -1 ∂t (AΣA ⊤ ))q + qµ ⊤ ∂t A ⊤ (AΣA ⊤ ) -1 (x -Aµ) - q 2 tr (x -Aµ)(x -Aµ) ⊤ ∂t (AΣA ⊤ ) -1 -∂t log J ∂t (AΣA ⊤ ) -1 = -(AΣA ⊤ ) -1 ∂t (AΣA ⊤ )(AΣA ⊤ ) -1 = -(AΣA ⊤ ) -1 (( ∂t A)A -1 ) -(( ∂t A)A -1 ) ⊤ (AΣA ⊤ ) -1 -A -⊤ Σ -1 ∂t ΣΣ -1 A -1 .With these results, the left and right sides of (28) become[x ⊤ • ∂t log A ⊤ • ∇ + ∂t + ∂t log J]q = -x ⊤ [( ∂t A)A -1 ] ⊤ (AΣA ⊤ ) -1 (x -Aµ)q - 1 2 tr((AΣA ⊤ ) -1 ∂t (AΣA ⊤ ))q +µ ⊤ ∂t A ⊤ (AΣA ⊤ ) -1 (x -Aµ)q - 1 2 tr (x -Aµ)(x -Aµ) ⊤ ∂t (AΣA ⊤ ) -1 q -∂t log |AA ⊤ |q +tr (AA ⊤ ) -1 ( ∂t A)A ⊤ + A( ∂t A ⊤ ) q = - q 2 tr ∂t (AΣA ⊤ )(AΣA ⊤ ) -1 + q 2 tr (AΣA ⊤ ) -1 (x -Aµ)(x -Aµ) ⊤ ∂t (AΣA ⊤ )(AΣA ⊤ ) -1 (A ĊA ⊤ ) kl ∂k ∂l q = -tr(A ĊA ⊤ (AΣA ⊤ ) -1 )q + tr (x -Aµ) ⊤ (AΣA ⊤ ) -1 (A ĊA ⊤ )(AΣA ⊤ ) -1 (x -Aµ) qand q(x, t) is a solution whenEQUATIONThus, for q to be a solution in the absence of rescaling (A = 1) requires Σ = Ċ, and combining this with (30) gives the additional constraintEQUATIONFinally, note that, under the assumed form of p t (x) given in (5) , when C(t) increases without bound, q(x, t) → N (0, ACA ⊤ (t)) asymptotically (under rescaling), and this distribution is stationary when Σ(t) = AΣA ⊤ → ACA ⊤ is time-independent and a solution to (31) .Here, we derive the form of the pfODE (6) in rescaled coordinates. Starting from the unscaled inflationary process (Appendix A.1) with f = 0 and GG ⊤ (x, t) = Ċ(t), substituting into (4) gives the pfODEEQUATIONAs in Appendix A.2, we again consider the rescaling transformation x = A(t) • x, t = t. To simplify the derivation, we start by parameterizing the particle trajectory using a worldline time τ such that dt = dτ while A remains a function of t. With this convention, the pfODE becomesEQUATIONEQUATIONEQUATIONTwo important things to note about this form: First, the score function ∇ x log p t (x) is calculated in the unscaled coordinates. In practice, this is the form we use when integrating the pfODE, though the transformation to the scaled coordinates is straightforward. Second, the rescaling has induced a second force due to the change of measure factor, and this force points inward toward the origin when A is a contraction. This overall attraction thus balances the repulsion from areas of high local density due to the negative score function, with the result that the asymptotic distribution is stabilized.More formally, recalling the comments at the conclusion of Appendix A.2, when C(t) grows without bound in (5), p t (x), the unscaled density, is asymptotically Gaussian with covariance C(t), and its rescaled form q(x, t) is a stationary solution of the corresponding rescaled Fokker-Planck Equation.In this case, we also haveEQUATIONwhere we have made use of (31) with Σ → C. That is, when the rescaling and flow are chosen such that the (rescaled) diffusion PDE has a stationary Gaussian solution, points on the (rescaled) flow ODE eventually stop moving.Here, we show that our pfODE in ( 6) is equivalent to the form proposed by [49] for isotropic C(t) and A(t). We begin by taking equation (6) and rewriting it such that our score term is computed with respect to the rescaled variable x:EQUATIONwhere we have made use of the transformation properties of the score function under the rescaling.If we then choose C(t) = c 2 (t)1 and A(t) = α(t)1 (i.e., isotropic noising and scaling schedules), this becomesEQUATIONwhere we have dropped tildes on x and t. This is exactly the same as the form given in Equation 4of [49] if we substitute α(t) → s(t), c(t) → σ(t).Here, we show the equivalence of our proposed un-scaled (32) and scaled (37) pfODEs to the unscaled and scaled ODEs obtained using the "Gaussian paths" flow matching formulation from [36] .Here, we will use the convention of the flow-matching literature in which t = 0 corresponds to the easily sampled distribution (e.g., Gaussian), while t = 1 corresponds to the target (data) distribution.In this setup, the flow x t = ψ t (x 0 ) is likewise specified by an ODE:EQUATIONwhere again, x 1 is a point in the data distribution and x 0 ∼ N (0, 1). In [36] , the authors show that choosingEQUATIONwith "dots" denoting time derivatives leads to a flowEQUATIONthat is, a conditionally linear transformation of the Gaussian sample x 0 .For our purposes, we can re-derive (42) for the general case where σ t (x 1 ) is no longer a scalar but a matrix-valued function of x 1 and time. That is, we rewrite (43) (equation 11 in [36] ) with a full covariance matrix Σ t (x 1 ):EQUATIONSimilarly, we can writeEQUATIONfrom which it is straightforward to show that ( 41) is again satisfied.This can be related to our pfODE (6) as follows: First, recall that, under the inflationary assumption (5) plus rescaling, our time-dependent conditional marginals areEQUATIONwhich is equivalent to (44) with µ t (x 1 ) = A t • x 1 , Σ t (x 1 ) = A t C t A ⊤ t .Note that, here again, we have reversed our time conventions from the main paper to follow the flow-matching literature: t = 0 is our inflated Gaussian and t = 1 is the data distribution. From these results, along with the constraint (31) required for inflationary flows to produce a stationary Gaussian solution asymptotically, we then have, substituting into (45):EQUATION)EQUATIONwhich is the pfODE (6) written in the rescaled form (39) . Thus, our inflationary flows are equivalent to a Gaussian paths flow matching approach for a particular choice of (matrix-valued) noise schedule and mean.Here, for simplicity of notation, denote the participation ratio (7) by R(Σ) and let Σ = diag(γ) in its eigenbasis, so thatEQUATION)and the change in PR under a change in covariance is given byEQUATION)Requiring that PR be preserved (dR = 0) then gives (8) .Now, we would like to consider conditions under which PR is not preserved (i.e., (8) does not hold). Assume we are given γ(t) (along with initial conditions γ(0)) and defineEQUATIONso thatEQUATIONby definition. Then we can rewrite (8) asEQUATIONIn the cases we consider, flows are expansive (d(log Tr(C 2 )) > 0), with the result that (56) drives R(γ) toward R(t). Thus, in cases where R(t) has an asymptotic value, the R(γ) should approach this value as well. In particular, for our dimension-reducing flows, we have γ = ρg ⊙ γ, givingEQUATIONwhere i = 1 . . . K are the dimensions with g i = g * and γ k (0) = γ 0k . That is, the asymptotic value of R(t) (and thus the asymptotic value of PR) is that of the covariance in which only the eigendimensions with g k = g * have been retained, as in (10) .Following an extensive set of experiments, the authors of [49] propose a set of preconditioning factors for improving the efficiency of denoiser training (11) that forms the core of score estimation. More specifically, they parameterize the denoiser network D θ (x; σ) asEQUATIONwhere F θ is the actual neural network being trained and c in , c out , c skip , and c noise are preconditioning factors. Using this parameterization of D θ (x; σ), they then re-write the original L 2 de-noising loss asEQUATION)where w(σ) is also a preconditioning factor, y is the original data sample, n is a noise sample and x = y + n. As detailed in [49] , these "factors" stabilize DBM training by:1. c in : Scaling inputs to unit variance across all dimensions, and for all noise/perturbation levels. This is essential for stable neural net training via gradient descent. 2. c out : Scaling the effective network output to unit variance across dimensions. 3. c skip : Compensating for c out , thus ensuring network errors are minimally amplified. The authors of [49] point out that this factor allows the network to choose whether to predict the target, its residual, or some value between the two. 4. w(σ): Uniformizing the weight given to different noise levels in the total loss. 5. c noise : Determining how noise levels should be sampled during training so that the trained network efficiently covers different noise levels. This is the conditioning noise input fed to the network along with the perturbed data. This quantity is determined empirically.In [49] , the authors propose optimal forms for all of these quantities based on these plausible first principles (cf. Table 1 and Appendix B.6 of that work). However, the forms proposed there rely strongly on the assumption that the noise schedule is isotropic, which does not hold for our inflationary schedules, which are diagonal but not proportional to the identity. Here, we derive analogous expressions for our setting.As in the text, assume we work in the eigenbasis of the initial data distribution Σ 0 and let C(t) = diag(γ(t)) be the noising schedule, such that the data covariance at time t is Σ(t) = Σ 0 + C(t). Assuming a noise-dependent weighting factor Λ(t) analagous to w(σ) above, we then rewrite (11) asEQUATIONThis clearly generalizes ( 59) by promoting all preconditioning factors either to matrices (C in , C out , C skip , Λ) or vectors (c noise ). We now derive forms for each of these preconditioning factors.The goal is to choose C in such that its application to the noised input y + n has unit covariance:EQUATIONMore explicitly, if W is the matrix whose columns are the eigenvectors of Σ 0 , thenEQUATIONwhere the square root is taken elementwise.B.1.2 C out , C skipWe begin by imposing the requirement that the target for the neural network F should have identity covariance:EQUATIONEQUATIONThis generalizes Equation 123 in Appendix B.6 of [49] .Again by analogy with [49] , we choose C skip to minimize the left-hand side of ( 70):EQUATIONEQUATIONwhich corresponds to Equation 131 in Appendix B.6 of [49] .Using ( 73) in ( 70) then allows us to solve for C out :EQUATIONEQUATIONOur goal in choosing Λ(t) is to equalize the loss across different noise levels (which correspond, via the noise schedule, to different times). Looking at the form of (62), we can see that this will be satisfied when Λ(t) is chosen to cancel the outermost factor of C outEQUATIONUsing these results, we now rewrite (62) using the preconditioning factors derived above:L(D θ ) = E t,y,n ∥Λ(t)C out F θ (C in (y + n); c noise ) -C -1 out (y -C skip (y + n)) ∥ 2 = E t,y,n ∥F θ (Σ -1 2 (t) • (y + n); c noise ) -Σ -1 0 + C -1 (t) 1 2 • (y -Σ 0 Σ -1 (t) • (y + n))∥ 2 .In practice, we precompute W and σ 2 0 via SVD and compute all relevant precoditioners in eigenspace using the forms given above. For c noise , we follow the same noise conditionining scheme used in the DDPM model [27] , sampling t uniformly from some interval t ∼ U[t min , t max ] and then setting c noise = (M -1)t, for some scalar hyperparameter M . We choose M = 1000, in agreement with [49, 27] . After this, as indicated above, our noise is sampled via n ∼ N (0, C(t)) with C(t) = Wdiag(γ(t))W ⊤ .As highlighted in main text, for constant end integration time T and ρ, the final scale ratio between preserved and compressed dimensions is dictated by the quantity g * -g i , which we called the inflation gap (IG). Higher inflation gaps (IGs) lead to more stringent exponential shrinkage towards zero in compressed dimensions (Tables 6, 11 ) and worse off generative performance (Table 2 ).In PR-Reducing experiments, we set ρ = 1 and constructed g by making all elements of g corresponding to preserved dimensions equal to 2 (i.e., g preserved = g max = 2) and all elements corresponding to compressed dimensions equal to g compressed = g min = g preserved -IG (Tables 5, 10 ). Of note, for PR-Preserving experiments, all elements of g are set to 1 (i.e., g = 1, IG = 0) and we chose ρ = 2, such that all dimensions are inflated to the same extent and we match exponential constant used for preserved dimensions in PR-Reducing experiments.Here we rewrite the pfODE (6) EQUATIONwe rewrite (6) asEQUATIONwhere in the last line we have expressed A(t) and ĊC -1 in their respective eigenspace (diagonal) representations, where the divisions are to be understood element-wise. For PR-Reducing schedules, this expression simplifies even further, since our scaling schedule becomes isotropic -i.e., A(t) = α(t)1.To integrate (83), we utilize either Euler's method for toy datasets and Heun's method (see Algorithm 1) for high-dimensional image datasets. The latter has been shown to provide better tradeoffs between number of neural function evaluations (NFEs) and image quality as assessed through FID scores in larger data sets [49] .In toy data examples, we chose a simple, linearly spaced (step size h = 10 -2 ) discretization scheme, integrating from t = 0 to t = t max when inflating and reversing these endpoints when generating data from the latent space. For higher-dimensional image datasets (CIFAR-10, AFHQv2), we instead discretized using t i = i N -1 (t maxϵ s ) + ϵ s when inflating, where t max is again the maximum time at which networks were trained to denoise and ϵ s = 10 -2 , similar to the standard discretization scheme for VP-ODEs [49, 30] (though we do not necessarily enforce t max = 1). When generating from latent space, this discretization is preserved but integration is performed in reverse.Toy models were trained using a smaller convolutional UNet architecture (ToyConvUNet) and our proposed preconditioning factors (Appendix B.1). For all toy datasets, we trained networks both Algorithm 1 Eigen-Basis pfODE Simulation using Heun's 2 nd order method 1: procedure HEUNSAMPLER(D θ (x, t), γ(t), α(t), W ⊤ , t i∈{0,...,N } ) 2:if running "generation" then ▷ Generate initial sample at t 0 3:x0 ∼ N (0, diag(α(t 0 ) ⊙ γ(t 0 ))) ▷ Sample from Gaussian latent space 4:else ▷ i.e., if running "inflation" 5:x 0 ∼ p data (x) ▷ Sample from target distribution 6: x0 = α(t 0 )(W ⊤ • x 0 ) ▷ Transformfor i ∈ {0, 1, ..., N -1} do: ▷ Solve equation (83) N times 9: di ← -1 2 α(t i ) ⊙ γ(ti) γ(ti) ⊙ D xi α(ti) , t i -xi α(ti)10:+ α(ti) α(ti) ⊙ xi ▷ Evaluate dx dt at t i 11: xi+1 ← xi + (t i+1 -t i ) di , x i+1 = xi+1 α(ti+1)▷ Take Euler step from t i to t i+112:d′ i ← -1 2 α(t i+1 ) ⊙ γ(ti+1) γ(ti+1) ⊙ D xi+1 α(ti+1) , t i+1 -xi+1 α(ti+1)13:+ α(ti+1) α(ti+1) ⊙ xi+1 ▷ Evaluate dx dt at t i+1 14: xi+1 ← xi + (t i+1 -t i ) 1 2 di + 1 2 d′ i ▷ Apply trapezoidal rule at t i+1 15:return xN ▷ Return Sample 16: end procedure by using original images as inputs (i.e., "image space basis") or by first transforming images to their PCA representation (i.e., "eigenbasis"). Networks trained using either base choice were able to produce qualitatively good generated samples, across all datasets. For all cases, we used a learning rate of 10 -5 , batch size of 8192, and exponential moving average half-life of 50 × 10 4 . For PR-Reducing schedules, we set ρ = 1 and constructed g as described in Appendix B.2 (Table 5 ). The only exceptions were networks used on mesh and HMC toy experiments (Appendices C.2.1, B.7), where we used instead g preserved = 1.15 across all preserved dimensions (circles, S-curve) and g compressed = 0.85 (circles), or g compressed = 0.70 (S-curve) -Table 5 , 2 nd and 6 th rows. This yields a softer effective compression (i.e., smaller IGs) and is needed to avoid numerical instability in these experiments.As explained in Appendix B.1, to construct our c noise preconditioning factor, we sampled t ∼ U(t min , t max ), with t min = 10 -7 across all simulations and t max equal to the values shown in Table 4 . In the same table, we also show training duration (in 10 6 images (Mimgs), as in [49] ), along with both the total number of dimensions (in the original data) and the number of dimensions preserved (in latent space) for each dataset and schedule combination. In Table 6 , we showcase latent space (i.e., end of "inflation") compressed dimension variances achieved for the different toy PR-Reducing experiments as a function of inflation gap (IG). As expected, higher IGs lead to more stringent shrinkage of compressed dimensions in latent space.For our image datasets (i.e., CIFAR-10 and AFHQv2), we utilized similar training hyperparameters to the ones proposed by [49] for the CIFAR-10 dataset, across all schedules explored (Table 7 ). 9 ). Here, training duration was determined for each schedule based on when computed Frechet Inception Distance (FID) scores [65] stopped improving.We also showcase in Table 10 the specific values used for elements of g corresponding to preserved vs. compressed dimensions at different inflation gaps.All networks were trained on the same DDPM++ architecture, as implemented in [49] and using our proposed preconditioning scheme and factors in the standard (e.g., image space) basis. No gradient clipping or mixed-precision training were used, and all networks were trained to perform unconditional generation. We run training in the image space basis (as opposed to in eigenbasis) because this option proved to be more stable in practice for non-toy datasets. Additionally, we estimate the eigendecomposition of the target datasets before training begins using 50K samples for CIFAR-10 and 15K samples for AFHQv2. Based on our experiments, any sample size above total number of dimensions works well for estimating the desired eigenbasis. Times utilized to construct conditioning noise inputs to networks (c noise (t)) were uniformly sampled (t ∼ U(t min , t max )), with t min = 10 -7 and t max = 15.01, across all experiments. For the AFHQv2 dataset, we chose to adopt a 32x32 resolution (instead of 64x64 as in [49] ) due to constraints on training time and GPU availability. Therefore, for our experiments, both datasets have a total of 3072 (i.e., 3x32x32) dimensions.Finally, training was performed in a distributed fashion using either 8 or 4 GPUs per each experiment (NVIDIA GeForce GTX TITAN X, RTX 2080) in a compute cluster setting. Generation (FID) and round-trip (MSE) experiments were performed on single GPU (NVIDIA RTX 3090, 4090, A5000, A6000). We report training duration in Mimgs and note that time needed to achieve 200Mimgs is approximately 2 days on 8GPUs (4 days on 4 GPUs) using hyperparameters shown in Tables 7, 8 , 9. This is in agreement with previous train times reported in [49] using an 8 GPU distributed training set up. For image datasets (CIFAR-10 and AFHQv2), we simulated full round-trips: integrating the pfODEs (6) forward in time to map original images into latent space and then backwards in time to reconstruct original samples. We run these round-trips for a set of 10K randomly sampled images, three times per each schedule investigated and compute pixel mean squared error between original and reconstructed images, averaged across the 10K samples. Values reported in Tables 1, 2 represent mean ±2 standard deviations of pixel MSE between these three different random seeds per each condition. For pfODE integration, we used the discretization schedule and Heun solver detailed above (Appendix B.3.2), with t max = 15.01, ϵ s = 10 -2 , and N = 118 for all conditions.For image datasets, we also computed Frechet Inception Distance (FID) scores [65] across 3 independent sets of 50K random samples, per each schedule investigated. Values reported in Tables 1, 2 represent mean ±2 standard deviations across these 3 sets of random samples per each condition. Here again, we used the discretization scheme and solver described in Appendix B.3.2 with t max = 15.01, ϵ s = 10 -2 , and N = 256 across all conditions. We chose N = 256 here (instead of 118) because this provided some reasonable trade-off between improving FID scores and reducing total compute time.To obtain our latent space random samples x(T ) at time t 0 = T (i.e, at the start of generation) we sample from a diagonal multivariate normal with either (1) all diagonal elements being 1 (for PR-Preserving schedule) or (2) all elements corresponding to preserved dimensions being 1 and all elements corresponding to compressed dimensions being equal to the same small value for a given inflation gap (see Table 11 ). For our CIFAR-10 comparison experiments against existing injective flow models, we used the same implementations for M-Flow [21] , Rectangular Flows [22] , and Canonical Manifold Flows [23] as in [23] . When training the comparison injective flows, we used the same hyper-parameters proposed in Appendix G.1 of [23] for the CIFAR-10 dataset. The only difference here is that we trained models with latent dimensions equal to d = [30, 40, 62] . Finally, comparison FID scores reported in Table 3 represent best score out of 3 independently generated sets, each with 10K samples. For our comparison models, inflation gap was fixed to IG = 1.02 while d was varied between 30, 40, and 62 and we utilized the same training hyper-parameters reported in Tables 7, 8 . FID scores were computed using the same discretization, solver, and general set up described above. Figure 8 : Traceplots (post-thinning) for 3 random chains for PR-Preserving and PR-Reducing schedules. A: Traceplots for 3 random PR-Preserving chains, after thinning by a factor of 5. "X axis" represents sample number and "Y axis" represents value of zeroth dimension of sample ( ŵ0 ).B: Same set up, only for 3 random PR-Reducing chains. Note that there is still some considerable correlation in the samples, even after thinning. Additionally, mixing is not particularly good. circles and 3D S-curve), Figure 10 . Here, we began by sampling 20K test points from a Gaussian latent space with appropriate diagonal covariance. For PR-Preserving schedules, this is simply a standard multivariate normal with either 2 or 3 dimensions. For PR-Reducing experiments, this diagonal covariance matrix contains 1's for dimensions being preserved and a smaller value (10 -2 for Circles, 2.5 × 10 -3 for S-curve) for dimensions being compressed.Next, we sampled uniformly from the surfaces of balls centered at zero and with linearly spaced Mahalanobis radii ranging from 0.5 to 3.5 (200 pts per ball). We then fit either a 2D alpha-shape (2D Circles) or a mesh (3D SCurve) to each one of these sets of points. These points thus represent "boundaries" that we use to assess coverage prior to and after integrating our pfODEs. We define the initial coverage of the boundary to be the set of points (out of the original 20K test points) that lie inside the boundary. We then integrate the pfODE backward in time (the "generation" direction)Alternative ScalingAlternative ScalingFigure 12 : Toy 3D → 2D dimension-reducing experiments with alternative scalings. Shown here are simulations of our 3D → 2D PR-Reducing pfODEs for 3D toy datasets (S-curve, Swirl) scaled either to unit variance across all 3 dimensions (first and third rows) or scaling the thickness dimension to 0.5, while leaving other dimensions scaled to 1 (second and fourth rows). Note that scaling all dimensions to 1 leads to some loss in original shape content when running generation (first and third rows, rightmost column). This is not the case when we make total variance contribution of the "thickness" dimension smaller (i.e., under the alternative scaling; second and fourth rows, rightmost column).z = sign(y)y 2 , which creates a curved (chair-like) shape in 3D (Figure 11 , bottom 2 rows). We then simulated our PR-Preserving pfODE for both embedding procedures and considering both the case in which no noise was added to the data or, alternatively, where some Gaussian noise is added to the initial distribution, giving it a small thickness. We used zero-mean Gaussian noise with σ of 0.2 and 0.5 for embedding types (1) and ( 2), respectively.As shown in Figure 11 , when no noise is added, our PR-Preserving pfODEs Gaussianize the original data points along the manifold plane (rows 1 and 3, rightmost columns). Alternatively, when noise is added and the manifold plane has some "thickness" the inflationary flows map original data into a lower-rank Gaussian (rows 3 and 4, rightmost columns). In both cases, the original PR is preserved (up to some small numerical error), as expected.For our 3D toy data PR-Reducing experiments, we tested how changing the relative scaling of different dimensions in the original datasets qualitatively changes generative performance. For the first experiment, we scaled all dimensions to variance 1 (Figure 12 , first and third rows). In this case, all dimensions contribute equally to total variance in the data. In contrast, for the second experiment (Figure 12 , second and fourth rows), we scaled the thickness dimension to variance 0.5 and all others to 1. In this case, the non-thickness dimensions together account for most of the total variance.We then trained neural networks on 3D S-curve and Swirl data constructed using these two different scaling choices and used these networks to simulate our PR-Reducing pfODEs (reduction from 3D → 2D) both forwards (Figure 12 left panels) and backwards (Figure 12 right panels) in time.Of note, the first scaling choice leads to generated samples that seem to loose some of the original shape content of the target dataset (first and third rows, rightmost columns). In contrast, scaling choice 2 is able to almost perfectly recover the original shapes (second and fourth rows, rightmost columns). This is because scaling the thickness dimension to 0.5 reduces the percent of total variance explained along that axis, and our PR reduction preferentially compresses in that direction, preserving most information orthogonal to it. By contrast, the first scaling choice spreads variance equally across all dimensions and, therefore, shape and thickness content of target distribution are more evenly mixed among different eigendimensions. As a result, compressing the last dimension in this case inevitably leads to loss of both shape and thickness content, as observed here.In Section 5 above, we considered the possibility that numerical errors in approximating the score function might result in errors in pfODE integration and thus miscalibration of our proposed inference procedure. There, we argued that if these score estimation errors can be modeled as white noise, integration using sufficiently small integration step sizes will maintain accuracy, as dictated by theorems on numerical integration of SDEs [59] . Here, we investigate the validity of this approximation for our trained score functions.As detailed in Appendices B.1 and B.3.1, we did not directly estimate scores but trained networks to estimate a denoiser ŷ = D θ (x, C(t)), where y are samples from the data and x = y + n are the noised samples with n ∼ N (0, C(t)). In this case, one can then compute scores for the noised distributions using:EQUATIONIn practice, however, this de-noised estimate contains some error ϵ = ŷy, which is the true residual error in our network estimates. Therefore, we rewrite our score expression as:EQUATIONwhere (ŷx) can be understood as the magnitude of the correction made by the denoiser at x [51] .Note that ϵ = 0 for the ideal denoiser (based on the true score function), but nonzero ϵ will result in errors in our pfODE integration.As argued above, these errors can be mitigated if they are uncorrelated across the data set, but this need not be true. To assess this in practice, we extracted estimation errors ϵ(x) across a large number of data samples (10K for 2D circles toys, 50K for CIFAR-10) and for networks trained on both PR-Preserving and select PR-Reducing schedules (PR-Reducing to 1D for circles at IG=2.0, and to 62D, 307D, and 615D for CIFAR-10, all at IG=1.02) and then computed cross-correlations for these errors along integration trajectories x(t):EQUATION)where ε is the mean residual across the entire data set. In practice, we use scaled correlations in which an entry R ij is normalized by σ i σ j the (zero-lag) variance of the residuals along the corresponding dimensions.Results of these calculations are plotted in Figure 13 , for the mean across diagonal elements of R.As the left panel of Figure 13 shows, residuals display negligible autocorrelation for networks trained to denoise toy data sets, while for CIFAR-10 (right panel), there is some cross-correlation at small time lags. This is likely due to the increased complexity of the denoising problem posed by a larger data set of natural images, in addition to the limited approximation capacity of the trained network. As a result, points nearby in data space make correlated denoising errors. Nevertheless, this small amount of autocorrelation does not seem to impact the accuracy of our round-trip experiments nor our ability to produce good-quality generated samples (Figures 5, 6 ; Table 1 ).Toy datasets were obtained from scikit-learn [93] and were de-meaned and standardized to unit variance prior to training models and running simulations. The only exceptions to this are the alternative 3D toy datasets detailed in Appendix C.2.3, where the third dimension was scaled to slightly smaller variance.For CIFAR-10 and AFHQv2 datasets, we apply the same preprocessing steps and use the same augmentation settings as those proposed for CIFAR-10 in [49] (cf. Appendix F.2), with the only change that we downsample the original AFHQv2 data to 32 × 32 instead of 64 × 64.Datasets:• CIFAR- Justification: As highlighted in Section 7, one of the main limitations of the proposed method lies in our choice of Participation Ratio (PR) as our dimensionality measure. This measure favors top principal components of the data when doing compression. Utilizing different (more complex) dimensionality metrics and noise and scaling schedules might yield pfODEs with more interesting compressive behavior and properties. We also note the need to train DBMs over much larger noise ranges than at present as a key limitation. Guidelines:• The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper. • The authors are encouraged to create a separate "Limitations" section in their paper.• The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be. • The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated. • The authors should reflect on the factors that influence the performance of the approach.For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon. • The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size. • If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness. • While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations. 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)? Answer: [NA] Justification: Although we proposed a new class of generative models, work presented here does not constitute a high risk for misuse (we do not release our pre-trained image generation models). We do not use scraped datasets. Guidelines:• The answer NA means that the paper poses no such risks.• Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters. • Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images. • We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort. • We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution. • For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.38th Conference on Neural Information Processing Systems (NeurIPS 2024). 
 
 Appendix:  


Decision: This paper presents a new class of generative models and associated inference algorithm that maps high-dimensional data to a lower-dimensional Gaussian via ODE integration. The main claimed contribution of the paper is that the method is well calibrated and identifiable, i.e., is uniquely defined on the low-dimensional space. 

Overall, the reviewers agree that the proposed approach is novel but raised several important concerns including: (1) Lack of comparisons with other approaches (including injective normalizing flows); (2) clarity regarding evaluation of calibrated uncertainties; and (3)  somewhat unintuitive results when more dimensions are retained (as the metrics get worse). The authors have, to some extent, addressed these issues satisfactorily in their comments by clarifying where their approach lies (unconditional generative modeling and learning structured representations) and provided additional results. However, the comparison with injective flows was not done due to time constraints, so I’d expect to see this in the final version of the paper. The authors have also clarified the relation of their work to flow matching by stating that it is essentially a “novel, dimension-reduction form of flow-matching”. 

One additional point I’d like to see discussed in the final version is concerned with calibration. Usually in the literature this refers to predictive calibration. However, how well the approximate posterior matches the true posterior is simply an evaluation of the posterior approximation method. While posterior comparisons across different models does not make sense, one can still evaluate (predictive) calibration, which the authors have done. I also agree that uncovering structured latent representations is an important problem, but in my view this has more to do with the underlying model than with the approximation method (while I recognise these two may be conflated sometimes). I think more clarity in the terminology and evaluation should be provided. 

 Given the positive comments and the recognition of the novelty of this work, I recommend acceptance.

Reviews: review: paper_summary: paper_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.
strength_weakness: strength_weakness: 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: 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.
limitations: limitations: The authors adequately discussed the limitations.
paper_summary: paper_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.
strength_weakness: strength_weakness: * 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: 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?
limitations: 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.
paper_summary: paper_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.
strength_weakness: strength_weakness: - 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: 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?
limitations: limitations: The authors openly discuss limitations of their method and future work.
paper_summary: paper_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.
strength_weakness: strength_weakness: ### 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: 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?
limitations: limitations: Limitations are properly addressed.

Comments: title: Reviewer Response
comment: We appreciate the reviewer's suggestions and are working to implement these additional comparisons.
We appreciate the reviewer providing these references. We did find several of these but missed some others. We are likewise eager to see the outcome of the injective flows comparisons.
Thanks again for your helpful comments. We are currently running these experiments and will report back if they are completed before the end of the discussion period.
Thank you for all the clarifications. In light of those, I increase my score to 5. I look forward to reading the updated version with improved clarity and additional results regarding model comparisons mentioned in the general rebuttal!
Thank you for your response. Without the additional experiments investigating the full spectrum of behavior as $d$ varies from $1$ to $D$, I will have to maintain my original score with the hopes that when this paper is accepted (either here or at another venue) that these experiments are included to tell the full store.
I thank the authors for their professional response and the clarifications. I also noticed that my references went missing in my original review, so I re-post them below. Having also read the other reviews, I will increase my score and would really like to see the comparisons with non-bijective flows.


[1] Draxler, F., Sorrenson, P., Zimmermann, L., Rousselot, A., & Köthe, U. (2024). Free-form flows: Make any architecture a normalizing flow. In International Conference on Artificial Intelligence and Statistics (pp. 2197-2205). PMLR.

[2] Sorrenson, P., Draxler, F., Rousselot, A., Hummerich, S., Zimmermann, L., & Köthe, U. (2024). Lifting architectural constraints of injective flows. In The Twelfth International Conference on Learning Representations.

[3] Kothari, K., Khorashadizadeh, A., de Hoop, M., & Dokmanić, I. (2021, December). Trumpets: Injective flows for inference and inverse problems. In Uncertainty in Artificial Intelligence (pp. 1269-1278). PMLR.

[4] Cranmer, K., Brehmer, J., & Louppe, G. (2020). The frontier of simulation-based inference. Proceedings of the National Academy of Sciences, 117(48), 30055-30062.

[5] Sharrock, L., Simons, J., Liu, S., & Beaumont, M. (2022). Sequential neural score estimation: Likelihood-free inference with conditional score based diffusion models. arXiv preprint arXiv:2210.04872.

[6] Schmitt, M., Pratz, V., Köthe, U., Bürkner, P. C., & Radev, S. T. (2023). Consistency Models for Scalable and Fast Simulation-Based Inference. arXiv preprint arXiv:2312.05440.