Paper: Exchangeable Dataset Amortization for Bayesian Posterior Inference (ICLR-2024)

Reviewer comment (Novelty):
The approach lacks novelty. Amortized inference has been studied extensively in the past, for example in the context of the variational autoencoder. Amortization on the dataset level is also not new and has been researched for years in the meta-learning community. The method is conceptually very similar to neural process (NP) approaches, with the difference being that inference is performed over the decoder parameters directly, resulting in no free parameters optimized for predictive performance. While this is an architectural difference, it does not require adaptations with respect to the posterior inference method, as both approaches simply optimize the ELBO with respect to the variational parameters.

Author response:
**Novelty** 

We would like to respectfully disagree with the reviewer on their point about novelty. While we agree that the methodologies used (eg. Variational Inference (VI), amortization on the dataset, transformers) do exist in current literature, it is important to note that its application for amortized Bayesian posterior estimation hadn’t been explored yet and this is precisely the void that we aimed to fill. It is true that in hindsight it might look like a simple application of existing concepts but this gap does indeed exist in literature and we feel that there is merit and contribution in addressing it since the class of parametric models with tractable likelihoods do span a significant volume of models used by practitioners. Thus, leveraging and combining existing methods to allow for faster and more efficient posterior estimation in such a class of models is important.

Additionally, as the reviewer points out, Neural Processes (NPs) are only applicable in predictive tasks and thus our proposed approach does tackle a different problem. In particular, our objective is not the same as NPs as we optimize the ELBO with respect to the variational parameters only (there are no other parameters in the likelihood model) whereas NPs optimize the ELBO with respect to both variational parameters as well as the parameters in the likelihood model. Though seemingly similar, it is important to note that the former is a direct application of the VI framework while the latter of Variational Expectation Maximization.

Finally, we would also like to point out that neither NPs nor Simulation-based Inference (SBI) approaches generally handle variable dimensional inputs in modeling the predictive tasks. Further, the two directions have mostly seen progress separately, without works comparing them under the common setting of probabilistic models with tractable likelihoods. Additionally, we also consider model mis-specification and OoD generalization based real world settings which have been relatively less explored in both SBIs and NPs.

Reviewer comment (Experiments):
The authors acknowledge similarities to NP-like methods but claim their method is new because NP methods are "predominantly designed for predictive modeling and thus cannot be used to provide useful information and uncertainty about model parameters." However, the paper largely focuses on posterior predictive performance. The only results studying the quality of the approximate posterior appear in Table 4 and Figures 4(c,d), which lack comparisons against non-amortized baselines. Additional experiments and baselines are needed to explore whether the method provides useful information and uncertainty about model parameters, and to assess the amortization gap introduced by the method.

Author response:
We want to thank the reviewer for the fruitful discussion and ideas mentioned. In particular, how to frame the presented novel method in the picture of other purposed methods (GP/NP/Meta-Learning) that utilize amortized inference for posterior prediction and parallel research streams like SBI that differ in the problem setup.

As mentioned by other reviewers, it is now brought to our attention that we need to focus on a more precise differentiation between existing methods and our proposed novel method. We want to thank the reviewer for highlighting methods from meta-learning and ideas for different aggregation schemes, which we will review carefully.

Furthermore, we thank the reviewer for clarifying the log marginal posterior predictive metric. In our setup, the likelihood is not learnable and only parameterizable by the posterior parameters $\theta$ sampled from the approximate posterior distribution $\theta_s \sim q_\varphi(\theta | \mathcal{D})$ of the whole given dataset, e.g., $\mathcal{D} = \\{(y_n, x_n)\\}$.
Further, since we do not use a context set for conditioning the marginal mentioned by the reviewer, this quantity would reduce to computing $\log p(y|x)$, which does not depend on $\theta$ and thus is constant irrespectively of forward or reverse KL.

Reviewer comment (Evaluation):
The paper does not justify why it uses L2 loss instead of evaluating the predictive log marginal likelihood, which is more typical for assessing the predictive performance of Bayesian models. Predictive log marginal likelihood would better measure the quality of predictive epistemic uncertainty estimates and, implicitly, the posterior approximation.

Author response:
Predictive log marginal likelihood: We are not quite clear what the reviewer means by this quantity, but we do add another metric definition in Appendix D (CNLL). An example of this metric in our proposed experiments is available in Tables 5 and 6. The reason we did not include this metric was because it was fairly well correlated with downstream $L_2$ or accuracy metrics, but for completeness we are happy to include it in the Appendix. We request the reviewer to let us know if they had a different metric in mind.

Reviewer comment (Evaluation):
It is unclear how much variance is introduced in the results due to the algorithm’s or network’s initialization. The authors should provide confidence intervals for the experimental results, as this is necessary to judge their statistical significance.

