Model-Informed Flows for Bayesian Inference

Thank you for your thorough review and valuable feedback on our manuscript. Below, we address each of your comments in turn. If anything remains unclear, we would be grateful for any further questions or clarifications.

Before addressing the individual reviews, we wish to clarify that MIF is not restricted to—nor does it rely on—full-rank Gaussian variational inference. Theorem 4 shows that MIF can recover the full-rank-Gaussian + VIP transformation, but Algorithm 1 makes clear that MIF itself is an autoregressive flow model, independent of any full-rank Gaussian assumption.

(Clarity 1) The notion of “prior information” needs further explanation. Does it mean that the parameters are informed by previous variables in the specified hierarchical model, or that the hierarchical model itself already includes neural network (NN) components? In the context of probabilistic programs, the model can be user-specified. Is the proposed approach applicable to generic hierarchical models?

By prior information we mean only the mean $\mu_i$ and scale $\sigma_i$ that define each conditional prior $p\left(z_i \mid z_{<i}\right)$. Our algorithm needs nothing else, so it works for any hierarchical model that can expose these two functions-whether the user writes them directly or we auto-extract them from frameworks such as PyMC or NumPyro.

We are willing to add text like below to the paper for clarity:

Prior information. For each latent variable $z_i$ we require two callables $\mu_i(z_{<i})$ and $\sigma_i(z_{<i})>0$ that output the mean and scale of its conditional prior $p(z_i \mid z_{<i}) = \mathcal N(\mu_i,\sigma_i^2)$. These are the only model‑specific inputs our flow uses.

LatentSpec = {
    "name": "z_i",                  # variable name
    "parents": ["z_j", ...],        # z_<i>
    "mean_fn": lambda parents: ..., # returns μ_i
    "scale_fn": lambda parents: ... # returns σ_i
}
model_spec = [LatentSpec, ...]      # one per latent

The callables can be hand‑written or auto‑parsed from frameworks such as PyMC or NumPyro; any hierarchical model that yields these $(\mu_i,\sigma_i)$ pairs is compatible with our Model‑Informed Flow.

(Clarity 2) It seems that MIF serves both as a way of specifying hierarchical priors and as a method for flexible variational posterior, but the topological order of variables may differ. Have you tried reversing the topological order (e.g., bottom-up vs top-down)?

Yes. We provide this experiment in Sec. 4.3: the row “MIF (w/o order)” in Table 2 trains the flow with the reverse (bottom‑up) ordering of latent variables. Performance drops on every dataset, confirming that matching the model’s natural (top‑down) order is important for capturing conditional structure efficiently.

(Significance 1 & Question 1)) The experiments are mostly limited to a few simple models and datasets, and the evaluation relies solely on ELBO. Could the authors evaluate the method’s ability to capture complex dependencies in the context of variational inference, and using a density estimation benchmark on, for example, the two-moons dataset? It would also be helpful to comment on the applicability of MIF to generic variational inference in probability programming, likelihood-based density estimation and sample generation. The experimental evaluation is somewhat limited. Is it possible to demonstrate that MIF offers general benefits, such as improving variational autoencoders, by evaluating on binarized MNIST (e.g., reporting marginal log-likelihood log p(x) using importance sampling)? Table 1 and Table 2 essentially use the same set of datasets, although the ablation is different.

Thank you for the suggestion. We’d like to clarify that MIF is not designed as a general-purpose density estimator but as a model-informed variational inference method for probabilistic programs. It cannot be applied to density estimation problems, as there is no prior or posterior.

In particular:

1. Different training objective. MIF is trained by maximizing$\mathbb{E}_{q(z)}[\log p(x,z) - \log q(z)],$ which requires an explicit latent-variable model $p(x,z)$. Pure density-estimation tasks (e.g., two-moons, binarized MNIST) optimize $\log p(x)$ directly and lack this programmatic structure.

