持续LLM升级循环：一种用于从稠密到稀疏LLM的预测器门控按组稀疏训练方案

# 持续大语言模型升级循环：一种用于稠密到稀疏大语言模型的预测器门控分块稀疏训练方案
来源：https://arxiv.org/html/2606.10722
Ruixuan Huang1  Jinyuan Shi2  Hantao Huang3,\*  Yifan Huang3  Ziyi Guan3  Hao Zeng2  Ian En\-Hsu Yen2,\*  Minghui Yu3,\*

###### 摘要

我们研究稠密到稀疏的持续训练，作为一种从稠密检查点构建通道稀疏大语言模型的方法。从 Qwen2.5-8B 稠密主干网络开始，我们在 32K 上下文中继续训练，并在 32K 阶段引入预测器门控稀疏 SwiGLU FFN。对于每个 token 和每一层，我们使用一个低秩预测器生成 FFN 通道路由 logits。然后应用分块 top-k 规则，在每个 64 通道块中保留 16 个通道，从而在 FFN 中间激活中实现 4 倍稀疏度。与事后稀疏推理方法不同，路由模块被放置在主要语言建模路径上，并在持续训练期间进行优化，使得稠密模型能够被升级循环为面向硬件的稀疏模型。我们报告了架构、训练方案、基准测试性能和训练经验教训。我们还识别出 RULER-CWE 上一个层级局部化的长上下文失败模式，并提出一种单层修复算法，显著改善了受影响的长度范围。

\{NoHyper\}††footnotetext:1 香港科技大学。2 Moffett AI。3 独立贡献者。\*通讯作者。

## 1 引言

