In-Context Learning for Full Bayesian Inference

Thank you for taking the time to read our manuscript and providing detailed feedback.

We address your questions and comments below. We use double quotation marks (") to mark parts directly quoted from the revised manuscript. In the manuscript itself, all added parts have the text color blue.

How to infer the posterior using the proposed method is not explicitly explained [...] how [continuous normalizing flows] are used in the proposed method

To address these points, we have added a new section "Implementing Flow Matching" (Section 3.5) to the manuscript that details how exactly the flow matching framework is used during training and how posterior samples are generated based on a pre-trained in-context learner.

Below is the added text quoted from the revised manuscript for your convenience:

“Implementing Flow Matching

During the training phase, a tuple $(\mathbf{z}\_1,\mathbf{x})$ is drawn from the distribution $P^{\mathbf{z},\mathbf{x}}$. Additionally, a time step $t \sim \mathcal{U}$ 0,1 $$ and a sample $\mathbf{z}\_0$ is drawn from the base distribution $P\_{\mathcal{B}}$, which is a standard Gaussian for all our applications. Subsequently, the ground-truth conditional flow $\psi(\mathbf{z}\_0|\mathbf{x}) = 1 - ( 1- \sigma_{min})t)\mathbf{z}\_0 + t \mathbf{z}\_1$ is computed, pushing forward $P_{\mathcal{B}}$ into $P^{\mathbf{z}|\mathbf{x}}$ up to time-point $t$. The transformer encoder processes $\mathbf{x}$ and the decoder takes the representation of the encoder into account in order to output $v\_{t,\mathbf{x}}^{\theta}(\psi(\mathbf{z}\_0|\mathbf{x}))$. This output should match the vector field that describes how the ground-truth flow $\psi(\mathbf{z}\_0|\mathbf{x})$ continues at time $t$. The discrepancy to the ground-truth vector field is measured with the MSE-loss in Equation 6. In the sampling phase, we are given $\mathbf{x}$ and the goal is to sample from $P^{\mathbf{z}|\mathbf{x}}$. To do so, first a vector $\mathbf{z}\_0 \sim P\_{\mathcal{B}}$ is drawn. The data $\mathbf{x}$ is passed through the encoder. The decoder defines a function that maps a time-point $t$ and a vector $\mathbf{\nu}$ onto a vector field: $(t, \mathbf{\nu}) \mapsto v_{t,\mathbf{x}}^{\theta}(\mathbf{\nu})$ taking $\mathbf{x}$ into account. This function is given to an ODE-solver in order to forward-solve the corresponding ODE with boundary conditions $0 \leq t \leq 1$.”

The new section provides a computational perspective on flow matching as it is used in conjunction with the introduced architecture.

