Multi-Task Bayesian In-Context Learning

📄 arXiv: 2606.20538v1 📥 PDF

作者: Qingyang Zhu, Eric Karl Oermann, Kyunghyun Cho

分类: cs.LG

发布日期: 2026-06-18

备注: ICML 2026

🔗 代码/项目: GITHUB


💡 一句话要点

提出多任务贝叶斯上下文学习以解决不确定性量化问题

🎯 匹配领域: 支柱八:物理动画 (Physics-based Animation)

关键词: 贝叶斯推断 上下文学习 多任务学习 不确定性量化 变换器模型

📋 核心要点

  1. 现有的贝叶斯推断方法在处理不确定性时面临计算复杂性和模型假设限制,导致鲁棒性不足。
  2. 论文提出了一种多任务上下文学习框架,通过将先验信息作为上下文数据集的前缀,增强了模型在不同先验下的适应能力。
  3. 实验结果表明,该方法在处理高维潜在结构和元分布外先验时,性能与理想贝叶斯预测器相当,且速度显著提升。

📝 摘要(中文)

贝叶斯预测推断为不确定性量化、数据效率和稳健泛化提供了原则性框架。然而,精确推断往往难以实现,且可扩展的近似方法可能计算开销大或需要限制性建模假设,导致预测性能下降。近期出现的基于先前数据拟合的上下文模型作为一种摊销替代方案,通过直接将数据集映射到预测分布,但现有方法与训练先验的支持紧密耦合,缺乏在测试时适应新先验的显式机制,导致在分布转移下的鲁棒性有限。本文提出了一种多任务上下文学习框架,用于摊销层次贝叶斯预测推断,明确将先验信息表示为上下文数据集的前缀。经过训练的变换器能够在不同先验家族间调整其预测。我们的方案在一系列逐渐增加难度的评估中表现出色,匹配了理想贝叶斯预测器,同时速度快几个数量级。我们还在实际的时空温度预测基准上展示了其实际相关性。

🔬 方法详解

问题定义:本文旨在解决贝叶斯预测推断中的不确定性量化和模型适应性问题。现有方法在面对分布转移时表现出鲁棒性不足,且计算开销较大。

核心思路:提出的框架通过将先验信息作为上下文数据集的前缀,利用变换器模型学习在不同先验间的适应性,从而实现更高效的预测推断。

技术框架:整体架构包括数据预处理、上下文信息编码、变换器模型训练和预测阶段。模型通过学习先验和目标任务的序列来进行适应性预测。

关键创新:最重要的创新在于引入了多任务上下文学习框架,使得模型能够在不同的先验分布下进行有效的预测,而不是仅依赖于训练时的先验支持。

关键设计:模型设计中使用了变换器结构,损失函数采用了适应性损失,参数设置上进行了优化以提高模型的学习效率和预测准确性。具体细节包括对高维潜在结构的处理和对元分布外先验的适应能力。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果显示,提出的方法在处理高维潜在结构和元分布外先验时,性能与理想贝叶斯预测器相当,且速度提升达几个数量级,显著优于现有基线方法,展现出良好的实用性和效率。

🎯 应用场景

该研究的潜在应用领域包括气象预测、金融市场分析和医疗数据分析等,能够在面对不确定性和数据分布变化时提供更为可靠的预测。其实际价值在于提升模型的适应性和效率,未来可能推动更多领域的智能决策支持系统的发展。

📄 摘要(原文)

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