Variational Routing: 校准 MoE Transformer 的可扩展贝叶斯框架¶

会议: ICML 2026
arXiv: 2603.09453
代码: 待确认
领域: 模型压缩 / LLM 效率 / AI 安全
关键词: 混合专家网络, 贝叶斯推断, 校准, 不确定性量化, 稀疏路由

一句话总结¶

提出变分路由框架 VMoER——通过对 MoE 层的路由决策进行变分推断而非权重推断，实现高效贝叶斯不确定性建模，在保持 <1% FLOPs 额外开销的同时将校准误差降低 94%、路由稳定性提升 38%。

研究背景与动机¶

领域现状：基础模型规模达到万亿参数，通过 MoE 稀疏专家路由实现高效扩展。然而当前路由机制采用确定性 Top-K 策略，在输入扰动下容易出现错误专家选择。

现有痛点：（1）确定性路由对输入噪声敏感，出现脆性失败；（2）预测高度过置信，校准误差大；（3）现有贝叶斯方法针对权重不确定性，计算开销大，不适用于万亿参数规模。

核心矛盾：在确保模型可靠部署的前提下，如何以最小计算成本为 MoE 模型注入不确定性感知能力。

本文目标：设计轻量级贝叶斯框架，直接对路由决策（而非权重）进行概率建模。

切入角度：将 MoE 路由重新表述为潜变量模型，观察到——（1）确定性路由隐含忽视了 logits→概率→选择的不确定性链条；（2）Top-K 操作本质上是多标签问题。

核心 idea：从权重空间转向决策空间进行变分推断——通过 amortised 推断直接对路由 logits 或温度参数进行概率建模，绕过高维权重后验的复杂性。

方法详解¶

整体框架¶

VMoER 包含两条互补推断路径——（1）Logit 空间推断：对路由 logits \(\mathbf{l}\) 应用变分高斯分布 \(q_\phi(\mathbf{l}|\mathbf{u})\)，显式建模专家间相关性；（2）选择空间推断：学习输入相关的温度参数 \(T_\phi(\mathbf{u})\) 动态调节决策边界，通过 Sample-K 而非 Top-K 实现随机化专家选择。

关键设计¶

变分高斯 Logit 路由（VGLR）:
- 功能：对路由 logits 进行 amortised 变分推断，通过全协方差建模显式捕捉专家相关性。
- 核心思路：采用居中先验 \(p(\mathbf{l}|\mathbf{u})=\mathcal{N}(\mathbf{l}_{det}, \mathbf{I})\)，其中 \(\mathbf{l}_{det}=\mathbf{u}\mathbf{W}_r\)。后验均值 \(\boldsymbol{\mu}_{post}(\mathbf{u})=\mathbf{l}_{det}+\Delta\boldsymbol{\mu}_\phi(\mathbf{u})\)，推断网络学习残差校正而非从零开始。Cholesky 因子化 \(\boldsymbol{\Sigma}_{post}=\mathbf{LL}^\top\) 参数化协方差（复杂度 \(O(N^2)\)，N≤64 可接受）。推断时 MC 采样平均化。
- 设计动机：权重空间推断间接通过线性投影传播参数噪声；直接建模路由决策变量（logits→概率）效率更高。全协方差超越 mean-field 假设捕捉专家间相关性。
变分温度缩放路由（VTSR）:
- 功能：学习输入相关温度参数 \(T_\phi(\mathbf{u})\) 动态调节 softmax 锐度，实现高效的单维度变分推断。
- 核心思路：约束变分族到 1D 流形——沿着由确定性 logits 与输入相关温度定义的轨迹 \(q_\phi(\mathbf{p}|\mathbf{u})=\text{Softmax}(\mathbf{l}_{det}/T_\phi(\mathbf{u}))\) 移动。通过 Gumbel-Softmax 进行 Sample-K 采样。KL 正则化简化为 Shannon 熵。
- 设计动机：VGLR 需多次采样导致推断延迟；VTSR 限制在标度参数空间，计算开销仅为 \(O(D_H)\)（<0.67% FLOPs）。
居中先验与残差学习:
- 功能：通过约束后验在确定性解附近进行，保证微调时预训练路由性能不丧失。
- 核心思路：后验不从零开始学习，而是学习残差 \(\Delta\boldsymbol{\mu}_\phi(\mathbf{u})\) 加到原 logits 上，使 KL 项自动围绕零进行 regularization。
- 设计动机：微调时路由往往陷入困境；居中先验提供稳定性。

训练策略¶

VGLR：\(\mathcal{L}_{ELBO}=\mathbb{E}_{q_\phi(\mathbf{l}|\mathbf{u})}[\log p(\mathbf{y}|\mathbf{l},\mathbf{u})]-\beta D_{KL}(q_\phi(\mathbf{l}|\mathbf{u})\|\mathcal{N}(\mathbf{0},\mathbf{I}))\)。VTSR：主要优化重构，通过代理损失 \(\mathcal{L}_{reg}=-\log T_\phi(\mathbf{u})\) 隐含推动温度朝向先验。

实验关键数据¶

主实验¶

数据集	模型	指标	MAP 基线	VGLR-MF	VGLR-FC	VTSR
OpenBookQA	Granite-3B	ECE ↓	0.252	0.026	0.015	0.052
OpenBookQA	Qwen-2.7B	ECE ↓	0.127	0.028	0.014	0.022
OpenBookQA	DeepSeek-16B	ECE ↓	0.168	0.067	0.054	0.060

消融¶

实验项	Granite ECE	Qwen ECE	发现
确定性 Top-K	0.252	0.127	基线过置信
固定温度缩放	0.107	0.102	跨模型不稳定（精度掉 3%）
VGLR-FC 全协方差	0.015	0.014	校准误差降 94%
噪声鲁棒性（σ=0.01）	Jaccard=0.532	Jaccard>0.612	VGLR 稳定性提升 38%
OoD 检测 AUROC	0.659（基线）	0.749（VGLR）	内部 logit 方差信号优于 gating 熵

关键发现¶

全协方差关键——显式建模相关性使校准显著改善。
VTSR 在准确率稳定性上优于全局固定温度。
内部推理不确定性为 OoD 检测提供比预测熵更强的信号。

亮点与洞察¶

概率生成模型视角：将 MoE 路由形式化为潜变量模型，将启发式负载均衡和辅助损失解释为隐含贝叶斯先验。
从权重空间转向决策空间：直接对路由 logits 或温度参数推断既捕捉必要不确定性又规避维数灾难。
双路径灵活设计：VGLR 最佳校准但推断延迟略高；VTSR 牺牲一点精度换取单过推理零额外采样成本。
可迁移构件：居中先验+残差学习、温度缩放 1D 流形简化可推广。

局限与展望¶

VTSR 训练不稳定——温度参数易陷入 collapse，需精心初始化。
评估仅限 MCQA next-token 预测任务，未涵盖长序列生成中的错误累积。
未评估更大规模——最大 DeepSeek-16B。
改进：稳定 VTSR 的变分目标；扩展到序列级不确定性；与权重空间贝叶斯方法混合。

评分¶

新颖性: ⭐⭐⭐⭐⭐ 首次系统将变分推断应用于 MoE 路由决策而非权重。
实验充分度: ⭐⭐⭐⭐ 3 SOTA 架构 + 多维评估；仅 MCQA 任务最大 16B。
写作质量: ⭐⭐⭐⭐⭐ 理论清晰、概率生成过程推导严谨。
价值: ⭐⭐⭐⭐⭐ 为万亿参数基础模型可靠部署指出高效路径。