Calibrating Scientific Foundation Models with Inference-Time Stochastic Attention¶

会议: NeurIPS 2026
arXiv: 2604.19530
领域: 自监督学习
关键词: 科学基础模型, 不确定性量化, 随机注意力, 推理时校准, 预测区间

一句话总结¶

针对科学基础模型输出确定性且缺乏可靠不确定性的缺陷，提出在推理时将 Softmax 注意力替换为多项分布采样的随机注意力（Stochastic Attention），无需重训练仅通过一维后验校准即可输出保均值、高锐度且校准严格的预测区间。

研究背景与动机¶

随着深度学习逐步渗透到天气预测、流体力学与时间序列分析等科学建模任务，基于 Transformer 的科学基础模型展现出了匹敌甚至超越传统数值计算工作流的高精度预测能力。然而，这类模型在本质上是机械不透明的确定性函数逼近器，缺乏可靠的不确定性度量。在科学预测的实际场景中，系统往往面临着异构数据、带噪监督、动力学方程不匹配以及跨分布偏移等严苛挑战；若缺乏经过统计校准的预测不确定性，即便单点预测极其精准，决策者也无法在具有高风险的实际控制与防灾调度中信任模型输出。

现有为深度模型赋予不确定性的方法在科学基础模型面前存在根本性权衡与矛盾。以深度集成（Deep Ensembles）、SWAG、MultiSWAG 以及变分优化（IVON）为代表的权重空间与集成方法，通常需要收集优化轨迹、训练多个独立模型或修改训练目标，随着模型参数扩展到数亿甚至数十亿量级，重复构建模型的计算成本变得完全不可承受；基于前向扰动的方法（如 MC Dropout）虽免去重新训练，但其随机性来自硬编码的特征掩码，无法自洽确定校准的扰动尺度；而修改注意力机制的方法（如分层随机注意力 HSA）往往侵入训练阶段，极易破坏预训练权重已学得的空间路由与物理规律表示。

本文的核心切入点在于挖掘标准注意力机制内在的概率结构：对于每个查询位置，Softmax 权重天然构成了键值 token 上的分类概率分布，确定性注意力输出本质上就是该分布下的精确数学期望。因此，无需修改任何模型参数，只需在推理时将该精确期望替换为其经验样本均值即可诱导预测分布。核心 idea：在推理阶段将注意力权重的 Softmax 精确期望替换为由单一浓度参数控制的归一化多项分布样本均值（Stochastic Attention），在理论上保均值并以 $1/\nu$ 速率收缩方差，进而通过匹配留出集残差尺度的单变量贝叶斯优化高效确定最优浓度，实现零重训练的高质量预测校准。

方法详解¶

整体框架¶

随机注意力（Stochastic Attention, SA）作用于冻结（Frozen）的主干 Transformer 网络中的每个注意力头，直接介入查询与键值交互的注意力瓶颈层。给定输入样本，算法在推理时保持模型权重完全不变，利用标准 Softmax 产生的概率向量作为多项分布的参数进行随机采样，并计算其加权均值作为注意力输出。通过多次随机前向传播（Monte Carlo passes），网络隐式生成了一组具有统计变异性的预测集成，进而构建经验预测分布。为了使输出预测区间的宽度与实际预测误差相匹配，SA 在留出校准集上最小化随机扰动范数与真实确定性残差范数之间的二阶偏差，通过单变量贝叶斯优化快速寻优浓度参数 $\nu$。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["输入数据 x"] --> B["注意力瓶颈概率解析<br/>保留 Softmax 分布作为类别分布参数"]
    B --> C["归一化多项分布采样<br/>抽取独立 token 样本并做经验平均"]
    C --> D["单变量残差尺度校准<br/>贝叶斯优化匹配留出集预测误差确定浓度"]
    D --> E["多轮前向蒙特卡洛预测集成<br/>输出校准 PIT 检验与高锐度预测区间"]

关键设计¶

1. 归一化多项分布采样：构建保均值且方差可控的内部扰动机制

