PADD：面向无路由教师指导MoE学生学习的路径对齐解压缩蒸馏

# PADD：面向无路由器教师的对齐路径解压缩蒸馏，用于指导MoE学生学习  
来源：https://arxiv.org/html/2606.10369  

###### 摘要  

随着大型语言模型（LLMs）不断扩展，在固定计算预算下增长模型容量变得越来越具有挑战性。我们提出路径对齐解压缩蒸馏（PADD），一个将无显式路由的密集教师知识蒸馏到混合专家（MoE）学生中，同时学习高质量路由策略的框架。PADD将知识蒸馏组织为两个阶段的四个步骤：初始化阶段（阶段I）通过教师神经元聚类和学生专家预热，在学生专家中构建多样化功能；训练阶段（阶段II–IV）将在线自适应蒸馏、路径精炼策略优化和奖励增强负载平衡集成在单一训练流水线中。在数学推理基准上的实验表明，PADD在相同推理成本下相比强基线取得了显著提升，并且MoE学生能够匹配甚至超越其密集教师。实验还展示了有效的教师到学生知识蒸馏以及稳定的路由行为。  

混合专家，知识蒸馏，密集到MoE迁移，强化学习，数学推理，专家路由  

## 1 引言  

参见图标题  
图1：路径对齐解压缩蒸馏（PADD）框架概览。PADD将密集到MoE知识蒸馏组织为四个阶段：阶段I执行神经元聚类和专家初始化，而阶段II–IV将在线自适应蒸馏、路径精炼策略优化和奖励增强负载平衡集成在统一的训练流水线中。  

随着大型语言模型（LLMs）的扩展，模型容量与有限计算预算之间的紧张关系日益加剧。当扩展到数千亿或万亿参数时，密集模型在训练吞吐量、推理延迟和内存带宽方面面临瓶颈（Kaplan等人，2020（https://arxiv.org/html/2606.10369#bib.bib1）；Hoffmann等人，2022（https://arxiv.org/html/2606.10369#bib.bib2））。混合专家（MoE）架构将参数解耦为稀疏激活的专家子网络，将容量与推理FLOPs解耦（Shazeer等人，2017（https://arxiv.org/html/2606.10369#bib.bib3）；Fedus等人，2022（https://arxiv.org/html/2606.10369#bib.bib4））。MoE可以将纠缠的密集表示解压缩为结构化的专家模块（Komatsuzaki等人，2023（https://arxiv.org/html/2606.10369#bib.bib6）；DeepSeek-AI等人，2024（https://arxiv.org/html/2606.10369#bib.bib7））。然而，大多数高性能模型仍然是密集的；从头训练MoE成本高昂，并且由于不相容的专家分解和路由策略，MoE到MoE蒸馏缺乏通用性（Dai等人，2022（https://arxiv.org/html/2606.10369#bib.bib9）；Zhang等人，2025（https://arxiv.org/html/2606.10369#bib.bib17））。此外，使用密集教师提供了灵活性：可以选择每个领域的最佳教师，向MoE学生提供任务专门化知识，而不增加推理成本。  
将密集模型转换为MoE面临一个根本性挑战：MoE依赖于路由决策，但密集模型缺乏显式路由。虽然可以执行稀疏升级以匹配参数形状（Komatsuzaki等人，2023（https://arxiv.org/html/2606.10369#bib.bib6）），但新的路由器没有来自密集教师激活的监督，必须从头学习。这导致路由器冷启动（Dai等人，2022（https://arxiv.org/html/2606.10369#bib.bib9））：早期训练无法区分句法标记和推理标记，导致随机噪声在专家间扩散（逻辑扩散）。传统蒸馏仅对齐输出（Hinton等人，2015（https://arxiv.org/html/2606.10369#bib.bib8）），无法传递内部处理偏好；离散路由跳跃破坏了思维链连续性，导致路径断裂（Zoph等人，2022（https://arxiv.org/html/2606.10369#bib.bib12））并使梯度不稳定。当MoE学生的容量（例如，每个标记的活动参数）远小于密集教师时，严重的能力差距阻止了对细粒度logits的吸收（Gu等人，2024（https://arxiv.org/html/2606.10369#bib.bib15））。然而，静态logits对齐是不够的，需要动态反馈来指导路由。强化学习（RL）原则上可以通过奖励提供这种反馈，并且on-policy蒸馏（例如，GRPO（Shao等人，2024（https://arxiv.org/html/2606.10369#bib.bib21）））将蒸馏与策略优化耦合，使得教师沿着学生的实际轨迹进行监督。然而，现有方法主要针对密集到密集或MoE到MoE蒸馏，无法弥合密集模型和MoE模型之间的结构不匹配。经典的MoE负载平衡（Shazeer等人，2017（https://arxiv.org/html/2606.10369#bib.bib3））仅控制激活频率，忽略专家质量，导致专家同质化，而诸如StableMoE、RSPO和R3（Dai等人，2022（https://arxiv.org/html/2606.10369#bib.bib9）；Zhang等人，2025（https://arxiv.org/html/2606.10369#bib.bib17）；Ma等人，2025（https://arxiv.org/html/2606.10369#bib.bib27））等方法稳定了已训练MoE模型的路由，但假设可用的专家结构，无法从密集教师中恢复路径级语义。因此，它们未能解决核心挑战：由架构不匹配引起的结构性缺陷。  
PADD针对与稀疏升级（Komatsuzaki等人，2023（https://arxiv.org/html/2606.10369#bib.bib6））和MoE化转换（Zhang等人，2022（https://arxiv.org/html/2606.10369#bib.bib46））互补的设置：我们将无路由器的密集教师的知识和路由迁移到已经预训练好的具有路由感知能力的MoE学生中，而不是从密集检查点构建新的MoE。目标是在固定推理预算下恢复教师的隐式模块结构并学习稳定的学生路由，而升级式流水线首先从密集权重构建专家结构，然后基本从头训练路由。  
图1（https://arxiv.org/html/2606.10369#S1.F1）总结了我们提出的路径对齐解压缩蒸馏（PADD）框架。PADD将蒸馏组织为两个阶段的四个步骤：初始化阶段（阶段I）对教师FFN神经元进行聚类以初始化学生专家，并预热专家形成不同的功能角色；训练阶段（阶段II–IV）将在线自适应蒸馏、路径精炼策略优化和奖励增强负载平衡结合在一个流水线中。这些步骤共同从密集教师中恢复了可学习的路径结构，并在固定推理成本下稳定了专家专门化。  
我们的主要贡献总结如下：  

