[论文解读] MiniViT: Compressing Vision Transformers with Weight Multiplexing

Jinnian Zhang, Houwen Peng+5 more|arXiv (Cornell University)|Apr 14, 2022

CCD and CMOS Imaging Sensors被引用 8

一句话总结

MiniViT 提出了一种用于视觉 Transformer 的新型压缩框架，采用权重复用技术——结合权重变换与知识蒸馏——在减少模型参数量的同时提升准确率。该方法在 Swin-B 上实现 48% 的参数压缩，并在 ImageNet 上将 Top-1 准确率提升 1.0%；同时将 DeiT-B 的参数量压缩 9.7 倍（从 8600 万降至 900 万），性能损失极小。

ABSTRACT

Vision Transformer (ViT) models have recently drawn much attention in computer vision due to their high model capability. However, ViT models suffer from huge number of parameters, restricting their applicability on devices with limited memory. To alleviate this problem, we propose MiniViT, a new compression framework, which achieves parameter reduction in vision transformers while retaining the same performance. The central idea of MiniViT is to multiplex the weights of consecutive transformer blocks. More specifically, we make the weights shared across layers, while imposing a transformation on the weights to increase diversity. Weight distillation over self-attention is also applied to transfer knowledge from large-scale ViT models to weight-multiplexed compact models. Comprehensive experiments demonstrate the efficacy of MiniViT, showing that it can reduce the size of the pre-trained Swin-B transformer by 48\%, while achieving an increase of 1.0\% in Top-1 accuracy on ImageNet. Moreover, using a single-layer of parameters, MiniViT is able to compress DeiT-B by 9.7 times from 86M to 9M parameters, without seriously compromising the performance. Finally, we verify the transferability of MiniViT by reporting its performance on downstream benchmarks. Code and models are available at here.

研究动机与目标

• 解决视觉 Transformer（ViTs）参数量过大的问题，以提升其在资源受限设备上的部署能力。
• 克服在 ViTs 中简单地跨层共享权重所导致的训练不稳定与性能下降问题。
• 开发一种适用于同构与分层 ViT 架构的通用压缩框架。
• 在高倍率压缩下仍保持模型性能，提升参数效率。
• 使压缩后的模型具备良好的迁移能力，适用于目标检测与分割等下游视觉任务。

提出的方法

• 提出权重复用技术，即在 Transformer 块之间共享权重的同时，对每一层应用特定的线性变换，以增强权重的多样性。
• 在每个共享层中，分别为多头自注意力（MSA）和多层感知机（MLP）模块应用独立的变换矩阵，避免权重更新完全相同。
• 将层归一化参数在各层之间解耦，以维持稳定的特征统计量，提升训练稳定性。
• 实施多级知识蒸馏——包括预测层、注意力层和隐藏状态蒸馏——以将大教师模型的知识迁移至紧凑的学生模型。
• 使用单一层参数集表示所有共享层，大幅减少模型大小，同时保持性能。
• 通过检测蒸馏对压缩后的模型进行微调，进一步提升其在下游任务中的性能。

实验结果

研究问题

• RQ1 能否在不引发训练不稳定或性能下降的前提下，有效将跨层权重共享应用于视觉 Transformer？
• RQ2 在 ViT 块之间共享权重时，哪些机制可稳定训练并保持性能？
• RQ3 多级知识蒸馏（预测层、注意力层与隐藏状态）在压缩 ViT 时，对保持准确率的效率如何？
• RQ4 在同构与分层 ViT 架构中，该方法可实现多大程度的参数压缩？
• RQ5 压缩后的 MiniViT 模型是否仍具备强迁移能力，适用于目标检测与分割等下游视觉任务？

主要发现

• MiniViT 在将预训练的 Swin-B 参数量减少 48% 的同时，使 ImageNet 上的 Top-1 准确率提升了 1.0%。
• 该方法将 DeiT-B 的参数量从 8600 万压缩至 900 万（压缩 9.7 倍），性能损失微小。
• 结合变换与蒸馏的权重复用技术可稳定训练并防止模型崩溃，而简单权重共享则无法实现。
• 注意力层与隐藏状态蒸馏的使用显著提升了性能，尤其在微调阶段结合检测蒸馏时效果更佳。
• MiniViT 在主干网络上实现了 57% 的参数压缩，同时在 COCO 目标检测任务上保持与基线模型相当的平均精度（mAP）。
• 压缩后的模型展现出强大的迁移能力，经微调后在下游检测与分割基准上达到最先进性能。