多任务贝叶斯In-Context学习

Multi-Task Bayesian In-Context Learning

精选理由

这篇论文提出多任务贝叶斯ICL框架,速度比传统方法快几个数量级,还能适应新先验,在温度预测上表现很好。

AI 摘要

该研究提出多任务上下文学习框架用于分层贝叶斯预测推理,将先验信息表示为上下文数据集的前缀。使用Transformer在序列先验和目标任务上训练,学会跨先验家族调整预测。在包含元分布外先验和高维潜在结构的评估中,该方法匹配oracle贝叶斯预测器,速度提升数个数量级。在真实世界时空温度预测基准上验证了实际效果。

原文 · arXiv cs.LG

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