💬 NLP

Adaptive Layer Selection for Layer-Wise Token Pruning in LLM Inference

本文提出了一种名为 ASL 的免训练方法，通过利用注意力分数排序的令牌方差来自适应地选择关键层进行令牌剪枝，从而在满足 KV 缓存预算的同时，显著提升了大语言模型在困难任务中的推理精度与性能平衡。

原作者： Rei Taniguchi, Yuyang Dong, Makoto Onizuka, Chuan Xiao

发布于 2026-04-17

📖 1 分钟阅读☕ 轻松阅读

原作者： Rei Taniguchi, Yuyang Dong, Makoto Onizuka, Chuan Xiao

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

这篇论文介绍了一种名为 ASL (Adaptive Selection Layer，自适应选择层) 的新方法，旨在解决大型语言模型（LLM）在处理超长文本时遇到的“内存爆炸”和“反应迟钝”的问题。

为了让你轻松理解，我们可以把整个 LLM 推理过程想象成一位正在阅读一本厚达 10 万页的百科全书的“超级侦探”。

1. 背景：侦探的困境（KV Cache 问题）

想象一下，这位侦探（LLM）正在读一本巨书。为了回答问题，他必须记住之前读过的所有关键信息。

KV Cache（键值缓存）：就是侦探脑子里的“笔记”。每读一页，他就要记下一笔。
问题：如果书有 10 万页，他的脑子（显存）根本装不下这么多笔记。笔记太多，不仅占地方，而且每次找信息都要翻遍所有笔记，速度极慢。

2. 现有的笨办法：固定剪枝（Layer-wise Token Pruning）

为了解决脑子装不下的问题，以前的方法（如 FastKV）想出了一个策略：“只记重点，扔掉废话”。

做法：他们规定，侦探读到第 15 页时，必须停下来，把前 15 页里最重要的 100 个词挑出来记在脑子里，剩下的几千个词全部扔掉。从第 16 页开始，只基于这 100 个词继续读。
缺点：这个“第 15 页”是死板规定的，不管书的内容是什么。
- 简单任务（比如问“书里第一章讲了什么”）：侦探在第 5 页就能找到答案。在第 15 页才做决定，虽然有点晚，但还能接受。
- 困难任务（比如“在 10 万页里找一句特定的暗语”）：侦探可能读到第 50 页还在迷茫，因为线索太隐蔽。如果强行在第 15 页就扔掉其他词，他可能把真正的线索也扔了，导致完全找不到答案。

这就好比：不管你是去超市买瓶水（简单），还是去大海捞针（困难），保安都规定你在进门的第 10 步就必须扔掉身上 90% 的行李。买水的人觉得没事，但捞针的人早就把针扔了。

3. 核心创新：ASL（自适应选择层）

这篇论文提出的 ASL，就像给侦探配了一个聪明的“直觉助手”。这个助手不规定死在第几页做决定，而是观察侦探的“注意力”变化。

核心原理：看“排名”变不变

助手会观察侦探对书中每个词的“关注度”（注意力分数）：

刚开始读（浅层）：侦探对书里的词都很模糊，觉得“这个好像有点用，那个也好像有点用”，大家的关注度排名忽高忽低，变化很大（方差大）。这时候扔东西太早，容易扔错。
读到深处（深层）：侦探终于看明白了！“啊！原来只有这几个词是关键的！”此时，关键词的排名非常稳定，不再乱跳（方差变小）。

ASL 的策略：

助手会一直盯着这个“排名稳定性”。
一旦它发现：“嘿，侦探的注意力已经稳定地集中在某几个词上了，排名不再剧烈波动了！”
立刻喊停：“就是现在！把不重要的词扔掉，只保留这几个核心词！”

比喻：

旧方法：不管你在听什么歌，都在第 30 秒切歌。
ASL：监听你的心跳。如果你听的是轻音乐，心跳平稳，第 10 秒就切；如果你听的是摇滚，心跳剧烈，那就等到第 60 秒心跳平稳了再切。它根据任务的“难度”自动调整时机。

4. ASL 带来的好处

