You Only Prune Once: Designing Calibration-Free Model Compression With Policy Learning
作者: Ayan Sengupta, Siddhant Chaudhary, Tanmoy Chakraborty
分类: cs.CL
发布日期: 2025-01-25 (更新: 2025-02-28)
💡 一句话要点
提出PruneNet，通过策略学习实现免校准数据集的模型压缩
🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture) 支柱九：具身大模型 (Embodied Foundation Models)
关键词: 模型压缩 剪枝 策略学习 免校准 大型语言模型
📋 核心要点
- 现有模型剪枝方法依赖校准数据集，限制了灵活性和可扩展性，在高压缩率下性能下降。
- PruneNet将剪枝视为策略学习，无需校准数据，基于模型内在属性学习剪枝策略。
- 实验表明，PruneNet在压缩LLaMA-2-7B时，性能优于现有方法，并在多任务理解中表现出鲁棒性。
📝 摘要（中文）
大型语言模型（LLMs）日益增长的规模给部署带来了重大挑战，因为它们对计算和内存的需求很高。目前的模型剪枝技术试图通过严重依赖外部校准数据集来确定要剪枝或压缩的参数来缓解这些问题，从而限制了它们在不同压缩率下的灵活性和可扩展性。此外，这些方法通常会导致严重的性能下降，尤其是在下游任务中，当受到更高的压缩率时。在本文中，我们提出了一种新的模型压缩方法PruneNet，它通过将模型剪枝重新定义为策略学习过程来解决这些限制。PruneNet将剪枝过程与模型架构分离，消除了对校准数据集的需求。它学习一种随机剪枝策略，仅基于内在模型属性来评估参数重要性，同时保留频谱结构以最小化信息损失。PruneNet可以在短短15分钟内压缩LLaMA-2-7B模型，在30%的压缩率下，保持其超过80%的zero-shot性能，优于现有方法，后者仅保持75%的性能。此外，在复杂的多任务语言理解任务中，PruneNet通过保持原始模型高达80%的性能来展示其鲁棒性，证明了它是传统结构化压缩技术的卓越替代方案。
🔬 方法详解
问题定义：论文旨在解决大型语言模型（LLMs）因体积庞大而难以部署的问题。现有剪枝方法依赖外部校准数据集来确定剪枝参数，这限制了其在不同压缩率下的泛化能力，并且在高压缩率下容易导致性能显著下降。
核心思路：PruneNet的核心思路是将模型剪枝过程建模为一个策略学习问题。通过学习一个随机剪枝策略，模型可以根据自身内在属性（例如参数的频谱结构）来评估参数的重要性，从而避免对外部校准数据的依赖。这种方法旨在保留模型的重要信息，同时实现高效的压缩。
技术框架：PruneNet的整体框架包含以下几个主要步骤：1) 初始化模型参数；2) 使用策略网络学习一个随机剪枝策略，该策略基于模型参数的内在属性来评估参数的重要性；3) 根据学习到的策略对模型进行剪枝；4) 对剪枝后的模型进行微调，以恢复性能。该框架的关键在于策略网络的训练，它旨在最大化剪枝后模型的性能，同时保持模型的频谱结构。
关键创新：PruneNet的关键创新在于其免校准数据集的剪枝方法。与传统方法不同，PruneNet不依赖外部数据来指导剪枝过程，而是通过学习一个策略来自动评估参数的重要性。此外，PruneNet还考虑了模型的频谱结构，以最小化剪枝过程中的信息损失。
关键设计：PruneNet使用一个策略网络来学习剪枝策略。该策略网络以模型参数作为输入，输出一个概率分布，表示每个参数被剪枝的概率。损失函数的设计至关重要，它需要平衡剪枝后的模型性能和模型的频谱结构。论文中可能使用了强化学习或模仿学习等方法来训练策略网络。具体的网络结构和损失函数细节需要参考论文正文。
🖼️ 关键图片
📊 实验亮点
PruneNet在LLaMA-2-7B模型上取得了显著的压缩效果。在30%的压缩率下，PruneNet保留了原始模型超过80%的zero-shot性能，优于现有方法（仅保留75%）。在复杂的多任务语言理解任务中，PruneNet也表现出良好的鲁棒性，保持了原始模型高达80%的性能，证明了其优越性。
🎯 应用场景
PruneNet适用于各种需要模型压缩的场景，例如在资源受限的设备上部署大型语言模型，或者在带宽有限的网络环境中进行模型推理。该方法可以降低模型的计算和存储成本，提高模型的推理速度，从而扩展LLM的应用范围。未来，该方法可以进一步应用于其他类型的深度学习模型，并与其他压缩技术相结合，以实现更高的压缩率和更好的性能。
📄 摘要（原文）
The ever-increasing size of large language models (LLMs) presents significant challenges for deployment due to their heavy computational and memory requirements. Current model pruning techniques attempt to alleviate these issues by relying heavily on external calibration datasets to determine which parameters to prune or compress, thus limiting their flexibility and scalability across different compression ratios. Moreover, these methods often cause severe performance degradation, particularly in downstream tasks, when subjected to higher compression rates. In this paper, we propose PruneNet, a novel model compression method that addresses these limitations by reformulating model pruning as a policy learning process. PruneNet decouples the pruning process from the model architecture, eliminating the need for calibration datasets. It learns a stochastic pruning policy to assess parameter importance solely based on intrinsic model properties while preserving the spectral structure to minimize information loss. PruneNet can compress the LLaMA-2-7B model in just 15 minutes, achieving over 80% retention of its zero-shot performance with a 30% compression ratio, outperforming existing methods that retain only 75% performance. Furthermore, on complex multitask language understanding tasks, PruneNet demonstrates its robustness by preserving up to 80% performance of the original model, proving itself a superior alternative to conventional structured compression techniques.