🤖 machine learning

MoE-nD: Per-Layer Mixture-of-Experts Routing for Multi-Axis KV Cache Compression

该论文提出了 MoE-nD 框架，通过混合专家机制为 Transformer 的不同层动态分配最优的令牌丢弃率与量化精度组合，在严格的全局显存预算下实现了比现有均匀压缩方法显著更低的精度损失和更高的推理效率。

原作者： Libo Sun, Peixiong He, Po-Wei Harn, Xiao Qin

发布于 2026-04-21

📖 1 分钟阅读☕ 轻松阅读

原作者： Libo Sun, Peixiong He, Po-Wei Harn, Xiao Qin

原始论文采用 CC BY 4.0 许可（http://creativecommons.org/licenses/by/4.0/）。 ✨ 这是对下方论文的AI生成解释。它不是由作者撰写或认可的。如需技术准确性，请参阅原始论文。阅读完整免责声明

这篇文章介绍了一种名为 MoE-nD 的新方法，旨在解决大型人工智能模型（LLM）在处理长文本时遇到的一个核心难题：“记忆”太占地方了。

为了让你轻松理解，我们可以把 AI 模型想象成一个正在写长篇小说的超级作家，而 KV Cache（键值缓存）就是作家手边的草稿纸。

1. 核心问题：草稿纸不够用了

当作家（AI）要写一个几万字的故事（长上下文）时，他需要不断回顾之前写过的内容。

现状：传统的做法是，无论故事写到哪一页，作家都使用同样大小的草稿纸，并且同样粗糙地记录所有细节。
后果：如果故事太长，草稿纸堆满了桌子（显存爆满），作家就写不下去了，或者不得不把桌子清空（导致模型崩溃或变慢）。

2. 旧方法的局限：一刀切

以前的压缩方法就像是一个只会用一种策略的管家：

策略 A：把草稿纸上的字写小一点（量化/降低精度）。
策略 B：把草稿纸撕掉一半，只留关键页（丢弃/蒸发）。
问题：管家对所有页面都使用完全相同的策略。比如，不管这一页是写“主角登场”的关键剧情，还是写“天气不错”的废话，管家都把它们撕掉同样的比例，或者都写得同样潦草。
结果：这导致要么关键剧情被撕坏了（模型变笨），要么废话占满了空间（压缩效果差）。

3. MoE-nD 的创意：智能的“分层管家”

MoE-nD 的核心思想是：每一页草稿纸的情况都不一样，需要“量体裁衣”。

作者发现：

开头几页（浅层）：可能主要是设定背景，稍微写潦草点（量化）或者撕掉一点（丢弃）都没关系。
中间几页（中层）：剧情转折，必须保留细节，不能乱撕。
结尾几页（深层）：逻辑推理最关键，必须字字珠玑，一点都不能丢。

MoE-nD 就像一个拥有“专家团”的智能管家：

专家团（MoE）：它把“撕纸”和“写小字”看作两个不同的专家技能。
智能路由（Routing）：在写故事之前，管家会先快速浏览一下（离线校准），然后给每一页草稿纸分配专属方案：
- 第 1 页：撕掉 50%，字写小点（省空间）。
- 第 10 页：只撕掉 10%，字保持清晰（保质量）。
- 第 20 页：完全不撕，但把字写得极小（极致压缩）。
全局预算：管家手里有一个总的“桌子面积”限制（显存上限），他会聪明地分配，确保在桌子不爆满的前提下，保留最重要的剧情。

4. 实验结果：既省地儿，又不糊涂

作者用这个新方法测试了 AI 写长故事和做数学题的能力：

惊人的压缩率：在保持 AI 智商（准确率）完全不下降的情况下，他们把原本需要 1.9 GB 的草稿纸空间，压缩到了 136 MB（相当于把 14 个书架的内容塞进了一个公文包）。
对比惨烈：其他旧方法（要么只撕纸，要么只写小字，且对所有页一视同仁）在同样的空间下，AI 的智商直接“掉线”，准确率暴跌。
为什么有效：因为旧方法像是一个笨拙的裁缝，给所有人都做同一码数的衣服；而 MoE-nD 是高级定制，给胖的做宽松版，给瘦的做修身版，每个人（每一层）都最舒服。

5. 什么时候这个方法不管用？

论文也诚实地指出，如果故事很短（比如只有几行字），或者给的空间非常充裕，这个“智能管家”就没必要大费周章去分配了。这时候，大家随便怎么压缩都行，因为草稿纸本来就不多。

总结

MoE-nD 就像是给 AI 的“记忆系统”装上了一个智能的、动态的资源管理器。它不再一刀切地压缩，而是根据每一层信息的重要程度，灵活决定是“撕掉”还是“写小”。

一句话比喻：
以前的压缩方法是把整本书的纸张都换成薄纸；
MoE-nD 的方法是把故事书里的“广告页”换成薄纸，把“精彩章节”换成厚纸，这样既省了书店的库存（显存），又没让读者（AI）错过任何精彩剧情。

MoE-nD 技术总结：基于每层混合专家路由的多轴 KV 缓存压缩

1. 研究背景与问题 (Problem)

在大语言模型（LLM）的长上下文推理中，KV 缓存（Key-Value Cache） 的显存占用已成为主要的性能瓶颈，甚至导致在消费级硬件上出现显存溢出（OOM）。现有的压缩方法通常针对 KV 张量的四个维度（Token 序列、精度、头维度、跨层共享）中的单一维度进行操作，且对所有 Transformer 层应用相同的压缩策略（均匀压缩）。

然而，论文通过实证观察发现：