简单任务不浪费：如果是简单的问答，侦探很快就能锁定重点，ASL 会早早地（比如第 10 层）就让他开始“精简笔记”，速度飞快。
困难任务不丢分：如果是很难的“大海捞针”任务，ASL 会等到侦探彻底想明白（比如第 25 层）再让他精简，准确率极高。
灵活适应：不需要人工去调参数，它自己就能根据任务难度决定“什么时候开始扔垃圾”。

5. 实验结果：真的好用吗？

作者让 ASL 和现有的“死板方法”（如 FastKV）在三个大考（InfiniteBench, RULER, NIAH）中 PK：

找针任务（NIAH）：在 25 万字的长文中找一句话。旧方法经常找丢，ASL 几乎全对。
复杂推理：在长文档中做数学题或代码调试。ASL 的准确率明显高于旧方法。
速度：虽然 ASL 因为要“观察”一下再决定，比死板方法稍微慢了一点点（就像多花了一秒思考），但它用这点微小的时间代价，换来了巨大的准确率提升。特别是在最难的任务上，它完胜对手。

总结

ASL 就像是一个“懂行”的图书管理员。
以前的管理员不管书多难，都机械地在第 15 页开始删书。
ASL 这个新管理员会观察读者（模型）的读书状态：“你看起来还在迷茫，先别删，继续读；哦，你眼神亮了，思路清晰了，好，现在可以开始删掉废话了！”

这种方法让大模型在处理超长文本时，既省内存，又不丢分，还能自动适应各种难度的任务，是长文本推理领域的一大进步。

这是一篇关于大语言模型（LLM）长上下文推理优化的技术论文总结。论文提出了一种名为 ASL (Adaptive Selection Layer) 的新方法，旨在解决现有层级 Token 剪枝（Layer-wise Token Pruning）技术在处理不同难度任务时的灵活性不足问题。

以下是该论文的详细技术总结：

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

KV Cache 瓶颈：随着 LLM 处理长上下文（如长文档分析、超长对话）的能力增强，Key-Value (KV) Cache 的显存占用成为推理的主要瓶颈。
现有方法的局限性：
- 当前的层级 Token 剪枝方法（如 FastKV, GemFilter 等）通常采用预定义的固定层（Fixed Selection Layer）来进行 Token 选择。
- 缺乏灵活性：这种固定设计无法适应不同任务的难度差异。
  - 简单任务（如普通问答 QA）：相关信息通常在浅层即可识别，过早剪枝（在中间层选择）效果良好。
  - 困难任务（如 KV 检索、多 Key NIAH）：上下文与查询的语义相似度高，早期难以定位关键信息。如果在固定层（如第 15 层）进行剪枝，会导致关键信息丢失，准确率急剧下降。
- 权衡困境：为了在困难任务上保持精度，现有方法不得不推迟剪枝层（增加显存）或增加 KV 预算，从而削弱了剪枝带来的内存和速度优势。

2. 核心方法论 (Methodology)

作者提出了 ASL (Adaptive Selection Layer)，一种**无需训练（Training-free）**的自适应方法，能够根据推理过程中的注意力模式动态决定 Token 剪枝的层。

核心思想

利用注意力分数排序后的 Token 秩（Token Ranks）的方差作为指标，来判断注意力是否已经稳定地聚焦于特定的 Token 子集。

具体步骤

监控秩方差：
- 从第 $L_{min}$ 层开始，计算最近 $L_{obs}$ 层中，按注意力分数排序的 Token 秩的方差。
- 为了平滑噪声并捕捉上下文中的连续区域，ASL 对注意力分数进行 1D 平均池化（类似 SnapKV）。
计算相对方差：
- 定义 $vartop(L) $为当前层及之前$ L_{obs}$ 层中 Top-K Token 的秩方差。
- 计算相对方差： $relative\_variance(L) = \frac{vartop(L)}{vartop(L_{min})}$ 。
自适应决策：
- 设定一个用户指定的阈值 $\tau$ （默认 0.3）。
- 当相对方差下降到 $\tau$ 以下时，意味着注意力分数已经稳定地聚焦在固定的 Token 子集上，此时判定当前层为最佳选择层（Selection Layer）。
- 一旦确定选择层，立即执行一次性的 Token 选择（One-shot selection），并将选中的 Token 传播到后续所有层。
与现有方法集成：
- ASL 主要在预填充（Prefilling）阶段运行，开销极小。
- 它可以与现有的 KV 压缩方法（如 SnapKV, GemFilter）无缝结合，以满足特定的 KV 预算要求。例如，结合 SnapKV 优化解码阶段，或结合 GemFilter 实现两遍策略。

