[論文レビュー] 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

ひとこと要約

本論文は、教師モデルと生徒モデルの間でパッチレベルの特徴多様体を一致させることで、ビジョントランスフォーマーを圧縮する細分化された多様体 distillation を提案する。損失項の分離により計算コストを著しく削減する。DeiT-Tiny を生徒モデルとして用いることで ImageNet-1k で 76.5% のトップ-1 精度を達成し、従来の distillation 法よりも +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.

研究の動機と目的

• ビジョントランスフォーマーの高い推論コストが、スマートフォンやスマートウォッチなどのエッジデバイスへの展開を制限する問題に対処すること。
• 主に CNN に特化した従来の知識 distillation 法が、ビジョントランスフォーマーのパッチレベル構造を十分に活用できないという制限を克服すること。
• 中間層の特徴関係を細分化された多様体ベースの方法で活用し、知識伝達を向上させる distillation 法の開発。
• 損失の分離とパッチマージを用いて、多様体マッチングの計算複雑性を低減しながら性能を維持すること。
• ソフトラベル distillation と固定深さの生徒アーキテクチャが、ビジョントランスフォーマー圧縮において優れた一般化性能をもたらすかを示すこと。

提案手法

• 本手法は、教師および生徒ネットワークのパッチレベル特徴表現を多様体空間としてモデル化し、各トランスフォーマー層におけるパッチ間関係を捉える。
• 計算複雑性を低減するため、三つの項（画像間、画像内、ランダムにサンプリングされた関係マッチング）から成る分離型多様体 distillation 損失を導入する。
• 画像間損失は、異なる画像間のパッチ間相対関係をマッチングする。一方、画像内損失は、同じ画像内での関係を維持する。
• ランダムにサンプリングされた損失項は、全ペアワイズ関係マップをパッチのサブセットで近似することで計算量を削減し、約 2 時間の速度向上を達成する。
• さらに、パッチマージ戦略を適用することで、多様体表現を単純化し、計算負荷を低減する。
• 生徒モデルは、教師のソフトラベルと中間特徴多様体をマッチングするように訓練され、一貫性のある distillation を保証するため、固定された深さとする。

実験結果

リサーチクエスチョン

• RQ1 ビジョントランスフォーマーにおけるパッチレベルの多様体関係は、知識 distillation のために効果的に活用可能か？
• RQ2 全多様体マッチングの高コストを、distillation の性能を損なわずにどのように低減できるか？
• RQ3 ソフトラベル distillation と固定深さの生徒訓練は、ハードラベルまたは可変深さのアプローチと比較して、より優れた一般化性能をもたらすか？
• RQ4 画像分類および下流タスクにおいて、細分化された多様体 distillation は、既存の distillation 法と比較して精度と効率の両面で優れているか？
• RQ5 損失重みやサンプリングサイズなどのハイパーパramータ設定（例：α=4.0, β=0.1, γ=0.2, K=192）は、提案された distillation フレームワークで最適なパフォーマンスをもたらすか？

主な発見

• 提案された細分化された多様体 distillation 法は、パrameter 数がたった 500 万の DeiT-Tiny を生徒モデルとして用い、ImageNet-1k で 76.5% のトップ-1 精度を達成し、従来の distillation 法よりも +2.0% の精度向上を達成した。
• CIFAR-10 および CIFAR-100 では、それぞれ +0.25% および +0.71% の転移学習性能向上を達成し、優れた一般化能力を示した。
• COCO 2017 におけるオブジェクト検出タスクでは、distilled Swin-Tiny モデルが 44.7% のボックス AP を達成し、distillation を行わなかった生徒モデルよりも +1.0% の向上を達成した。
• アブレーションスタディにより、ソフトラベル distillation と固定深さの生徒訓練が性能向上に不可欠であることが確認され、これらのコンponent を除去すると顕著な精度低下が生じた。
• 分離損失設計により、全多様体マッチングと比較して計算複雑性が約 2 時間のオーダーで低減され、効率的な訓練が可能になった。
• ハイパーパramータのアブレーションにより、デフォルト設定（α=4.0, β=0.1, γ=0.2, K=192）が最適なパフォーマンスをもたらすことが判明したが、さらなるチューニングで結果の向上が期待できる。