Author response:
**Variance due to Initialization**

We agree with the reviewer that it would be nice to have an estimate of variance due to initialization in the posterior estimation. We would like to point out, however, that all our metrics are evaluated over 100 datasets, each with 25 samples from the approximate variational distribution which already provides an estimate of the reliability of the model. We omit the standard deviation in this setup so as to improve readability in the tables, but we will open-source our code for the community to freely use the approach. The experiments are computationally extensive to run since we test on a wide variety of probabilistic models but we would be happy to attempt to provide some estimates. Is there a particular metric / setup that the reviewer is interested in seeing the effect of initialization on?

Reviewer comment (Experiments):
A non-amortized version of the proposed method should be added as a baseline in order to assess the amortization gap — the approximation error introduced solely by amortization.

Author response:
**Non Amortized Baseline** 

We thank the reviewer for this point and note that we do consider two non amortized baselines, optimization and MCMC. These baselines in particular show that the amortized VI setup is not as strong as optimization in predictive modeling but does compare significantly favorably to MCMC. We also decided to opt out of KDE estimates of MCMC in Figure 4 just to make the figures clearer as we already have the true posterior on it.

Reviewer comment (Experiments):
The authors state that normalizing flows do not increase approximation accuracy due to the mode-seeking behavior of the reverse KL objective. A discussion and/or comparison with recent natural-gradient based methods (e.g., those using expressive Gaussian mixture approximations and objective terms to prevent mode collapse) would be of interest.

Author response:
**Related Work**

We thank the reviewer for pointing out some of the relevant works on VI approaches leveraging a Gaussian Mixture distribution as the approximating distribution as well as different context averaging methods like Bayesian Context Averaging. We believe that empirical comparisons against these approaches is a contribution on its own and thus leave it as future work. However, we do include a discussion on them in the main text as well as the Appendix in the revised draft.

Reviewer comment (Experiments):
The authors argue that deep sets are inferior to transformer-based architectures due to naive sum/mean-based aggregation. Since Bayesian aggregation methods have been proposed (e.g., [4]) to address this specific issue, it would be interesting to compare Bayesian aggregation with transformer-based aggregation.

Author response:
**Related Work**

We thank the reviewer for pointing out some of the relevant works on VI approaches leveraging a Gaussian Mixture distribution as the approximating distribution as well as different context averaging methods like Bayesian Context Averaging. We believe that empirical comparisons against these approaches is a contribution on its own and thus leave it as future work. However, we do include a discussion on them in the main text as well as the Appendix in the revised draft.

Reviewer comment (Presentation):
There are figures all over the place while they all belong to section 4. It is very confusing to see experimental figures in the middle of the introduction. In addition, when reading the experiment section, the reader has to jump back to the introduction to see the figure.

Author response:
We apologize for the figure positioning and will try to improve it where possible.

Reviewer comment (Writing):
Equation (8) seems to be very similar to equation (9) from prior work. Would it be easier to start from there, explain what $\chi$ is and say that you use the reverse KL instead? I feel that previous explanations dilute the message and make things hard to follow while equation (9) is straightforward to understand.

Author response:
**Starting from Equation 8-9** 

We argue that the reason Equation 9 seems familiar from prior work is because it is the general starting point of defining the optimization problem for Simulation-based Inference (SBI) approaches. However, we rely on the reverse KL / Variational Inference (VI) framework where the derivation leads to the well studied ELBO objective. Thus, we discuss the connections to ELBO and amortization more clearly in the start to provide a complete story as well as the exact loss formulation, which is defined in Equation 7. Mathematically, equations 7 and 8 describe an equivalent optimization problem.

Reviewer comment (Presentation):
Section 4.3 seems to be full of methodological elements while being in the middle of the experiments section. I think grouping all the methodological elements together would help clarify the contributions.

Author response:
**Section 4.3** 

We provide each section of the experiments with its own set of methodological details and defer the readers to the Appendix for additional tangential details (eg. the choice of mis-specification explored). This is because we want each experimental section to be self-contained about the experiments discussed there. We will take the reviewer’s concerns into account and update Section 4.3 in the next revision to clarify some of the concerns regarding the setup.

Reviewer comment (Writing):
In section 4.3 it is said that "In contrast, we can leverage our proposed reverse KL approach to train an amortized inference model to predict the posterior over the assumed probabilistic model’s parameters by directly using the available unpaired data during training." It is not clear to me how this is done while this seems to be a contribution of the paper. It would be worth to expand on this more.

