Calibrating Scientific Foundation Models with Inference-Time Stochastic Attention

📄 arXiv: 2604.19530v1 📥 PDF

作者: Akash Yadav, Taiwo A. Adebiyi, Ruda Zhang

分类: cs.LG, cs.CE, stat.ML

发布日期: 2026-04-21

💡 一句话要点

提出基于随机注意力的科学基础模型校准方法，提升预测不确定性

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 科学基础模型 不确定性校准 随机注意力 时间序列预测 天气预报 推理时修改 集成学习

📋 核心要点

  1. 现有科学基础模型输出确定性结果，缺乏对预测不确定性的有效校准，限制了其在高风险场景的应用。
  2. 论文提出随机注意力机制，通过在推理时随机化注意力权重，生成预测集成，无需重新训练模型。
  3. 实验表明，该方法在天气和时间序列预测等任务上，能有效提升模型校准能力，并获得更清晰的预测区间。

📝 摘要（中文）

基于Transformer的科学基础模型越来越多地应用于高风险场景，但现有架构提供确定性输出，对校准的预测不确定性支持有限。我们提出了一种轻量级的推理时修改方法——随机注意力，通过用受单个浓度参数控制的归一化多项式样本替换softmax权重来随机化注意力，从而生成预测集成而无需重新训练。为了设置这个参数，我们引入了一个校准目标，使随机注意力输出与目标匹配，从而产生一个高效的单变量事后调整问题。我们在天气和时间序列预测的两个科学基础模型以及一个额外的回归任务上评估了这种机制。在针对不确定性感知基准的测试中，我们发现随机注意力实现了最强的原生校准和最清晰的预测区间，同时具有相当的覆盖率，并且只需要几分钟的事后调整，而有竞争力的基准需要数天的重新训练。

🔬 方法详解

问题定义：现有基于Transformer的科学基础模型在进行预测时，通常输出的是确定性的结果，无法提供关于预测不确定性的有效信息。这在高风险的科学应用场景中是一个严重的问题，因为决策者需要了解预测结果的可信度，以便做出更明智的判断。现有方法要么需要耗时的模型重训练，要么校准效果不佳。

核心思路：论文的核心思路是在推理阶段引入随机性，通过对注意力机制的softmax权重进行随机采样，生成多个不同的预测结果，从而形成一个预测集成。这种方法的核心在于，通过引入随机性，可以模拟模型预测的不确定性，并利用集成学习的思想，提高预测的鲁棒性和校准性。

技术框架：该方法主要包含以下几个步骤：1) 在推理阶段，对于Transformer模型的每个注意力层，不再使用标准的softmax函数计算注意力权重，而是使用一个参数化的多项式分布进行采样。2) 通过调整多项式分布的浓度参数，控制随机性的强度。3) 对同一个输入进行多次采样，得到多个不同的预测结果。4) 将这些预测结果进行集成，例如取平均值或使用其他集成方法，得到最终的预测结果和不确定性估计。5) 使用一个校准目标函数，调整浓度参数，使得模型的预测结果与真实值之间的校准误差最小化。

关键创新：该方法最重要的创新点在于，它是一种轻量级的推理时修改方法，无需对模型进行重新训练，即可实现预测不确定性的校准。与现有方法相比，该方法具有更高的效率和更低的计算成本。此外，通过引入可调节的浓度参数，可以灵活地控制随机性的强度，从而更好地适应不同的任务和数据集。

关键设计：关键设计包括：1) 使用多项式分布对注意力权重进行采样，而不是直接添加噪声。2) 引入浓度参数来控制随机性的强度，并通过校准目标函数进行优化。3) 使用集成学习的方法，将多个预测结果进行融合，得到最终的预测结果和不确定性估计。校准目标函数的设计至关重要，需要能够有效地衡量模型的校准误差，并指导浓度参数的优化。

📊 实验亮点

实验结果表明，随机注意力机制在天气和时间序列预测任务上，能够显著提升模型的校准能力，并获得更清晰的预测区间。与现有不确定性感知基线方法相比，该方法在实现相当覆盖率的同时，只需要几分钟的事后调整，而基线方法需要数天的重新训练。这表明该方法具有更高的效率和更低的计算成本。

🎯 应用场景

该研究成果可广泛应用于需要高可靠性和不确定性估计的科学领域，例如天气预报、气候建模、金融时间序列分析、医疗诊断等。通过提供校准的预测不确定性，该方法可以帮助决策者更好地理解预测结果的可靠性，从而做出更明智的决策，并降低风险。未来，该方法可以进一步扩展到其他类型的科学模型和任务中。

📄 摘要（原文）

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 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 two scientific foundation models for weather and timeseries forecasting along with an additional regression task. Across benchmarks against uncertainty-aware baselines, we find that Stochastic Attention achieves the strongest native calibration and the sharpest prediction intervals at comparable coverage, while requiring only minutes of post-hoc tuning versus days of retraining for competitive baselines.