Multi-Task Bayesian In-Context Learning

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

• The paper introduces a novel MT-BICL framework that augments transformers with multi-task prefixes for flexible, test-time prior adaptation.
• The paper demonstrates efficient amortized inference with negligible KL divergence and robust out-of-distribution performance across diverse prior distributions.
• The paper bridges traditional hierarchical Bayesian inference and neural sequence modeling, enabling uncertainty-aware predictions through explicit hierarchical conditioning.

Multi-Task Bayesian In-Context Learning: Bridging Hierarchical Bayesian Inference and Amortized Neural Sequence Modeling

Problem Formulation and Prior Work

Multi-task Bayesian in-context learning (MT-BICL) extends the in-context learning (ICL) paradigm by introducing flexible, explicit control over priors in hierarchical Bayesian inference using neural sequence models. While the Bayesian framework offers principled uncertainty quantification through posterior predictive distributions (PPDs), exact inference is infeasible in complex latent variable models, and scalable approximations, such as MCMC or variational inference, either incur substantial computational overhead or require restrictive modeling assumptions. Recent amortized approaches—e.g., Prior-Data Fitted Networks (PFNs)—circumvent some computational bottlenecks by meta-training neural models directly to map datasets to probabilistic predictions. However, such methods hard-code a fixed prior into the model parameters, precluding adaptation to novel priors at test time, which hampers robustness, particularly under distribution shift.

Explicit Prior Conditioning via Multi-Task Prefixes

MT-BICL proposes augmenting ICL architectures (specifically, decoder-only transformers) by prefixing the target task context with multiple auxiliary "prior" datasets, each corresponding to synthetic tasks drawn from a shared, but possibly episode-specific, prior. This construction enables amortized hierarchical Bayesian inference: the model can infer both the episode-level prior and the task-specific quantities solely from data, without additional parameter updates or retraining. Changing the prefix datasets at test time effectively alters the induced prior, providing a direct interface for controlling the predictor's prior sensitivity. This approach bridges the gap between meta-trained neural forecasting and explicit hierarchical Bayesian models, allowing adaptive predictive inference under variable prior assumptions.

Framework and Training Objective

The method involves hierarchical episode generation: for each meta-training episode, hyperparameters are sampled from a meta-prior, used to instantiate $K$ auxiliary "prior" datasets and one target dataset. The transformer is trained to maximize the log-likelihood of the target dataset labels, conditioned on both the target context and the concatenated prior prefix (i.e., the $K$ data blocks). This structure mirrors the conditional independence and marginalization properties of hierarchical Bayesian models, and allows the blackbox function implemented by the transformer to implicitly perform marginalized inference over the latent hierarchy. The framework is demonstrated on both conjugate (linear regression with Gaussian priors) and non-conjugate (logistic regression) likelihood families, using robust probabilistic metrics (KL and TV divergence) to the oracle Bayesian predictive.

Empirical Results

The MT-BICL model exhibits a number of strong empirical properties:

• Test-Time Prior Adaptation: Explicit conditioning on the prior-prefix datasets enables robustness to test-time prior shifts, both in- and out-of-meta-distribution (OoMD). The model's posterior predictive closely tracks the true hierarchical Bayesian PPD, achieving negligible KL divergence to the oracle across target context sizes and prior types.
• Robust Out-of-Distribution Generalization: Under systematically varied OoMD shifts—such as heavy-tailed Student’s t priors (including regimes with undefined variance or mean) and high-dimensional, flow-transformed priors (spiral flows)—the MT-BICL model demonstrates robust adaptation, with generalization bounded by the diversity present in the training meta-prior. Generalization degrades only when the test prior diverges drastically from training support (e.g., extreme heavy-tailedness), in line with the theoretical limitations of amortized inference.
• Mechanistic Prior Sensitivity: MT-BICL mechanistically interprets the prefix datasets as encoding prior information, not as additional evidence—a hypothesis confirmed via controlled prior adaptation and baseline comparisons with pooled-evidence Bayesian models. The neural model exhibits correct shrinkage behavior with increasing $K$ and appropriate posterior contraction in the presence of growing prior evidence, quantitatively matching hierarchical Bayesian expectations.
• Inference Efficiency: MT-BICL requires only a single forward pass for inference, yielding orders-of-magnitude lower inference latency relative to MCMC or SVI baselines, especially in models with structured/high-dimensional latent spaces, with essentially no loss in predictive calibration.
• Real-World Applicability: On spatiotemporal temperature prediction (ERA5), the prefix-conditioning mechanism improves negative log-likelihood (NLL) and MSE in the IID regime and demonstrates resilience to temporal distribution shift. Sequential prior prefixes offer advantages over set-based alternatives when correlations between training and test periods are maintained, whereas permutation-invariant architectures become advantageous under severe out-of-distribution shift.

Limitations

Despite strong robustness, the architecture imposes quadratic computational cost with respect to sequence length due to the transformer's self-attention. Non-permutation invariance of standard decoder-only transformers is a potential misalignment with the exchangeability inherent in Bayesian inference, although empirical investigations indicate that sensitivity to observation order is small in practice for the studied regimes. The generalization of amortized inference to new regimes is fundamentally limited by the diversity and coverage of the meta-training support; extrapolation under severe distributional shift remains challenging, particularly when the test priors possess statistical properties (e.g., number of moments) not observed during training.

Implications and Future Directions

MT-BICL enables scalable, data-driven amortized inference with explicit, controllable prior representations, closing the gap between fast neural predictors and rigorous Bayesian conditioning. The proposed mechanism supports test-time prior adaptation in resource-constrained and safety-critical applications, such as personalized medicine or adaptive trial planning, where integrating heterogeneous external data as priors is operationally required. The separation of prior information into prefix datasets generalizes to arbitrary prior families, as long as sufficient coverage and representational capacity is provided during meta-training.

Future work should address scaling challenges for extremely long context lengths, explore more complex forms of prior structure (e.g., combinatorial, graph-structured), implement more flexible permutation-equivariant transformer variants, and investigate meta-training strategies for improved extrapolation to out-of-meta-support priors. Further analysis of the intersection between in-context meta-learning and Bayesian nonparametrics is warranted, particularly in applications with unlimited or evolving prior structure.

Conclusion

Multi-Task Bayesian In-Context Learning introduces a tractable, principled mechanism for combining neural amortized inference with explicit prior adaptation through multi-dataset prefixing. This enables both computational and statistical efficiency in Bayesian predictive inference, aligning transformer-based in-context learning models more closely with the desiderata of hierarchical Bayesian reasoning and robust adaptation under prior uncertainty. The demonstrated empirical performance underscores its utility as a general-purpose module for uncertainty-aware prediction with user-controllable priors, complementing and extending meta-inference paradigms.