2. Model-informed architecture. Our shift/scale/translation networks consume “prior inputs” $f_i, g_i$ derived from the known generative model. Without a specified probabilistic program, those modules have no signal to condition on, and MIF’s inductive bias disappears.

3. Workload scale. Probabilistic inference is inherently more expensive than density estimation, and state-of-the-art PPL benchmarks typically involve only tens to low-hundreds of latent dimensions. That is why we focused on six canonical hierarchical models (8 Schools, Credit, Funnel, Radon, Movielens, IRT) in Table 1 and 2 and 3 more models (Seeds, Sonar, Ionosphere) in Table 3, which capture challenging posterior geometries yet reflect real-world PPL use cases.

While it would be interesting in future work to adapt MIF for pure density estimation (e.g., by re-casting it as a normalizing flow or integrating it into a VAE encoder), those extensions lie outside our current scope. We will clarify these points in the revised manuscript.

(Question 2) With a limited-capacity conditioning network (e.g., a one-layer affine transformation), does MIF capture posterior dependencies better than alternative flow-based approaches such as coupling flows, MAF, or IAF?

Yes. Even with a minimal conditioning network (a single affine layer), MIF still captures posterior dependencies more effectively than other methods. Tables 1 and 2 demonstrate that this low-capacity MIF matches—or exceeds—the accuracy of those alternatives, which require deeper, higher-capacity networks to reach similar performance.

(Qeustion 3) Equation 2.2 and 2.3, 2.8, the likelihood term does not depend on z_i and the joint density factorizes between z and x? The way of writing it is confusing. The parent variables pi(x_j) are the z s?

In equation 2.3, the likelihood term would be $\prod_{j=1} p\left(x_j \mid \pi\left(x_j\right)\right)$. As the reviewer suggests, here the set of parent variable $\pi\left(x_j\right)$ would include the latent variables $z_j$. Our reason for writing the model this way is that we wish to consider models where arbitrary variables could be observed. So, latent variables might be conditioned on both observed and latent variables, and similarly for observed variables. For this reason, we settled on our (admittedly subtle) notation where the parent sets $\pi(z_i)$ and $\pi(x_i)$ include both latent and observed variables. We will clarify this in the revised version of the paper.

(Question 4) Line 135: What is the quadratic form in Equation 3.4?

We believe the reviewer is referring to our comment about "the quadratic form in Equation (3.4)" on line 135. We apologize for this somewhat ambiguous phrasing. By “quadratic form” we mean the product of $g_i(\pi(z_i))$ and $f_i(\pi(z_i))$, which introduces a quadratic dependence on $z$ that an affine function of $z$ cannot capture; we will clarify this wording in the revision.

Thank you for your thorough review and valuable feedback on our manuscript. Below, we address each of your comments in turn. If anything remains unclear, we would be grateful for any further questions or clarifications.

(Weakness 1) The inherent sequential nature of MIF limits it's scalability and practicality. This has been recognised by the authors but it'd would been nice to see them hint towards resolutions for this major limitation.

Funnel geometries are notoriously difficult for any inference method, and while MIF's sequential updates do add some non-parallelisable cost, the overhead is modest in practice (see the runtime table below). Looking forward, we plan to exploit conditional independencies among latent variables-thereby reducing the $D$-dependence of forward autoregressiveness while preserving the MIF-derived benefits. For example, in a tree-structured hierarchy $p(\theta, z, x)$ with global latent $\theta$, all local factors $q(z|\theta)$ can be sampled in parallel. We leave a full exploration of these parallel schemes to future work.

(Weakness 2) The evaluations, though seemingly thorough, lack error bars or standard deviations. The numbers presented in table 1, for example, are not very conclusive as there are minute differences b/w MF-VIP and FR-VIP. Similarly I cannot ascertain the value of different components of MIF since all the numbers are very close to each other. This, in my opinion, is another major limitation.

