arXiv:2604.19530cs.LGcs.CE2026-04

让科学大模型输出带可信度,推理时随机化注意力即可

Calibrating Scientific Foundation Models with Inference-Time Stochastic Attention

论文配图:Calibrating Scientific Foundation Models with Inference-Time Stochastic Attention
图 1 · 摘自论文原文
  • 用随机采样替换注意力权重,不重训就能生成预测集合
  • 在气象与时间序列任务中,校准效果最佳且区间更精确
  • 仅需调一个参数,计算成本比现有方法低近1000倍

基于Transformer的科学基础模型在高风险场景中应用日益广泛,但现有架构输出确定性结果,难以提供可校准的预测不确定性。本文提出Stochastic Attention,一种轻量级的推理时修改方法:通过将softmax权重替换为受单一浓度参数控制的归一化多项式采样,实现注意力的随机化,从而生成无需重训练的预测集成。为确定该参数,引入校准目标,使随机注意力输出匹配预期分布,转化为高效的单变量后处理优化问题。我们在天气预测、时间序列建模及多个回归任务上评估该机制。相比其他具备不确定性感知能力的基线模型,样本平均随机注意力在原生校准性能上表现最优,且预测区间更紧凑,在相同校准水平下表现更优;其适应成本比次优基线低近三个数量级。

原文摘要 · Abstract (English)

Transformer-based scientific foundation models are increasingly deployed in high-stakes settings, but current architectures give deterministic outputs and provide limited support for calibrated predictive uncertainty. We propose Stochastic Attention, a sample average lightweight inference-time modification that randomizes attention by replacing softmax weights with normalized multinomial samples controlled by a single concentration parameter, and produces predictive ensembles without retraining. To set this parameter, we introduce a calibration objective that matches the stochastic attention output with the target, yielding an efficient univariate post-hoc tuning problem. We evaluate this mechanism on scientific foundation models for weather and time-series forecasting, as well as several regression tasks. Across benchmarks against uncertainty-aware baselines, we find that Sample Average Stochastic Attention achieves the strongest native calibration and the sharpest prediction intervals at comparable calibration, with adaptation costs nearly three orders of magnitude lower than the next-best baseline.

不确定性建模注意力机制推理优化

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。