ZipLM: Inference-Aware Structured Pruning of Language Models

type

status

date

slug

summary

1 简介

1.1 背景

模型压缩技术可以分为剪枝、量化、蒸馏。本文关注结构化压缩，目标是通过删除一些子组件实现减少模型大小，比如在权重矩阵中删除一些行或列。

相较于对单个权重的非结构化剪枝，结构化剪枝的优势在于它重构了模型维度形状，使得无需特殊硬件支持即可在任何硬件上减少计算。

但结构化剪枝也带来了重要挑战：

模型对于结构化压缩高敏感，大多数方法需要逐渐压缩，包括重训练设计，来恢复模型精度

结构化压缩比只知识蒸馏更复杂

大多数现有技术无法提供加速保证：模型被剪枝到一个固定稀疏度或 FLOPS 目标，然后在目标推理环境中进行评估，如果没有达到目标环境规格，整个过程需要从头再。

1.2 ZipLM

实现最有性能，无需重训练

通过推理感知算法实现，在每一个剪枝步骤中成功平衡损失-运行时间

考虑了运行时间，避免删除组件操作带来过大开销

提供了加速保证

1.3 贡献

新型结构化剪枝方法，该方法统一了过去工作中研究的重点部分，包括：权重大小、激活值影响、移除线性冗余结构。同时也考虑了局部（layer wise）和全局的相关性，称之为”推理感知“，确保在任意给定配置下都能获得期望的延迟或吞吐量

该方法通过一种新颖的逐层的 token 级别的蒸馏实现，能够在小数据集上持续提高准确率，同时不需要手动进行层级匹配，规避了之前结构化剪枝技术的局限性

第一个在后训练/单次压缩和渐进剪枝环境中都能达到最优结果的方法，且适用于 encoder 和 decoder 语言模型，无需任何修改

实用 & 高效

2 技术路线

ZipLM算法通过在权重矩阵的行上分解目标函数，并独立解决每行的稀疏线性回归问题，以找到要从网络中删除的最佳权重和结构，并补偿删除引起的误差。该算法利用显著性得分和基于Hessian矩阵的权重更新。它一次删除一个结构，并相对于剩余结构重新计算逆Hessian。该算法旨在处理结构化剪枝，其中考虑了多个块内部和之间的相关性。它在计算效率上具有较低的复杂性，并且与现有方法相比具有较低的复杂性。该算法还扩展到具有推理感知性的情况，考虑了批处理大小、序列长度和目标硬件上的加速度。该算法可以优化损失与加速度的权衡，并允许不同的现实世界指标优化。该文本进一步描述了通过延迟表集成运行时约束和使用结构化剪枝搜索的SPDY方法。

3 实验

ZipLM算法在单次运行中为一组所需加速度生成整个压缩模型系列，并在相同一组超参数下实现了实用和高效的结果。它与非结构化剪枝和量化兼容，即使在基于CPU的环境中也能够实现最先进的结果。该算法被证明与稀疏的动态形式和权重/激活的较低位宽表示兼容，为像商用CPU这样的边缘部署环境提供了更高的压缩比。

ZipLM在计算成本的一小部分下实现了卓越的结果，使其成为生成更小、更快和高度准确模型的一种经济有效方法。它优于所有先前的BERTbase蒸馏和剪枝技术，并通过简单剪枝基线BERTlarge模型，与经过大幅优化的MobileBERT模型的性能相匹配。在压缩GPT2时，ZipLM优于DistilGPT2，同时体积减小60%，速度提升30%。该算法为压缩模型提供了加速度保证，确保在任何给定配置中实现所需的延迟或吞吐量。

4 方法详述

4.1 ZipLM 结构化剪枝算法（局部相关性）

大多数结构化剪枝都基于以下关于重要性的假设之一或之二：

更低的或者更平均的权重大小更容易被修剪

输入激活值较小的可以被移除且损失不大