Thank you for raising this point. We agree that quantifying variability is crucial, so we've added standard deviations to Table 1 in this rebuttal (based on five independent runs) and will include the full error bars in the final manuscript. As in the manuscript, for each run, we perform 100,000 optimization iterations using 256 Monte Carlo samples to approximate the ELBO during training, then evaluate the final ELBO with 100,000 fresh samples. We also sweep six learning rates (log-spaced from $10^{-1}$ to $10^{-6}$ ) and report the best result.

We will add these to the revised version of the paper.

(Weakness 3) For Figure 2, why have the authors not included any other variational approximations here? There are several complex and highly-expressive variational approximations that could've been considered here as baselines. I can read that they intended to compare this with IAF (Line 204) but this result is not present.

We apologize for the confusion but MIF with $\epsilon$-conditioning is in fact just an inverse autoregressive flow-by removing forward dependence on previous latents. We will make sure to clarify this in the revised version of the paper. The sole aim of Figure 2 is to isolate and compare $z$ -conditioning versus $\epsilon$-conditioning under varying network capacity. As expected, $z$ -conditioning ensures accurate sampling in funnel-like geometries even with simple (affine) nets, while $\epsilon$-conditioning only matches that performance once the nets become expressive enough to implicitly learn the same dependencies. We therefore restrict Figure 2 to these two closest algorithmic variants; broader comparisons against other highly expressive VI families are provided later (e.g., Table 3 and the review-paper benchmarks).

(Weakness 4 & Limitation 1) There are no runtime comparisons. Since IAF and FAF have difference in their ability to capture VIP but the one's faster than the other, the authors should have considered an accuracy to runtime comparison for these two methods. / The authors should have commented on the runtime performance. It is important that they elucidate to the reader the differences between the runtimes/latency of MIF vs IAF

Apologies for not including runtime comparisons earlier. As Algorithm 1 shows, vanilla MIF and its $\epsilon$-conditioning variant (i.e., IAF) require essentially the same total FLOPs, but vanilla MIF executes them sequentially-incurring an $O(D)$ wall-clock-whereas IAF can launch those operations in parallel and achieve $O(1)$ wall-clock per sample. This reflects an explicit trade-off: full MIF guarantees accurate sampling in funnel-like geometries even with shallow conditioning networks, while $\epsilon$-conditioning achieves constant-time complexity at the cost of needing higher-capacity nets to implicitly recover those dependencies. In response to your suggestion, we have now measured wall-clock runtimes for both variants on our benchmark models and will include a comprehensive runtime comparison in the final manuscript. Below is the runtime table where $D$ is the number of latent variable:

Unfortunately, we cannot provide figures as part of this discussion. However, for a full accuracy to runtime comparison, we will provide another version of Figure 2 with the same set of results, but with the x-axis measuring the runtime for each method. This will still show some benefits of using MIF.

(Minor 1) Why did the authors choose to label the y-axis in Figure 2 as 'Best Elbo'? The notion of 'best' is very ambiguous and inconsistent from the rest of the paper where the authors quote Negative ELBO.

Thank you for pointing out the inconsistency in our labeling. We used "BestELBO" to emphasize that, at each iteration, the curve shows the maximum ELBO achieved over our learning-rate sweep, rather than the negative ELBO that we minimize elsewhere. We will clarify this in the revised figure.

Thank you for your thorough review and valuable feedback on our manuscript. Below, we address each of your comments in turn. If anything remains unclear, we would be grateful for any further questions or clarifications.

The experiment section mostly contains comparisons to mean-field Gaussian VI, full-rank Gaussian VI, and VIPs. However, in the related work section, several other related algorithms are brought up (i.e. ASVI, Cascading Flows, EMF). However, these are not included in the comparison. Can the authors comment on, if any, the difference in performance that they expect relative to MIF?

