针对大模型长文本推理,提出分层动态压缩方案提升内存效率。
MoE-nD: Per-Layer Mixture-of-Experts Routing for Multi-Axis KV Cache Compression

- 按层独立分配丢弃率与量化位数,实现混合专家路由优化
- 在14倍压缩下保持基线精度,优于所有对比方法
- 适用于长上下文推理任务,尤其对复杂推理题效果显著
KV缓存是长上下文大模型推理的主要瓶颈。现有压缩方法仅在四维KV张量的单一维度上操作——如序列丢弃(token eviction)、量化(precision)、低秩投影(head dimension)或跨层共享,且对各层采用统一策略。我们发现这种同质性导致性能损失:不同层对压缩操作的响应差异显著,最优的逐层丢弃率与量化组合并非均匀分布。为此提出MoE-nD,一个分层混合专家框架,可在全局内存预算下为每层选择独立的(丢弃率, K位, V位)配置。通过离线校准的贪心求解器确定最小预测质量损失的路由策略;推理时通过单一注意力补丁联合应用各层异构的丢弃与量化。在LongBench-v1的4个任务子集(16k输入,每任务n=50,适配推理模型协议)上,MoE-nD的异构变体以14倍压缩(136MB)达到与未压缩基线(1.9GB)相当的性能,而其他测试压缩基线(1d、2d_uniform、2d)在相近或更小内存下均未超过8/100得分。该优势在AIME推理基准上持续显现,相较最强逐层量化基线提升6至27分(共八种配置)。两个未见提升案例(MATH-500和LongBench的TREC)有明确成因(输入过短,求解器默认保留率keep=1.0),清晰界定分层丢弃路由的有效范围。
原文摘要 · Abstract (English)
KV cache memory is the dominant bottleneck for long-context LLM inference. Existing compression methods each act on a single axis of the four-dimensional KV tensor -- token eviction (sequence), quantization (precision), low-rank projection (head dimension), or cross-layer sharing -- but apply the same recipe to every layer. We show that this homogeneity leaves accuracy on the table: different layers respond very differently to each compression operation, and the optimal per-layer mix of eviction and quantization is far from uniform. We propose MoE-nD, a mixture-of-experts framework that routes each layer to its own (eviction-ratio, K-bits, V-bits) tuple under a global memory budget. An offline-calibrated greedy solver chooses the routing that minimizes predicted quality loss; at inference time, per-layer heterogeneous eviction and quantization are applied jointly through a single attention patch. On a 4-task subset of LongBench-v1 (16k inputs, n=50 per task, adapted reasoning-model protocol; see section Experiments), MoE-nD's hetero variant matches our uncompressed 1.9~GB baseline at 14x compression (136~MB) while every other compressed baseline we tested (1d, 2d_uniform, 2d) at comparable or smaller memory stays under 8/100. The gains hold on AIME reasoning benchmarks (+6 to +27 pts over the strongest per-layer-quantization baseline across eight configurations). Two null results -- MATH-500 and LongBench's TREC -- share a principled cause (short inputs, solver picks keep=1.0 on most layers), cleanly characterizing when per-layer eviction routing has headroom to help.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。