Multi-Task Bayesian In-Context Learning

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 learning networks 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 heavy-tailed 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.

Many AI systems can learn from examples, but they often struggle when the assumptions behind a problem change. Bayesian methods handle this by combining prior knowledge with new evidence, but they are often too slow for practical use. We propose Multi-Task Bayesian In-Context Learning, a transformer-based method that represents prior knowledge using extra related datasets placed in the model’s context. By changing these prior datasets, users can change the model’s assumptions at test time without retraining. Our experiments show that this approach closely matches strong Bayesian methods while being much faster. This could help build AI systems that adapt quickly, use related data effectively, and provide more reliable uncertainty estimates.