# Multi-Task Bayesian In-Context Learning

arXiv: https://arxiv.org/abs/2606.20538 · HTML: https://mlsanity.com/paper/2606.20538/

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Builds on 30 of its 38 references that are in the map.
- First posted on arXiv 18 Jun 2026.

## Builds on (most cited first, 30 of 30)

- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [Model-Agnostic Meta-Learning for Fast Adaptation of Deep Networks](https://mlsanity.com/paper/1703.03400/) (2017)
- [RoFormer: Enhanced Transformer with Rotary Position Embedding](https://mlsanity.com/paper/2104.09864/) (2021)
- [Neural Ordinary Differential Equations](https://mlsanity.com/paper/1806.07366/) (2018)
- [Variational Inference: A Review for Statisticians](https://mlsanity.com/paper/1601.00670/) (2016)
- [An Explanation of In-context Learning as Implicit Bayesian Inference](https://mlsanity.com/paper/2111.02080/) (2021)
- [Transformers learn in-context by gradient descent](https://mlsanity.com/paper/2212.07677/) (2022)
- [What Can Transformers Learn In-Context? A Case Study of Simple Function Classes](https://mlsanity.com/paper/2208.01066/) (2022)
- [What learning algorithm is in-context learning? Investigations with linear models](https://mlsanity.com/paper/2211.15661/) (2022)
- [TabPFN: A Transformer That Solves Small Tabular Classification Problems in a Second](https://mlsanity.com/paper/2207.01848/) (2022)
- [The frontier of simulation-based inference](https://mlsanity.com/paper/1911.01429/) (2019)
- [Pyro: Deep Universal Probabilistic Programming](https://mlsanity.com/paper/1810.09538/) (2018)
- [Transformers Can Do Bayesian Inference](https://mlsanity.com/paper/2112.10510/) (2021)
- [Neural Processes](https://mlsanity.com/paper/1807.01622/) (2018)
- [Many-Shot In-Context Learning](https://mlsanity.com/paper/2404.11018/) (2024)
- [Benchmarking Simulation-Based Inference](https://mlsanity.com/paper/2101.04653/) (2021)
- [Pretraining task diversity and the emergence of non-Bayesian in-context learning for regression](https://mlsanity.com/paper/2306.15063/) (2023)
- [Towards a Neural Statistician](https://mlsanity.com/paper/1606.02185/) (2016)
- [Meta-Learning Probabilistic Inference For Prediction](https://mlsanity.com/paper/1805.09921/) (2018)
- [General-Purpose In-Context Learning by Meta-Learning Transformers](https://mlsanity.com/paper/2212.04458/) (2022)
- [All-in-one simulation-based inference](https://mlsanity.com/paper/2404.09636/) (2024)
- [Meta-in-context learning in large language models](https://mlsanity.com/paper/2305.12907/) (2023)
- [The Gaussian Neural Process](https://mlsanity.com/paper/2101.03606/) (2021)
- [Can Transformers Learn Full Bayesian Inference in Context?](https://mlsanity.com/paper/2501.16825/) (2025)
- [Amortized Probabilistic Conditioning for Optimization, Simulation and Inference](https://mlsanity.com/paper/2410.15320/) (2024)
- [Position: The Future of Bayesian Prediction Is Prior-Fitted](https://mlsanity.com/paper/2505.23947/) (2025)
- [Does learning the right latent variables necessarily improve in-context learning?](https://mlsanity.com/paper/2405.19162/) (2024)
- [Distribution Transformers: Fast Approximate Bayesian Inference With On-The-Fly Prior Adaptation](https://mlsanity.com/paper/2502.02463/) (2025)
- [Episodic Multi-Task Learning with Heterogeneous Neural Processes](https://mlsanity.com/paper/2310.18713/) (2023)
- [Supervised Contrastive Block Disentanglement](https://mlsanity.com/paper/2502.07281/) (2025)

## Built on by (most cited first, 1 of 1)

- [DUALVIEW: Leveraging Broad and Output-Specific Evidence for LLM-based Multi-Output Text Regression](https://mlsanity.com/paper/2610.19206/) (2026)

## Abstract (arXiv)

Bayesian predictive inference provides a principled framework for uncertainty quantification, data efficiency, and robust generalization. However, exact inference is often intractable, and scalable approximations may remain computationally expensive or require restrictive modeling assumptions that degrade predictive performance. Prior-Data Fitted and in-context models have recently emerged as an amortized alternative by learning to map datasets directly to predictive distributions, but existing approaches are tightly coupled to the support of the training prior and lack explicit mechanisms for adapting to new priors at test time, resulting in limited robustness under distribution shift. We introduce a multi-task in-context learning framework for amortized hierarchical Bayesian predictive inference that explicitly represents prior information as a prefix of in-context datasets. A transformer trained on sequences of prior and target tasks learns to adapt its predictions across families of priors. On a suite of evaluations with increasing difficulty, including out-of-meta-distribution priors and priors with high-dimensional latent structures, our method matches oracle Bayesian predictors while being orders of magnitude faster. We further demonstrate its practical relevance on a real-world spatiotemporal temperature prediction benchmark. Code is available at https://github.com/martianmartina/multi-task-bayesian-icl/.