层间异质性（Layer Heterogeneity）：不同的 Transformer 层对不同的压缩操作（如 Token 丢弃 vs. 量化）的敏感度差异巨大。例如，在激进丢弃 90% 的 Token 时，不同层的精度损失差异可达 689 倍。
单一策略的局限性：现有的均匀压缩策略（Uniform Compression）无法适应这种层间差异，导致在追求高压缩比时，精度损失过大；或者为了保精度而浪费了压缩空间。
缺乏联合路由：目前尚无工作能在全局显存预算下，针对每一层联合优化“丢弃率（Eviction Ratio）”和"K/V 量化位数（Quantization Bits）”这两个轴。

2. 核心方法论 (Methodology)

论文提出了 MoE-nD（Mixture-of-Experts for n-Dimensional Routing），一种基于混合专家（MoE）思想的框架，旨在为每一层动态分配最优的压缩配置。

2.1 核心设计

专家定义：将不同的压缩操作（不同的 Token 保留率、K 量化位数、V 量化位数）视为“专家”。
路由机制：引入一个离线校准的贪心求解器（Offline-calibrated Greedy Solver）。
- 离线阶段：使用一个轻量级的代理指标（基于 27 个 Token 的推理提示，计算压缩后注意力输出的 L2 误差）来构建每层的敏感度表（Sensitivity Table）。
- 求解过程：在满足全局显存预算 $M$ 的约束下，求解器迭代地选择“单位显存节省带来的精度提升最大”的层进行升级（即选择更激进的压缩配置），直到预算耗尽。
- 输出：为每一层 $\ell$ 分配一个独特的配置三元组 $(keep\_ratio_\ell, kbits_\ell, vbits_\ell)$ 。

2.2 推理实现

异构注意力补丁（Heterogeneous Attention Patch）：由于不同层保留的 Token 数量不同，导致序列长度不一致。MoE-nD 修改了注意力内核，为每一层维护独立的 Token 位置数组 $p^{(\ell)}$ ，并在 RoPE（旋转位置编码）反旋转时使用该层的位置信息，而非全局位置。
触发机制：为了防止某些层（保留率=1.0）永远不触发丢弃导致成本失控，设置了基于生成步数的触发器（每生成 128 个 Token 触发一次丢弃检查）。

3. 关键贡献 (Key Contributions)

首次联合路由：首次提出了在每层粒度上联合路由“序列轴（Eviction）”和“精度轴（Quantization）”的方法。
实证发现：揭示了层间对压缩操作的敏感度存在巨大差异（特别是激进丢弃时），且不同层对“丢弃”和“量化”的偏好截然不同（有的层丢弃损失小，有的层量化损失小）。
消融实验结论：通过对比实验（ $2d_{uniform} \to 2d \to 2d_{hetero}$ ）证明，每层丢弃路由（Per-layer Eviction Routing） 是性能提升的关键驱动力，而每层量化路由带来的收益微乎其微。
明确适用边界：通过两个“零结果”（MATH-500 和 TREC）明确了方法的适用场景：仅在输入长度较长或预算较紧、迫使求解器进行非平凡路由分配时有效；在短文本场景下，求解器会自动选择保留所有 Token，退化为均匀策略。

4. 实验结果 (Results)

实验基于 DeepSeek-R1-Distill-Qwen-7B 模型，在 LongBench-v1（4 个任务子集）、AIME-24/25 和 MATH-500 上进行评估。

4.1 LongBench-v1 (长上下文)

极致压缩：MoE-nD 的异构变体（ $2d_{hetero}$ ）在 14 倍压缩（显存从 1.9GB 降至 136MB）的情况下，其平均准确率（12.0）完全匹配未压缩基线（11.5），且误差在统计显著性范围内。
对比优势：在相同或更小的显存预算下，其他基线方法（包括单轴丢弃、均匀双轴压缩、仅每层量化）的准确率均低于 8/100，损失高达 33%-64%。
具体数据：在 136MB 显存下， $2d_{hetero}$ 得分为 12.0，而仅做每层量化的 $2d$ 仅为 5.9，证明了“每层丢弃路由”的关键作用。

4.2 AIME 推理基准

在 AIME-24 和 AIME-25 上，MoE-nD 在所有预算设置下均优于非异构的双轴基线。
在极紧预算（如 AIME-25, b=64）下，MoE-nD 达到 30% 准确率，而其他方法接近 0%。
有趣的是，在部分推理任务中，激进的 Token 丢弃甚至能提升性能（去除干扰性中间 Token），MoE-nD 能利用这一特性。

4.3 零结果分析 (Null Results)

在 MATH-500（短提示）和 TREC（短提示）上，MoE-nD 未表现出优势。
原因：由于输入短且预算宽松，贪心求解器自动为绝大多数层选择了 keep=1.0（不丢弃），导致异构路由退化为均匀策略。这验证了该方法在“短上下文”场景下无需复杂路由的结论。

5. 意义与影响 (Significance)

打破均匀假设：挑战了 LLM 压缩中“所有层一视同仁”的传统假设，证明了自适应、非均匀的压缩策略能显著释放显存潜力。
高效推理：使得在消费级硬件（如单张 H200 甚至更低端显卡）上运行超长上下文（16k+）的推理成为可能，且几乎不损失精度。
方法论模板：为未来的模型压缩研究提供了一个新范式，即通过校准的“路由器”在多个正交的压缩维度上进行联合优化，而非单一维度的堆叠。
工程落地：虽然目前尚未完全兼容 FlashAttention（存在 O(T^2) 的中间显存开销），但其提出的异构注意力补丁机制为未来实现提供了明确的技术路径。

总结：MoE-nD 通过智能地“因材施教”（为每一层定制压缩策略），在保持推理精度的前提下，将 KV 缓存压缩了 14 倍，解决了长上下文推理中的显存瓶颈问题，并清晰地界定了该技术的有效边界。