引言

我们正处于大语言模型 (LLMs) 的“黄金时代”。从 LLaMA 到 GPT-4，这些模型展现出的推理、编程和创意写作能力在十年前是无法想象的。然而，这种智能伴随着巨大的代价: 计算成本。

对 GPT-175B 这样的模型进行推理需要数百 GB 的 GPU 显存和大规模的并行处理能力。主要的瓶颈在于自回归解码 (autoregressive decoding) 过程。因为 LLM 是逐个 token 生成文本的——第 \(n\) 个 token 的生成依赖于之前所有的 \(n-1\) 个 token——所以计算负载随着序列长度线性增加。每一个 token 的生成都需要完整地通过数十亿个参数。

为了解决这个问题，研究人员长期以来一直提出“早退 (early-exit) ”或“层跳过 (layer-skipping) ”策略。逻辑很简单: 并非每个 token 都需要动用模型的全部脑力。如果模型只是预测单词 “the” 或一个标点符号，它当然可以跳过几层，对吧？

虽然这听起来很直观，但它引入了一个关键问题: 键值 (Key-Value, KV) 缓存管理不善 。 当你为了节省时间而跳过层时，你破坏了 Transformer 的记忆机制，通常会导致幻觉或“Token 崩溃 (token collapse) ” (即模型开始无休止地重复自己) 。

今天，我们将深入探讨一篇提出巧妙解决方案的研究论文。作者没有跳过整个层，而是引入了 FFN-SkipLLM 。 这种方法策略性地只跳过前馈网络 (FFN) 模块——这是层中计算成本最高的部分——同时保持注意力 (Attention) 机制的完整性。结果是模型速度更快，保持了智能，并且至关重要的是，避免了困扰其他压缩方法的幻觉问题。

传统层跳过的问题

要理解为什么需要 FFN-SkipLLM，我们首先需要了解之前的尝试为何失败。

在标准的 Transformer 架构中，每一层处理信息并将其传递给下一层。为了加速推理，像 SkipDecode 或 ShortGPT 这样的方法尝试跳过整个层。如果一个 token 很容易预测，这些方法可能会在第 10 层而不是第 32 层退出网络。

然而，LLM 依赖 KV 缓存来存储先前 token 的上下文。如果当前 token 跳过了第 15 层，它就不会为该层生成 Key 或 Value 状态。当下一个 token 试图关注第 15 层的上一个 token 时，它会在记忆中发现一个“空洞”。

之前的工作试图通过复制其他层的状态或重新计算缺失值来修补这个问题，但这些都是权宜之计。正如这篇论文的作者所发现的，这些方法通常会导致知识密集型任务的性能严重下降。

AI 回答的对比。SkipDecode 和 ShortGPT 未能识别出印度总理，而 FFN-SkipLLM 成功了。

如图 1 所示，当被问及“谁是印度总理？”时，传统的跳过方法 (SkipDecode 和 ShortGPT) 惨遭失败。它们要么产生幻觉说该职位已被废除，要么退化成不知所云的乱语。相比之下, FFN-SkipLLM (绿色气泡) 尽管跳过了大约 25% 的计算，仍能检索到正确的信息。

背景: Transformer 层的解剖

为什么 FFN-SkipLLM 能成功而其他方法却失败了？答案在于 Transformer 层本身的架构。

单个 Transformer 层由两个主要模块组成:

  1. 多头注意力 (Multi-Head Attention, MHA) : 这允许 token 关注序列中的其他 token。这也是 KV 缓存所在的地方。
  2. 前馈网络 (Feed-Forward Network, FFN) : 这处理由注意力机制收集的信息。它通常被描述为网络的“记忆”或“知识”存储。

至关重要的是，这两个模块在大小上并不相等。

表格显示参数计数。FFN 权重明显大于注意力权重。

如表 1 以 LLaMa-7B 为例所示，前馈网络占据了该层大约三分之二的参数。相比之下，注意力机制相对轻量级。

这引出了论文的核心假设: 如果我们只跳过沉重的 FFN 模块，但保持轻量级的注意力模块运行，我们就可以在不破坏 KV 缓存的情况下节省大量计算资源。

核心方法: FFN-SkipLLM

研究人员提出了一种细粒度的策略: 输入自适应前馈跳过 (Input-Adaptive Feed-Forward Skipping) 。 与其丢弃整个层，他们执行注意力模块 (正确更新 KV 缓存) ，但如果有条件地认为 FFN 模块是冗余的，则跳过它。

1. FFN 模块的冗余性

我们如何知道一个 FFN 模块是否冗余？作者分析了进入 FFN 模块的输入张量与离开它的输出张量之间的余弦相似度 (cosine similarity) 。

如果余弦相似度很高 (接近 1.0) ，这意味着 FFN 模块并没有怎么改变表示——它实际上什么也没做。如果相似度很低，说明 FFN 模块进行了显著的转换。

各层的余弦相似度图表。中间层显示出高相似度，表明冗余。

图 2 揭示了关于 LLM 行为的三个关键见解:

  1. 高相似度: FFN 的输入和输出之间通常存在很高的相似度，表明存在固有的冗余。
  2. “冷”区域 (The “Cold” Regions) : 最初的几层和最后的几层 (标记为红色) 具有较低的相似度。这些是“冷区域”。这里的 FFN 正在做繁重的工作, 绝不应该被跳过。
  3. 冗余的中间部分: 中间层 (黄色区域) 显示出单调增加的相似度。这是跳过的最佳区域。