确定性注意力机制中，查询矩阵 $\mathbf{Q}\in\mathbb{R}^{n_q\times d}$ 与键矩阵 $\mathbf{K}\in\mathbb{R}^{n_k\times d}$ 计算点积缩放后，第 $t$ 个位置的归一化权重 $\boldsymbol{\pi}_t = \operatorname{softmax}(\mathbf{s}_t)\in\Delta^{n_k-1}$ 定义了离散概率单纯形。标准注意力输出 $\mathbf{o}_t = \boldsymbol{\pi}_t^\top \mathbf{V}$ 是值向量在 $\boldsymbol{\pi}_t$ 下的理论期望。为在不改变模型固有预测倾向的前提下引入随机性，SA 定义整数浓度参数 $\nu\in\mathbb{N}$，从多项分布中抽取计数值向量 $\mathbf{W}_t \sim \operatorname{Multinomial}(\nu, \boldsymbol{\pi}_t)$，并将其归一化构造随机注意力权重 $\widetilde{\boldsymbol{\pi}}_t = \frac{1}{\nu}\mathbf{W}_t$。该过程等价于独立抽取 $\nu$ 个类别变量 $z_t^{(m)}\overset{\text{iid}}{\sim} \operatorname{Categorical}(\boldsymbol{\pi}_t)$，并将对应的值向量做经验算术平均： $$ \widetilde{\mathbf{o}}t = \widetilde{\boldsymbol{\pi}}t^\top \mathbf{V} = \frac{1}{\nu}\sum}^\nu \mathbf{v $$ 该设计具有两个极为优良的数学性质。首先是}均值保持性与确定性渐进恢复（Mean Preservation and Deterministic Recovery）：在给定 $(\boldsymbol{\pi}_t, \mathbf{V})$ 条件下，$\mathbb{E}[\widetilde{\boldsymbol{\pi}}_t \mid \boldsymbol{\pi}_t] = \boldsymbol{\pi}_t$ 且 $\mathbb{E}[\widetilde{\mathbf{o}}_t \mid \boldsymbol{\pi}_t, \mathbf{V}] = \mathbf{o}_t$；当 $\nu \to \infty$ 时，依据强大数定律，随机输出几乎必然收敛到确定性注意力输出 $\mathbf{o}_t$，保证了预训练特征路由中心的稳定。其次是方差可缩放性（Covariance Scaling）：随机权重的协方差严格服从 $\operatorname{Cov}(\widetilde{\boldsymbol{\pi}}_t \mid \boldsymbol{\pi}_t) = \frac{1}{\nu}(\operatorname{diag}(\boldsymbol{\pi}_t) - \boldsymbol{\pi}_t\boldsymbol{\pi}_t^\top)$，传递到值空间后的协方差为： $$ \operatorname{Cov}(\widetilde{\mathbf{o}}_t^\top \mid \boldsymbol{\pi}_t, \mathbf{V}) = \frac{1}{\nu}\mathbf{V}^\top \left(\operatorname{diag}(\boldsymbol{\pi}_t) - \boldsymbol{\pi}_t\boldsymbol{\pi}_t^\top\right)\mathbf{V} $$ 这表明模型输出空间的扰动幅度完全由标量参数 $\nu$ 沿规范的 $1/\nu$ 速率平滑调控，将高维注意力随机化优雅退化为一个单一浓度控制维度。

2. 单变量残差尺度校准：基于留出集方差匹配的后验参数寻优

为了摒弃启发式或经验试错设定噪声参数的弊端，SA 引入了严谨的后验校准准则。在留出校准集 $\mathcal{D}_{\text{cal}} = \{(\mathbf{x}_i, \mathbf{y}_i)\}_{i=1}^{N_{\text{cal}}}$ 上，令 $f_{\boldsymbol{\theta}}(\mathbf{x})$ 表示冻结主干的确定性输出，$f_{\boldsymbol{\theta}, \nu}(\mathbf{x})$ 表示启用随机注意力后的单次前向随机输出。定义单样本诱导随机偏移范数 $R_\nu(\mathbf{x}) = \|f_{\boldsymbol{\theta}, \nu}(\mathbf{x}) - f_{\boldsymbol{\theta}}(\mathbf{x})\|_2$ 与确定性残差范数 $Z(\mathbf{x}, \mathbf{y}) = \|\mathbf{y} - f_{\boldsymbol{\theta}}(\mathbf{x})\|_2$。校准目标函数定义为二者均方误差的最小化： $$ \mathcal{L}{\text{SA}}(\nu) := \mathbb{E}}, \mathbf{y})\sim\mathcal{D{\text{cal}}}\left[\operatorname{Var}(R\nu(\mathbf{x})\mid\mathbf{x},\mathbf{y}) + \left(\mathbb{E}[R_\nu(\mathbf{x})\mid\mathbf{x},\mathbf{y}] - Z(\mathbf{x},\mathbf{y})\right)^2\right] $$ 这一损失函数具有清晰的统计解释：平方偏置项强行拉齐注意力诱导的扰动期望幅度与实际任务标签残差幅度，而方差项则惩罚不必要的随机抖动。在实现中，所有网络层共享单一标量 $\nu$，使得搜索空间收敛为一维整数域 $\Xi\subset\mathbb{N}$。利用带噪声评估的不确定性贝叶斯优化（Bayesian Optimization under Uncertainty），构建对数线性代理模型 $\ln \widehat{s}(\nu) = a\ln\nu + \ln b + \epsilon z$ 来拟合平均单样本扰动尺度 $\widehat{s}(\nu)$，并与目标残差尺度 $\widehat{s}_0$ 快速对齐。全流程仅需 1–3 分钟的前向抽样即可解析求出全局最优整数浓度 $\nu^\star$。

