[论文解读] Learning Efficient Vision Transformers via Fine-Grained Manifold Distillation

Zhiwei Hao, Jianyuan Guo+6 more|arXiv (Cornell University)|Jul 3, 2021

Advanced Neural Network Applications被引用 23

一句话总结

本文提出细粒度流形蒸馏方法，通过匹配教师模型与学生模型在图像块级别的特征流形，压缩视觉Transformer模型，同时通过解耦损失项显著降低计算成本。该方法在ImageNet-1k上实现76.5%的top-1准确率，使用DeiT-Tiny学生模型，相比先前蒸馏方法提升+2.0%准确率，同时在迁移学习和目标检测任务中保持优异泛化能力。

ABSTRACT

In the past few years, transformers have achieved promising performances on various computer vision tasks. Unfortunately, the immense inference overhead of most existing vision transformers withholds their from being deployed on edge devices such as cell phones and smart watches. Knowledge distillation is a widely used paradigm for compressing cumbersome architectures via transferring information to a compact student. However, most of them are designed for convolutional neural networks (CNNs), which do not fully investigate the character of vision transformer (ViT). In this paper, we utilize the patch-level information and propose a fine-grained manifold distillation method. Specifically, we train a tiny student model to match a pre-trained teacher model in the patch-level manifold space. Then, we decouple the manifold matching loss into three terms with careful design to further reduce the computational costs for the patch relationship. Equipped with the proposed method, a DeiT-Tiny model containing 5M parameters achieves 76.5% top-1 accuracy on ImageNet-1k, which is +2.0% higher than previous distillation approaches. Transfer learning results on other classification benchmarks and downstream vision tasks also demonstrate the superiority of our method over the state-of-the-art algorithms.

研究动机与目标

• 为解决视觉Transformer模型推理成本过高的问题，以限制其在智能手机、智能手表等边缘设备上的部署。
• 克服现有知识蒸馏方法的局限性，这些方法主要针对CNN设计，未能有效利用视觉Transformer的图像块级结构。
• 开发一种基于中间层特征关系的细粒度、流形化蒸馏方法，以提升知识迁移效果。
• 通过损失解耦与图像块合并策略，在保持性能的同时降低流形匹配的计算复杂度。
• 展示软标签蒸馏与固定深度学生模型架构在视觉Transformer压缩中的优越性。

提出的方法

• 该方法将教师与学生网络的图像块级特征表示建模为流形空间，捕捉每个Transformer层中图像块之间的相互关系。
• 提出一种由三项组成、解耦的流形蒸馏损失：图像间关系匹配、图像内关系匹配以及随机采样关系匹配，以降低计算复杂度。
• 图像间损失用于匹配不同图像之间图像块的相对关系，而图像内损失则用于保持同一图像内部的图像块关系。
• 随机采样损失项通过使用部分图像块近似完整成对关系图，减少计算量，实现约两个数量级的速度提升。
• 采用图像块合并策略进一步简化流形表示并降低计算负载。
• 学生模型被训练以匹配教师模型的软标签与中间特征流形，且采用固定深度以确保一致的蒸馏过程。

实验结果

研究问题

• RQ1 能否有效利用视觉Transformer中的图像块级流形关系进行知识蒸馏，以提升模型压缩效果？
• RQ2 如何在不牺牲蒸馏性能的前提下，降低完整流形匹配的高计算成本？
• RQ3 软标签蒸馏与固定深度学生训练是否相比硬标签或可变深度方法带来更好的泛化能力？
• RQ4 在图像分类与下游任务中，细粒度流形蒸馏与现有蒸馏方法相比，在准确率与效率方面表现如何？
• RQ5 在所提出的蒸馏框架中，哪些超参数设置（如损失权重、采样大小）能获得最佳性能？

主要发现

• 所提出的细粒度流形蒸馏方法在使用仅500万参数的DeiT-Tiny学生模型时，于ImageNet-1k上实现76.5%的top-1准确率，相比先前蒸馏方法提升+2.0%准确率。
• 在CIFAR-10与CIFAR-100上，该方法分别将迁移学习性能提升+0.25%与+0.71%，展现出优越的泛化能力。
• 在COCO 2017目标检测任务中，蒸馏后的Swin-Tiny模型达到44.7%的box AP，相比未使用蒸馏的学生模型提升+1.0%。
• 消融实验表明，软标签蒸馏与固定深度学生训练对性能增益至关重要，若移除这些组件将导致准确率显著下降。
• 解耦损失设计使计算复杂度相比完整流形匹配降低约两个数量级，从而实现高效训练。
• 超参数消融分析显示，默认设置（α=4.0, β=0.1, γ=0.2, K=192）可获得最优性能，尽管进一步调优可能带来更优结果。