Exchangeable Dataset Amortization for Bayesian Posterior Inference

We thank the reviewer for their valuable comments. Our response aims to address the raised concerns, and we are available for further discussions to resolve any outstanding comments or ambiguities.

Contribution and Metrics

We refer the reviewer to the common shared response which provides a discussion about the novelty of our work as well as the different contributions made. It also refers to the exact formulae used to define the metrics that we consider.

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.

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.

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.

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?

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.

We hope that our response clarifies the reviewer's concerns and we would be happy to answer any additional questions and concerns.

We sincerely appreciate the reviewer's valuable comments. In our response, we aim to address and clarify the majority of the concerns raised and would be happy to engage in additional discussions should there be any remaining unresolved comments.

Contribution and Metrics

We refer the reviewer to the common shared response which provides a discussion about the novelty of our work as well as the different contributions made. It also refers to the exact formulae used to define the metrics that we consider.

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.

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.

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.

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.

Figure Positioning

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

We hope that our response clarifies the reviewer's concerns and we would be happy to answer any additional questions and concerns.

We thank the reviewer for their valuable comments and hope that our response addresses and clarifies the raised concerns. We are open to further discussions to resolve any remaining comments or ambiguities.

Contribution and Metrics

We refer the reviewer to the common shared response which provides a discussion about the novelty of our work and the different contributions made as well as references to the exact formulae used to define the metrics that we consider.

Literature Gap

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.

Variable Dimension Handling

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.

GMM’s Variable Dimension Handling

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.

Ablation Limitations

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.

Choice of MCMC

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.

Figure Positioning

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

[1] Liu, Sulin, Xingyuan Sun, Peter J. Ramadge, and Ryan P. Adams. "Task-agnostic amortized inference of gaussian process hyperparameters." Advances in Neural Information Processing Systems 33 (2020): 21440-21452.

We hope that our response clarifies the reviewer's concerns and we would be happy to answer any additional questions and concerns.

Specifically, the tractability of Equation (9) hinges on the condition that the dataset $\mathcal{D}$ is sampled in accordance with the assumed probabilistic model $p(\mathcal{D} | \mathbf{\theta})$. In contrast, Equation (8) remains tractable even without this strict condition, enabling training with data from diverse sources. As is the case with most novel work, our proposed method indeed does share mathematical and experimental similarities to current existing works (eg. the concept of amortization, variational inference and transformers have been in the literature for a while now).

Moreover, it is crucial to emphasize that to the best of our knowledge, none of the existing methods within SBI or NPs demonstrate the capability to effectively handle inputs of variable dimensions while our method can, which in itself is a novel contribution. Additionally, we also test for OoD generalization of such models not only through synthetic model mis-specification settings but also through their ability to transfer to real-world regression and classification problems. We believe that we are the first ones to highlight the differences in the predictive ability of Forward (SBI) and Reverse KL amortization approaches, especially under different kinds of distribution shifts.

Metrics

We thank the reviewer for bringing this up. To make things clearer and more precise, we have added Section D in the Appendix that now provides the exact metrics that we use. We would like to clarify that we use the loss and accuracy metrics on the predictive distribution and not the approximate posterior, exactly because expected $L_2$ loss in the parameter space is not suited for problems with multimodal posteriors as Reviewer iJJB pointed out. Thus, a good proxy for understanding if we capture a reasonable posterior is to see if it leads to good predictive performance. Additionally, to assess the quality of the posterior we provide some symmetric KL divergence based metrics where the true posterior is available in closed form.

We hope that this addresses the reviewers' concerns regarding the novelty and contributions of this work, as well as clarifies the exact metrics that we use.