3. 主要贡献 (Key Contributions)

问题识别：揭示了现有层级剪枝方法中“固定选择层”的僵化性，证明了任务难度（语义相似度）显著影响最佳剪枝层的选择。
提出 ASL：设计了一种基于注意力秩方差监控的自适应方法，无需微调即可根据任务难度动态调整剪枝时机。
灵活集成：展示了 ASL 如何与现有的 KV 缓存缩减方法（SnapKV, GemFilter）结合，同时优化预填充和解码阶段，并满足用户指定的 KV 预算。
性能验证：在 InfiniteBench、RULER 和 NIAH 三个基准测试（最长 256k 上下文）上，证明了 ASL 在困难任务上显著优于最先进（SOTA）的层级剪枝方法，同时保持了可比的解码速度和显存节省。

4. 实验结果 (Results)

实验使用了 Llama-3.1-8B-UL 和 Qwen2.5-7B 模型，在 InfiniteBench、RULER 和 NIAH 数据集上进行评估。

准确率提升（Accuracy）：
- 困难任务：在 KV 检索（KV Retrieval）和 NIAH（Haystack 中的针）等困难任务上，ASL 的表现显著优于 FastKV 和 GemFilter。例如，在 RULER 的 128k 上下文设置下，ASL 的准确率明显高于固定层剪枝方法。
- 全量 KV 对比：当在剪枝前使用全量 KV 时，ASL 在困难任务上的表现接近全量 KV 水平，而 FastKV 则大幅下降。
- 平均表现：在 InfiniteBench 上，ASL（特别是 ASL_2pass 版本）在 Llama-3.1-8B-UL 上取得了最高的平均分。
效率权衡（Efficiency Trade-off）：
- TTFT (首字延迟)：由于 ASL 在困难任务上会选择更深的层进行剪枝，其 TTFT 略高于 FastKV（因为需要计算更多层的全注意力），但远低于未剪枝的全量 KV。
- TPOT (每字延迟) 与吞吐量：一旦进入解码阶段，ASL 与其他剪枝方法的解码速度（TPOT）基本一致，因为后续层处理的 Token 数量相同。
- 显存占用：ASL 引入的额外显存开销（用于存储秩缓存）微乎其微，与 FastKV 等方法的显存占用几乎相同。
参数敏感性：
- 相对方差阈值 $\tau$ 对性能有影响。实验表明 $\tau=0.3$ 在准确率和 TTFT 之间取得了较好的平衡。

5. 意义与影响 (Significance)

打破“一刀切”的剪枝策略：ASL 证明了在 LLM 推理中，根据任务动态调整计算策略是可行的且必要的。它解决了固定层剪枝在复杂检索任务中失效的痛点。
无需训练的即插即用：作为一种无需微调（Training-free）的方法，ASL 可以立即应用于现有的开源或闭源 LLM，降低了部署门槛。
精度与效率的更好平衡：对于需要处理超长上下文且对精度要求较高的场景（如法律文档检索、复杂代码分析），ASL 提供了一种比固定层剪枝更优的解决方案，能够在不牺牲太多推理速度的前提下，显著提升长文本理解的准确性。
未来方向：论文指出未来可以探索将 ASL 集成到多轮剪枝（Progressive）或分组剪枝（Grouped）的方法中，以进一步挖掘潜力。

总结：ASL 通过观察注意力机制的收敛特性，动态决定何时“放手”剪枝，从而在保持 LLM 长上下文推理能力的同时，实现了更智能、更高效的显存管理。