通过有限专家库实现通信高效的专家路由

# 通过有限专家库实现通信高效 MoE 的专家路由

来源: https://arxiv.org/html/2605.05278

###### 摘要

资源高效的机器学习越来越多地使用稀疏混合专家（MoE）架构，其中门控机制既作为学习组件，也作为控制计算、通信和准确性的路由接口。受有限速率 MoE 门控解释的启发，我们将门控视为随机信道，并利用 $I(X;T)$ 来量化分配给所选专家的路由信息。为了使相关信息量在合成示例之外也可行，我们构建了一个基于预训练 CNN 专家和离散数据依赖选择规则的有限库 MNIST 实验。由于所选模型属于有限候选集，算法互信息 $I(S;W)$ 允许通过经验后验 $q(W|S)$ 采用闭式离散熵估计器。通过扫描数据依赖参数 $\alpha$，我们观察到 $\widehat{I}(S;W)$ 单调跟踪泛化间隙，而 Xu-Raginsky 界表现出预期的松驰性。我们还与均匀联合界基线进行了比较，并引入了 $I(X;T)$ 的经验估计器以及 Blahut-Arimoto 过程，以便在专家库上追踪准确率-速率曲线。所提出的框架为分析资源感知的 MoE 推理系统提供了实用工具，并将 $I(X;T)$ 和 $D(R_g)$ 解释为高效专家路由的设计代理。

## I 引言