近似其他结构的线性组合的结构大多数是冗余的

ZipLM 联合考虑了这三个方面

4.1.1 问题描述

从逐层应用结构化压缩的想法开始，我们希望每一层能够尽可能多的保留输出特征。这种设置在后训练量化和非结构化剪枝中很常见。

首先，给定一个小校准集，让其通过整个网络，可以获得每一层的”参考“输入和输出。

然后，对于每一层，给定校准输入和该层的原始权重，我们的目标是找到遵守压缩约束且使得输出最近似原始输出的压缩权重，误差通过平方误差矩阵测量。

如果假设输入和权重矩阵有着恰当的矩阵形式，问题公式化为：

这一目标可以沿着矩阵行进行分解，变成一组稀疏线性回归问题。如果将逐行的问题看作独立的，那么就变成了过去的工作。由于我们要进行结构化剪枝，这些问题变得依赖：我们期望在所有行中修剪相同的权重索引（一列）。因此，找到找到最优权重等于找到：

要删除的所需形状的最佳结构。假设其应用在所有行，对应的剪枝掩码为，被剪枝处的值为1，其他为0。

对剩余权重进行相应的更新，以最佳方式补偿由于删除引起的误差

4.1.2 显着性分数和权重更新

令为 L2 最小化问题的海森矩阵，该矩阵与权重无关。

将定义为第 i 行中掩码下的权重子集，并通过定义对应与掩码下的条目的逆海森矩阵的子矩阵。

然后，可以得到最优的掩码和权重更新：

通过扩展求解 L2 最小化问题的 Optimal Brain Surgeon 公式获得同时覆盖所有行的结果。

海森矩阵的逆的子选择被所有行共享。

此外，由于通常只考虑相同大小的非重叠集合，额外求逆的开销为。

由于掩码中的结构数量通常很小，例如注意力头通常由 64 列组成，求逆的总开销很低。

4.1.3

仅通过 $4.1.2 中的方程求得要剪枝的结构统一了权重大小和激活值影响，但仍忽略了结构之间的任何相关性。下面通过“单次剪枝”来解决这个问题：始终应用更新并对于剩余结构完全重新计算。

例如，如果存在两个冗余结构和，先删除并更新以补偿这个删除，此时不再容易被剪枝。如果没有这种一次删除一个结构，这两个结构都会被错误地删除。

新问题又出现了，执行这个策略需要在每一步对剩余结构对应的海森矩阵的逆进行开销的重计算，这一步很慢。

可以通过一步高斯消除直接逆向删除与相对应的行和列来避免开销过大，逐块应用直到覆盖整个结构：

只需要时间。

伪代码

利用这样一个事实：权重矩阵和海森矩阵的逆中的被剪枝的值不会影响任何后续计算，因此即使他们不全为零也可以忽略不计。但是最后必须通过与整体掩码相乘来明确地修剪他们。

4.1.4 被修剪的结构

对于 transformers，考虑三种类型：

丢弃注意力头

缩小全连接网络层的扩展中间维度

删除残差部分

我们通过在注意力块的输出矩阵中删除个连续列以及在全连接网络的第二个线性层中删除各个列来实现这一点。一旦这些列结构被清零，前面层中的相应行就饿可以安全地删除而不会改变任何输出。

重要的是，通过剪枝（例如 FC2 层中的列，而不是 FC1中的等效行），我们可以使用 ZipLM 剪枝器通过海森矩阵信息利用输入相关性。

4.2 推理感知的结构化剪枝（全局相关性）

4.2.1 动机

推理感知的结构化剪枝的主要好处在于，剪枝决策并不纯粹由显着性分数来指导，而是由与模型中每个组件的删除相关的损失与加速权衡来指导。

先前的方法仅关注修剪，直到达到特定的稀疏阈值，而不考虑与压缩阈值相对应的现实世界加速，这在设置之间可能有很大差异。例如，CoFi 生成的 95% 稀疏 BERT 在 V100 GPU 上加速了 12 倍，但在 A100 GPU 上只有 5 倍。