2. 算法

基于这些观察，FFN-SkipLLM 在 token 生成期间采用了一种简单但有效的算法:

  1. 预热阶段 (Warm-Up Phase) : 由于“注意力汇聚 (Attention Sink) ”现象 (即最初的几个 token 充当注意力机制的锚点) ，序列的前几个 token (例如前 25-30 个 token) 由完整模型处理。不允许跳过。这稳定了生成过程。
  2. 识别区域: 模型识别“冷启动” (前几层) 和“冷结束” (后几层) 。这些层中的 FFN 总是被执行。
  3. 自适应跳过: 对于中间的层，模型计算 FFN 输入与简化传递输出之间的余弦相似度。如果相似度超过特定阈值，模型预测该 FFN 是冗余的。
  4. 贪婪跳过 (Greedy Skip) : 由于冗余往往在中间层单调增加，如果模型决定跳过一个 FFN，它可以贪婪地跳过接下来的 \(k\) 个 FFN 模块，进一步节省计算资源。

这种方法的妙处: 因为注意力模块从不被跳过，KV 缓存始终得到完美维护。模型始终拥有完整的过去“记忆”，即使它跳过了某些步骤的“处理” (FFN) 。

实验与结果

作者对 FFN-SkipLLM 进行了严格的“知识密集型”评估。与仅检查流畅性的标准基准测试不同，这些测试确定模型是否真的记住了事实并遵循指令。

1. 基于事实的问答

此任务测试模型回忆特定实体和属性的能力 (例如，“谁写了《1984》？”) 。

表格比较 Factoid-QA 上的性能。FFN-SkipLLM 即使在跳过的情况下仍保持高准确率。

表 2 将 FFN-SkipLLM 与完整模型及其他跳过方法进行了比较。

  • 完整模型 (Full Model) : 79.02% 准确率。
  • SkipDecode: 降至 73.33%。
  • ShortGPT: 降至 70.49%。
  • FFN-SkipLLM: 保持 78.89% 的准确率。

即使跳过了约 25% 的 FFN 模块，性能下降也可以忽略不计 (不到 0.2%) 。这证实了跳过的 FFN 模块对于回忆事实确实是冗余的。

2. 上下文摘要

摘要任务测试模型处理长上下文和生成连贯文本的能力。研究人员使用 GPT-4 来评判压缩模型生成的摘要质量。

条形图显示摘要性能。FFN-SkipLLM 在连贯性和流畅性方面匹配或超过基线。

图 3 显示了在连贯性、一致性、流畅性和相关性方面的结果。

  • 红线 (Ours) 始终徘徊在顶部附近，即使被剪枝的 FFN 百分比增加 (x 轴向右移动) 。
  • 有趣的是，在约 10-12% 的跳过率下，FFN-SkipLLM 有时会获得比完整模型更好的分数。作者推测，去除冗余处理可能会减少信号中的“过度思考”或噪声。

3. 多轮对话

现实世界的使用通常涉及来回对话。作者使用 MT-Bench 数据集评估了编程、数学和角色扮演等类别的性能。

图表显示 8 个类别的性能。FFN-SkipLLM 在编程和数学方面表现稳健。

图 4 突显了该方法的稳健性。

  • 在像编程 (Coding) 和数学 (Fermi) 这样的复杂类别中，FFN-SkipLLM (红线) 在高达 25% 的跳过率下仍非常接近完整模型 (蓝线) 。
  • 随机跳过基线 (橙线) 几乎立即崩溃，证明了你不能随意跳过 FFN——你必须针对冗余的那些。

4. 效率与加速

这项研究的最终目标是速度。跳过 FFN 真的能转化为现实世界的收益吗？

表格显示 FLOPs 减少和吞吐量。吞吐量随着 FLOPs 减少而增加。

表 4 显示了 FLOPs (浮点运算次数) 的减少和吞吐量 (每秒 token 数) 的增加。

  • 在 50% 的跳过率下，该模型达到了 12.45 tokens/second 的吞吐量，而密集模型为 9.08。
  • FLOPs 的减少是巨大的，在 50% 的场景中每个 token 节省了高达约 43 亿次运算。

虽然计算余弦相似度有一点小的开销，但随着跳过率的增加，跳过巨大的 FFN 模块所带来的节省远远超过了成本。

结论与启示

FFN-SkipLLM 代表了我们思考模型压缩方式的转变。与其将 Transformer 层视为一个必须整体保留或丢弃的原子单元，这项工作鼓励我们深入层的内部。

通过认识到:

  1. FFN 很重 (占 2/3 的计算量) ，
  2. FFN 在中间层通常是冗余的 , 以及
  3. 注意力对于记忆 (KV 缓存) 至关重要 ,

作者设计了一种方法，规避了之前“层丢弃”技术带来的灾难性失败。

这块“隐藏的宝石”表明，未来的 LLM 架构在设计时可能会考虑到模块化，也许只有在输入复杂性需要时才动态分配 FFN 容量。对于学生和研究人员来说，FFN-SkipLLM 是一个完美的例子，说明了分析网络的基本属性 (如余弦相似度和参数分布) 如何能够带来简单、优雅且高效的优化策略。