MoE 架构通过门控机制组合专用预测器，该机制要么对专家进行软加权，要么进行离散选择 [1](https://arxiv.org/html/2605.05278#bib.bib1), [2](https://arxiv.org/html/2605.05278#bib.bib2)。稀疏 MoE 模型目前广泛用于大型神经网络架构中，因为对于每个输入仅激活一小部分专家，从而实现了可扩展且资源高效的机器学习和人工智能 [3](https://arxiv.org/html/2605.05278#bib.bib3)。这种资源感知的推理机制在通信和计算受限的信号处理系统中尤其相关，包括航空航天应用，其中板载处理、带宽、延迟和能量预算均有限。一个互补的理论问题是：当专家路由受到通信、隐私或压缩约束时，应如何解释门控机制。

最近的一种信息论表述将门控视为随机信道 $X \to T$，并使用门控互信息 $R_g = I(X;T)$ 作为操作信息速率。在这种观点下，$R_g$ 控制有多少输入信息到达专家库，而学习算法对样本的依赖性则由 $I(S;\Theta)$ [4](https://arxiv.org/html/2605.05278#bib.bib4) 衡量，其中 $\Theta$ 收集完整的 MoE 参数。由此产生的风险分解将速率失真项 $D(R_g)$ 与信息论泛化罚项相结合。这种观点与经典的信息论泛化界 [5](https://arxiv.org/html/2605.05278#bib.bib5), [6](https://arxiv.org/html/2605.05278#bib.bib6)、通信受限学习 [7](https://arxiv.org/html/2605.05278#bib.bib7), [8](https://arxiv.org/html/2605.05278#bib.bib8)、具有信息约束的分层决策 [9](https://arxiv.org/html/2605.05278#bib.bib9), [10](https://arxiv.org/html/2605.05278#bib.bib10), [11](https://arxiv.org/html/2605.05278#bib.bib11), [12](https://arxiv.org/html/2605.05278#bib.bib12) 以及 MoE 模型的风险分析 [13](https://arxiv.org/html/2605.05278#bib.bib13), [14](https://arxiv.org/html/2605.05278#bib.bib14) 相关。这也符合边缘推理、分布式传感、自主平台以及航空或航天信号处理流水线中资源高效 AI 系统的设计，在这些系统中，路由决策可能不仅代表神经网络选择，还代表通信和计算决策。

一个自然的担忧是，合成实验虽然有助于隔离信息速率效应，但本身并不能证明该框架如何应用于标准深度学习基准。困难不在于概念层面，而在于统计和计算层面：对于具有约 $10^5$ 个连续值参数的 CNN，参数变量是高维的，随机训练样本在权重上诱导连续后验，且 $I(S;\Theta)$、$I(X;T)$ 和 $D(R_g)$ 通常没有闭式解。因此，任何面向基准的扩展都必须构造得使这些量保持可直接测量或有意义地可近似。

本文提供了这样的扩展。我们考虑 MNIST 上的一组有限预训练 CNN 专家库，以及一个离散选择规则，该规则将训练样本 $S$ 映射到候选模型上的后验分布 $q(W|S)$，其中 $W$ 现在表示选定的候选者索引（这是对上述更一般 $\Theta$ 的刻意限制）。这种设置足够丰富以与基准相关，但又足够有限以允许对 $I(S;W)$ 进行精确的基于熵的评估。与先前的工作通过 PAC-Bayes 或噪声迭代分析为随机梯度方法估计 $I(S;\Theta)$ 不同 [15](https://arxiv.org/html/2605.05278#bib.bib15), [16](https://arxiv.org/html/2605.05278#bib.bib16)，我们的有限库构造使得 $I(S;W)$ 可以从闭式离散熵计算，代价是将 $W$ 限制为预训练候选集。我们此外开发了离散 MoE 门控下 $I(X;T)$ 的插件估计器，并在同一专家库实例化相应的 Blahut-Arimoto 速率失真求解器。从系统角度来看，这提供了一种可行的方式来研究如何利用路由信息作为资源感知 MoE 推理的设计代理。

我们的贡献如下：

1. 我们在 MNIST 上引入了一种有限库 CNN 协议，将深度模型选择重构为资源高效 MoE 推理的可行信息论实验。
2. 我们推导了由实践中使用的 $\alpha$-混合选择规则诱导的 $I(S;W)$ 的蒙特卡洛估计器，以及 Miller-Madow 偏差分析和自助法置信区间。
3. 我们将所得界与有限库上的均匀联合界基线进行比较，并明确报告松驰因子。
4. 我们通过固定专家库上的 Blahut-Arimoto 求解器，经验实例化输入依赖的有限速率路由扩展，产生经验速率失真曲线 $\widehat{D}(\rho)$。

**组织结构：** 本文其余部分组织如下。第二部分 [II](https://arxiv.org/html/2605.05278#S2) 回顾有限速率 MoE 框架和 Xu-Raginsky 泛化界。第三部分 [III](https://arxiv.org/html/2605.05278#S3) 介绍 MNIST 上的有限库 CNN 协议和 $\alpha$-混合选择规则。第四部分 [IV](https://arxiv.org/html/2605.05278#S4) 推导 $I(S;W)$ 的离散熵估计器，分析其偏差和方差，并与联合界基线进行比较。第五部分 [V](https://arxiv.org/html/2605.05278#S5) 通过 Blahut-Arimoto 求解器将构造扩展到输入依赖的有限速率路由。第六部分 [VI](https://arxiv.org/html/2605.05278#S6) 报告实验结果，第七部分 [VII](https://arxiv.org/html/2605.05278#S7) 和第八部分 [VIII](https://arxiv.org/html/2605.05278#S8) 讨论局限性并得出结论。

## II 有限速率 MoE 背景

令 $(X,Y) \sim \mathcal{D}$ 表示输入-标签对，并考虑 $n$ 个专家 $\{h_g(\cdot; W_g)\}_{g=1}^n$。门控将输入 $x$ 映射到专家上的概率向量，并采样路由变量 $T \in [n]$。令 $\Theta=(W_{\mathrm{gate}}, W_{\mathrm{exp}})$，其中 $W_{\mathrm{exp}}=(W_1, \ldots, W_n)$ 收集完整的 MoE 参数。我们全程假设损失取值在 $[0,1]$ 内，这是 Xu-Raginsky 界所要求的。总体风险和实证风险分别为

$$
R(\Theta)=\mathbb{E}_{(X,Y)\sim\mathcal{D}}\mathbb{E}_{T\sim P_{W_{\mathrm{gate}}}(\cdot|X)}\left[\ell\left(h_T(X;W_T),Y\right)\right], \tag{1}
$$

$$
R_S(\Theta)=\frac{1}{m}\sum_{j=1}^m\mathbb{E}_{T\sim P_{W_{\mathrm{gate}}}(\cdot|x_j)}\left[\ell\left(h_T(x_j;W_T),y_j\right)\right], \tag{2}
$$

其中 $S=\{(x_j,y_j)\}_{j=1}^m \sim \mathcal{D}^m$。门控被解释为信道 $P(T|X)$，其实现的门控速率为

$$
R_g \triangleq I(X;T). \tag{3}
$$

对应的单字母速率失真目标为

$$
D(R_g) \triangleq \inf_{P(T|X): I(X;T)\leq R_g} \mathbb{E}\left[\ell\left(h_T(X;W_T),Y\right)\right], \tag{4}
$$

其中在计算下确界时专家库是固定的。在高维深度网络中，(4) 很少被精确求解；尽管如此，它作为正则化门控的设计代理仍是有用的。

信息论泛化组件由 $I(S;\Theta)$ 控制。将 Xu-Raginsky 界 [5](https://arxiv.org/html/2605.05278#bib.bib5) 专门应用于 MoE 规则给出

$$
\left|\mathbb{E}[R(\Theta)]-\mathbb{E}[R_S(\Theta)]\right| \leq \sqrt{\frac{2}{m}I(S;\Theta)}. \tag{5}
$$

如果学习到的门控在实现速率 $R_g$ 下对速率失真目标是经验近最优的，即，

$$
\mathbb{E}[R_S(\Theta)] \leq D(R_g) + \delta_m, \tag{6}
$$

那么结合 (5) 和 (6) 可得

$$
\mathbb{E}[R(\Theta)] \leq D(R_g) + \delta_m + \sqrt{\frac{2}{m}I(S;\Theta)}. \tag{7}
$$

方程 (7) 将表达性/通信项 $D(R_g)$ 与估计项 $I(S;\Theta)$ 分开。我们强调，条件 (6) 是一个设计假设，*并未* 在我们的有限库实验中得到验证；该实验隔离了估计项。

## III MNIST 上的有限库 CNN 协议

为了使 $I(S;\Theta)$ 可行，我们将 $\Theta$ 限制为有限的、预训练的候选集，使得所选候选者的索引扮演学习参数的角色。稍微滥用符号，我们在下面用 $W$ 表示这个取索引值的随机变量。

### III-A 候选库

基准扩展使用一组有限个 $R$ 预训练 CNN 分类器

$$
\mathcal{W}_R \triangleq \{W_1, \ldots, W_R\}. \tag{8}
$$

在实现中，每个候选者使用相同的紧凑 CNN 架构，包括两个卷积层，分别具有 16 和 32 个滤波器，每个层后接最大池化，一个具有 64 个隐藏单元的全连接层，以及一个 10 类 softmax 输出。每个候选者在独立抽取的 10,000 张 MNIST 训练图像子集上（跨候选者有放回采样）预训练一个 epoch，使用学习率为 $10^{-3}$ 的 Adam 优化器和独立的随机初始化。该过程产生的候选者在 MNIST 测试集上的测试准确率在 0.89-0.92 之间，成对预测不一致率约为 0.06-0.10，表明存在非平凡多样性。默认配置总结于表 I。

对于大小为 $m$ 的样本 $S$，候选者 $r$ 的经验 0-1 误差为

$$
\widehat{R}_S(W_r) \triangleq \frac{1}{m}\sum_{j=1}^m \mathbf{1}\left\{\arg\max_{c\in\{0,\ldots,9\}} f_c(x_j; W_r) \neq y_j\right\}, \tag{9}
$$

其中 $f_c(x; W_r)$ 是候选者 $r$ 分配给类别 $c$ 的 softmax 分数。

**表 I:** MNIST 有限库 CNN 结果。熵和互信息以 nat 为单位报告。
**表 II:** 混合参数 $\alpha$ 对信息-泛化项的影响。

### III-B 数据依赖的离散选择

给定 $S$，经验风险最小化候选者为

$$
r^\star(S) \triangleq \mathop{\mathrm{arg\,min}}_{1\leq r\leq R} \widehat{R}_S(W_r), \tag{10}
$$

其中平局（在 $m=256$ 下很少出现）通过选择最小索引来打破。代码中使用的随机化学习规则是有限库上的 $\alpha$-混合后验：

$$
q_\alpha(r|S) \triangleq \frac{1-\alpha}{R} + \alpha \, \mathbf{1}\{r=r^\star(S)\}, \quad r\in\{1,\ldots,R\}. \tag{11}
$$

然后选定索引 $W$ 根据以下分布采样：

$$
W \,|\, S \sim q_\alpha(\cdot|S). \tag{12}
$$

这种构造在 $\alpha=0$ 时的样本无关学习器和 $\alpha=1$ 时的经验风险最小化器之间进行插值。增加 $\alpha$ 会增加 $W$ 对 $S$ 的依赖性，从而增加 $I(S;W)$。

### III-C 对库构建的敏感性

为了评估鲁棒性，我们变化了 $R\in\{10,25,50\}$ 和每个候选者的预训练子集大小 $\{5\text{k}, 10\text{k}, 20\text{k}\}$。表 II 中报告的趋势（$\widehat{I}(S;W)$ 随 $\alpha$ 单调增加）在所有配置中均保持；$\widehat{I}(S;W)$ 的绝对值大致按 $\log R$ 缩放，这与熵上限 $H(W) \leq \log R$ 一致。默认值 $R=25$ 是在信息多样性和蒙特卡洛估计成本之间取得平衡而选择的。样本大小 $m=256$ 的选择使得界代理 $\sqrt{2I/m}$ 与观察到的间隙保持在同一数量级，从而允许有意义的并排比较。

> **图 1:** 有限库 MNIST 协议。小样本 $S$ 用于对预训练 CNN 候选者评分，形成后验 $q_\alpha(r|S)$，采样模型，并估计信息-泛化项。

## IV 在有限库中估计 $I(S;W)$

由于 $W$ 在有限集合中取值，$I(S;W)$ 可以直接估计。精确恒等式为

$$
I(S;W) = H(W) - H(W|S). \tag{13}
$$

令 $S_1, \ldots, S_M$ 为独立的蒙特卡洛训练样本，并写为

$$
q_{i,r} \triangleq q_\alpha(r|S_i). \tag{14}
$$

经验边...