VTrans: Accelerating Transformer Compression with Variational Information Bottleneck based Pruning
作者: Oshin Dutta, Ritvik Gupta, Sumeet Agarwal
分类: cs.LG
发布日期: 2024-06-07 (更新: 2024-06-11)
💡 一句话要点
VTrans：基于变分信息瓶颈的Transformer剪枝加速压缩
🎯 匹配领域: 支柱五：交互与反应 (Interaction & Reaction)
关键词: Transformer剪枝 变分信息瓶颈 模型压缩 自然语言处理 资源受限设备
📋 核心要点
- 现有Transformer剪枝方法通常忽略嵌入层，导致模型参数冗余，且压缩过程耗时。
- VTrans利用变分信息瓶颈（VIB）指导迭代剪枝，压缩包括嵌入层在内的所有结构组件，保留关键权重。
- 实验表明，VTrans在BERT、RoBERTa、GPT-2和LLaMA-2-7B上实现了更高的压缩率和更快的压缩速度，性能损失小。
📝 摘要（中文）
近年来，压缩大型预训练Transformer模型以适应资源受限设备的需求日益增长。然而，传统的剪枝方法通常保持嵌入层不变，导致模型过度参数化。此外，为了保持剪枝模型的性能，它们需要使用大型数据集进行大量的压缩时间。为了解决这些挑战，我们提出了VTrans，一个由变分信息瓶颈（VIB）原则指导的迭代剪枝框架。我们的方法压缩所有结构组件，包括嵌入、注意力头和层，使用VIB训练的掩码。这种方法仅保留每层中的必要权重，确保符合指定的模型大小或计算约束。值得注意的是，我们的方法实现了比现有最先进方法高达70%的压缩率，包括任务无关和任务特定的方法。我们进一步提出了我们方法的更快变体：Fast-VTrans仅利用3%的数据，以及Faster-VTrans，一种时间高效的替代方案，涉及VIB掩码的独占微调，与以前的方法相比，将压缩速度提高了高达25倍，且性能损失最小。在BERT、ROBERTa和GPT-2模型上的大量实验证实了我们方法的有效性。此外，我们的方法展示了压缩大型模型（如LLaMA-2-7B）的可扩展性，与以前的剪枝方法相比，实现了卓越的性能。此外，我们使用基于注意力的探测来定性地评估模型冗余并解释我们方法的效率。值得注意的是，我们的方法将未剪枝模型中对特殊和当前token具有高注意力的头作为剪枝的首要候选对象，而保留的头被观察到更多地关注于任务关键的关键词。
🔬 方法详解
问题定义：现有Transformer模型压缩方法，特别是剪枝方法，通常忽略嵌入层的压缩，导致模型仍然存在大量的参数冗余。此外，为了在剪枝后保持模型性能，需要使用大量数据进行耗时的微调，压缩效率较低。
核心思路：VTrans的核心思路是利用变分信息瓶颈（VIB）原则来指导Transformer模型的剪枝过程。VIB旨在找到输入数据的信息瓶颈，只保留对预测目标最有用的信息，从而实现模型压缩。通过VIB，VTrans可以自适应地确定哪些权重、注意力头和层对于保持模型性能至关重要，并剪掉冗余的部分。
技术框架：VTrans是一个迭代剪枝框架，主要包含以下几个阶段： 1. VIB掩码训练：使用VIB原则训练掩码，用于指示哪些权重、注意力头和层应该被保留或剪掉。 2. 模型剪枝：根据训练好的掩码，对Transformer模型进行剪枝，移除冗余的权重、注意力头和层。 3. 模型微调：对剪枝后的模型进行微调，以恢复因剪枝而损失的性能。 4. 迭代：重复以上步骤，直到达到目标模型大小或计算约束。
关键创新：VTrans的关键创新在于： 1. 统一的剪枝框架：VTrans可以同时剪枝Transformer模型的嵌入层、注意力头和层，实现更彻底的压缩。 2. VIB指导的剪枝：利用VIB原则自适应地确定哪些权重是重要的，避免了手动选择或启发式规则的局限性。 3. 快速压缩变体：提出了Fast-VTrans和Faster-VTrans，通过减少数据使用量或仅微调VIB掩码，显著加速了压缩过程。
关键设计：VTrans的关键设计包括： 1. VIB损失函数：使用VIB损失函数来训练掩码，该损失函数平衡了模型复杂度和预测精度。 2. 掩码参数化：使用可学习的参数来表示掩码，使得掩码可以根据数据进行自适应调整。 3. 迭代剪枝策略：采用迭代剪枝策略，逐步减小模型大小，避免一次性剪枝导致性能大幅下降。
📊 实验亮点
VTrans在BERT、RoBERTa和GPT-2等模型上实现了比现有方法高达70%的压缩率。Fast-VTrans仅使用3%的数据即可达到与全量数据训练相近的性能。Faster-VTrans通过仅微调VIB掩码，将压缩速度提高了高达25倍，且性能损失最小。在LLaMA-2-7B上的实验表明，VTrans在压缩大型模型方面具有良好的可扩展性，并优于其他剪枝方法。
🎯 应用场景
VTrans适用于各种需要部署大型Transformer模型的资源受限场景，例如移动设备、嵌入式系统和边缘计算。它可以显著减小模型大小和计算复杂度，降低部署成本，提高推理速度，并支持在算力有限的设备上运行复杂的自然语言处理任务。该技术还有助于降低云端服务器的计算负载，实现更高效的云计算。
📄 摘要（原文）
In recent years, there has been a growing emphasis on compressing large pre-trained transformer models for resource-constrained devices. However, traditional pruning methods often leave the embedding layer untouched, leading to model over-parameterization. Additionally, they require extensive compression time with large datasets to maintain performance in pruned models. To address these challenges, we propose VTrans, an iterative pruning framework guided by the Variational Information Bottleneck (VIB) principle. Our method compresses all structural components, including embeddings, attention heads, and layers using VIB-trained masks. This approach retains only essential weights in each layer, ensuring compliance with specified model size or computational constraints. Notably, our method achieves upto 70% more compression than prior state-of-the-art approaches, both task-agnostic and task-specific. We further propose faster variants of our method: Fast-VTrans utilizing only 3% of the data and Faster-VTrans, a time efficient alternative that involves exclusive finetuning of VIB masks, accelerating compression by upto 25 times with minimal performance loss compared to previous methods. Extensive experiments on BERT, ROBERTa, and GPT-2 models substantiate the efficacy of our method. Moreover, our method demonstrates scalability in compressing large models such as LLaMA-2-7B, achieving superior performance compared to previous pruning methods. Additionally, we use attention-based probing to qualitatively assess model redundancy and interpret the efficiency of our approach. Notably, our method considers heads with high attention to special and current tokens in un-pruned model as foremost candidates for pruning while retained heads are observed to attend more to task-critical keywords.