Multi-Task Bayesian In-Context Learning

2026-06-18

AI Summary

Traditional Bayesian inference is computationally expensive and difficult to scale, while existing in-context learning approaches lack the flexibility to adapt to new priors at test time, resulting in poor robustness under distributional shift. This work proposes a multitask in-context learning framework that explicitly encodes prior information as a prefix in the Transformer input sequence, enabling adaptive prediction across diverse prior families. By integrating hierarchical Bayesian reasoning with sequence modeling, the method supports efficient generalization to both unseen priors and high-dimensional latent-structured priors. Experiments demonstrate that the approach achieves predictive performance comparable to an ideal Bayesian inference engine while accelerating inference by several orders of magnitude, with practical efficacy validated on real-world spatiotemporal temperature forecasting tasks.

Abstract

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/.

Problem

• Bayesian inference
• in-context learning
• distribution shift
• prior adaptation
• amortized inference

Innovation

• Multi-Task In-Context Learning
• Amortized Bayesian Inference
• Prior Adaptation
• Transformer-based Prediction
• Distribution Shift Robustness