Please refer to Table 3, this presents a broad set of comparisons: Gaussian Mean‑Field VI (MF, [1]), Gaussian Mixture Model VI (GMMVI, [2]), Sequential Monte Carlo (SMC, [3]), Annealed Flow Transport (AFT, [4]), Flow Annealed Importance‑Sampling Bootstrap (FAB, [5]), Continual Repeated AFT (CRAFT, [6]), and Uncorrected Hamiltonian Annealing (UHA, [7]). We chose these baselines because they are the most closely related to our approach and were identified as top performers in the benchmark study of [8]. While we believe this suite of experiments already demonstrates the significance of our algorithmic contributions, if the reviewer would like to see comparisons to any other particular algorithms, we would be happy to include these (even during the discussion period, if possible).

In Algorithm 1, we see that MIF involves training three neural networks. This could be quite computationally heavy. Can the authors comment on the relative computational cost between MIF and other competing methods?

In practice, adding the third conditioning network has minimal impact on runtime. It introduces only an $O(D)$-sized parameter block-negligible relative to the overall model-and the dominant cost remains evaluating the log-density, which all autoregressive flows share. As a result, per-iteration runtimes for MIF are effectively unchanged compared to two-network variants or other flow-based VI methods. Below, we present the wall-clock comparison between MIF with two versus three conditioning networks, demonstrating that the extra network adds negligible overhead while enhancing expressiveness.

We will add these measurements to the appendix of the paper.

The MIF structure seems rather flexible, and I wonder if it could be helpful for problems with other interesting structure/geometry (i.e. multi-modal posteriors). Can the authors comment on this a bit on what they think?

Thank you for the question. MIF is inherently flexible: it extends autoregressive flows with an architecture already proven to sample effectively in funnel-like geometries, and because autoregressive flows capture complex, multimodal posterior shapes we expect MIF to perform well on other challenging structures too. Looking ahead, we plan to study tree-structured hierarchical models where independent branches can be sampled in parallel, thereby reducing—or even removing—the sequential sampling cost and further speeding up inference.

[1] C. M. Bishop and N. M. Nasrabadi. Pattern recognition and machine learning, volume 4. Springer, 2006.

[2] O. Arenz, P. Dahlinger, Z. Ye, M. Volpp, and G. Neumann. A unified perspective on natural gradient variational inference with gaussian mixture models. arXiv preprint arXiv:2209.11533, 2022.

[3] P. Del Moral, A. Doucet, and A. Jasra. Sequential monte carlo samplers. Journal of the Royal Statistical Society Series B: Statistical Methodology, 68(3):411-436, 2006.

[4] M. Arbel, A. Matthews, and A. Doucet. Annealed flow transport monte carlo. In International Conference on Machine Learning, pages 318-330. PMLR, 2021.

[5] L. I. Midgley, V. Stimper, G. N. Simm, B. Schölkopf, and J. M. Hernández-Lobato. Flow annealed importance sampling bootstrap. arXiv preprint arXiv:2208.01893, 2022.

[6] A. Matthews, M. Arbel, D. J. Rezende, and A. Doucet. Continual repeated annealed flow transport monte carlo. In International Conference on Machine Learning, pages 15196-15219. PMLR, 2022.

[7] A. Thin, N. Kotelevskii, A. Doucet, A. Durmus, E. Moulines, and M. Panov. Monte carlo variational auto-encoders. In International Conference on Machine Learning, pages 1024710257. PMLR, 2021.

[8] D. Blessing, X. Jia, J. Esslinger, F. Vargas, and G. Neumann. Beyond elbos: A large-scale evaluation of variational methods for sampling. arXiv preprint arXiv:2406.07423, 2024.

Thank you for your thorough review and valuable feedback on our manuscript. Below, we address each of your comments in turn. If anything remains unclear, we would be grateful for any further questions or clarifications.

(Weakness 1) it’s unclear if the full MIF is currently a practical model family due to its sampling time scaling in variable dimensionality compared to the parallel IAF-like families

We have measured the computational cost compared to IAF, please see the following table below.