- • 我们提出路径对齐解压缩蒸馏（PADD），一个统一的密集到MoE蒸馏框架，组织为两个阶段的四个步骤：初始化阶段（阶段I）用于专家设置，训练阶段（阶段II–IV）集成了在线自适应蒸馏、路径精炼策略优化和奖励增强负载平衡。PADD系统性地解决了路由器冷启动、能力差距、路径断裂和专家同质化问题，使得MoE学生在相同推理成本下能够匹配甚至超越密集教师。  
- • 通过分析密集教师FFN中的神经元激活模式并进行聚类，我们构建了学生专家初始化，并使用预热训练赋予学生专家与教师中不同神经元组对应的不同功能角色，从而在源头上缓解了路由器冷启动和专家同质性。  
- • 我们提出了一种在线自适应蒸馏机制，沿着学生实际路由路径提供平滑且可吸收的语义监督，并引入路径精炼组相对策略优化（PR-GRPO），利用路由偏移抑制由离散路由引起的梯度不稳定性，显著提高了MoE强化学习的稳定性。  
- • 我们提出奖励增强负载平衡，联合建模专家激活频率和性能质量，增加高质量专家的激活概率，缓解传统负载平衡引发的“专家同质化”现象，并增强MoE专家之间的长期分工。  

## 2 方法论  

在密集到MoE蒸馏中，密集教师缺乏显式路由，因此传统离线蒸馏无法指导MoE学生学习路由决策。PADD将知识蒸馏与强化学习相结合，通过两个阶段的四个步骤实现：初始化阶段（阶段I）执行初始化和预热；训练阶段（阶段II–IV）将在线自适应蒸馏、路径精炼策略优化和奖励增强负载平衡集成在统一的训练流水线中。我们将数据集\(\mathcal{D}\)划分为四个不重叠的子集：\(\mathcal{D}_A\)用于阶段I中的激活统计和聚类，\(\mathcal{D}_B\)用于阶段I中的专家预热，\(\mathcal{D}_C\)用于阶段II–IV中的主训练，\(\mathcal{D}_D\)用于评估。我们使用下标\(\mathrm{T}\)和\(\mathrm{S}\)区分教师和学生变量（例如，\(p_{\mathrm{T}}\)，\(A_{i,\mathrm{S}}\)）。在开始PADD多阶段密集到MoE蒸馏训练之前，我们对预训练的密集教师应用标准GRPO，以学习可以有效地蒸馏给MoE学生的任务特定推理策略。  

