VTrans：基于变分信息瓶颈的 Transformer 压缩方法
机器学习
2024-06-13 v2
摘要
近年来，随着对资源受限设备上大型预训练变换器模型压缩的日益关注，传统剪枝方法往往不处理嵌入层，导致模型过度参数化。此外，这些方法还需要大数据集以维持剪枝后模型的性能。为此，我们提出了 VTrans，一种基于变分信息瓶颈（VIB）原则引导的迭代剪枝框架。我们的方法使用 VIB 训练的掩码压缩所有结构组件，包括嵌入层、注意力头和层数。该方法确保仅保留每一层中必需的权重，从而满足指定的模型大小或计算限制。值得注意的是，我们的方法在任务无关和任务特定方面均实现了比以前最先进方法高达 70% 的压缩比。我们进一步提出了我们方法的快速变体：Fast-VTrans 仅使用 3% 的数据，Faster-VTrans 是一种高效替代方案，仅对 VIB 掩码进行微调，使压缩速度提升最高可达 25 倍，同时与以前的方法相比性能损失最小。对 BERT、ROBERTa 和 GPT-2 模型的大量实验证实了我们方法的有效性。此外，我们的方法在压缩大型模型（如 LLaMA-2-7B）方面表现出良好的可扩展性，性能优于以前的剪枝方法。我们还使用基于注意力的探测来定性评估模型冗余性和我们方法的效率。值得注意的是，我们的方法将注意力高度集中在特殊和当前记号的注意力头视为主要的剪枝候选人，而保留的注意力头则观察到更关注任务关键词。
引用
@article{arxiv.2406.05276,
title = {VTrans: Accelerating Transformer Compression with Variational Information Bottleneck based Pruning},
author = {Oshin Dutta and Ritvik Gupta and Sumeet Agarwal},
journal= {arXiv preprint arXiv:2406.05276},
year = {2024}
}