使用现有方法，如果现实世界的时序无法满足推理要求，则必须使用不同的稀疏值重复整个过程，直到实现目标加速，这既耗时又容易出错。

推理感知的另一个优势是它可以针对不同的现实世界指标（例如延迟或吞吐量）进行优化。

4.2.2 运行时间感知

通过延迟表为目标推理环境继承运行时约束，其中我们记录运行注意力块修剪 0~-1 个头的时间，和全连接层缩小中间维度（通过缩小因子）的时间，以 10% 的相对步长知道稀疏度到达了 ≈99%。这用来快速估计不同的每层稀疏配置的运行时间。

4.2.3 寻找最优稀疏配置

目标是找到一种每层稀疏配置，在最大化准确性的同时满足一定的加速约束。

一种流行范例是生成大量跨层具有不同稀疏度分布的修剪模型，然后选择一个满足目标约束且精度最高的模型。

为了使其在计算上可行，剪枝成本低廉且准确至关重要。

ZipLM 独立处理每一层，这使得可以预先计算每个层具有不同稀疏度的多个修剪版本的数据库。利用该算法的一次一次的性质，可以在一次运行中生成整个数据库。

虽然我们的算法与用于查找分层轮廓的各种搜索方法兼容，但我们采用了最近的 SPDY 方法。

4.2.3 结构化 SPDY 搜索

SPDY 方法专为非结构化剪枝而设计，并在不同稀疏度级别的每层敏感度之前分配二次方。这在我们的结构化修剪场景中是无效的，例如它表明删除完整层仅比将其修剪到 99% 稀疏度稍微困难一些。因此，使用标准 SPDY 会导致算法探索大量次优配置，严重浪费计算资源。

为了缓解这个问题，对于结构化稀疏性，我们引入了更好的先验作为剪枝引起的相对分层平方误差，定义为，对于完全丢弃的层，其值为 1。

此外，原始的 SPDY 方法使用收缩邻域搜索，这在结构化压缩的运行时间和解决方案质量方面具有很高的方差。因此，我们执行固定数量的 1000 个步骤，随机突变期望 10% 的分层敏感度系数。最后，我们注意到，通过此过程评估的任何候选者实际上都实现了目标加速，从而显着减少了搜索时间。

4.3 分层令牌蒸馏

对于结构化剪枝，通常应用分层蒸馏目标来传输中间表示。然而，结构化剪枝会产生相对于固定教师架构的兼容性问题，导致大多数方法开发定制的蒸馏策略。

一种流行方法通过将教师层子集静态或动态映射到学生层子集来解决该问题。它们的主要限制是手动层选择，其中做出“最佳”选择需要评估所有可能的组合，这可能非常昂贵。另一个限制是中间层之间的形状匹配，这是通过引入附加到学生输出的可学习线性变换矩阵来解决的。

我们通过利用 ZipLM 保留隐藏维度大小的事实以不同的方式解决这些挑战，并建议在整个模型中使用中间令牌表示的蒸馏。由此产生的最小化目标由三个组成部分组成：

其中和分别代表学生和教师模型，是输入，是与任务相关的损失（例如文本分类的交叉熵），是输出 logits 之间的 KL 散度，而是我们的 token 级别的蒸馏损失。在连续的 Transformer 层之间传递的隐藏张量具有恒定形状，其中代表批量大小，代表序列长度，代表模型架构定义的隐藏大小。该张量可以解释为向量的集合，每个向量携带输入标记的中间模型表示。损失定义为对应于输入序列中每个非填充标记的向量之间的欧几里德距离 Δ，在所有未剪枝层上取平均值。形式上，对于第 k 层，它被定义为:

其中 P 代表填充标记集。该公式鼓励学生模型为每个标记生成与教师模型生成的向量表示类似的向量表示。