### 2.1 阶段I：基于神经元聚类的专家初始化与预热对齐  

阶段I使用不重叠的数据集\(\mathcal{D}_A\)和\(\mathcal{D}_B\)分为两步。专家预热仅使用\(\mathcal{D}_B\)，不重复使用\(\mathcal{D}_A\)中的样本，从而防止聚类统计信息泄漏到专家拟合中。第一步对教师的FFN执行激活统计和神经元聚类，以构建学生专家的目标功能结构。第二步使用冻结的路由器对学生专家执行预热对齐，形成与教师隐式结构对齐的初始功能差异。  

**激活统计与聚类。**  
设\(L_{\mathrm{T}}\)和\(L_{\mathrm{S}}\)分别为教师和学生的层数。教师是密集的；每一层的前馈网络（FFN）有两个线性层：\(W_1 \in \mathbb{R}^{d_{\text{ff},\mathrm{T}} \times d}\)（其中\(d\)是隐藏维度）和\(W_2\)，它们之间有一个激活函数。学生层\(l\)对应于教师层\(\lfloor l \cdot L_{\mathrm{T}} / L_{\mathrm{S}} \rfloor\)。无论哪一层更多，这都成立，从而为每个学生层提供一个唯一的教师层。学生是MoE，每层有\(N\)个专家；每个专家是一个2层FFN。我们提取教师FFN的\(W_1\)，其中行\(k\)（\(w_k \in \mathbb{R}^d\)）是第\(k\)个神经元的权重向量。当\(d_{\text{ff},\mathrm{T}} \neq d_{\text{ff},\mathrm{S}}\)时，我们通过均匀采样将教师神经元调整到\(d_{\text{ff},\mathrm{S}}\)，将调整后的维度记为\(d_{\text{ff}}\)。密集模型缺乏显式路由，但其FFN神经元在相似输入下会共同激活，揭示了隐式模块结构（Qiu等人，2024（https://arxiv.org/html/2606.10369#bib.bib16））。我们对\(w_k\)执行基数约束的K-Means聚类，将教师神经元划分为\(N\)个簇\(C_j\)，每个簇对应学生专家\(E_{j,\mathrm{S}}\)：  

\[
\min_{C} \sum_{j=1}^{N} \sum_{k \in C_j} \| w_k - \mu_j \|^2, \quad \text{s.t.} \ |C_j| = \frac{d_{\text{ff}}}{N}
\]

(1)  
其中\(\mu_j \in \mathbb{R}^d\)是簇\(C_j\)的质心。获得\(C_j\)后，我们计算簇\(C_j\)在数据集\(\mathcal{D}_A\)上的激活分布\(p_{j,\mathrm{T}}\)，这将成为学生专家\(j\)的学习目标。更具体地说，对于每个样本\(x \in \mathcal{D}_A\)，我们将其前向传播到对应的教师FFN层，并记录神经元\(k\)的第一线性层输出（FFN中间激活）\(h_k^{(x)} \in \mathbb{R}\)。我们将簇\(C_j\)中的神经元在\(\mathcal{D}_A\)的所有样本上取平均，得到\(\bar{h}_j\)，然后对\(j=1,\dots,N\)应用softmax得到\(p_{j,\mathrm{T}}\)：  

\[
\bar{h}_j = \frac{1}{|\mathcal{D}_A|} \sum_{x \in \mathcal{D}_A} \frac{1}{|C_j|} \sum_{k \in C_j} h_k^{(x)},
\]

(2)  
\[
p_{j,\mathrm{T}} = \frac{\exp(\bar{h}_j / \xi)}{\sum_{j'=1}^{N} \exp(\bar{h}_{j'} / \xi)}, \quad \xi > 0,
\]

其中\(\xi\)是控制softmax锐度的温度参数。分布\(p_{j,\mathrm{T}}\)通过\(\mathcal{L}_{\text{init}}\)指导第二步中的学生专家学习。此步骤仅执行统计，不训练学生。  