3. 多轮前向蒙特卡洛预测集成：兼顾校准充分性与锐度最大化

在确定最优操作点 $\nu^\star$ 后，推断阶段对测试样本 $\mathbf{x}$ 并行执行 $M$ 次随机前向传播，得到 Monte Carlo 样本预测集合 $\{\widehat{\mathbf{y}}^{(m)}(\mathbf{x})\}_{m=1}^M$。对于标量或网格场坐标输出，通过这些样本的经验累积分布函数直接计算任意置信度下的对称预测区间（如 95% 预测区间 PI-95）以及概率积分变换（Probability Integral Transform, PIT）诊断值。与事后对标量方差进行简单缩放（如温度缩放、保序回归）不同，SA 的扰动在网络底层特征交互中非线性传播，使得不同样本根据其注意力集中程度自然获得具有异方差性（Heteroscedastic）的置信边界，在实现 PIT 分布极度趋近于均匀分布（充分校准）的同时，保持了极度紧凑的区间宽度（极高锐度）。

损失函数 / 训练策略¶

SA 在方法本质上是完全无需重训练的（Zero Retraining）。冻结基础模型的所有预训练参数 $\boldsymbol{\theta}$，仅在推理时替换注意力加权计算。超参数优化目标完全退化为式 (9) 所示的标量校准损失 $\mathcal{L}_{\text{SA}}(\nu)$。在具体贝叶斯优化过程中，仅需设置采样前向预算 $M$（通常取 10–20 次）以及校准批次数 $B$ 即可在单卡数分钟内收敛。

实验关键数据¶

主实验¶

论文在三大基准上系统评估了 SA：全球天气预报（ClimaX 72小时 $Z_{500}$ 预测）、时间序列预测（TimesFM 在 ETT 基准上的少样本适配）以及 8 个 UCI 标准工程回归任务（使用 FT-Transformer）。对比基线涵盖权重后验近似（SWAG、MultiSWAG）、变分学习（IVON）、前向掩码扰动（MC Dropout、Contextual Dropout）以及同属注意力扰动但侵入训练的 HSA。在 ClimaX 评估中，对比了空间异常相关系数（ACC）、未经修正的原始 PIT 距均匀分布的 Wasserstein-1 距离（Native $W_1$）、温度缩放后校准距离（Scaled $W_1$）以及 95% 预测区间平均宽度（PI-95）。

方法	ACC ↑	原始 $W_1$ ↓	缩放后 $W_1$ ↓	PI-95 区间宽度 ↓	训练/适配额外开销
SA (本文, $\nu=4$)	0.9022	0.0560	0.0560	1051.9	~2 分钟 (单卡推理)
Contextual Dropout	0.9599	0.1048	0.0501	1080.9	训练期变分拟合
MultiSWAG (10 seeds)	0.9648	0.2188	0.0534	1075.5	~12 天 (独立训练10次)
SWAG	0.9648	0.2195	0.0475	1167.1	~29 小时 (轨迹收集)
IVON	0.9381	0.2132	0.0558	1422.0	训练期二阶变分优化
HSA	0.7476	0.1787	0.0367	2882.8	训练架构破坏 (精度下降22%)

在 UCI 8 个标准回归任务上，以 FT-Transformer 为主干的测试表明，在经过统一覆盖率校准后，SA 在 7 个数据集上取得了最窄的置信区间宽度，其余基线区间宽度普遍比 SA 宽出 1.2× 至 3.0×。

消融实验¶

针对注意力浓度参数 $\nu$ 的操作点选取，论文在 ClimaX 上深入探索了“校准优先”（Calibration-first, 基于式 (9) 选取 $\nu=4$）与“打分优先”（Score-oriented, 直接最小化连续概率排序列分 CRPS 选取 $\nu=25$）的特性解耦。同时，论文消融了 SWAG 模型的方差尺度（Scale 1 至 64），验证了基线的欠离散绝非超参数未对齐的假象。

操作配置 / 尺度选择	验证集 RMSE ↓	原始 $W_1$ 偏差 ↓	95% 覆盖率 (%) ↑	预测区间宽度 ↓	CRPS 打分 ↓
SA ($\nu=4$, 校准优先)	-	0.056	95.2%	1051.9	221.68
SA ($\nu=25$, 打分优先)	-	0.085	87.4%	782.4	135.42
SWAG (Scale 1)	265.73	0.215	8.9%	47.4	-
SWAG (Scale 16, 折中最佳)	263.86	0.174	32.0%	171.3	-
SWAG (Scale 64, 校准最佳)	267.44	0.116	60.4%	379.6	-