We note that MIF imposes little extra computational cost in terms of flops, but the sequential nature of sampling does impose some limitations on parallelization. However, as the above table shows, this is usually quite modest. Given the large improvements in accuracy, we believe there are many situations where it would be worth paying this cost.

Moreover, the primary goal of our work is to expose the previously unexplored relationship between flow‑based models and variational inferred parameters (VIP), and to introduce MIF as a family that captures all the benefits revealed by this connection. As discussed in Sections 4.4 and 5, there is a clear trade‑off between MIF and IAF‑like families:

• When network capacity is limited (e.g., few hidden units), MIF often outperforms IAF‑style flows, as shown in Figure 2.
• When ample capacity is available and fast sampling is a priority, the $\epsilon$‑conditioned variant of MIF (our IAF‑like family) is preferable.

Moreover, even if one opts for an IAF‑style architecture, it remains beneficial to incorporate the other key components from our VIP analysis into the variational family.

(Question 1) It seems from the formulation in Definition 3 that one could alternatively absorb the (scaled) translation term into the shift function $m$, by making $m$ dependent on $\epsilon$ in addition to $z$. Is there any reason why an explicit translation function is preferred?

An affine FAF (i.e. one where both m and log s are affine functions, see Definition 2) cannot capture the composite mapping $T = T_{\mathrm{VIP}} \circ T_A$ (Equation 3.3) without an explicit translation term (please see Section 3, just before Definition 3). Concretely, even if we allow the shift function $m$ to depend on both $\epsilon$ and $z$, it still cannot reproduce the quadratic dependence on functions of $z$ required by Equation 3.4. By introducing a separate translation function $t$, we can match that quadratic form exactly and thus express $T = T_{\mathrm{VIP}} \circ T_A$ as an affine generalized FAF, as proven in Theorem 4. Although it is possible to omit the translation term when there are no constraints on the shift and log‑scale functions, our experiments show that including it consistently improves performance.

(Question 2) In Figure 2, the Funnel plot seems to show that the performance of MIF and MIF ( $\epsilon$-cond) are similar for larger hidden sizes, but Table 2 shows a significant difference. Is this because Table 2 uses smaller hidden sizes?

Apologies for the confusion—Table 2 was indeed run with smaller hidden sizes than Figure 2, which explains the larger performance gap.

(Question 3) It's a little unclear to me if the conditioning on prior distributions is a novel contribution of this paper or if it has previously been used in other works on NFs for inference in probabilistic programs/hierarchical Bayesian models. Could the authors clarify this?

The idea of conditioning on prior distributions has indeed been used in the VIP paper, albeit only in concert with a mean-field Gaussian variational distribution (Section 2.4). As we stated above, the main goal of our work is to analyze the previously unexplored relationship between flow‑based models and VIP, which we found helps address sampling from posteriors with challenging curvature. Previous works on normalizing flows for inference have typically focused on creating more expressive variational families and improving efficiency. Our method is novel in that we explicitly incorporate not only prior information but also the other key components from our analysis to build a variational family that is proven to perform well on posteriors with funnel‑like geometries. We will expand the Related Work section in the final version of the manuscript to compare in detail with prior NF‑based inference methods and highlight our distinct contributions.

(Question 4) Could the authors provide a citation for the term "forward autoregressive flow"? It appears to be essentially the masked autoregressive flow of (Papamakarios et al. 2017).

Apologies for the omission. Indeed, by "forward autoregressive flow", we mean to refer to the work of Papamakarios et al. (2017). While they refer to this simply as an "autoregressive flow" we add the term "forward" to contrast more clearly with "inverse autoregressive flows". To make this explicit, we will add this clarification after Definition 2 in the revised version. Also, please note that our Definition 2 coincides with the autoregressive flow formulation given in Section 3.2 of Papamakarios et al. (2017).

G. Papamakarios, T. Pavlakou, and I. Murray. Masked autoregressive flow for density estimation. Advances in neural information processing systems, 30, 2017.