新增: DeepSeek 全部 33 篇官方论文(原文 PDF 走 Git LFS + 中文精读笔记与索引)
This commit is contained in:
ylhan
@@ -0,0 +1,98 @@
|
||||
# mHC：流形约束超连接
|
||||
|
||||
> **原文标题**：mHC: Manifold-Constrained Hyper-Connections
|
||||
> **arXiv**：[2512.24880](https://arxiv.org/abs/2512.24880) ｜ **发布**：2025-12-31（v1），2026-01-05 修订（v2）
|
||||
> **作者/机构**：Zhenda Xie、Yixuan Wei、Huanqi Cao 等 20 位作者，DeepSeek-AI（深度求索）
|
||||
> **原文 PDF**：[01-mHC-2512.24880.pdf](./01-mHC-2512.24880.pdf)
|
||||
|
||||
## 一、摘要（全文翻译）
|
||||
|
||||
近来，以超连接（Hyper-Connections, HC）为代表的研究，通过扩展残差流宽度并多样化连接模式，对过去十年间已被普遍采用的残差连接范式进行了拓展。这种多样化虽然带来了可观的性能收益，却从根本上破坏了残差连接固有的恒等映射（identity mapping）性质，进而导致严重的训练不稳定与可扩展性受限，并且还额外引入了显著的内存访问开销。为应对这些挑战，我们提出了流形约束超连接（Manifold-Constrained Hyper-Connections, mHC），这是一个通用框架：它将 HC 的残差连接空间投影到某个特定流形之上，以恢复恒等映射性质，同时配合严格的基础设施优化来保证效率。实证实验表明，mHC 在大规模训练中是有效的，能够带来切实的性能提升与更优的可扩展性。我们期待 mHC 作为 HC 的一种灵活而实用的扩展，将有助于更深入地理解拓扑架构设计，并为基础模型的演进提示有前景的方向。
|
||||
|
||||
## 二、核心贡献
|
||||
|
||||
1. **定位 HC 的失效机理**：论文指出 HC 中可学习的残差映射 $\mathcal{H}_l^{\mathrm{res}}$ 不受约束，多层展开后的复合映射 $\prod_{i=1}^{L-l}\mathcal{H}_{L-i}^{\mathrm{res}}$ 无法保持特征全局均值，导致信号无界放大或衰减；实测该复合映射的「最大增益幅值」（Amax Gain Magnitude）峰值高达 3000，而理想值应为 1。
|
||||
2. **提出 mHC 通用框架**：把 HC 的残差连接空间投影到双随机矩阵流形（伯克霍夫多胞形，Birkhoff polytope）上，既恢复恒等映射的守恒性质，又保留残差流之间的信息交互，避免了「退回恒等矩阵即牺牲表达力」的两难。
|
||||
3. **给出三条严格理论性质**：谱范数有界（$\|\mathcal{H}_l^{\mathrm{res}}\|_2 \le 1$）保证非扩张、双随机矩阵对乘法封闭保证全深度稳定、伯克霍夫多胞形作为置换矩阵的凸包赋予「置换的凸组合」这一几何解释；并对输入/输出映射 $\mathcal{H}_l^{\mathrm{pre}}$、$\mathcal{H}_l^{\mathrm{post}}$ 施加非负约束以避免正负系数组合引起信号抵消。
|
||||
4. **配套系统级优化**：通过内核融合、TileLang 混合精度算子、选择性重计算与 DualPipe 通信重叠，把扩展率 $n=4$ 时的大规模训练额外时间开销压到仅 6.7%。
|
||||
5. **大规模实证验证**：在 3B/9B/27B 的 MoE 模型上，27B 于 8 个下游基准全面超过 Baseline、多数超过 HC（BBH +2.1、DROP +2.3），最终 loss 相对 Baseline 降低 0.021，且计算扩展曲线上的优势随算力提升基本保持。
|
||||
|
||||
## 三、方法与技术要点
|
||||
|
||||
### 3.1 预备：HC 的公式化与动态/静态映射
|
||||
|
||||
标准残差连接的单元结构（图 1a）为 $\mathbf{x}_{l+1} = \mathbf{x}_l + \mathcal{F}(\mathbf{x}_l, \mathcal{W}_l)$，递归展开得 $\mathbf{x}_L = \mathbf{x}_l + \sum_{i=l}^{L-1}\mathcal{F}(\mathbf{x}_i, \mathcal{W}_i)$，其中 $\mathbf{x}_l$ 这一项即所谓恒等映射，它保证浅层信号不加修改地直达深层。HC（图 1b）把残差流宽度由 $C$ 扩到 $n\times C$，单层传播变为
|
||||
|
||||
$$\mathbf{x}_{l+1} = \mathcal{H}_l^{\mathrm{res}}\mathbf{x}_l + \mathcal{H}_l^{\mathrm{post}\top}\mathcal{F}(\mathcal{H}_l^{\mathrm{pre}}\mathbf{x}_l, \mathcal{W}_l),$$
|
||||
|
||||
其中 $\mathcal{H}_l^{\mathrm{res}}\in\mathbb{R}^{n\times n}$ 在残差流内部混合特征，$\mathcal{H}_l^{\mathrm{pre}}\in\mathbb{R}^{1\times n}$ 把 $nC$ 维残差流读成 $C$ 维层输入，$\mathcal{H}_l^{\mathrm{post}}\in\mathbb{R}^{1\times n}$ 把层输出写回残差流。由于 $n$（典型取 4）远小于 $C$，这些额外映射的 FLOPs 开销可忽略，HC 因而把「残差流的信息容量」从与计算量强相关的层输入维度中解耦，成为继模型 FLOPs 与数据量之外的第三条扩展轴。三组映射均由依赖输入的动态部分（经 RMSNorm、tanh 门控与初始化为小值的门控因子 $\alpha$）与静态部分（可学习偏置 $\mathbf{b}$）相加构成。消融（表 1）显示，仅启用 $\mathcal{H}_l^{\mathrm{res}}$ 即得 $-0.022$ 的损失差，三项全开为 $-0.027$，说明残差流内部的信息交换是收益主要来源，也正是 mHC 的着力点。
|
||||
|
||||
### 3.2 两个痛点：数值不稳定与系统开销
|
||||
|
||||
（1）**数值不稳定**。HC 递归展开后为 $\mathbf{x}_L = \left(\prod_{i=1}^{L-l}\mathcal{H}_{L-i}^{\mathrm{res}}\right)\mathbf{x}_l + \sum_{i=l}^{L-1}\left(\prod_{j=1}^{L-1-i}\mathcal{H}_{L-j}^{\mathrm{res}}\right)\mathcal{H}_i^{\mathrm{post}\top}\mathcal{F}(\mathcal{H}_i^{\mathrm{pre}}\mathbf{x}_i, \mathcal{W}_i)$。与标准残差相比，复合映射偏离恒等且不保持特征全局均值，前向与反向都会出现信号爆炸或消失。论文用复合映射的行和最大值（前向）与列和最大值（反向）两个指标刻画，称 Amax Gain Magnitude：27B 模型上该指标峰值约 3000，且 HC 在约 12k step 出现意外 loss 飙升，与梯度范数的不稳定高度相关。（2）**系统开销**。HC 虽在 FLOPs 上廉价，却撞上「内存墙」：每 token 的残差流维护 I/O 从读 $2C$／写 $C$ 膨胀到读 $(5n+1)C+n^2+2n$／写 $(3n+1)C+n^2+2n$，约为 $n$ 倍；三个映射的中间激活还需为反向保留，显存大增、常需梯度检查点；流水线并行的通信量亦放大 $n$ 倍，气泡更大、吞吐下降。
|
||||
|
||||
### 3.3 流形约束超连接：把残差映射投影到双随机矩阵流形
|
||||
|
||||
mHC 的核心（图 1c）是把 $\mathcal{H}_l^{\mathrm{res}}$ 约束到特定流形上。作者指出，原始恒等映射靠 $\mathcal{H}_l^{\mathrm{res}}=\mathbf{I}$ 换来稳定，却完全禁止残差流内部的信息交换，浪费了多流架构的潜力；需要一个「既跨层稳定、又允许流间交互」的中间地带。他们选择双随机矩阵流形（伯克霍夫多胞形）：
|
||||
|
||||
$$\mathcal{P}_{\mathcal{M}^{\mathrm{res}}}(\mathcal{H}_l^{\mathrm{res}}) \coloneq \left\{\mathcal{H}_l^{\mathrm{res}}\in\mathbb{R}^{n\times n}\mid \mathcal{H}_l^{\mathrm{res}}\mathbf{1}_n=\mathbf{1}_n,\ \mathbf{1}_n^{\top}\mathcal{H}_l^{\mathrm{res}}=\mathbf{1}_n^{\top},\ \mathcal{H}_l^{\mathrm{res}}\geqslant 0\right\}.$$
|
||||
|
||||
$n=1$ 时该条件退化为标量 1，恰好恢复原始恒等映射。选择双随机性带来三条性质：**范数保持**——谱范数 $\|\mathcal{H}_l^{\mathrm{res}}\|_2\le 1$，映射非扩张，缓解梯度爆炸；**复合封闭**——双随机矩阵对乘法封闭，故任意深度的复合映射仍是双随机的，稳定性贯穿整个深度；**伯克霍夫多胞形的几何解释**——该集合是置换矩阵的凸包，残差映射相当于置换的凸组合，反复施加会单调地增强流间信息混合，构成稳健的特征融合机制。此外，输入/输出映射施加 Sigmoid 非负约束（可视作另一类特殊流形投影），防止正负系数组合导致信号抵消。
|
||||
|
||||
### 3.4 参数化与 Sinkhorn-Knopp 投影
|
||||
|
||||
为保留完整上下文，mHC 先把隐状态展平成 $\vec{\mathbf{x}}_l=\mathrm{vec}(\mathbf{x}_l)\in\mathbb{R}^{1\times nC}$，再做动态映射（投影矩阵 $\varphi_l^{\mathrm{pre}},\varphi_l^{\mathrm{post}}\in\mathbb{R}^{nC\times n}$、$\varphi_l^{\mathrm{res}}\in\mathbb{R}^{nC\times n^2}$，后者经 $\mathrm{mat}(\cdot)$ 重塑为 $n\times n$），最后经约束得到最终映射：$\mathcal{H}_l^{\mathrm{pre}}=\sigma(\tilde{\mathcal{H}}_l^{\mathrm{pre}})$，$\mathcal{H}_l^{\mathrm{post}}=2\sigma(\tilde{\mathcal{H}}_l^{\mathrm{post}})$，$\mathcal{H}_l^{\mathrm{res}}=\text{Sinkhorn-Knopp}(\tilde{\mathcal{H}}_l^{\mathrm{res}})$。Sinkhorn-Knopp 先用指数函数使元素全为正，再交替对行、列归一化：以 $\mathbf{M}^{(0)}=\exp(\tilde{\mathcal{H}}_l^{\mathrm{res}})$ 为起点迭代 $\mathbf{M}^{(t)}=\mathcal{T}_r(\mathcal{T}_c(\mathbf{M}^{(t-1)}))$，$t_{\max}\to\infty$ 时收敛到双随机矩阵；实践中取 $t_{\max}=20$。
|
||||
|
||||
### 3.5 高效基础设施设计
|
||||
|
||||
**内核融合**：RMSNorm 用于高维 $\vec{\mathbf{x}}_l\in\mathbb{R}^{1\times nC}$ 时延迟显著，作者把「除以范数」重排到矩阵乘法之后（等价），并以混合精度（$\varphi_l$: tfloat32，$\vec{\mathbf{x}}_l$: bfloat16）与共享访存融合缓解带宽瓶颈。三类专用内核分别计算三个映射，并借助矩阵乘单元把对 $\vec{\mathbf{x}}_l$ 的两次扫描融合、反向的两次矩阵乘合并为一个内核；Sinkhorn-Knopp 迭代在单内核内完成，其反向为自研内核，在片上重算中间结果并遍历整个迭代；应用侧把 $\mathcal{H}_l^{\mathrm{post}}$、$\mathcal{H}_l^{\mathrm{res}}$ 与残差合并融合，使该内核读取元素数由 $(3n+1)C$ 降到 $(n+1)C$、写入由 $3nC$ 降到 $nC$。**重计算**：前向后丢弃 mHC 内核的中间激活，反向时重算（跳过沉重的层函数 $\mathcal{F}$），连续 $L_r$ 层只需保存首层输入 $\mathbf{x}_{l_0}$；最优块长由最小化总显存给出 $L_r^{*}=\arg\min_{L_r}\left[nC\times\lceil L/L_r\rceil+(n+2)C\times L_r\right]\approx\sqrt{nL/(n+2)}$，且与流水线阶段边界对齐。**DualPipe 通信重叠**：扩展 DualPipe 调度，把 MLP 层的 $\mathcal{F}_{\mathrm{post,res}}$ 内核放到专用高优先级计算流上执行，并对注意力层的长时操作不使用持久化内核，以便抢占式调度；各阶段首激活 $\mathbf{x}_{l_0}$ 已在本地缓存，故重计算与流水线通信依赖解耦。
|
||||
|
||||
## 四、主要实验结果
|
||||
|
||||
实验基于 DeepSeek-V3 风格的 MoE 架构训练四个模型变体，HC 与 mHC 的扩展率均取 $n=4$：3B、9B、27B（数据量与参数量成比例）为主结果对象，外加一个在 1T token 固定语料上训练的 3B 模型用于观察 token 扩展行为。
|
||||
|
||||
| 基准 / 任务（指标，样本数） | 本模型 27B w/ mHC | 对比：27B w/ HC | 对比：27B Baseline | 说明 |
|
||||
|---|---|---|---|---|
|
||||
| BBH（EM，3-shot） | 51.0 | 48.9 | 43.8 | mHC 相对 HC 再提升 2.1 |
|
||||
| DROP（F1，3-shot） | 53.9 | 51.6 | 47.0 | 相对 HC 再提升 2.3 |
|
||||
| GSM8K（EM，8-shot） | 53.8 | 53.2 | 46.7 | 略优于 HC |
|
||||
| HellaSwag（Acc.，10-shot） | 74.7 | 74.3 | 73.7 | 小幅领先 |
|
||||
| MATH（EM，4-shot） | 26.0 | 26.4 | 22.0 | 唯一低于 HC 的基准（−0.4） |
|
||||
| MMLU（Acc.，5-shot） | 63.4 | 63.0 | 59.0 | 小幅领先 |
|
||||
| PIQA（Acc.，0-shot） | 80.5 | 79.9 | 78.5 | 小幅领先 |
|
||||
| TriviaQA（EM，5-shot） | 57.6 | 56.3 | 54.3 | 领先 HC 1.3 |
|
||||
|
||||
| 稳定性 / 扩展性 / 系统指标 | 结果 | 说明 |
|
||||
|---|---|---|
|
||||
| 27B 最终 loss（相对 Baseline） | −0.021 | mHC 消除了 HC 在第 12k step 附近的 loss 飙升 |
|
||||
| 复合映射最大增益幅值 | 约 1.6（HC 近 3000） | 降低约三个数量级 |
|
||||
| 计算扩展（3B→9B→27B） | 相对 loss 优势随算力提升基本保持，仅轻微衰减 | 计算最优配置点 |
|
||||
| 大规模训练额外时间开销（$n=4$） | 6.7% | 内部大规模训练实测 |
|
||||
|
||||
mHC 在 8 个下游基准上全面超过 Baseline，并在其中 7 个上超过 HC：推理类任务提升最明显（BBH +2.1、DROP +2.3），仅在 MATH 上略低于 HC（26.0 对 26.4）。稳定性方面，mHC 的 loss 与梯度范数曲线均与 Baseline 同样平稳，无 HC 那样的突变。传播稳定性上，单层映射因 Sinkhorn-Knopp 只迭代 20 次而使反向梯度增益略偏离 1，复合偏差增大但仍有界（最大约 1.6），相较 HC 的近 3000 降低约三个数量级；图 8 还显示 HC 在最大增益较大时其余元素也普遍偏大（所有传播路径都失稳），而 mHC 始终稳定。
|
||||
|
||||
## 五、结论、局限与影响
|
||||
|
||||
- **结论**：论文识别出 HC 在扩展残差流宽度、多样化连接的同时，因连接不受约束而导致信号发散、破坏跨层信号能量守恒，进而引发训练不稳定并限制可扩展性。mHC 把残差连接空间投影到特定流形，用 Sinkhorn-Knopp 对残差映射施加双随机约束，将信号传播转化为特征的凸组合；实验证实它恢复了恒等映射性质，可稳定地大规模训练且可扩展性优于 HC，而基础设施优化使这些收益的计算开销几乎可忽略。
|
||||
- **局限性**：论文未设置专门的局限性章节。作者明确提及的近似性只有一处：Sinkhorn-Knopp 需限制迭代次数（取 $t_{\max}=20$）以换取计算效率，因此得到的是近似解，单层映射的反向梯度增益会略偏离 1，复合映射的偏差最大约 1.6。其余方面（例如该约束对表达力上限、对不同 $n$ 与不同模型规模的普适边界）论文未明确说明。
|
||||
- **影响**：mHC 被定位为 HC 范式的泛化扩展，作者主张该框架可容纳面向不同学习目标的各种流形约束探索，并期待后续对几何约束的研究能在「可塑性」与「稳定性」间取得更好权衡。论文还希望 mHC 重新激发社区对宏观架构（macro-design）设计的兴趣，通过加深对拓扑结构如何影响优化与表示学习的理解，为下一代基础架构的演进提供新路径。
|
||||
|
||||
## 六、关键术语对照表
|
||||
|
||||
| English | 中文 |
|
||||
|---|---|
|
||||
| Hyper-Connections (HC) | 超连接 |
|
||||
| Manifold-Constrained Hyper-Connections (mHC) | 流形约束超连接 |
|
||||
| Residual Connection | 残差连接 |
|
||||
| Residual Stream | 残差流 |
|
||||
| Identity Mapping | 恒等映射 |
|
||||
| Expansion Rate ($n$) | 扩展率 |
|
||||
| Doubly Stochastic Matrix | 双随机矩阵 |
|
||||
| Birkhoff Polytope | 伯克霍夫多胞形 |
|
||||
| Sinkhorn-Knopp Algorithm | Sinkhorn-Knopp 算法（交替行列归一化） |
|
||||
| Amax Gain Magnitude | 最大增益幅值 |
|
||||
| Spectral Norm | 谱范数 |
|
||||
| Kernel Fusion | 内核融合 |
|
||||
| Recomputation (Gradient Checkpointing) | 重计算（梯度检查点） |
|
||||
| DualPipe / Pipeline Parallelism | DualPipe 调度 / 流水线并行 |
|
||||
| Macro Design / Micro Design | 宏观设计 / 微观设计 |
|
||||
Binary file not shown.
@@ -0,0 +1,143 @@
|
||||
# Engram：条件记忆——面向大语言模型的可扩展查找稀疏轴
|
||||
|
||||
> **原文标题**：Conditional Memory via Scalable Lookup: A New Axis of Sparsity for Large Language Models
|
||||
> **arXiv**：[2601.07372](https://arxiv.org/abs/2601.07372) ｜ **发布**：12 Jan 2026（最后修订 12 Jul 2026, v2）
|
||||
> **作者/机构**：Xin Cheng、Rui Tian、Wangding Zeng、Damai Dai 等；北京大学（Peking University）与 DeepSeek-AI（深度求索）联合完成，张辉帅、赵东岩为通讯作者
|
||||
> **原文 PDF**：[02-Engram-条件记忆-2601.07372.pdf](./02-Engram-条件记忆-2601.07372.pdf)
|
||||
|
||||
## 一、摘要（全文翻译）
|
||||
|
||||
混合专家（MoE）通过条件计算扩展容量，但 Transformer 缺少原生的知识查找原语，被迫以计算低效地模拟检索。为此我们提出**条件记忆（conditional memory）**这一互补的稀疏轴，并以 **Engram** 实例化——该模块将经典 N-gram 嵌入现代化，实现 $\mathcal{O}(1)$ 查找。通过形式化**稀疏容量分配（Sparsity Allocation）**问题，我们揭示一条 U 形缩放定律，以优化神经计算（MoE）与静态记忆（Engram）的权衡。据此我们把 Engram 扩展至 27B 参数，在严格等参数、等 FLOPs 条件下超越 MoE 基线。值得注意的是，记忆模块本应有助于知识检索（MMLU +3.4；CMMLU +4.0），我们却观察到通用推理（BBH +5.0；ARC-Challenge +3.7）与代码/数学（HumanEval +3.0；MATH +2.4）的提升更大。机制分析表明，Engram 把骨干浅层从静态重建中解放出来，有效加深网络以进行复杂推理。此外，把局部依赖委派给查表后，它释放注意力容量处理全局上下文，大幅提升长上下文检索（Multi-Query NIAH：84.2 → 97.0）。最后，Engram 确立了"基础设施感知的效率"：确定性寻址支持运行时从主机内存预取，开销可忽略不计。我们设想条件记忆将成为下一代稀疏模型不可或缺的建模原语。代码见：https://github.com/deepseek-ai/Engram 。
|
||||
|
||||
## 二、核心贡献
|
||||
|
||||
1. **提出「条件记忆」新稀疏轴**：MoE 只解决了条件计算，而语言中大量命名实体与公式化短语本质局部、静态、高度模式化，本应由廉价查表承担。Engram 把 N-gram 嵌入重造为一等建模原语，以确定性哈希完成 $\mathcal{O}(1)$ 检索，避免 Transformer 用多层注意力与前馈网络在运行时"重建静态查找表"。
|
||||
2. **形式化稀疏容量分配并发现 U 形缩放定律**：在总参数与 FLOPs 固定的前提下定义分配比例 $\rho$（非激活参数中分给 MoE 专家的比例），验证损失对 $\rho$ 呈 U 形；纯 MoE（$\rho=100\%$）并非最优，把约 20%–25% 稀疏预算划给 Engram 最好，最优位置 $\rho\approx75\%\text{--}80\%$ 跨算力规模稳定。
|
||||
3. **27B 规模超越同参同算力 MoE**：Engram-27B 把路由专家由 72 削至 55，释放参数转为 5.7B 记忆表，总参数 26.7B、激活参数 3.8B 不变，知识、推理、阅读理解、代码数学全谱系提升。
|
||||
4. **证明增益源于「有效深度」增加**：LogitLens 与 CKA 显示 Engram 浅层表示对齐 MoE 更深的层（第 5 层约等于基线第 12 层），早期层 KL 散度系统性更小，显式查找让网络更快完成特征组合。
|
||||
5. **长上下文与系统效率双重收益**：局部依赖交给查表后注意力可专注全局，Multi-Query NIAH 由 84.2 升至 97.0；确定性寻址使 100B 参数表可卸载到主机内存异步预取，吞吐损失最高 2.8%。
|
||||
|
||||
## 三、方法与技术要点
|
||||
|
||||
### 3.1 总体结构：检索 + 融合两阶段
|
||||
|
||||
给定输入序列 $X=(x_1,\dots,x_T)$ 与第 $\ell$ 层隐状态 $\mathbf{H}^{(\ell)}\in\mathbb{R}^{T\times d}$，Engram 对每个位置 $t$ 先**检索**（由局部上下文确定性取回静态嵌入），再**融合**（用当前隐状态调制并做轻量卷积）。模块以残差注入：$\mathbf{H}^{(\ell)}\leftarrow\mathbf{H}^{(\ell)}+\mathbf{Y}$，随后接标准 Attention 与 MoE；标准输入嵌入与 un-embedding 不变，模块只插入部分层。
|
||||
|
||||
### 3.2 分词器压缩与多头哈希检索
|
||||
|
||||
子词分词器优先保证无损重建，会把语义等价的词面（如 `Apple` 与 `␣apple`）分到不同 ID，稀释语义密度。为此预计算满射函数 $\mathcal{P}:V\to V'$，按规范化文本等价（NFKC、小写化等）把原始 token ID 折叠为规范 ID，在 128k 词表上使有效词表缩小 23%（附录 C 为 23.43%）。
|
||||
|
||||
对位置 $t$ 做映射后得到后缀 N-gram $g_{t,n}=(x'_{t-n+1},\dots,x'_t)$。直接参数化全部 N-gram 组合空间不可行，故用哈希方案并设 $K$ 个哈希头缓解冲突：
|
||||
|
||||
$$z_{t,n,k}\triangleq\varphi_{n,k}(g_{t,n}),\quad \mathbf{e}_{t,n,k}=\mathbf{E}_{n,k}[z_{t,n,k}]$$
|
||||
|
||||
实践中 $\varphi_{n,k}$ 用轻量乘法-XOR 哈希实现。记忆向量由各阶各头结果拼接而成：$\mathbf{e}_t\triangleq\|_{n=2}^{N}\|_{k=1}^{K}\mathbf{e}_{t,n,k}$，表 $\mathbf{E}_{n,k}$ 取质数大小 $M_{n,k}$ 以降低冲突。
|
||||
|
||||
### 3.3 上下文感知门控
|
||||
|
||||
检索到的 $\mathbf{e}_t$ 是上下文无关的先验，静态性使其缺乏适应性，还可能被哈希冲突或一词多义污染。Engram 借用注意力范式：把已聚合全局上下文的 $\mathbf{h}_t$ 作 Query，把 $\mathbf{e}_t$ 同时投影为 Key 与 Value，
|
||||
|
||||
$$\mathbf{k}_t=\mathbf{W}_K\mathbf{e}_t,\quad \mathbf{v}_t=\mathbf{W}_V\mathbf{e}_t$$
|
||||
|
||||
为梯度稳定，对 Query 与 Key 先做 RMSNorm 再算标量门：
|
||||
|
||||
$$\alpha_t=\sigma\!\left(\frac{\text{RMSNorm}(\mathbf{h}_t)^{\top}\text{RMSNorm}(\mathbf{k}_t)}{\sqrt{d}}\right),\qquad \tilde{\mathbf{v}}_t=\alpha_t\cdot\mathbf{v}_t$$
|
||||
|
||||
这强制语义对齐：若记忆与上下文矛盾，门值趋零以抑制噪声。随后用深度可分离因果卷积扩大感受野（核宽 $w=4$、膨胀取最大 N-gram 阶数、SiLU 激活）：
|
||||
|
||||
$$\mathbf{Y}=\text{SiLU}\left(\text{Conv1D}\left(\text{RMSNorm}(\tilde{\mathbf{V}})\right)\right)+\tilde{\mathbf{V}}$$
|
||||
|
||||
### 3.4 与多分支架构的集成
|
||||
|
||||
默认骨干采用多分支结构（残差流扩展为 $M$ 条并行分支、连接权重可学习，即 Manifold-Constrained Hyper-Connections，$M=4$）。为兼顾效率与表达力采用参数共享：稀疏嵌入表与 Value 投影 $\mathbf{W}_V$ 全分支共享，每分支各有独立 Key 投影 $\{\mathbf{W}_K^{(m)}\}_{m=1}^{M}$，以获得分支特异门控：
|
||||
|
||||
$$\alpha_t^{(m)}=\sigma\!\left(\frac{\text{RMSNorm}(\mathbf{h}_t^{(m)})^{\top}\text{RMSNorm}(\mathbf{W}_K^{(m)}\mathbf{e}_t)}{\sqrt{d}}\right),\quad \mathbf{u}_t^{(m)}=\alpha_t^{(m)}\cdot(\mathbf{W}_V\mathbf{e}_t)$$
|
||||
|
||||
这样 $1$ 个 $\mathbf{W}_V$ 与 $M$ 个 $\mathbf{W}_K^{(m)}$ 可融合成单次稠密 FP8 矩阵乘，提升 GPU 算力利用率。
|
||||
|
||||
### 3.5 计算与存储解耦的系统设计
|
||||
|
||||
MoE 路由依赖运行时隐状态，而 Engram 索引仅取决于输入 token ID，天然可预测。训练时按模型并行把大表分片到各 GPU，前向用 All-to-All 收集激活行、反向分发梯度，记忆容量随加速器数量线性增长。推理时把表卸载到主机内存，因索引在前向之前已知，主机可经 PCIe 异步预取，用前序块的计算掩盖传输延迟。因此放置层是硬件-算法协同设计：越深可掩盖延迟的窗口越长，但第 6.2 节消融显示建模性能偏好早期介入。N-gram 还服从 Zipf 分布，少量模式占绝大多数访问，可构建多级缓存层次（热条目放 HBM/DRAM，长尾放 NVMe SSD）。
|
||||
|
||||
### 3.6 稀疏容量分配定律与无限记忆
|
||||
|
||||
定义三个指标：总可训练参数 $P_{\mathrm{tot}}$（不含词表嵌入与 LM head）、每 token 激活参数 $P_{\mathrm{act}}$（决定 FLOPs）、非激活参数 $P_{\mathrm{sparse}}\triangleq P_{\mathrm{tot}}-P_{\mathrm{act}}$。分配比例：
|
||||
|
||||
$$P_{\mathrm{MoE}}^{(\mathrm{sparse})}=\rho\,P_{\mathrm{sparse}},\qquad P_{\mathrm{Engram}}=(1-\rho)\,P_{\mathrm{sparse}}$$
|
||||
|
||||
$\rho=1$ 即纯 MoE；$\rho<1$ 则减少路由专家、把释放参数转为 Engram 槽位。作者在 $C=2\times10^{20}$（$P_{\mathrm{tot}}\approx5.7$B、$P_{\mathrm{act}}=568$M、106 专家）与 $C=6\times10^{20}$（$P_{\mathrm{tot}}\approx9.9$B、$P_{\mathrm{act}}=993$M、99 专家）两档、稀疏度 $P_{\mathrm{tot}}/P_{\mathrm{act}}\approx10$ 下扫描 $\rho$：$C=6\times10^{20}$ 时验证损失从 $\rho=100\%$ 的 1.7248 降到最优 $\rho\approx80\%$ 的 1.7109（$\Delta=0.0139$）；即使 MoE 份额压到 $\rho\approx40\%$（5.7B 模型 46 专家、9.9B 模型 43 专家）仍与纯 MoE 相当。U 形两端各有病因：MoE 独大时只能靠深度低效重建静态模式；Engram 独大时丧失条件计算能力，动态推理受损。无限记忆档固定 3B MoE 骨干（$P_{\mathrm{act}}=568$M、100B tokens），槽位数从 $2.58\times10^5$ 扫到 $1.0\times10^7$（最多约 13B 参数），验证损失在 log 空间线性下降，同等预算下比把 N-gram 嵌入与词表嵌入平均的 OverEncoding 潜力更大。
|
||||
|
||||
### 3.7 大规模预训练配置
|
||||
|
||||
四个模型（Dense-4B、MoE-27B、Engram-27B、Engram-40B）共用数据课程与激活参数规模：262B tokens、DeepSeek-V3 分词器（128k 词表）、30 层 Transformer、隐层 2560、每块含 32 头 MLA、经 mHC（扩张率 4）连接 FFN、Muon 优化器。MoE-27B 为 2 共享 + 72 路由专家（top-6）；Engram-27B 改为 2 + 55（top-6，$\rho=74.3\%$），省下的参数变成 5.7B 记忆，模块置于第 2、15 层，最大 N-gram 阶数 3、哈希头 8 个、维度 1280；嵌入参数用 Adam、学习率放大 5 倍、无权重衰减，卷积零初始化以保持恒等映射。Engram-40B 保持同一骨干与算力预算，仅把记忆扩到 18.5B（总 39.5B）。
|
||||
|
||||
### 3.8 消融与机制分析要点
|
||||
|
||||
3B MoE 骨干（12 层、0.56B 激活参数、100B tokens，基线验证损失 1.808）上加 1.6B 记忆，参考配置用 {2,3}-gram 置于第 2、6 层得 1.768。层扫描显示单模块放第 2 层最优（1.770）：一次注意力已足以提供有上下文的 Query，同时仍早到可替代底层局部聚合；同一预算拆成两个模块置于第 2、6 层更好，兼顾早期介入与后期门控精度。消融显示分支特异融合、上下文感知门控、分词器压缩贡献最大，去掉深度可分离卷积仅轻微退化，固定预算下容量分给 4-gram 略次优。推理期完全抑制记忆输出后，知识类仅保留 29%–44%（TriviaQA 29%），阅读理解类保留 81%–93%（C3 93%），印证记忆是参数化知识的主要仓库、上下文任务靠骨干注意力。门控激活集中在命名实体、套语及中文成语与历史人物（如"四大发明"、"张仲景"），跨语言一致。
|
||||
|
||||
## 四、主要实验结果
|
||||
|
||||
| 基准 / 任务 | 本模型（Engram-27B） | 对比对象（MoE-27B） | 说明 |
|
||||
|---|---|---|---|
|
||||
| Pile (loss) | 1.960 | 2.091 | 语言建模，越低越好 |
|
||||
| 验证集 (loss) | 1.634 | 1.768 | 同分布验证集 |
|
||||
| MMLU (Acc., 5-shot) | 60.4 | 57.4 | +3.0（摘要与引言口径 +3.4） |
|
||||
| MMLU-Pro (Acc.) | 30.1 | 28.3 | +1.8 |
|
||||
| CMMLU (Acc.) | 61.9 | 57.9 | +4.0 |
|
||||
| C-Eval (Acc.) | 62.7 | 58.0 | +4.7 |
|
||||
| AGIEval (Acc., 0-shot) | 41.8 | 38.6 | +3.2 |
|
||||
| ARC-Challenge (Acc.) | 73.8 | 70.1 | +3.7 |
|
||||
| BBH (EM, 3-shot) | 55.9 | 50.9 | +5.0 |
|
||||
| CCPM (Acc., 0-shot) | 87.1 | 79.6 | +7.5，知识类增幅最大 |
|
||||
| PIQA (Acc.) | 73.5 | 71.9 | +1.6 |
|
||||
| TriviaQA (EM) | 50.7 | 48.8 | +1.9 |
|
||||
| DROP (F1, 1-shot) | 59.0 | 55.7 | +3.3 |
|
||||
| RACE-High (Acc.) | 78.2 | 75.4 | +2.8 |
|
||||
| HumanEval (Pass@1) | 40.8 | 37.8 | +3.0 |
|
||||
| MBPP (Pass@1) | 48.2 | 46.6 | +1.6 |
|
||||
| GSM8K (EM, 8-shot) | 60.6 | 58.4 | +2.2 |
|
||||
| MGSM (EM) | 49.4 | 46.8 | +2.6 |
|
||||
| MATH (EM, 4-shot) | 30.7 | 28.3 | +2.4 |
|
||||
| CruxEval-i (EM) | 32.2 | 30.7 | +1.5 |
|
||||
| RULER 32k Multi-Query NIAH (Acc.) | 97.0 | 84.2 | 等预训练损失设置（46k vs 50k） |
|
||||
| RULER 32k Variable Tracking (Acc.) | 87.2 | 77.0 | 同上，等损失设置 |
|
||||
| RULER 32k Multi-Values (Acc.) | 95.5 | 92.7 | 同上 |
|
||||
| LongPPL Code (PPL ↓) | 2.45 | 2.49 | 等损失设置，略优 |
|
||||
| LongPPL L-CoT (PPL ↓) | 13.59 | 14.16 | 等损失设置 |
|
||||
| 吞吐 4B-Dense + 100B Engram（CPU 卸载） | 8,858.28 tok/s | 8,903.62 tok/s（基线） | H800，512 序列，损失 0.5% |
|
||||
| 吞吐 8B-Dense + 100B Engram | 6,140.02 tok/s | 6,315.52 tok/s（基线） | 损失 2.8%，为最差情形 |
|
||||
|
||||
三种稀疏架构在等训练算力下全面优于 iso-FLOPs 的 Dense-4B，说明稀疏本身的缩放优势；Engram-27B 在同参同算力下对 MoE-27B 的领先不只出现在知识密集任务（MMLU +3.0、CMMLU +4.0、MMLU-Pro +1.8），在通用推理（BBH +5.0、ARC-Challenge +3.7、DROP +3.3）与代码数学（HumanEval +3.0、MATH +2.4、GSM8K +2.2）上增益更大。Engram-40B 进一步降低损失并在多数基准上更好，但未在每个任务上严格压过 Engram-27B，作者归因于欠训练——两者差距在训练末期仍在扩大。长上下文中，等预训练损失（46k vs 50k）下 Engram-27B 在复杂检索上大幅领先，等算力（50k）下全面提升；即使用仅约 82% 预训练算力的 41k 检查点，也能在 LongPPL 上追平、在 RULER 上超越完全训练的 MoE-27B。系统中，100B 记忆表全放主机 DRAM、每次检索都走 PCIe 的保守设定下吞吐损失仅 2.8%。
|
||||
|
||||
## 五、结论、局限与影响
|
||||
|
||||
- **结论**：论文把条件记忆确立为与 MoE 条件计算互补的稀疏轴。Engram 以分词器压缩、多头哈希检索、上下文感知门控与多分支集成，把经典 N-gram 嵌入现代化为 $\mathcal{O}(1)$ 的可扩展查找原语。由此形式化的稀疏容量分配给出稳定的 U 形定律：纯 MoE 次优，把 20%–25% 稀疏预算分给记忆最优。27B 规模上 Engram 在同参同算力下全面超越 MoE 基线；机制分析表明它把静态重建从浅层卸载从而"加深"网络，并使注意力得以专注全局上下文；确定性寻址让海量参数表可卸载到主机内存而几乎不损吞吐。
|
||||
- **局限性**：论文未设独立的"局限"章节，作者自述的相关限制包括：（1）插入位置受建模效果与系统延迟双重约束，早期注入最优但牺牲门控精度；（2）固定 1.6B 预算下 4-gram 略次优，更大记忆规模下是否有利未做结论；（3）第 6.3 节推理期抑制消融造成训练-推理不一致，故只分析信噪比最高的知识类与阅读理解类；（4）Engram-40B 未在每个任务上严格占优，作者判断为欠训练；（5）系统实验用稠密骨干的 nano-vLLM 原型与全量 CPU 卸载，属保守基线。其余方面（如超大规模下哈希冲突率、多语言覆盖面）论文未明确说明。
|
||||
- **影响**：该工作把"算法-系统协同设计"提升为设计原则：以往 N-gram 嵌入方案都在第 0 层、使访存与计算串行化，中层注入才能实现通信-计算重叠，并借助 Zipf 分布利用多级存储层次；这为绕过 GPU 显存上限、激进扩张参数规模提供了可落地路径。
|
||||
|
||||
## 六、关键术语对照表
|
||||
|
||||
| English | 中文 |
|
||||
|---|---|
|
||||
| Conditional Memory | 条件记忆 |
|
||||
| Conditional Computation | 条件计算 |
|
||||
| Engram | Engram（记忆印迹模块，论文未给中文译名） |
|
||||
| Sparsity Allocation | 稀疏容量分配 |
|
||||
| Mixture-of-Experts (MoE) | 混合专家 |
|
||||
| Iso-parameter / Iso-FLOPs | 等参数 / 等算力（同 FLOPs） |
|
||||
| N-gram Embedding | N-gram 嵌入 |
|
||||
| Hashed N-gram | 哈希 N-gram |
|
||||
| Multi-Head Hashing | 多头哈希 |
|
||||
| Tokenizer Compression | 分词器压缩 |
|
||||
| Context-aware Gating | 上下文感知门控 |
|
||||
| Multi-branch Architecture | 多分支架构 |
|
||||
| Manifold-Constrained Hyper-Connections (mHC) | 流形约束超连接 |
|
||||
| Multi-head Latent Attention (MLA) | 多头潜在注意力 |
|
||||
| Effective Depth | 有效深度 |
|
||||
| LogitLens | LogitLens（逐层逻辑回归探针） |
|
||||
| Centered Kernel Alignment (CKA) | 中心核对齐 |
|
||||
| Zipfian Distribution | 齐夫分布 |
|
||||
| Multi-Level Cache Hierarchy | 多级缓存层次 |
|
||||
| Prefetch-and-Overlap | 预取与重叠 |
|
||||
| Needle-in-a-Haystack (NIAH) | 大海捞针检索 |
|
||||
| LongPPL / RULER | 长上下文困惑度 / RULER 长上下文基准 |
|
||||
Binary file not shown.
@@ -0,0 +1,114 @@
|
||||
# NSA：原生可训练、硬件对齐的稀疏注意力
|
||||
|
||||
> **原文标题**：Native Sparse Attention: Hardware-Aligned and Natively Trainable Sparse Attention
|
||||
> **arXiv**：[2502.11089](https://arxiv.org/abs/2502.11089) ｜ **发布**：16 Feb 2025（最后修订 27 Feb 2025, v2）
|
||||
> **作者/机构**：Jingyang Yuan、Huazuo Gao、Damai Dai、Junyu Luo、Liang Zhao、Zhengyan Zhang、Zhenda Xie、Y. X. Wei、Lean Wang、Zhiping Xiao、Yuqing Wang、Chong Ruan、Ming Zhang、Wenfeng Liang、Wangding Zeng；DeepSeek-AI（深度求索）联合北京大学（Peking University, PKU-Anker LLM Lab）与华盛顿大学（University of Washington）
|
||||
> **原文 PDF**：[03-NSA-原生稀疏注意力-2502.11089.pdf](./03-NSA-原生稀疏注意力-2502.11089.pdf)
|
||||
|
||||
## 一、摘要（全文翻译）
|
||||
|
||||
长上下文建模对下一代语言模型至关重要，然而标准注意力机制的高昂计算开销带来了严峻的计算挑战。稀疏注意力为在保持模型能力的同时提升效率提供了颇具前景的方向。我们提出 NSA——一种原生可训练的稀疏注意力（Natively trainable Sparse Attention）机制，它把算法创新与硬件对齐（hardware-aligned）的优化结合起来，以实现高效的长上下文建模。NSA 采用动态分层稀疏策略，将粗粒度的 token 压缩与细粒度的 token 选择相结合，从而同时保持全局上下文感知能力与局部精度。我们的方法以两项关键创新推进了稀疏注意力的设计：（1）借助算术强度（arithmetic intensity）平衡的算法设计，并配合面向现代硬件的实现优化，我们取得了显著的加速；（2）我们实现了端到端训练，在不牺牲模型性能的前提下减少了预训练算力开销。如图 1 所示，实验表明：用 NSA 预训练的模型在通用基准、长上下文任务与指令式推理上均保持或超越了全注意力（Full Attention）模型。同时，在 64k 长度序列上，NSA 在解码、前向传播与反向传播各阶段相比全注意力都取得了显著加速，验证了其效率贯穿整个模型生命周期。
|
||||
|
||||
## 二、核心贡献
|
||||
|
||||
1. **提出原生可训练的分层稀疏注意力框架 NSA**：把每个查询对应的 key-value 重映射为三条并行分支——粗粒度 token 压缩（compression）、细粒度 token 选择（selection）、局部滑动窗口（sliding window），再用可学习门控加权融合，在 $N_t \ll t$ 的高稀疏度下同时保住全局感知与局部精度。
|
||||
2. **以算术强度平衡为准则做硬件对齐设计**：区分「训练/预填充是算力受限、解码是访存受限」两种优化目标；实现上按 GQA 分组加载查询、组内共享稀疏 KV 块，用 Triton 写专用核，使计算削减真正转化为延迟下降，覆盖预填充、解码与训练三个阶段。
|
||||
3. **稀疏模式在训练中原生学习，而非事后套用**：压缩与选择算子可反向传播，无需聚类、哈希等离散操作；27B 总参数 / 3B 激活参数模型的预训练损失全程低于全注意力基线。
|
||||
4. **性能不降反升**：通用基准 9 项中 7 项超过全注意力（平均 0.456 vs 0.443），LongBench 平均分 0.469 为全部方法最高（高出全注意力 +0.032、高出 Exact-Top +0.046）。
|
||||
5. **全生命周期加速与长上下文能力增强**：64k 序列下解码最高 11.6×、前向 9.0×、反向 6.0× 加速，且序列越长优势越大；64k 大海捞针（Needle-in-a-Haystack）检索全位置满分，AIME 24 上 NSA-R 相比全注意力基线在 8k/16k 生成预算下分别高出 +0.075 / +0.054。
|
||||
|
||||
## 三、方法与技术要点
|
||||
|
||||
### 3.1 为何要「原生」稀疏：对现有稀疏注意力的再思考（论文第 2 节）
|
||||
|
||||
**高效推理的幻觉（2.1 节）。** 一是**阶段受限**：H2O 只在解码阶段稀疏，预填充仍需昂贵的注意力图计算与索引构建；MInference 只做预填充稀疏。二者至少有一个阶段的计算量与全注意力相当，在书本摘要（预填充主导）或长思维链（解码主导）等负载下都无法全程加速。二是**与现代架构不兼容**：Quest 让每个头独立选择 KV 子集，在多头注意力（MHA）下计算与访存都稀疏；但在分组查询注意力（GQA）等共享 KV 架构下，同组各头实际需要的访存量是各自选择的**并集**——算得少了，但读得没少。
|
||||
|
||||
**可训练稀疏的神话（2.2 节）。** 事后（post-hoc）套用稀疏迫使模型偏离预训练轨迹：已有工作表明 top 20% 的注意力仅覆盖总注意力分数的 70%，预训练形成的「检索头」在剪枝下非常脆弱。训练侧有两类硬伤：ClusterKV 的 k-means 聚类、MagicPIG 的 SimHash 选择等**离散操作**切断计算图，梯度无法流经 token 选择；HashAttention 的 token 粒度选择虽理论可训，却要从 KV cache 加载大量彼此离散的 token，非连续访存令 FlashAttention 的分块计算失效，被迫退回低硬件利用率。
|
||||
|
||||
### 3.2 总体框架：三路重映射与门控融合（3.1、3.2 节）
|
||||
|
||||
标准注意力为 $\mathbf{o}_t=\operatorname{Attn}(\mathbf{q}_t,\mathbf{k}_{:t},\mathbf{v}_{:t})$，其中 $\alpha_{t,i}=e^{\mathbf{q}_t^{\top}\mathbf{k}_i/\sqrt{d_k}}$。NSA 的核心思路是把 $\mathbf{k}_{:t},\mathbf{v}_{:t}$ 换成更紧凑、信息密度更高的表示：
|
||||
|
||||
$$\tilde{K}_t=f_K(\mathbf{q}_t,\mathbf{k}_{:t},\mathbf{v}_{:t}),\qquad \tilde{V}_t=f_V(\mathbf{q}_t,\mathbf{k}_{:t},\mathbf{v}_{:t})$$
|
||||
|
||||
$$\mathbf{o}^{*}_t=\operatorname{Attn}(\mathbf{q}_t,\tilde{K}_t,\tilde{V}_t)=\sum_{c\in\mathcal{C}}g_t^{c}\cdot\operatorname{Attn}(\mathbf{q}_t,\tilde{K}_t^{c},\tilde{V}_t^{c})$$
|
||||
|
||||
映射集合仅取三种：$\mathcal{C}=\{\text{cmp},\text{slc},\text{win}\}$；门控 $g_t^{c}\in[0,1]$ 由输入特征经 MLP 加 sigmoid 得到。重映射后总长度 $N_t=\sum_{c\in\mathcal{C}}\operatorname{size}[\tilde{K}^{c}_t]$，靠 $N_t\ll t$ 维持高稀疏比。
|
||||
|
||||
### 3.3 三条分支的算法设计（3.3 节）
|
||||
|
||||
**Token 压缩（3.3.1）。** 把连续的 key/value 分块聚合为块级表示：$\tilde{K}^{\text{cmp}}_t=\{\varphi(\mathbf{k}_{id+1:id+l})\mid 0\leqslant i\leqslant\lfloor\frac{t-l}{d}\rfloor\}$，其中 $l$ 为块长、$d$ 为块间滑动步长，$\varphi$ 是带块内位置编码的可学习 MLP；通常取 $d<l$ 以缓解信息碎片化。该分支代价低，负责粗粒度全局扫描。
|
||||
|
||||
**Token 选择（3.3.2）。** 只靠压缩表示会丢失细粒度信息，故再挑出最重要的若干连续块。选块而非选 token 有两个理由：GPU 对连续块访问的吞吐远高于随机索引读，块计算也能充分利用 Tensor Core（FlashAttention 的基本原则）；且注意力分数具有空间连续性，相邻 key 的重要性往往相近（已有工作与本文可视化均证实）。重要性分数**几乎不额外花钱**——直接复用压缩分支产生的中间注意力分数：
|
||||
|
||||
$$\mathbf{p}_t^{\text{cmp}}=\operatorname{Softmax}\left(\mathbf{q}_t^{T}\tilde{K}_t^{\text{cmp}}\right)$$
|
||||
|
||||
当压缩块与选择块切分一致（$l'=l=d$）时，选择块重要性分数 $\mathbf{p}_t^{\text{slc}}$ 就等于 $\mathbf{p}_t^{\text{cmp}}$；切分不一致时按空间关系求和聚合：
|
||||
|
||||
$$\mathbf{p}_t^{\text{slc}}[j]=\sum_{m=0}^{\frac{l'}{d}-1}\sum_{n=0}^{\frac{l}{d}-1}\mathbf{p}_t^{\text{cmp}}\left[\frac{l'}{d}j-m-n\right]$$
|
||||
|
||||
对 GQA/MQA 这类跨头共享 KV 的模型，组内必须做**一致的块选择**，否则解码时访存量退化为并集，故先把组内各头分数相加：${\mathbf{p}_t^{\text{slc}}}'=\sum_{h=1}^{H}\mathbf{p}_t^{\text{slc},(h)}$。随后取 top-$n$ 块：
|
||||
|
||||
$$\mathcal{I}_t=\{i\mid\operatorname{rank}({\mathbf{p}_t^{\text{slc}}}'[i])\leqslant n\},\qquad \tilde{K}^{\text{slc}}_t=\operatorname{Cat}\left[\{\mathbf{k}_{il'+1:(i+1)l'}\mid i\in\mathcal{I}_t\}\right]$$
|
||||
|
||||
由于分数来自可微的注意力计算（而非聚类或哈希），选择路径的梯度可自然回流。
|
||||
|
||||
**滑动窗口（3.3.3）。** 局部模式学得最快，容易「短路」学习过程，使模型忽视压缩与选择分支。因此专门划出滑动窗口分支 $\tilde{K}^{\text{win}}_t=\mathbf{k}_{t-w:t},\ \tilde{V}^{\text{win}}_t=\mathbf{v}_{t-w:t}$ 显式承担局部上下文，让另两条分支专心学习各自特征。为进一步阻断捷径学习，三条分支各自持有**独立的 key/value 投影**，仅在门控处汇总——以极小开销换取「局部识别」与「长程识别」之间的梯度隔离。
|
||||
|
||||
### 3.4 硬件对齐的核设计（3.4 节）
|
||||
|
||||
压缩与滑窗分支可直接复用 FlashAttention-2 核，难点在稀疏选择分支。若沿用 FlashAttention 把时间上连续的查询块载入 SRAM 的策略，块内不同查询可能指向互不相交的 KV 块，访存低效。NSA 的关键改动是**更换查询分组方式**：对每个查询位置，把同一 GQA 组内**所有查询头**（共享同一组稀疏 KV 块）一起载入 SRAM。具体三点：
|
||||
|
||||
1. **组为中心的加载**：内层循环加载该组在位置 $t$ 的全部头查询 $Q\in\mathbb{R}^{[h,d_k]}$ 及其共享的稀疏块索引 $\mathcal{I}_t$；
|
||||
2. **共享 KV 抓取**：按 $\mathcal{I}_t$ 顺序把连续 KV 块 $K\in\mathbb{R}^{[B_k,d_k]},V\in\mathbb{R}^{[B_k,d_v]}$ 载入 SRAM，$B_k$ 为满足 $B_k\mid l'$ 的核块大小；
|
||||
3. **网格上的外层循环**：内层长度（正比于所选块数 $n$）对不同查询块几乎相同，故把查询/输出循环交给 Triton 的 grid 调度器，简化并优化核实现。
|
||||
|
||||
组内共享消除冗余 KV 搬运，并把负载均衡到各流式多处理器（SM），逼近最优算术强度。
|
||||
|
||||
## 四、主要实验结果
|
||||
|
||||
**预训练配置**：GQA + MoE 骨干，27B 总参数 / 3B 激活参数，30 层，隐层 2560；GQA 分 4 组、共 64 个头，$d_q=d_k=192$、$d_v=128$；MoE 采用 DeepSeekMoE，72 个路由专家 + 2 个共享专家，top-k 取 6，首层 MoE 换成 SwiGLU 形式的 MLP 以保证稳定。NSA 超参：压缩块长 $l=32$、滑动步长 $d=16$、选择块长 $l'=64$、选择块数 $n=16$（其中固定激活首个块与 2 个局部块）、窗口 $w=512$。全注意力与稀疏模型均在 8k 长度文本上预训练 270B token（引言处写作 260B token），随后在 32k 文本上用 YaRN 做继续训练与监督微调，两者都训到完全收敛。
|
||||
|
||||
| 基准 / 任务 | 本模型（NSA） | 对比对象 | 说明 |
|
||||
|---|---|---|---|
|
||||
| 通用基准平均（9 项） | 0.456 | 全注意力 0.443 | 9 项中 7 项领先 |
|
||||
| MMLU / MMLU-PRO / CMMLU | 0.565 / 0.286 / 0.587 | 0.567 / 0.279 / 0.576 | 知识类基本持平或略优 |
|
||||
| BBH / GSM8K / MATH / DROP | 0.521 / 0.520 / 0.264 / 0.545 | 0.497 / 0.486 / 0.263 / 0.503 | 推理类提升明显（DROP +0.042、GSM8K +0.034） |
|
||||
| MBPP / HumanEval | 0.466 / 0.348 | 0.482 / 0.335 | 代码类一降一升 |
|
||||
| LongBench 平均 | 0.469 | 全注意力 0.437；Exact-Top 0.423；Quest 0.392；InfLLM 0.383；H2O 0.303 | 全部方法最高 |
|
||||
| HPQ / 2Wiki（多跳 QA） | 0.437 / 0.356 | 0.350 / 0.305 | 分别 +0.087 / +0.051 |
|
||||
| LCC（代码理解） | 0.232 | 0.163 | +0.069 |
|
||||
| PassR-en（段落检索） | 0.905 | 0.830 | +0.075 |
|
||||
| 64k 大海捞针检索 | 全位置满分 | — | 分层稀疏同时覆盖全局扫描与局部精确定位 |
|
||||
| AIME 24（8k / 16k 生成预算） | NSA-R 0.121 / 0.146 | Full Attention-R 0.046 / 0.092 | 8k 下 +0.075，16k 下 +0.054 |
|
||||
| 训练加速（64k，Triton 同后端） | 前向 9.0× / 反向 6.0× | FlashAttention-2 | 序列越长优势越大 |
|
||||
| 解码访存量（8k→64k，等价 token 数） | 2048 / 2560 / 3584 / 5632 | 全注意力 8192 / 16384 / 32768 / 65536 | 预期加速 4× / 6.4× / 9.1× / 11.6× |
|
||||
|
||||
通用基准上 NSA 在 9 项中 7 项超过全注意力，推理类提升最突出，说明稀疏预训练迫使模型聚焦最关键信息、过滤无关通路的噪声；仅 MBPP（−0.016）与 MMLU（−0.002）落后。长上下文侧，在与基线同等稀疏度（每查询激活 2560 token，含前 128 个 token 与 512 个局部 token，对应 NSA 在 32k 序列下的平均激活量）下取得最高 LongBench 平均分，多跳问答与代码理解提升最大，表明原生稀疏使各组件同步适配。AIME 24 上 NSA-R 在两个预算下均优于全注意力基线，16k 预算下仍有提升（0.121→0.146），表明扩展推理深度未出现灾难性遗忘。
|
||||
|
||||
## 五、结论、局限与影响
|
||||
|
||||
- **结论**：NSA 把分层 token 压缩与块级 token 选择整合进可训练架构，在保持全注意力性能的同时实现训练与推理双向加速：通用基准与基线相当，长上下文与推理能力更强，并伴随可测量的延迟下降。
|
||||
- **局限性**：论文未设专门的局限性章节。可辨识的自述局限有两点：通用基准上 NSA「可能无法在较短序列上充分发挥效率优势」；注意力分数块状聚类的确切性质「仍需进一步研究」。此外，长上下文对比中作者排除了 LongBench 上所有模型得分都过低的若干子集（其无法提供有意义的比较）。
|
||||
- **影响**：该工作把「稀疏注意力能否既快又可训」从推理期工程推进到预训练架构层面；其算术强度协同设计范式与压缩+选择+滑窗的三分支结构，成为后续 DeepSeek 系模型（如 DeepSeek-V3.2 中的 DSA 稀疏注意力）的直接前身；长上下文能力不再必须付出全注意力的二次方代价。
|
||||
|
||||
## 六、关键术语对照表
|
||||
|
||||
| English | 中文 |
|
||||
|---|---|
|
||||
| Native Sparse Attention (NSA) | 原生稀疏注意力 |
|
||||
| Full Attention | 全注意力 |
|
||||
| Arithmetic Intensity | 算术强度 |
|
||||
| Token Compression | token 压缩（块级语义聚合） |
|
||||
| Token Selection | token 选择 |
|
||||
| Sliding Window | 滑动窗口 |
|
||||
| Blockwise Selection | 块级选择 |
|
||||
| Importance Score | 重要性分数 |
|
||||
| Grouped-Query Attention (GQA) | 分组查询注意力 |
|
||||
| Multi-Query Attention (MQA) | 多查询注意力 |
|
||||
| Multi-Head Attention (MHA) | 多头注意力 |
|
||||
| KV-cache Eviction | KV 缓存驱逐 |
|
||||
| FlashAttention | FlashAttention（分块精确注意力实现） |
|
||||
| Tensor Core | 张量核心 |
|
||||
| Gating Mechanism | 门控机制 |
|
||||
| Needle-in-a-Haystack | 大海捞针（长上下文检索测试） |
|
||||
Binary file not shown.
@@ -0,0 +1,113 @@
|
||||
# 面向专家混合模型的无辅助损失负载均衡策略
|
||||
|
||||
> **原文标题**：Auxiliary-Loss-Free Load Balancing Strategy for Mixture-of-Experts
|
||||
> **arXiv**：[2408.15664](https://arxiv.org/abs/2408.15664) ｜ **发布**：28 Aug 2024
|
||||
> **作者/机构**：Lean Wang、Huazuo Gao、Chenggang Zhao、Xu Sun、Damai Dai（DeepSeek-AI；部分作者兼属北京大学计算机学院多媒体信息处理国家重点实验室）
|
||||
> **原文 PDF**：[04-Aux-Loss-Free-Load-Balancing-2408.15664.pdf](./04-Aux-Loss-Free-Load-Balancing-2408.15664.pdf)
|
||||
|
||||
## 一、摘要（全文翻译）
|
||||
|
||||
对于专家混合（Mixture-of-Experts, MoE）模型而言，专家负载不均衡会导致路由崩塌（routing collapse）或计算开销增加。现有方法通常采用辅助损失（auxiliary loss）来鼓励负载均衡，但较大的辅助损失会在训练中引入不可忽视的干扰梯度，从而损害模型性能。为了在控制负载均衡的同时不在训练中产生不需要的梯度，我们提出了 Loss-Free Balancing，其特色是一种无辅助损失的负载均衡策略。具体而言，在 top-K 路由决策之前，Loss-Free Balancing 会先对每个专家的路由分数施加一个专家级别的偏置（expert-wise bias）。通过根据每个专家近期的负载动态更新其偏置，Loss-Free Balancing 能够持续维持专家负载的均衡分布。此外，由于 Loss-Free Balancing 不产生任何干扰梯度，它也提升了 MoE 训练所能带来的模型性能上限。我们在参数量最高 3B、训练 token 数最高 200B 的 MoE 模型上验证了 Loss-Free Balancing 的性能。实验结果表明，与传统由辅助损失控制的负载均衡策略相比，Loss-Free Balancing 同时取得了更好的性能和更好的负载均衡。
|
||||
|
||||
## 二、核心贡献
|
||||
|
||||
1. **提出无辅助损失的负载均衡策略（Loss-Free Balancing）**：把负载均衡从「损失函数上的正则项」改为「路由分数上的可调偏置」，从根源消除辅助损失引入的干扰梯度，破解均衡与性能不可兼得的矛盾。
|
||||
2. **给出可迭代的专家偏置更新算法**：在 top-K 选择前为每个专家引入偏置 $b_i$，每个 batch 后按「重负载降偏置、轻负载升偏置」更新，并实证确定更新率 $u=0.001$。
|
||||
3. **打破「均衡 vs. 性能」的两难**：论文在 $\alpha\in\{0,10^{-4},10^{-3},10^{-2}\}$ 上展示辅助损失系数越小越不均衡、越大性能越差（Figure 2），本方法在两方面均优于该基线。
|
||||
4. **均衡度呈量级改善**：1B 模型验证困惑度由 9.56 降至 9.50、$\text{MaxVio}_{\text{global}}$ 由 0.72 降至 0.04；3B 模型困惑度由 7.97 降至 7.92、$\text{MaxVio}_{\text{global}}$ 由 0.52 降至 0.04。
|
||||
5. **与专家并行（Expert Parallelism）天然兼容**：computation batch 越大均衡越好，而辅助损失方法在大 batch 下基本停滞。
|
||||
6. **揭示 Expert Choice（EC）的未来 token 泄漏**：理论上每 token 泄漏量可超过 $K\log_2\frac{1-R}{R}$ bit，并有实验佐证；本方法遵守因果约束。
|
||||
|
||||
## 三、方法与技术要点
|
||||
|
||||
### 3.1 背景：MoE 层与辅助损失的困境
|
||||
|
||||
MoE 层用多个专家 FFN 替换标准 Transformer 的 MLP 层，对第 $t$ 个 token 的输入 $\mathbf{u}_t$，输出为
|
||||
|
||||
$$\mathbf{h}_t=\mathbf{u}_t+\sum_{i=1}^{N}g_{i,t}\operatorname{FFN}_i(\mathbf{u}_t),\quad s_{i,t}=G(\mathbf{u}_t^{\top}\mathbf{e}_i),$$
|
||||
|
||||
其中 $g_{i,t}$ 仅在 $s_{i,t}$ 落入 top-K 时为 $s_{i,t}$，否则为 0。负载不均衡的两类危害是路由崩塌（只选少数专家，其余专家训练不足）与多设备下的计算瓶颈。为此主流做法引入辅助损失：
|
||||
|
||||
$$\mathcal{L}_{\text{Balance}}=\alpha\sum_{i=1}^{N}f_i P_i,\quad f_i=\frac{N}{KT}\sum_{t=1}^{T}\mathbbm{1}(\text{Token }t\text{ selects Expert }i),\quad P_i=\frac{1}{T}\sum_{t=1}^{T}s_{i,t}.$$
|
||||
|
||||
痛点是：它本质上是额外正则项，会与语言建模目标产生冲突梯度。$\alpha$ 太小则负载不均衡甚至路由崩塌，太大则明显拖累性能——即论文所称「负载均衡与模型性能之间的两难」。
|
||||
|
||||
### 3.2 核心设计：偏置门控分数（biased gating score）
|
||||
|
||||
Loss-Free Balancing 不改损失函数，而是改路由的判定依据：为每个专家加一个偏置 $b_i$，用**带偏置的分数**决定 top-K 归属，但**最终加权输出仍用原始分数**：
|
||||
|
||||
$$g_{i,t}=\begin{cases}s_{i,t},& s_{i,t}+b_i\in\operatorname{Topk}\left(\{s_{j,t}+b_j\mid 1\le j\le N\},K\right),\\ 0,&\text{otherwise.}\end{cases}$$
|
||||
|
||||
关键在于：$b_i$ 只影响「选谁」、不进入 $g_{i,t}$，因此既不改变专家输出的加权方式，也不向反向传播注入与语言建模目标无关的梯度；均衡由偏置数值承担，性能由主损失主导，两者解耦。
|
||||
|
||||
### 3.3 偏置更新算法
|
||||
|
||||
偏置按「重负载压低、轻负载抬高」的原则迭代更新。对每个 batch：先按上式正常前向与训练，然后统计各专家被分配的 token 数 $c_i$ 及其均值 $\overline{c_i}$，得到负载违背误差 $e_i=\overline{c_i}-c_i$，再更新
|
||||
|
||||
$$b_i \leftarrow b_i + u\cdot\operatorname{sign}(e_i).$$
|
||||
|
||||
作者强调只使用**历史** batch 的负载信息：用当前序列的负载会破坏语言建模的因果约束，导致未来 token 信息泄漏。这一「事后更新」性质使本方法兼具控制力与因果安全性。
|
||||
|
||||
### 3.4 与其他负载均衡方法的对比
|
||||
|
||||
论文从三个维度刻画各方案：是否均衡、是否引入干扰梯度、是否泄漏未来 token。
|
||||
|
||||
| 负载均衡方法 | 专家负载均衡 | 干扰梯度 | 未来 token 泄漏 |
|
||||
|---|---|---|---|
|
||||
| 辅助损失控制（强） | 均衡 | 强 | 无泄漏 |
|
||||
| 辅助损失控制（弱） | 不均衡 | 弱 | 无泄漏 |
|
||||
| Expert Choice | 均衡 | 无 | 有泄漏 |
|
||||
| **Loss-Free（本文）** | **均衡** | **无** | **无泄漏** |
|
||||
|
||||
EC 让每个专家被分配完全相同数量的 token 以实现完美均衡，但每个 token 的目标专家取决于同序列/批次中的未来 token，破坏因果约束；辅助损失方法则始终需权衡均衡与性能。本文方法是唯一同时满足三者的方案。
|
||||
|
||||
### 3.5 算法变体的实证选择
|
||||
|
||||
论文对偏置更新算法做了消融，据此选定主实验配置：
|
||||
|
||||
- **更新率**：$u=0.0001$ 收敛太慢（训练早期均衡差），$u=0.01$ 令 $b_i$ 后期剧烈波动、反而恶化均衡，$u=0.001$ 兼顾两者。
|
||||
- **更新规则**：换成 $b_i+u\cdot e_i$（违背更严重的专家偏置变化更快）虽略微改善均衡，但困惑度都更差，故保留 sign 版本。
|
||||
- **加性 vs. 乘性偏置**：乘性形式 $s_{i,t}*b_i$ 参与 top-K（初始化为 1），均衡相近而性能略差，故采用加性偏置。
|
||||
- **softmax 门控**：softmax 的归一化使两个专家间的分数差受其他专家影响，偏置更难调，此时改用 $b_i+u\cdot e_i$ 变体且 $u=10^{-3}$。
|
||||
|
||||
## 四、主要实验结果
|
||||
|
||||
| 基准 / 任务 | 本模型 | 对比对象 | 说明 |
|
||||
|---|---|---|---|
|
||||
| 1B MoE，验证困惑度 | 9.50（Loss-Free） | 9.56（辅助损失控制，$\alpha$=0.001） | 1B 模型训练 100B tokens |
|
||||
| 1B MoE，$\text{MaxVio}_{\text{global}}$ | 0.04 | 0.72 | 均衡度提升约 18 倍 |
|
||||
| 3B MoE，验证困惑度 | 7.92（Loss-Free） | 7.97（辅助损失控制） | 3B 模型训练 200B tokens |
|
||||
| 3B MoE，$\text{MaxVio}_{\text{global}}$ | 0.04 | 0.52 | 均衡度提升约 13 倍 |
|
||||
| 1B，更新规则消融（$u$=0.001） | 9.50 / 0.044（sign 版本） | 9.53 / 0.028、9.51 / 0.036、9.51 / 0.040（$u\cdot e_i$ 变体） | sign 版本性能最优 |
|
||||
| 1B，偏置形式消融 | 9.50 / 0.044（加性） | 9.52 / 0.041、9.52 / 0.036、9.54 / 0.048（乘性 Bias） | 加性偏置性能最好 |
|
||||
| 1B，softmax 门控 | 9.599 / 0.027（Loss-Free） | 9.604 / 0.937（Loss-Controlled，$\alpha$=0.0003） | 困惑度略优，均衡显著更优 |
|
||||
| Expert Choice 泄漏实验（2B MoE / 100B tokens） | chunk=512 出现约 10% 异常 loss 下降 | chunk=2048 / 8192 无此现象；打乱 token 后异常消失 | 证明未来 token 泄漏真实存在 |
|
||||
| 理论泄漏量（9 层 MoE，16 专家，平均激活 2） | 每 token 超过 50 bit | — | 足以让每个 token 推断其后继 token 的身份 |
|
||||
|
||||
在 1B 与 3B 两个规模上，本方法**同时**取得更低困惑度与更好的全局负载均衡：困惑度领先 0.05–0.06，而 $\text{MaxVio}_{\text{global}}$ 从 0.72/0.52 降到 0.04，属量级差异。训练曲线（Figure 3）与 computation batch 扫描（Figure 5）显示其均衡在绝大部分训练步中更优、且 batch 越大优势越明显，而辅助损失方法在大 batch 下基本不再改善。论文未报告下游任务评测。
|
||||
|
||||
## 五、结论、局限与影响
|
||||
|
||||
- **结论**：论文提出 Loss-Free Balancing——一种不引入辅助损失梯度的 MoE 负载均衡控制方法，解决了辅助损失在强制均衡时引入额外梯度、损害性能的问题。在 1B/3B MoE 上，相比辅助损失训练，本方法同时取得更好的性能与负载均衡。
|
||||
- **局限性**：论文未设单独的局限性章节，可归纳的自述边界为：（1）实验规模仅到 3B 参数、最多 200B tokens；（2）主实验固定用 sigmoid 门控（其基线优于 softmax），softmax 场景需改用 $u\cdot e_i$ 变体并单独调参；（3）3B 模型训练 token 数前后不一致——摘要、正文与附录 B 写 200B，第 6 节结论写 300B；（4）未报告下游任务指标，仅以验证困惑度与 $\text{MaxVio}$ 衡量。
|
||||
- **影响**：该策略已被 DeepSeek-V3 采用，成为大规模 MoE 训练中替代辅助损失的实用方案：它把负载均衡从「损失层面」移到「路由层面」，解除均衡强度与性能的绑定；又因与专家并行天然兼容，并行规模越大优势越显著。论文对 EC 泄漏的分析也警示了 MoE 路由的因果安全性。
|
||||
|
||||
## 六、关键术语对照表
|
||||
|
||||
| English | 中文 |
|
||||
|---|---|
|
||||
| Mixture-of-Experts (MoE) | 专家混合 |
|
||||
| Loss-Free Balancing | 无辅助损失负载均衡 |
|
||||
| Auxiliary Loss / Loss-Controlled | 辅助损失 / 辅助损失控制 |
|
||||
| Interference Gradient | 干扰梯度 |
|
||||
| Routing Collapse | 路由崩塌 |
|
||||
| Expert-wise Bias | 专家级偏置 |
|
||||
| Biased Gating Score | 偏置门控分数 |
|
||||
| Top-K Routing | Top-K 路由 |
|
||||
| Bias Update Rate | 偏置更新率 |
|
||||
| Maximal Violation (MaxVio) | 最大违背度 |
|
||||
| Expert Choice (EC) | 专家选择路由 |
|
||||
| Future Token Leakage | 未来 token 泄漏 |
|
||||
| Expert Parallelism | 专家并行 |
|
||||
| Computation Batch | 计算批次 |
|
||||
| DeepSeekMoE | 深度求索 MoE 架构 |
|
||||
Binary file not shown.
Reference in New Issue
Block a user