Multi-Task Bayesian In-Context Learning

• URL: http://arxiv.org/abs/2606.20538v1
• Date: Thu, 18 Jun 2026 17:50:10 GMT
• Title: Multi-Task Bayesian In-Context Learning
• Abstract summary: We introduce a multi-task in-context learning framework for amortized hierarchical Bayesian predictive inference.<n>A transformer trained on sequences of prior and target tasks learns to adapt its predictions across families of priors.<n>Our method matches oracle predictors while being orders of magnitude faster than existing approaches.
• Score: 41.16196712546687
• License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/
• 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/.

Related papers

• Guidance for Prior Change via Density Ratio Estimation [5.018167017476858]
  PriorGuide was introduced as an inference-time guidance method.<n>We propose an unbiased test-time guidance framework that leverages Density Ratio Estimation (DRE) to learn a score guidance term.<n>Our method matches or outperforms PriorGuide on C2ST and MMD in most tasks.
  arXiv  Detail & Related papers  (2026-08-22T02:17:41Z)
• Class Adaptive Conformal Training [32.3021967024497]
  Conformal Prediction offers principled framework for uncertainty quantification.<n>Class Adaptive Conformal Training (CaCT) learns to shape prediction sets class-conditionally without making distributional assumptions.<n>CaCT consistently outperforms prior conformal training methods.
  arXiv  Detail & Related papers  (2026-01-14T14:41:23Z)
• Goal-Oriented Time-Series Forecasting: Foundation Framework Design [11.999600538978044]
  We propose a training methodology that enables forecasting models to adapt their focus to application-specific regions of interest without retraining.<n>The approach partitions the prediction space into fine-grained segments during training, which are dynamically reweighted and aggregated to emphasize the target range specified by the application.<n> Experiments on standard benchmarks and a newly collected wireless communication dataset demonstrate that our method not only improves forecast accuracy within regions of interest but also yields measurable gains in downstream task performance.
  arXiv  Detail & Related papers  (2025-04-24T12:34:43Z)
• Bayesian Test-Time Adaptation for Vision-Language Models [51.93247610195295]
  Test-time adaptation with pre-trained vision-language models, such as CLIP, aims to adapt the model to new, potentially out-of-distribution test data.<n>We propose a novel approach, textbfBayesian textbfClass textbfAdaptation (BCA), which in addition to continuously updating class embeddings to adapt likelihood, also uses the posterior of incoming samples to continuously update the prior for each class embedding.
  arXiv  Detail & Related papers  (2025-03-12T10:42:11Z)
• Relational Conformal Prediction for Correlated Time Series [56.59852921638328]
  We address the problem of uncertainty quantification in time series by exploiting correlated sequences.<n>We propose a novel distribution-free approach based on conformal prediction framework and quantile regression.<n>Our approach provides accurate coverage and achieves state-of-the-art uncertainty quantification in relevant benchmarks.
  arXiv  Detail & Related papers  (2025-02-13T16:12:17Z)
• When Rigidity Hurts: Soft Consistency Regularization for Probabilistic Hierarchical Time Series Forecasting [69.30930115236228]
  Probabilistic hierarchical time-series forecasting is an important variant of time-series forecasting. Most methods focus on point predictions and do not provide well-calibrated probabilistic forecasts distributions. We propose PROFHiT, a fully probabilistic hierarchical forecasting model that jointly models forecast distribution of entire hierarchy.
  arXiv  Detail & Related papers  (2023-10-17T20:30:16Z)
• Boosted Control Functions: Distribution generalization and invariance in confounded models [10.503777692702952]
  We introduce a strong notion of invariance that allows for distribution generalization even in the presence of nonlinear, non-identifiable structural functions.<n>We propose the ControlTwicing algorithm to estimate the Boosted Control Function (BCF) using flexible machine-learning techniques.
  arXiv  Detail & Related papers  (2023-10-09T15:43:46Z)
• When Rigidity Hurts: Soft Consistency Regularization for Probabilistic Hierarchical Time Series Forecasting [69.30930115236228]
  Probabilistic hierarchical time-series forecasting is an important variant of time-series forecasting. Most methods focus on point predictions and do not provide well-calibrated probabilistic forecasts distributions. We propose PROFHiT, a fully probabilistic hierarchical forecasting model that jointly models forecast distribution of entire hierarchy.
  arXiv  Detail & Related papers  (2022-06-16T06:13:53Z)
• Few-shot Conformal Prediction with Auxiliary Tasks [29.034390810078172]
  We develop a novel approach to conformal prediction when the target task has limited data available for training. We obtain substantially tighter prediction sets while maintaining desirable marginal guarantees by casting conformal prediction as a meta-learning paradigm. We demonstrate the effectiveness of this approach across a number of few-shot classification and regression tasks in natural language processing, computer vision, and computational chemistry for drug discovery.
  arXiv  Detail & Related papers  (2021-02-17T17:46:57Z)
• Robust Validation: Confident Predictions Even When Distributions Shift [19.327409270934474]
  We describe procedures for robust predictive inference, where a model provides uncertainty estimates on its predictions rather than point predictions. We present a method that produces prediction sets (almost exactly) giving the right coverage level for any test distribution in an $f$-divergence ball around the training population. An essential component of our methodology is to estimate the amount of expected future data shift and build robustness to it.
  arXiv  Detail & Related papers  (2020-08-10T17:09:16Z)
• Video Prediction via Example Guidance [156.08546987158616]
  In video prediction tasks, one major challenge is to capture the multi-modal nature of future contents and dynamics. In this work, we propose a simple yet effective framework that can efficiently predict plausible future states.
  arXiv  Detail & Related papers  (2020-07-03T14:57:24Z)
• A machine learning approach for forecasting hierarchical time series [4.157415305926584]
  We propose a machine learning approach for forecasting hierarchical time series. Forecast reconciliation is the process of adjusting forecasts to make them coherent across the hierarchy. We exploit the ability of a deep neural network to extract information capturing the structure of the hierarchy.
  arXiv  Detail & Related papers  (2020-05-31T22:26:16Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.