[论文解读] ZipLM: Inference-Aware Structured Pruning of Language Models

Eldar Kurtic, Elias Frantar+1 more|arXiv (Cornell University)|Feb 7, 2023

Topic Modeling被引用 7

一句话总结

ZipLM 是一种新颖的推理感知结构化剪枝框架，通过迭代地根据损失-推理时间权衡移除模型组件，在 BERT 和 GPT 模型的单次剪枝与渐进式剪枝设置下，实现了最先进的准确率-加速比权衡。该方法可在任何环境中保证目标推理加速，并在单次运行中以极低的计算成本生成一整套压缩后的模型。

ABSTRACT

The breakthrough performance of large language models (LLMs) comes with major computational footprints and high deployment costs. In this paper, we progress towards resolving this problem by proposing a novel structured compression approach for LLMs, called ZipLM. ZipLM achieves state-of-the-art accuracy-vs-speedup, while matching a set of desired target runtime speedups in any given inference environment. Specifically, given a model, a dataset, an inference environment, as well as a set of speedup targets, ZipLM iteratively identifies and removes components with the worst loss-runtime trade-off. Unlike prior methods that specialize in either the post-training/one-shot or the gradual compression setting, and only for specific families of models such as BERT (encoder) or GPT (decoder), ZipLM produces state-of-the-art compressed models across all these settings. Furthermore, ZipLM achieves superior results for a fraction of the computational cost relative to prior distillation and pruning techniques, making it a cost-effective approach for generating an entire family of smaller, faster, and highly accurate models, guaranteed to meet the desired inference specifications. In particular, ZipLM outperforms all prior BERT-base distillation and pruning techniques, such as CoFi, MiniLM, and TinyBERT. Moreover, it matches the performance of the heavily optimized MobileBERT model, obtained via extensive architecture search, by simply pruning the baseline BERT-large model. When compressing GPT2, ZipLM outperforms DistilGPT2 while being 60% smaller and 30% faster. Our code is available at: https://github.com/IST-DASLab/ZipLM.

研究动机与目标

• 为解决大语言模型（LLMs）在资源受限环境中的高计算成本和部署障碍。
• 克服现有结构化剪枝和蒸馏方法的局限性，这些方法通常无法保证运行时加速，或需要大量微调和超参数调优。
• 将单次剪枝与渐进式剪枝统一为一个框架，无需修改即可适用于编码器（BERT）和解码器（GPT）架构。
• 提供一种实用、高效且推理感知的方法，通过一次优化运行生成满足特定目标加速比的完整压缩模型家族。

提出的方法

• ZipLM 使用一种新颖的显著性准则，结合权重大小、激活影响以及线性冗余检测，同时考虑层间局部与全局相关性。
• 它引入了一种推理感知剪枝策略，评估每个组件对模型准确率和运行时的影响，优先移除能实现最佳损失-运行时间权衡的组件。
• 该方法采用逐层的 token 级蒸馏技术，在不需手动匹配层的情况下提升小数据集上的准确率。
• 它通过统一的一组超参数，在单次运行中实现端到端的多目标加速压缩（例如 2x、5x、10x），确保一致的性能和速度保障。
• ZipLM 兼容非结构化剪枝与量化，可在基于 CPU 的推理环境中实现进一步压缩和性能提升。

Figure 1 : Illustration of the ZipLM pipeline: 1) inference specifications, 2) runtime benchmarking of candidates for pruning, 3) gradual structured pruning until all speedup targets are met.

实验结果

研究问题

• RQ1 结构化剪枝是否能在不修改架构的前提下，在单次剪枝与渐进式剪枝设置下均实现最先进的准确率-加速比权衡？
• RQ2 推理感知剪枝策略是否能保证在不同硬件环境中实现目标运行时加速，而无需重复微调？
• RQ3 与基于蒸馏的方法及先前剪枝技术相比，该方法在 BERT 和 GPT 模型上的效率和准确率表现如何？
• RQ4 在极端压缩比（例如 50 倍）下，结构化剪枝的可扩展性如何，是否会导致模型崩溃？

主要发现

• 在 BERT-base 上，ZipLM 超越了所有先前的蒸馏与剪枝方法，包括 CoFi、MiniLM 和 TinyBERT，并在不进行架构搜索的情况下，实现了与 MobileBERT 相当的性能。
• 在 GPT2 上，ZipLM 生成的模型比 DistilGPT2 小 60%，快 30%，同时性能更优。
• 在 BERT-base 上，ZipLM 仅使用 4,096 个校准样本，即在 1.5 倍加速下达到 87.6 的 F1 分数，在 2.0 倍加速下达到 84.7 的 F1 分数，优于 Kwon 等人在 2,048 个样本下的表现。
• 该方法仅需 115 次总训练轮次，即可在 14 个加速目标（2x–15x）下完成 BERT-base 的压缩，效率比 CoFi 高 4.87 倍（CoFi 需要 560 次轮次）。
• 在 CPU 上，ZipLM 实现了 13 倍加速并完全恢复准确率，最大压缩下达到 50 倍加速，相比先前的复合流水线方法性能高出 3–10 倍。
• 缩放定律表明，剪枝后的模型遵循线性准确率-加速关系：BERT-large 的 F1 ≈ 92.1 − 0.3×加速比，BERT-base 的 F1 ≈ 90.3 − 0.6×加速比，表明大模型具有更高的鲁棒性。