关键发现¶

天然具备极佳校准度：在所有任务中，所有基线在原生状态下均呈现出严重的 U 型 PIT 直方图（严重欠离散、过分自信），在 TimesFM 上覆盖率低至 15.6%–20.8%；而 SA 在原生状态下即可达到 61.5%（TimesFM）与 95.2%（ClimaX），其 Native $W_1$ 与经外部温度缩放后的基线持平或更优。
训练期随机化的脆弱性：同在注意力层引入随机性的 HSA，因训练期破坏了预训练模型的空间注意力分配，在 ClimaX 气象大模型上的相关系数从 0.96 骤降至 0.74，预测区间被迫拉宽至 SA 的 2.74 倍；SA 凭借保均值特性完美避开了这一灾难。
算力开销断崖式领先：在百兆级参数的大模型上，MultiSWAG 需要约 12 天的累计训练机时，而 SA 仅需 1–3 分钟评估单参数后验，且单次前向时延仅微增 11%（$\nu=4$ 时仅 0.457 秒 vs 原生 0.424 秒）。

亮点与洞察¶

将注意力权重重构为期望并用蒙特卡洛经验均值离散化：该视角极为精巧，直接将神经网络内部的 Softmax 概率分布转化为天然的随机数发生源，省去了外挂噪声生成器或额外采样分支的复杂开销。
校准优先与评分优先在同一单参数曲线上平滑演进：揭示了 CRPS 等评分指标容易被极度收敛的确定性预测（缩减预测误差项）掩盖不良校准的机理，并证明通过调节单一标量 $\nu$ 即可自由切换高校准与高敏锐度操作点。
可低成本复用到任意冻结的预训练 Transformer 基础模型中：无论是针对自回归时间序列预测、空间视觉网格气象大模型还是表格特征模型，该方法即插即用，特别契合大规模科学前沿计算无法承受重训练的真实痛点。

局限与展望¶

全网共享全局单一浓度参数的表达力受限：目前全网络各层、各注意力头均采用同一标量 $\nu$，虽然保证了一维优化的极简与稳健，但无法刻画深浅层不同语义空间所蕴含的不同不确定度；未来可拓展为逐层（Layer-wise）或逐头（Head-wise）的多变量校准。
校准损失目前局限于连续回归型残差：目前的优化目标针对连续欧氏距离残差构建，对于分类、离散生成任务或受约束的李群流形空间，需要设计专用的残差泛函来进行概率匹配。

评分¶

新颖性: ⭐⭐⭐⭐⭐ 巧妙重构注意力为期望并引入推理时归一化多项抽样，理论形式优美。
实验充分度: ⭐⭐⭐⭐⭐ 跨越时间序列、全球气象与 8 个回归基准，四维指标覆盖全面，消融深入。
写作质量: ⭐⭐⭐⭐⭐ 统计概念清晰，数学推导严格，问题阐述极具说服力。
价值: ⭐⭐⭐⭐⭐ 极大地降低了前沿大模型获得可靠可信预测不确定性的成本，应用落地前景广阔。

方法	ACC ↑	原始 \(W_1\) ↓	缩放后 \(W_1\) ↓	PI-95 区间宽度 ↓	训练/适配额外开销
SA (本文, \(\nu=4\))	0.9022	0.0560	0.0560	1051.9	~2 分钟 (单卡推理)
Contextual Dropout	0.9599	0.1048	0.0501	1080.9	训练期变分拟合
MultiSWAG (10 seeds)	0.9648	0.2188	0.0534	1075.5	~12 天 (独立训练10次)
SWAG	0.9648	0.2195	0.0475	1167.1	~29 小时 (轨迹收集)
IVON	0.9381	0.2132	0.0558	1422.0	训练期二阶变分优化
HSA	0.7476	0.1787	0.0367	2882.8	训练架构破坏 (精度下降22%)

操作配置 / 尺度选择	验证集 RMSE ↓	原始 \(W_1\) 偏差 ↓	95% 覆盖率 (%) ↑	预测区间宽度 ↓	CRPS 打分 ↓
SA (\(\nu=4\), 校准优先)	-	0.056	95.2%	1051.9	221.68
SA (\(\nu=25\), 打分优先)	-	0.085	87.4%	782.4	135.42
SWAG (Scale 1)	265.73	0.215	8.9%	47.4	-
SWAG (Scale 16, 折中最佳)	263.86	0.174	32.0%	171.3	-
SWAG (Scale 64, 校准最佳)	267.44	0.116	60.4%	379.6	-