**专家预热。**  
我们通过将簇质心\(\mu_j\)映射到每一层路由器的Linear权重，初始化学生路由器权重。然后，我们在\(\mathcal{D}_B\)上使用冻结的路由器对学生专家进行预热，以防止早期路由不稳定，并确保所有专家在学习功能差异之前接收相等的训练信号。路由固定为均匀分布（每个专家以概率\(1/N\)激活），因此我们只训练专家网络。设\((x,y)\)为输入-输出对，\(t\)为标记位置，\(\pi_{\mathrm{S}}\)为学生策略，\(p_{\mathrm{T}}\)和\(p_{\mathrm{S}}\)分别为给定上下文时教师和学生的下一标记分布。预热损失包括语言建模损失\(\mathcal{L}_{\text{LM}}\)、知识蒸馏损失\(\mathcal{L}_{\text{KD}}\)和初始化损失\(\mathcal{L}_{\text{init}}\)：  

\[
\mathcal{L}_{\text{LM}} = -\mathbb{E}_{(x,y) \sim \mathcal{D}_B} \sum_{t=1}^{|y|} \log \pi_{\mathrm{S}}(y_t | x, y_{<t})
\]

(3)  
\[
\mathcal{L}_{\text{KD}} = -\mathbb{E}_{(x,y) \sim \mathcal{D}_B} \sum_{t=1}^{|y|} \sum_{v \in \mathcal{V}} p_{\mathrm{T}}(v | x, y_{<t}) \log p_{\mathrm{S}}(v | x, y_{<t})
\]

(4)  
\[
\mathcal{L}_{\text{init}} = -\sum_{j=1}^{N} \mathbb{E}_{x \in \mathcal{D}_B} \left[ p_{j,\mathrm{T}} \log \frac{\exp(s_{j,\mathrm{S}}(x) / \tau)}{\sum_{j'=1}^{N} \exp(s_{j',\mathrm{S}}(x) / \tau)} \right]
\]

(5)  
其中\(\mathcal{V}\)是词表，\(s_{j,\mathrm{S}}(x)\)是学生专家\(j\)的softmax前logits，\(\tau\)是温度。总体上，预热损失为：  

\[
\mathcal{L}_{\text{warmup}} = \mathcal{L}_{\text{LM}} + \lambda_{\text{KD}} \mathcal{L}_{\text{KD}} + \lambda_{\text{init}} \mathcal{L}_{\text{init}}
\]

其中\(\lambda_{\text{KD}}, \lambda_{\text{init}} > 0\)是超参数。  

### 2.2 阶段II：在线自适应蒸馏  

阶段II启动训练流水线，开始联合训练路由器。路由器现在被训练；但是，为了确保学生路由器在早期训练中不从劣势专家开始，我们引入了在线自适应蒸馏。设\(A_{i,\mathrm{S}}\)和\(A_{i,\mathrm{T}}\)分别为学生和教师对专家\(i\)的输出。在向前传播中，当\(A_{i,\mathrm{S}} > 0\)时，监督加强（较低温度，更自信的信号）；当\(A_{i,\mathrm{S}} < 0\)时，偏差被纠正（较高温度，更多探索）。这防止了过拟合到不正确的专家路径。  

### 2.3 阶段III：PR-GRPO路径精炼策略优化  

在反向传播中，为缓解由路由偏移引起的策略梯度不稳定性，我们引入了路径精炼组相对策略优化（PR-GRPO）。学生有\(L_{\mathrm{S}}\)个MoE层，每层有一个路由器和\(N\)个专家。路由器是一个单线性层，将隐藏状态\(h \in \mathbb{R}^d\)映射到\(N\)维logits；经过Softmax后得到\(G_{\theta,\mathrm{S}}(x_t) \in \mathbb{R}^N\)。每层通过Top-\(K\)选择独立路由，并通过相应的2层FFN进行前向传播。对于多层，我们在计算偏移之前聚合路由输出。设\(\pi_{\theta,\mathrm{S}}\)为学生策略，\(\theta_{\text{old}}\)为上次更新时的参数，\(x_t\)为时间\(t\)的输入。路由偏移\(\Gamma_{i,t,\mathrm{S}}\)测量相对于前一步的路由器决策变化：  

\[
\Gamma_{i,t,\mathrm{S}} = \| G_{\theta,\mathrm{S}}(x_t) - G_{\theta_{\text{old}},\mathrm{S}}(x_t) \|_2
\]

(7)  
设\(a_t\)为标记动作，\(s_t\)