大语言模型 (LLM) 的成本日益受到参数大小和每个 token 计算量的主导。这在使用仅解码器 Transformer 模型[26 (https://arxiv.org/html/2606.10722#bib.bib13)]时尤为明显，其中每个 token 都要通过每一层中的密集前馈网络 (FFN)。在基于 SwiGLU 的现代模型[24 (https://arxiv.org/html/2606.10722#bib.bib14)]中，FFN 将隐藏状态扩展到更大的中间维度，应用逐元素门控，并将结果投影回模型维度。在长上下文训练和长提示预填充期间，即使对于任何特定 token 只有一部分中间通道是重要的，这种计算也会应用于大的 token 块[20 (https://arxiv.org/html/2606.10722#bib.bib2)]。因此，减少这种逐 token 的 FFN 成本是使长上下文 LLM 更高效的核心要求。

条件计算为实现这一目标提供了一种自然途径。混合专家 (MoE) 模型通过将每个 token 路由到少量专家 FFN 来减少计算，但它们引入了独立的专家参数池、专家级路由、负载均衡问题，并且通常需要全对全通信[17 (https://arxiv.org/html/2606.10722#bib.bib15),7 (https://arxiv.org/html/2606.10722#bib.bib16)]。一种更细粒度的替代方案是将稠密 FFN 的中间通道视为神经元级别的计算单元，并为每个 token 仅激活其中一部分。它保留了原始的稠密 FFN 参数化，同时在每一层内部暴露了条件计算，可以将稠密 FFN 转变为一个稀疏的子 FFN，其活跃组件随 token 而变化。

然而，并非所有通道稀疏性都能带来可用的加速。上下文稀疏方法表明，LLM 计算高度依赖于输入[20 (https://arxiv.org/html/2606.10722#bib.bib2),1 (https://arxiv.org/html/2606.10722#bib.bib3),35 (https://arxiv.org/html/2606.10722#bib.bib4)]。通道级方法[30 (https://arxiv.org/html/2606.10722#bib.bib6)]表明，许多 FFN 通道可以基于逐 token 基础被抑制。但是，如果稀疏掩码源自原生 FFN 激活，那么路由信号可能来得太晚，无法跳过主导的门控和上投影。为了使通道稀疏性成为面向硬件的计算路径，稀疏掩码必须在主 FFN 通道计算被具体化之前可用。

本文研究稠密到稀疏的持续训练过程，旨在将稠密 LLM 转换为通道级稀疏模型。我们引入一个轻量级低秩预测器，将每个 token 隐藏状态映射到 FFN 中间通道上的 logits。该预测器在 FFN 计算之前产生路由信号。然后我们施加固定的逐块 top-k 规则，在每个 64 通道块中保留 16 个活跃通道，使活跃 FFN 中间宽度减少 4 倍。由此产生的稀疏路径通过前向传播中的硬稀疏执行和通过 top-k 边界的软代理梯度进行训练。我们在 Qwen2.5-8B 骨干模型架构[31 (https://arxiv.org/html/2606.10722#bib.bib12)]中实现了这个想法。该模型首先在 8K 上下文长度下从头开始训练，然后在 32K 上下文长度下继续训练。稀疏性在 32K 阶段的后期引入，反映了一种实际的升级循环场景[16 (https://arxiv.org/html/2606.10722#bib.bib9)]。

我们的实验将得到的稀疏模型与 (1) 一个稠密伴生运行和 (2) 一个基于密集激活幅度的朴素 4 倍稀疏基准线 (Dense-4x) 进行比较。在一系列广泛评估基准上，我们的稀疏模型一致地比朴素 Dense-4x 基准线更接近稠密基准线，这表明通道稀疏性可以作为语言建模路径的一部分进行训练，而不仅仅是作为一种训练后压缩启发式方法应用。

本文揭示了稀疏升级循环并非无损。我们重点介绍了在训练过程中出现的两种有启发性的失败模式。首先，通道级路由不能直接继承 MoE 风格的平衡。其次，长上下文评估暴露了在广泛短上下文基准中几乎不可见的失败。在 RULER-CWE[12 (https://arxiv.org/html/2606.10722#bib.bib11)]上，稀疏模型在较短长度下表现良好，但出现了尖锐的 12K-16K 断崖。我们提出一种稠密回退干预方法，将大部分这种失败定位到单个层和最终的问题-回答后缀。

#### 贡献。

总之，本文做出以下贡献：

- •**预测器门控的通道稀疏模式。** 我们引入一种预测器门控的稀疏模式，其中低秩逐 token 预测器在 FFN 通道计算之前产生通道路由 logits，并且硬性的逐块 top-k 规则在每个 64 通道块中保留 16 个通道。当受专用内核或稀疏硬件支持时，它可以跳过非活跃的门控、上投影和下投影通道。
- •**稠密到稀疏的持续训练方案。** 我们提出一个稠密到稀疏的持续训练方案，将在 8K 训练并在 32K 继续训练的 Qwen2.5-8B 稠密模型转换为预测器门控的 4 倍稀疏 FFN 模型，性能下降有限。
- •**训练经验教训。** 我们报告了两个训练教训。首先，MoE 风格的平衡偏差不能干净地转移到逐块通道路由。其次，RULER-CWE 揭示了一种层级局部化的长上下文失败，可以通过单层干预进行部分修复。

## 2 相关工作

#### 上下文稀疏性。

先前的工作表明，LLM 计算是依赖输入的，并且通常可以被稀疏化。Dejavu 预测 token 依赖的注意头集和 MLP 神经元的稀疏集以实现高效推理[20 (https://arxiv.org/html/2606.10722#bib.bib2)]，而 ShadowLLM 通过更强的重要性标准和面向部署的预测器改进了基于预测器的上下文稀疏性[1 (https://arxiv.org/html/2606.10722#bib.bib3)]。Sirius 强调无训练的上下文稀疏性在复杂生成任务上可能很脆弱，并引入了纠正机制以恢复质量[35 (https://arxiv.org/html/2606.10722#bib.bib4)]。本文更侧重于训练，其中稀疏门控被插入到主 SwiGLU FFN 计算中，并在持续训练期间进行优化。

#### 通道稀疏 FFN。

混合通道 (MoC) 通过为每个 token 从原生门控投影中选择 top-k 通道，在 SwiGLU FFN 内部引入持久通道稀疏性[30 (https://arxiv.org/html/2606.10722#bib.bib6)]。然而，稀疏决策仍然源自原生门控输出，因此在活跃集已知之前必须计算门控投影。SparkTransformer 针对 FFN 和注意力提出了一种新的稀疏架构，使用 top-k 激活掩码和通过重新分配现有维度形成的低成本预测器[32 (https://arxiv.org/html/2606.10722#bib.bib7)]。该预测器是 Spark 重新设计参数化的内部结果。得到的稀疏设计并不直接与现有的 SwiGLU 架构兼容。

#### MoE 和稠密到稀疏转换。

专家级条件计算将 token 路由到少量完整的 FFN 专家。相关的稠密到稀疏方法包括 MoEfication，它将 FFN 参数划分为功能专家并构建专家路由器[33 (https://arxiv.org/html/2606.10722#bib.bib8)]，以及 Aran 等人，它从稠密检查点初始化稀疏 MoE 模型[16 (https://arxiv.org/html/2606.10722#bib.bib9)]。我们在此使用的通道级视图在精神上与 MoE 类似，因为 FFN 中间通道表现得像细粒度的条件单元，但它不引入独立的专家参数池或专家级全对全调度。它保留了稠密 SwiGLU 参数化，并执行逐 token 的通道选择。

参见标题图1：我们提出的预测器门控稀疏 SwiGLU FFN。一个低秩预测器将 token 隐藏状态映射到 FFN 通道路由 logits。

## 3 方法

本节描述了我们在稠密到稀疏持续训练过程中使用的预测器门控稀疏 FFN。我们从稠密 SwiGLU 计算开始，然后介绍提出的低秩预测器、逐块 top-k 路由，以及保持稀疏路径可训练的技术细节。

### 3.1 稠密 SwiGLU FFN

对于层 \(l\) 和 token 位置 \(t\)，设 FFN 输入为 \(x_{l,t} \in \mathbb{R}^{d}\)，其中 \(d\) 是隐藏维度。我们用 \(d_i\) 表示稠密 FFN 中间维度。一个标准的 SwiGLU FFN 计算：
\[
\begin{aligned}
u_{l,t} &= W^{g}_{l} x_{l,t}, \\
v_{l,t} &= W^{u}_{l} x_{l,t},
\end{aligned}
\tag{1}
\]
\[
m^{\mathrm{dense}}_{l,t} = \mathrm{SiLU}(u_{l,t}) \odot v_{l,t} = u_{l,t} \odot \sigma(u_{l,t}) \odot v_{l,t},
\tag{2}
\]
\[
y^{\mathrm{dense}}_{l,t} = W^{o}_{l} m^{\mathrm{dense}}_{l,t}.
\tag{3}
\]
其中 \(u_{l,t}, v_{l,t}, m^{\mathrm{dense}}_{l,t} \in \mathbb{R}^{d_i}\)。在我们的 8B 规模训练配置中，\(d=4096\) 且 \(d_i=14336\)。

### 3.2 预测器门控通道稀疏性

我们引入一个逐层预测器 \(P_l\)，将 token 隐藏状态映射到 FFN 中间通道上的 logits：
\[
p_{l,t} = P_l(x_{l,t}) \in \mathbb{R}^{d_i}.
\tag{4}
\]
稀疏门控 \(g_{l,t}\) 由这些 logits 构建，并替换稠密 SwiGLU 门控 \(\sigma(u_{l,t})\)：
\[
\begin{aligned}
m_{l,t} &= u_{l,t} \odot g_{l,t} \odot v_{l,t}, \\
y_{l,t} &= W^{o}_{l} m_{l,t}.
\end{aligned}
\tag{5}
\]
因此，预测器直接调制 FFN 中间激活，并从语言建模损失接收梯度。

预测器是一个小型的两层低秩模块：
\[
P_l(x) = W^{p,2}_l W^{p,1}_l x,
\tag{7}
\]
维度为：
\[
W^{p,1}_l: \mathbb{R}^{d} \to \mathbb{R}^{448}, \quad W^{p,2}_l: \mathbb{R}^{448} \to \mathbb{R}^{d_i}.
\tag{8}
\]
隐藏大小 448 是 \(d_i/32\)。预测器输出还包括一个块级仿射校准，具有在每个 64 通道块内共享的可学习尺度和偏置参数。

### 3.3 带软代理梯度的逐块硬 Top-k

FFN 中间维度被划分为固定大小的块，每个块大小 \(C_{\mathrm{bank}}=64\)。对于每个 token、层和块，我们保留 \(k=16\) 个通道：
\[
\begin{aligned}
p_{l,t} &= [p_{l,t}^{(1)}, \ldots, p_{l,t}^{(G)}], \quad p_{l,t}^{(b)} \in \mathbb{R}^{C_{\mathrm{bank}}}, \\
S_{l,t}^{(b)} &= \operatorname{topk}(p_{l,t}^{(b)}, k=16).
\end{aligned}
\tag{9}
\]
门控为：
\[
g_{l,t,i} = 
\begin{cases}
\sigma(p_{l,t,i}), & i \in S_{l,t}^{(b)}, \\
0, & \text{否则}.
\end{cases}
\tag{11}
\]
对于 \(d_i=14336\)，这给出 \(G=d_i/C_{\mathrm{bank}}=224\) 个块，每个 token 每层保留 \(d_s=Gk=3584\) 个通道。这里，\(d_s\) 是图1 (https://arxiv.org/html/2606.10722#S2.F1) 中所示的稀疏中间维度。

添加预测器和离散的 top-k 算子改变了 FFN 路由规则，但端到端训练路径仍然是一个标准语言建模计算图，带有掩码的自定义梯度。在前向传播中，每个块首先对预测器 logits 进行归一化，选择前 16 个通道，并构建一个二元掩码 \(h_{l,t}\)。然后 FFN 中间激活为：
\[
m_{l,t} = u_{l,t} \odot v_{l,t} \odot \sigma(p_{l,t}) \odot h_{l,t},
\tag{12}
\]
因此，在损失所见的模型路径中，非活跃通道恰好为零。

在反向传播中，通过 \(W^{o}_l, W^{g}_l, W^{u}_l\) 以及所选中间激活的梯度遵循所选掩码下的普通自动微分。不可微的 top-k 选择由软 top-k 代理处理：在每个块的阈值附近，硬指示器被赋予一个类 Sigmoid 的局部梯度。因此，语言建模损失仍然可以更新路由边界附近的预测器 logits 和预测器参数。

### 3.4 硬件影响

我们的训练和验证在 GPU 上作为逐块稀疏 FFN 语义的功能仿真进行。模型在前向路径中使用硬稀疏门控进行训练，但 GPU 实现仍然在应用 top-k 掩码之前具体化了稠密的门控和上投影 \(W^{g}x\) 和 \(W^{u}x\)。因此，在 GPU 上实际加速的部分主要是下游投影 \(W^{o}m\)，其中零化的中间激活可以被聚集并仅与选定的行/列相乘。因此，我们的实验验证了稀疏语义和模型质量，而不是展示实测的 4 倍端到端 GPU 挂钟加速。

该算法原理比我们的 GPU 实现所暴露的更强大。对于任何非活跃通道 \(i\) 且 \(h_{l,t,i}=0\)，项
\[
u_{l,t,i} v_{l,t,i} \sigma(p_{l,t,i}) h_{l,t,i}
\tag{13}
\]
对 FFN 输出的贡献恰好为零。因此，专用稀疏硬件可以将预测器生成的掩码置于 FFN 投影之前，并避免在所有三个 FFN 阶段计算非活跃通道：门控投影 \(W^{g}\)、上投影 \(W^{u}\) 和下投影 \(W^{o}\)。在 64:16 的块规则下，对于每个 token 和层，理想的稀疏执行形状从稠密中间宽度 \(d_i\) 缩减为选定的稀疏宽度 \(d_s = d_i/4\)，同时保持相同的数学稀疏 FFN 路径。因此，预期的效率目标