Author response:
To clarify, we would like to point the reviewer to the differences in Equations 8 and 9. Equation 9 is only tractable for training if $\chi$ defines sampling according to the probabilistic model $p(\mathcal{D} | \mathbf{\theta})$, since then we observe the pairs $\{(\mathcal{D}, \mathbf{\theta})\}$. However, in general, we often only observe a stream of data $\{\mathcal{D}\}$ from some data generating function $\chi’$ coming from an unknown probabilistic model, without its corresponding parameters $\mathbf{\theta}$. In this setting of model mis-specification, we find that when we train both the proposed method and SBI baseline on $\chi$ and evaluate on $\chi’$, the proposed method generalizes better.

Additionally, even though we don’t know the underlying probabilistic model, we can still model the observations relatively well. Practitioners often define a probabilistic model to explain this stream of data; which could be different from the ground-truth underlying probabilistic model and hence wrong. However, one could still model the data well by finding a good set of parameters of this ill-specified probabilistic model. A caveat of SBI approaches, which are generally framed as a Forward KL optimization problem, is that they cannot be trained using data obtained from $\chi’$ and thus would not be able to leverage diverse data coming from different and unknown probabilistic models. In contrast, our proposed method can leverage such data during training, and we show in Table 3 that leveraging such data improves performance even further.

More succinctly, this difference arises from swapping $\chi$ with $\chi’$ in Equations 8 and 9 such that $\chi’$ does not follow $p(\mathcal{D} | \mathbf{\theta})$ while $\chi$ does. This swap makes Equation 9 infeasible to optimize, while Equation 8 still remains feasible. We hope that this clarifies the reviewer’s doubts regarding how this setting is formulated.

As an example, we can consider the probabilistic model as a nonlinear Bayesian Neural Network model. It can only be trained under the SBI framework if $\chi$ represents nonlinear data sampled according to this model. However, we might receive a stream of data of interest that might follow a linear relationship. This data cannot be part of the training framework in the SBI setup, but can be in our proposed method.

Reviewer comment (Novelty):
I think the following paper should be discussed in the related works. It addresses the problem of amortized Bayesian inference for datasets of different cardinality. It exploits the fact that the scores of each individual observation can be composed to produce the score of the joint observations. This joint score can then be used to efficiently produce samples from the posterior distribution.  
Geffner, T., Papamakarios, G., & Mnih, A. (2023). Compositional Score Modeling for Simulation-Based Inference.

Author response:
**Related Work** 

We have added a discussion about the work on score modeling for SBI in the Appendix. We were not aware of this work but as we point out, it is still different from the proposed approach since we are not leveraging the SBI framework for amortized posterior estimation and instead propose a complementary research stream and show benefits of performing amortized VI. However, we provide a small discussion on the similarities in the Appendix A of the revised draft.

We thank the reviewer for their response. We would like to politely disagree with the reviewer as we do mention how we tackle a different problem than the NP framework; as well as we do perform a direct comparison with SBI, which is the Forward KL approach. Details of how we are different is also present in the last paragraph of Section 2. Additionally, we believe that GPs are only tangentially related to the work in so much as how amortization has been used to learn the kernel function for GPs. We additionally provide discussion with all such related work in Appendix A as well as describe how we are different in more detail.

*"I have to disagree because this is precisely what is done in Geffner, T., Papamakarios, G., & Mnih, A. (2023). Compositional Score Modeling for Simulation-Based Inference. It constructs a score model that can handle inputs of variable dimensions and be used to sample from the posterior."*

We think that maybe the reviewer is confusing by what we mean by dimensionality. Both our model and Geffner et. al can handle variable number of data-points / observations when modeling the posterior distribution. In our work, we provide an additional flexibility by allowing for modeling observations lying on variable dimensional spaces. That is, the same network can be used to model the posterior over weight vectors for 2-D regression as well as 5-D regression. Geffner et. al on the other hand don't look at this axis of variable dimensionality, and instead look at amortizing over variable number of data-points.

Thank you again for your thoughtful review. We made a significant effort to address your feedback and include additional experiments, and would appreciate it if you would consider raising your score in light of our response. Please let us know if there are further questions we can address.

Reviewer comment (Novelty):
*Literature Gap*: The paper seems to omit relevant literature on amortized GP hyperparameters [1][2][3]. And GP belongs to the framework considered here because it has closed form likelihoods given model parameters. Specifically, [1] produces amortized inference for point estimate of the posterior for GP hyperparameters given a dataset. The neural network architecture proposed in [1] also makes use of transformer for permutation equivariance. [1][2][3] also generalize to unseen datasets, with the same meta-learning flavor. This paper broadens the perspective by considering general probabilistic models and considering a distribution rather than a point estimate. But the basic idea and architecture choices share the same spirit, which decreases the novelty in the methodological contribution.

