一文详解视觉Transformer模型压缩和加速策略
一文详解视觉Transformer模型压缩和加速策略
视觉Transformer（ViT）在计算机视觉领域实现了革命性突破，但其高计算和内存需求限制了实际应用。本研究系统评估了量化、低秩近似、知识蒸馏和剪枝四种主要的模型压缩技术，以优化ViT在资源受限环境中的性能。
1. 引言
Transformer架构由Vaswani等人引入，以其强大的数据建模能力、可伸缩性以及出色的捕捉长距离依赖性的能力，彻底改变了机器学习领域。最初在自然语言处理（NLP）领域取得突破后，Transformer开始应用于计算机视觉任务，如图像分类和目标检测，并在多模态情感分析和医疗保健领域展现出巨大潜力。
然而，ViT的高性能伴随着庞大的参数集，数量达到数亿个，这在推理过程中产生显著的内存和计算开销。因此，开发和改进ViTs的模型压缩技术变得至关重要，尤其是在工业应用中。
2. 相关工作
2.1 量化
量化是促进神经网络高效推理的关键技术。通过将网络转换为低比特表示，可以减少计算需求和内存占用。量化感知训练（QAT）通过在每次梯度更新后重新量化模型参数，平衡了精度保留和量化效益。
2.2 低秩近似
ViT中的自注意力机制具有二次计算复杂度，Chen等人发现其注意力矩阵具有低秩特性。基于此，研究者开发了Nyström、Performer和Linformer等方法，并探索了与稀疏注意力机制的结合。
2.3 知识蒸馏
知识蒸馏通过教师模型的软标签训练紧凑的学生模型。Touvron等人提出的蒸馏token进一步优化了这一过程，通过自注意力机制参与优化蒸馏过程。
2.4 剪枝
剪枝通过减少模型维度复杂性来简化架构。研究表明，经过剪枝的模型有时甚至能超过原始模型的性能。
3. 方法论
3.1 量化
- 后训练量化（PTQ）：直接调整权重，但可能导致显著的精度下降。
- 量化感知训练（QAT）：通过在浮点表示的模型上执行前向和后向过程，并在每次梯度更新后重新量化参数，使模型适应量化变化。
- 学习步长量化（LSQ）：改进了量化器的配置过程，优化了量化间隔。
- DIFFQ：引入可微分的模型压缩技术，使用伪量化噪声实现完全可微的量化过程。
3.2 知识蒸馏
知识蒸馏技术包括软蒸馏和硬蒸馏，通过最小化教师模型和学生模型的软化对数输出之间的Kullback-Leibler（KL）散度来传输知识。
3.3 剪枝
剪枝通过调整隐藏层之间的权重核维度来减少参数数量。目标是在保持准确性的同时减少模型复杂性。
3.4 低秩近似
利用自注意力的低秩性质进行低秩近似，旨在保持准确性的同时显著降低计算复杂度。
4. 实验
实验在Tesla V100-SXM2 16GB GPU上进行，使用PyTorch作为主要平台，数据集限于CIFAR-10和CIFAR-100。
4.1 实验设置
主要评估指标包括模型大小、推理速度和准确性。最佳压缩技术应在准确性上几乎无影响，同时显著降低模型大小并提升推理速度。
4.2 不同模型压缩方法的比较
- 量化和剪枝：量化技术，特别是动态量化，将模型大小减小至25倍。剪枝率0.1导致显著准确性降低。
- 知识蒸馏：DeiT基础模型在保持准确性的同时，推理速度几乎是标准ViT的两倍。DeiT微型配置在CIFAR-10上实现95.43%的准确率，速度提升四倍，仅占原始模型大小的6%。
- 低秩近似：Nyströmformer技术在准确性和速度之间实现微妙平衡，受landmark数量选择影响。
4.3 混合方法的探索
结合量化和知识蒸馏的混合方法显示出优化潜力。DeiT基础模型与动态量化结合，推理速度提升超过两倍，模型大小减小到原始尺寸的四分之一。
5. 结论
本研究系统评估了四种主要的模型压缩技术，发现后训练量化和知识蒸馏在减小模型大小和加快推理时间方面效果显著。结合量化和知识蒸馏的混合方法在保持可接受的性能降低的同时，实现了显著的效率提升。这些发现为优化Vision Transformers的操作效率提供了重要参考。
论文链接：https://arxiv.org/pdf/2404.10407
本文原文来自公众号“极市平台”，仅用于学术分享，侵权删。