部署大型语言模型 (LLM) 的竞赛实际上是一场对抗两个约束条件的战斗: 计算成本和延迟。我们希望模型巨大、能力强且聪明 (如 Llama-2 或 OPT-175B) ,但也希望它们能在我们的设备上即时运行,而不需要一座核电站来供电。

解决这个问题的一个流行技术是稀疏化 (Sparsification) ——简单地关闭神经网络的一部分以节省工作量。然而,随机删除神经元会让模型变“笨”。行业标准已经转向*上下文稀疏性 (Contextual Sparsity) *,即模型根据当前的输入,动态决定自身的哪些部分是必要的。

在这篇文章中,我们将深入探讨 ShadowLLM , 这是一篇重新定义我们如何处理动态剪枝的研究论文。研究人员指出了当前最先进方法 (如 DejaVu) 的两个主要缺陷: 它们使用了低效的预测机制,且对“神经元重要性”的定义过于简单。ShadowLLM 提出了一种统一的预测器和基于梯度的标准,相比现有方法实现了 15% 的准确率提升和 20% 的速度提升。

体量巨大的问题

LLM 由数十亿个参数组成。对于生成的每一个 token,模型通常会激活所有这些参数。这是低效的,因为对于任何特定任务——比如回答关于历史的问题——模型并不需要它那些与编写 Python 代码或创作诗歌相关的神经元。

静态稀疏性与上下文稀疏性

有两种方法可以为这些模型“瘦身”:

  1. 静态稀疏性 (Static Sparsity) : 永久删除看起来不重要的注意力头 (Attention Heads) 或神经元。虽然速度快,但这通常会降低模型执行“上下文学习 (In-context Learning) ” (根据提示适应新任务) 的能力。
  2. 上下文稀疏性 (Contextual Sparsity) : 保留所有权重,但对于每个特定的输入 (上下文) ,动态预测哪些神经元可以被忽略。

为了保持准确性,上下文稀疏性是更优的方法。如下图所示,不同的输入 (“Hamburger Cheese Needs…” 与 “A Quantitative Approach…”) 在网络中激活的路径完全不同。

上下文稀疏性图解,展示了不同输入对应不同路径。 图 2: 上下文稀疏性根据上下文 (输入) 本身来修剪神经元和注意力头。训练一个预测器来根据输入 token 动态预测稀疏模式,可以提高模型质量。

挑战在于“上下文”。如果注意力头的重要性随输入发生剧烈变化,我们需要一种机制来准确且快速地预测这种变化。

研究人员分析了 OPT-1.3B 模型中注意力头的排名方差 (Rank Variance) 。 如果一个头总是很重要,它的排名方差就会很低。然而,数据显示出巨大的方差,尤其是在早期和晚期层中。这证明了创建一个静态的“好神经元”列表是不可能的;模型必须动态适应。

散点图显示注意力头之间的高排名方差。 图 3: 具有较高排名方差的注意力头表明其对上下文的依赖性更强。这种方差在模型的早期和晚期层中最为明显,因此必须要进行动态修剪。

ShadowLLM 方法

ShadowLLM 改进了之前的最先进框架 DejaVu 。 虽然 DejaVu 成功实现了上下文稀疏性,但它遭遇了两个瓶颈:

  1. 逐层预测: 它在每一层都使用一个预测器来猜测下一层的稀疏性。这种“走走停停”的交通方式造成了延迟开销。
  2. 基于幅度的标准: 它假设如果激活值很大 (高幅度) ,该神经元就是重要的。正如我们将看到的,“响亮”的神经元并不总是聪明的。

ShadowLLM 引入了一种精简的架构和一种在数学上更优越的判断神经元价值的方法。

1. 统一预测器

ShadowLLM 不会在每一层都停下来问“我接下来该修剪什么?”,而是在最开始使用单一的预测器。

该方法获取第一个注意力层的输出,并将其输入到一个统一的预测器中。这个预测器会一次性预测所有后续层的稀疏模式。

对比 ShadowLLM 的单一预测器与 DejaVu 的多预测器的图解。 图 4: (1) 使用单个预测器对整个 LLM 进行建模可提高模型性能 (左图) ,而 (2) 在评估神经元修剪标准时利用基于梯度的信息可提高模型质量。

这种架构转变对性能有深远的影响。在 DejaVu 中,系统必须不断启动异步内核。在 ShadowLLM 中,预测的繁重工作是预先完成的。如下表所示,这使得各种模型尺寸下预测所需的浮点运算次数 (FLOPs) 减少了近 20%。

表格显示 ShadowLLM 对比 DejaVu 的 FLOPs 减少量。 表 1: 相比 DejaVu,ShadowLLM 显著降低了预测器的计算开销。

2. 全局与局部修剪

因为 ShadowLLM 一次性预测整个模型的状态,它解锁了全局修剪 (Global Pruning)

  • 局部修剪 (逐层) : 强制每一层都达到 (例如) 50% 的稀疏度。
  • 全局修剪: 强制整个模型达到 50% 的稀疏度。这意味着第 5 层可能只有 10% 的稀疏度 (非常重要的层) ,而第 10 层可能有 90% 的稀疏度 (不重要的层) 。

研究人员发现,全局修剪始终优于局部修剪,因为它尊重了并非所有层都生而平等这一事实。