Author response:
We thank the reviewer for bringing these related lines of work to our attention. While they share similarities in the sense that they amortize the inference process for the kernel function of Gaussian Processes (GPs), it is crucial to delineate the distinctions in our approach. Unlike these methodologies, our framework does not involve the estimation of kernel functions for GPs. Instead, we present a more comprehensive framework for Bayesian posterior estimation based on amortized inference, emphasizing the modeling of the entire posterior distribution rather than relying on point estimates. However, since our approach does share some of the underlying similarities (connections to meta learning, use of transformer architecture, etc.), we include a discussion on the mentioned approaches in Appendix A of the revised manuscript.

Reviewer comment (Theory):
*Variable Dimension Handling*: The method of managing variable-dimensions through masking might be limited. It wastes GPU memory and is not equivariant w.r.t. dimensions. There could be potential benefits in exploring the neural network employed by [1] if each dimension has its own parameters, such as in linear regression and GMM.

Author response:
We agree with the reviewer about predicting feature-specific parameters in modeling the posterior distribution. We had actually tried such an approach for linear regression where we leveraged attention across features as well as across observations. However, not only was it too compute intensive because of such interleaved attention operations but it also did not perform comparably to our current approach. 

While we do agree with the reviewer’s point that using a parameter for each feature dimension allows for leveraging the existing permutation equivariance in the feature space for a class of probabilistic models, we argue that such a setting can be fairly limited. That is, it is conceptually straightforward to learn such a system for linear regression or Gaussian Mixture models but learning such a general system that exploits all the equivariances present in more complex models (eg. Bayesian Neural Nets) is far less obvious. In particular, leveraging the methodology explored in [1] would not work for models with deep learning based probabilistic models, for which it is not clear if the parameters corresponding to some intermediate layers should be tied to certain feature dimensions, and if so, which? In particular, explicitly modeling all the equivariances present in a deep neural network is likely intractable and hence while we agree that it is a very interesting direction, we feel that it is an orthogonal line of research which would deserve a separate exploration.

Reviewer comment (Reproducibility):
*Clarity on GMM*: It is unclear how the proposed approach would manage variable number of mixtures in the case of Gaussian Mixture Models (GMM).

Author response:
We would like to clarify that although our methodology accommodates the modeling of Gaussian Mixture Model (GMM) tasks with varying feature dimensions, it does not currently support variable numbers of mixtures within the same model. We acknowledge any lack of clarity on this matter and will make necessary revisions in the draft to mitigate potential confusion.

Reviewer comment (Experiments):
*Ablation Limitations*: While the mention of an ablation study is commendable, it would be beneficial to see a more comprehensive study, for instance, sweeping across dimensions ranging from 1-100D, to see how the approach extrapolates on dimensions different than the training data.

Author response:
We would also like to clarify that for variable dimensional experiments, we choose the dimensions 1-100 uniformly so even if we were to do a sweep over these dimensions to see how the performance varies, it remains “in-distribution” testing. That being said, we take the reviewer’s point under consideration and showcase the trends over different dimensional setups in Figures 5-7 over different KL-based optimization strategies, architecture choices and variational assumptions and hope that it solves the reviewer’s concerns. For OoD testing, we refer the reviewer to the model mis-specification and tabular experiments.

Reviewer comment (Experiments):
*Choice of MCMC*: What was the motivation behind choosing Langevin instead of HMC? Further, is the paper referring to stochastic gradient Langevin dynamics? Maybe HMC type algorithms such as NUTS [4] should also be considered, since they have shown to be performing well empirically and given the number of datapoints is not too large that needs stochastic gradient.

Author response:
We apologize for not providing more clarity on this. We actually use Langevin Dynamics (full batch).  We also experimented with NUTS sampler but found it to be quite slow in nonlinear Bayesian Neural Network parameter estimation, especially since for each metric we would need to run the MCMC algorithm 100 times as we consider an average over 100 different datasets for each experiment.

Reviewer comment (Presentation):
*Figure Positioning*: The positions of the figures within the paper were not specified. Could the authors provide more clarity on this aspect?

Author response:
We apologize for the figure positioning and will try to improve it where possible.

Reviewer comment (Experiments):
*Comprehensive Ablation Study*: Would the authors consider conducting an ablation that sweeps on dimensionality from 1-100D?

Author response:
We would also like to clarify that for variable dimensional experiments, we choose the dimensions 1-100 uniformly so even if we were to do a sweep over these dimensions to see how the performance varies, it remains “in-distribution” testing. That being said, we take the reviewer’s point under consideration and showcase the trends over different dimensional setups in Figures 5-7 over different KL-based optimization strategies, architecture choices and variational assumptions and hope that it solves the reviewer’s concerns. For OoD testing, we refer the reviewer to the model mis-specification and tabular experiments.