To my understanding, in the ICL framework, the (foundation) model, typically, Transformer, learns algorithms implicitly through the pre-training (e.g., Garg's work). In this work, however, the model seems designed for inference and pre-training on many related tasks. I would call such a framework meta-learning, rather than ICL.

We fully agree that ICL is a special case of meta-learning and included an additional explanation in the revised manuscript in the first paragraph of Section 2:

“ICL is a special case of meta-learning [1] characterized by using a large pre-trained model to learn from a context dataset without explicitly updating task-specific parameters. [...]”

We use the more specific term in-context learning instead of meta-learning to emphasize that our approach yields posterior samples directly based on a context set processed by a large transformer network without requiring additional parameter updates. Taking seminal work on meta-learning into account, e.g. [2,3,4], we believe that the term in-context learning is useful to contrast the underlying process of our approach.

Thank you again for your constructive feedback. Please let us know if our response addresses your questions and concerns, and if there is anything else we need to clarify.

References

[1] Vilalta, Ricardo, and Youssef Drissi. "A perspective view and survey of meta-learning." Artificial intelligence review 18 (2002): 77-95.

[2] Finn, Chelsea, et al. "Model-agnostic meta-learning for fast adaptation of deep networks." International conference on machine learning. PMLR, 2017.

[3] Sung, Flood, et al. "Learning to compare: Relation network for few-shot learning." Proceedings of the IEEE conference on computer vision and pattern recognition. 2018.

[4] Vinyals, Oriol, et al. "Matching networks for one shot learning." Advances in neural information processing systems 29 (2016).

Dear Reviewer prH2,

We are pleased to hear that our clarifications helped you understand our work. We would also like to thank you for inquiring about the role of the transformer in our paper, which is indeed very insightful.

This raises another question: why is a Transformer necessary in your work?

There are several conceptual reasons why transformers are central to our approach:

First, transformers are a commonly used baseline architecture working extremely well across a broad variety of fields such as computer vision, NLP, or even tabular data [1,2,3]. In addition, the papers introducing the PFN methodology [4,5], a central pillar of our approach, heavily rely on the transformer. Therefore, using a transformer is a natural and reasonable choice for our proposal, whose central point is to extend [4,5] and demonstrate the feasibility of in-context learning for full Bayesian inference. Furthermore, the notion of in-context learning itself, which we want to extend to full Bayesian inference, is closely connected to the transformer architecture [6,7,8]. This is confirmed by findings suggesting the attention mechanism to play a crucial role for in-context learning [9,10].

Ablation Study

Intrigued by your question, we conducted another ablation experiment to study your question from an empirical point-of-view and investigate in which cases the transformer architecture is indeed a crucial component. For the ablation, we exchange the transformer-encoder in our approach with an MLP using skip connections and batch normalization and ensure that the transformer encoder and the MLP have approximately the same number of parameters. We then ran the previous experiments using the different models (GLMs, FA, GMMs) in three different scenarios each; every scenario has 50 synthetic and 17 real-world datasets.

In summary, the results confirm our initial hypothesis that a transformer encoder is important for complex conditioning on the input data.

For the GLM scenarios, this is especially pronounced with substantial performance drops across all scenarios when using the MLP baseline instead of the transformer. In the case of FA and GMMs, the improvement of the transformer-encoder is smaller, albeit consistent. In particular, there is no case where the MLP performance is above one standard error of the transformer’s performance. For FA and GMMs, the gap between the MLP and the transformer encoder is, for example, especially noticeable in scenarios 1 and 3 of the FA cases on the real-world data and in scenarios 1 and 2 for the GMM cases.

If you have any remaining concerns or questions, please let us know.

References

[1] Dosovitskiy, Alexey, et al. "An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale." International Conference on Learning Representations. 2020.

[2] Dubey, Abhimanyu, et al. "The llama 3 herd of models." arXiv preprint arXiv:2407.21783 (2024).

[3] Gorishniy, Yury, et al. "Revisiting deep learning models for tabular data." Advances in Neural Information Processing Systems 34 (2021): 18932-18943.

[4] Müller, Samuel, et al. "Transformers Can Do Bayesian Inference." International Conference on Learning Representations.

[5] Hollmann, Noah, et al. "TabPFN: A Transformer That Solves Small Tabular Classification Problems in a Second." The Eleventh International Conference on Learning Representations.

[6] Panwar, Madhur, Kabir Ahuja, and Navin Goyal. "In-Context Learning through the Bayesian Prism." The Twelfth International Conference on Learning Representations.

[7] Dong, Qingxiu, et al. "A survey on in-context learning." arXiv preprint arXiv:2301.00234 (2022).

[8] Garg, Shivam, et al. "What can transformers learn in-context? a case study of simple function classes." Advances in Neural Information Processing Systems 35 (2022): 30583-30598.

[9] Akyürek, Ekin, et al. "In-Context Language Learning: Architectures and Algorithms." Forty-first International Conference on Machine Learning.

[10] Lee, Ivan, Nan Jiang, and Taylor Berg-Kirkpatrick. "Is attention required for ICL? Exploring the Relationship Between Model Architecture and In-Context Learning Ability." The Twelfth International Conference on Learning Representations.

Thank you for taking the time to read our manuscript and providing detailed feedback.

We address your questions and comments below. We use double quotation marks (") to mark parts directly quoted from the revised manuscript. In the manuscript itself, all added parts have the text color blue.

While ICL may yield accurate approximations of posterior distributions, it requires a high computational investment in training.

In-context learning is inherently expensive during the pre-training phase and our method is no exception—the arguably most prominent in-context learners, Large Language Models, are currently even trained on more than ten thousand GPUs in order to achieve state-of-the-art performance [1]. We have revised the limitations section in section 5 of the manuscript to emphasize this aspect even more clearly. It is noteworthy, however, that when practically applied to a new dataset, our approach can sample the posterior distribution very quickly. To substantiate this claim, we have conducted additional runtime comparisons of all methods. While some VI methods are faster than ICL, at the cost of inferior performance, our approach is consistently faster than HMC (for FA even achieving a factor 8 of speed-up).

For your convenience, we also include the results (in Appendix E in the revised manuscript) below.

Thank you again for your constructive feedback. Please let us know if our response addresses your questions and concerns, and if there is anything else we need to clarify.

the rationale for specific architectural choices, such as the reliance on continuous normalizing flows over other methods, is underexplored.

While exploring other architectural choices is an interesting avenue, we would like to note that [2] introduces theoretical results backing the choice of continuous normalizing flows over, for instance, various variants of Diffusion methods. They also empirically find flow matching to be a more robust and efficient objective. [3] confirms these results across several simulation-based inference tasks that are closely related to those in our paper.

To again substantiate our claims, we added additional empirical results to the manuscript by replacing flow matching based on optimal transport (OT) paths in our approach with a diffusion objective based on variance preserving (VP) diffusion probability paths from [4]. When comparing samples based on the VP diffusion probability paths and our initial flow matching implementation with OT transport paths, we find that our original approach consistently yields superior samples.

Thank you again for your constructive feedback. Please let us know if our response addresses your questions and concerns, and if there is anything else we need to clarify.

References

[1] Dubey, Abhimanyu, et al. "The llama 3 herd of models." arXiv preprint arXiv:2407.21783 (2024).

[2] Lipman, Yaron, et al. "Flow matching for generative modeling." arXiv preprint arXiv:2210.02747 (2022).

[3] Wildberger, Jonas, et al. "Flow matching for scalable simulation-based inference." Advances in Neural Information Processing Systems 36 (2024).

[4] Song, Yang, et al. "Score-based generative modeling through stochastic differential equations." arXiv preprint arXiv:2011.13456 (2020).

For your convenience, we also include the results when comparing samples based on the VP diffusion probability paths and our initial flow matching implementation with OT transport paths (in Appendix C in the revised manuscript) below.

GLMs: Comparison of the OT flow matching and the VP diffusion objective on 50 synthetic and 17 real-world datasets for three different scenarios. All results within two standard errors of the best average result for each scenario are marked in bold.

FA: Comparison of the OT flow matching and the VP diffusion objective on 50 synthetic and 17 real-world datasets for three different scenarios. All results within two standard errors of the best average result for each scenario are marked in bold.

GMMs: Comparison of the OT flow matching and the VP diffusion objective on 50 synthetic and 17 real-world datasets for three different scenarios. All results within two standard errors of the best average result for each scenario are marked in bold.

Dear Reviewer rWy9,

Thank you again for providing highly valuable feedback on our manuscript . We are pleased that you point out that we propose an “innovative framework for Bayesian inference”, and include “strong empirical comparisons”.

As the discussion period will end in approximately 72 hours, we would like to ask if our answers so far have successfully addressed your concerns and answered your questions.

In summary, we addressed the points raised by you as follows: We added a table comparing the inference time of our ICL method with that of all other approaches, which shows that ICL is consistently faster than HMC during deployment. Additionally, we conducted an ablation study using a diffusion objective to address your concern regarding the reliance on continuous normalizing flows. Our new empirical results confirm the original choice of using flow matching with optimal transport paths.

Thank you again for your time. If there are any remaining concerns, please let us know. In case we sufficiently addressed your questions, we would greatly appreciate it if you consider raising the score.

Dear Reviewer rWy9,

We would like to kindly remind you that the extended rebuttal period will end in less than 72 hours.

We would be happy to hear if we have addressed your concerns so far and if you have further questions. In case we have sufficiently addressed your questions, we would greatly appreciate it if you consider raising the score.

Thank you again for your time.

Thank you for taking the time to read our manuscript and providing detailed feedback.

We address your questions and comments below. We use double quotation marks (") to mark parts directly quoted from the revised manuscript. In the manuscript itself, all added parts have the text color blue.

Why do the authors consider extending the ICL capability to fully Bayesian inference? What is the advantage of extending the ICL capability to full Bayesian inference beyond Bayesian inference?

In contrast to previous literature such as [1], that focuses on the posterior predictive distributions, we tackle a much more complex and challenging problem, namely learning the posterior of latent variables. In contrast to [1], this posterior is multivariate and thus much more difficult to obtain. Except for cases where the model is linear in the parameters and a conjugate prior is chosen, obtaining such a posterior distribution of latent variables usually requires MCMC methods and a careful choice of the sampling procedure. While variational and local approximations exist, these do not provide nominal coverage guarantees. Hence neither [1] nor variational approaches would be able to, e.g., provide a 95%-credible interval for the effect of the covariates on medical charges in a hospital. Full Bayesian inference is a fundamental approach in probabilistic modeling which is basically always desirable, even though often infeasible.

To emphasize why full Bayesian inference is a central paradigm, we added some extra content in Section 1.1 and marked it blue in the updated manuscript. We think that this addition will provide valuable context especially to readers less familiar with Bayesian inference.

Below is the added text quoted from the revised manuscript for your convenience:

“Full Bayesian inference is thus fundamental for a wide range of applications in Bayesian statistics and probabilistic machine learning, even though it cannot always be achieved.”

To strengthen the novelty of this work, it would be better if the authors could demonstrate that the ICL capability of being good at posterior approximation, can be effective for improving the performance of Bayesian optimization or few-shot learning as conducted in [1].

Our scope is fundamentally different from that of Mueller et al. [1]. While [1] shows that transformers can learn univariate and discrete posterior predictive distributions, we target continuous multivariate distribution posteriors of latent variables. This implies that the tasks considered in [1], for instance on Bayesian Optimization, are not well-suited to benchmark our approach.

Thank you again for your constructive feedback. Please let us know if our response addresses your questions and concerns, and if there is anything else we need to clarify.

References

[1] Müller, Samuel, et al. "Transformers Can Do Bayesian Inference." International Conference on Learning Representations.

If ICL uses other distribution match methods such as VAE or Score matching instead of the continuous normalization flow (CNF), could the ICL still approximate the posterior distribution accurately? It seems less clear how important the type of loss in Eq. (5) is in the sense of accurate approximation for ICL. It would be better to conduct the ablation study comparing the loss of Eq. (5) with the loss of other distribution matching methods, and thus explaining why the loss of Eq. (5) is necessary.

We would like to highlight that [2] provides theoretical justification for the use of continuous normalizing flows and empirically demonstrates that flow matching is a more robust and efficient objective compared to various variants of Diffusion objectives. Similarly, [3] corroborates these findings across various simulation-based inference tasks closely related to those discussed in our paper.

To further support our claims, we incorporated additional comprehensive empirical results across three selected scenarios for GLMs, GMMs, and FA with 50 synthetic and 17 real-world datasets each into the manuscript by replacing the flow matching objective with optimal transport (OT) paths from our original approach with variance preserving (VP) diffusion probability paths as proposed in [4]. A comparison of samples generated using VP diffusion probability paths and our original flow matching implementation with OT transport paths reveals that our initial approach consistently produces higher-quality samples.

Thank you again for your constructive feedback. Please let us know if our response addresses your questions and concerns, and if there is anything else we need to clarify.

References

[2] Lipman, Yaron, et al. "Flow matching for generative modeling." arXiv preprint arXiv:2210.02747 (2022).

[3] Wildberger, Jonas, et al. "Flow matching for scalable simulation-based inference." Advances in Neural Information Processing Systems 36 (2024).

[4] Song, Yang, et al. "Score-based generative modeling through stochastic differential equations." arXiv preprint arXiv:2011.13456 (2020).

In Tables 1, 2, 3, and 4, how complex models are used for the baseline of VI and Laplace approximation (LA)? Could you compare the mode size used for LA, VI, and ICL?

Since our approach operates in-context and is thus a meta-learner (learning multiple tasks at the same time whereas LA, VI, etc. only learn on one specific dataset), it is impossible to directly compare the numbers of parameters. From a theoretical perspective, all benchmarked methods (variational inference, Hamiltonian Monte Carlo, and our in-context learning approach) implement inference for the same models (generalized linear models, factor analysis, Gaussian mixture models), which therefore always have the same number of parameters. While these other methods thus require solving a separate optimization problem for every single task, our method amortizes this computation via the meta-learned ICL model. In order to provide some comparison between the methods, we have however performed a runtime benchmark. The benchmark highlights the efficiency of our ICL approach during the deployment phase. While some VI methods are faster than ICL, at the cost of inferior performance, our approach is consistently faster than HMC (for FA even achieving a factor 8 of speed-up).

For your convenience, we also include the results (in Appendix E in the revised manuscript) below.

Thank you again for your constructive feedback. Please let us know if our response addresses your questions and concerns, and if there is anything else we need to clarify.

Dear Reviewer nf5T,

Thank you again for providing highly valuable feedback on our manuscript. We are pleased that you find our manuscript “well written”, and point out that our ICl approach can yield the posterior distribution on complex real-world datasets.

As the discussion period will end in approximately 72 hours, we would like to ask if our answers so far have successfully addressed your concerns and answered your questions.

In summary, we addressed the points raised by you as follows: We explain why extending ICL to full Bayesian inference is a valuable contribution. We also clarify this point further in the main body of the manuscript. Additionally, we conducted an ablation study using a diffusion objective to address your concern regarding the reliance on continuous normalizing flows, which confirms our original choice of using flow matching with optimal transport paths.

To compare the computational expense for running ICL, LA, VI and HMC in practical applications, we have included the runtimes of all methods in the paper and find that ICL is consistently faster than HMC during inference.

Thank you again for your time. If there are any remaining concerns, please let us know. In case we sufficiently addressed your questions, we would greatly appreciate it if you consider raising the score.

Dear Reviewer nf5T,

We would like to kindly remind you that the extended rebuttal period will end in less than 72 hours.

We would be happy to hear if we have addressed your concerns so far and if you have further questions. In case we have sufficiently addressed your questions, we would greatly appreciate it if you consider raising the score.

Thank you again for your time.

Thank you for taking the time to read our manuscript and providing detailed feedback.

We address your questions and comments below. We use double quotation marks (") to mark parts directly quoted from the revised manuscript. In the manuscript itself, all added parts have the text color blue.

there is insufficient discussion regarding concerns and constraints related to model specification. I believe that this limitations are not fully addressed in the most cases.

Regarding issues with model specification, we would like to clarify that the structure of the Bayesian model we want to learn in-context directly determines the distribution of synthetic data we have to use to fit the in-context learner. Our goal is to demonstrate that established GLM, GMM, and FA approaches with standard modeling assumptions and commonly used priors can be fitted in-context. Therefore, the generative processes are fixed for our analysis. Changing this process would result in models not estimating the correct posterior distribution and not allowing the access to a ground truth or gold standard model (as given by the analytic solution or the HMC in our experiments). In particular, this would mean that we cannot judge whether one approach works better or worse than another approach, as no reference distribution is available.

We thank the reviewer for the comment and fully acknowledge that discussing the issue of model (mis-)specification is indeed very valuable. We, therefore, extended Section 3.3 to include a discussion regarding model misspecification for FA and GMMs.

Below is the added text for your convenience:

“To be able to access a reference or ground truth distribution, the data generating processes in our experiments need to match the structure of the GLM, FA and GMM approaches. While the generative processes of FA and GMMs directly prescribe how all parts of the data are generated, this can potentially cause a discrepancy between synthetically generated and real-world datasets. However, our empirical results (Section 4.1) demonstrate that the in-context learner can generalize to real-world data despite the discrepancy to the simulated datasets.”

We further added a discussion of model misspecification and the associated limitations in the discussion part (Section 5) of our manuscript. Below is the added text quoted from the revised manuscript for your convenience:

“Even though our experiments show that ICL works well despite being trained on data that is potentially very different from real-world data, the approach will only be as flexible as the data and model structures it was trained on. As a result, ICL might fail if the model, which implies the synthetic data generation, is severely misspecified. However, this is the same limitation as when misspecifying the hypothesis space of, e.g., a deep neural network or other machine learning approaches, effectively providing the model with the wrong inductive bias.”

Although the paper briefly mentions synthetic data, a more thorough exploration of the design of synthetic datasets is essential.

Key aspects of this point have been answered in our response above. The structure of the standard GLM, GMM and FA approaches directly determines how the synthetic data has to be generated, where GLMs are somewhat of an exception. As noted above, we will make this more clear in a revised manuscript version and thank the reviewer for bringing up this important point.

Thank you again for your constructive feedback. Please let us know if our response addresses your questions and concerns, and if there is anything else we need to clarify.

I appreciate the updates in the revised manuscript, particularly the inclusion of discussions the possibility of misspecification and the more detailed comparisons between gradient methods and the proposed approach. The presented results appear promising and provide meaningful statements considering the standard of ICLR 2025.

However, there remains a significant gap in addressing the generalizability of the proposed method to novel data. Specifically, the authors should demonstrate how the method performs when there is a slight increase in the discrepancy between training and test data, and whether such shifts result in a notable drop in performance. This would provide more clarity on the robustness of the method.

Additionally, I am concerned about the reproducibility. While the updated experimental sections seem valuable for readers who explore the method, implementation does not still reflect the updates in the rebuttal phase.

Depending on the updates, I am inclined to raise my score from 6 to 8. I continue to monitor their responses.

SGLD vs. ICL: Evaluation on 50 synthetic and 17 real-world datasets for four different GMM scenarios. All results within two standard errors of the best average result for each scenario are marked in bold.

For your convenience, we also include the results of benchmarking SGLD below. (Please also refer to Appendix F in the revised manuscript) .

SGLD vs. ICL: Evaluation on 50 synthetic and 17 real-world datasets for six different GLM scenarios. All results within two standard errors of the best average result for each scenario are marked in bold.