对比全局与局部修剪性能的图表。 图 6: 全局修剪 (蓝线) 优于局部修剪策略 (橙线) 。全局修剪适应了不同层重要性的变化。

重新思考修剪标准: 超越幅度

这就论文最重要的贡献可能在于对“重要性”的数学分析。

在之前的工作中,重要性由 L2 范数 (L2Norm) (即幅度) 定义。如果一个神经元输出一个很大的数字,它就会被保留。ShadowLLM 主张我们应该关注 敏感度 (Sensitivity) : 如果我们移除这个神经元,损失函数会发生多大变化?

为了计算这一点,作者利用了梯度信息。他们探索了几个标准:

  • L2Norm (基准) : 激活幅度 (\(||A||_2\))。
  • GradNorm: 梯度的幅度 (\(||\frac{\partial \mathcal{L}}{\partial A}||_2\))。
  • PlainAct: 激活和梯度的组合 (\(||A \cdot \frac{\partial \mathcal{L}}{\partial A}||_1\))。这衡量了激活实际上对梯度的影响程度。
  • Fisher: 二阶近似。

胜出者: PlainAct

研究人员发现 PlainAct (基于梯度的激活) 是最佳选择。相比简单的幅度,它能更好地捕捉模型的真实敏感度。

有趣的是,一些标准如 GRASP (使用 Hessian 矩阵近似) 虽然理论上很强,但在实践中却失败了。为什么? 可预测性。

我们需要训练一个小型神经网络 (预测器) 在运行时猜测这些值。如果修剪标准噪声太大或有大量异常值 (如 GRASP) ,预测器就无法学习它。PlainAct 既是一个准确的重要性指标,又容易被神经网络学习。

柱状图显示预测器在不同标准上的表现。 图 14: PlainAct 是一个很好的修剪标准,并且易于学习 (高 Spearman 相关性) 。GRASP 有许多异常值,使得预测变得困难。

实验结果

当在 OPT 和 Llama-2 模型上测试时,统一预测器与 PlainAct 标准的结合产生了令人印象深刻的结果。

准确率与困惑度

当与 DejaVu 风格的预测器 (使用局部修剪和 L2Norm) 进行对比时,ShadowLLM 在相同的稀疏度水平下始终实现更低的困惑度 (Perplexity,即更好的语言建模能力) 。

对比 ShadowLLM 与 DejaVu 在困惑度上的表现。 图 7: DejaVu 风格预测器 (L2Norm) 与 ShadowLLM (PlainAct) 的比较。随着稀疏度增加,ShadowLLM 保持了更低的困惑度。

这种优势延伸到了零样本 (Zero-shot) 下游任务 (如推理或阅读理解) 。在七个不同的评估任务中,ShadowLLM 展示了更强的准确率-稀疏度权衡。

显示 7 个任务中准确率提升的图表。 图 8: 在七个下游评估任务中,ShadowLLM 相比 DejaVu 展现了一致的准确率提升。

延迟与速度

如果在这种情况下模型没有变快,那么准确率就毫无意义。通过消除逐层预测器的开销,ShadowLLM 实现了显著的加速。

对于 OPT-30B 模型,ShadowLLM 在所有生成长度下都比 DejaVu 更快。虽然静态修剪 (固定稀疏度) 自然是最快的,但它缺乏上下文方法的准确性。ShadowLLM 弥补了这一差距,提供了接近静态的速度和动态的智能。

显示生成时间与生成长度关系的图表。 图 11 (右) : OPT-30B 上的平均生成时间。ShadowLLM (橙色) 始终快于 DejaVu (蓝色) 。

当我们看大局时——将准确率与延迟绘制在一起——我们可以看到 ShadowLLM (蓝色星星) 在帕累托前沿 (Pareto frontier) 上占据主导地位。对于任何给定的延迟预算,ShadowLLM 都比 DejaVu 风格的方法提供更高的准确率。

准确率与延迟的散点图。 图 1: 相比 DejaVu,ShadowLLM 实现了更好的准确率与延迟权衡。

规模扩展

性能提升不仅限于小模型。随着模型尺寸的增加 (从 1.3B 到 66B 参数) ,ShadowLLM 和 DejaVu 之间的差距进一步拉大。在 OPT-66B 模型上,生成时间的减少是巨大的。

柱状图对比不同模型尺寸的生成时间。 图 10: 不同模型尺寸下的平均单次推理时间。ShadowLLM 相比 DejaVu 提供了持续的速度提升。

结论

ShadowLLM 代表了上下文稀疏性技术的成熟。作者证明,要有效地优化大型语言模型,我们必须关注两个不同的优化层面:

  1. 预测架构: 从细粒度的逐层预测器转向整体的、早期阶段的预测器,这减少了延迟并实现了全局修剪策略。
  2. 重要性的定义: 从简单的基于幅度的启发式方法转向基于梯度的指标 (PlainAct) ,确保我们保留的是真正影响模型输出的神经元。

通过结合这些见解,ShadowLLM 相比现有技术实现了 15% 的准确率提升和 20% 的延迟降低。随着 LLM 规模持续增长,像 ShadowLLM 这样允许我们智能地“跳过”不必要计算部分的技术,对于高效的实际部署将至关重要。