[論文レビュー] TinyViT: Fast Pretraining Distillation for Small Vision Transformers

Kan Wu, Jinnian Zhang+5 more|arXiv (Cornell University)|Jul 21, 2022

Advanced Neural Network Applications被引用数 22

ひとこと要約

TinyViTは、大規模な事前学習済み教師から知識を転移させることで小型のビジョン・トランスフォーマーを訓練する高速な事前学習蒸留フレームワークを導入し、はるかに少ないパラメータでImageNetおよび下流タスクで高い性能を達成します。

ABSTRACT

Vision transformer (ViT) recently has drawn great attention in computer vision due to its remarkable model capability. However, most prevailing ViT models suffer from huge number of parameters, restricting their applicability on devices with limited resources. To alleviate this issue, we propose TinyViT, a new family of tiny and efficient small vision transformers pretrained on large-scale datasets with our proposed fast distillation framework. The central idea is to transfer knowledge from large pretrained models to small ones, while enabling small models to get the dividends of massive pretraining data. More specifically, we apply distillation during pretraining for knowledge transfer. The logits of large teacher models are sparsified and stored in disk in advance to save the memory cost and computation overheads. The tiny student transformers are automatically scaled down from a large pretrained model with computation and parameter constraints. Comprehensive experiments demonstrate the efficacy of TinyViT. It achieves a top-1 accuracy of 84.8% on ImageNet-1k with only 21M parameters, being comparable to Swin-B pretrained on ImageNet-21k while using 4.2 times fewer parameters. Moreover, increasing image resolutions, TinyViT can reach 86.5% accuracy, being slightly better than Swin-L while using only 11% parameters. Last but not the least, we demonstrate a good transfer ability of TinyViT on various downstream tasks. Code and models are available at https://github.com/microsoft/Cream/tree/main/TinyViT.

研究の動機と目的

• リソースが限られたデバイス向けに効率的な小型ビジョン・トランスフォーマーの開発を促進する。
• 蒸留を通じて小型ViTが大規模事前学習データの利点を得られるようにする。
• 事前学習蒸留のための訓練メモリと計算コストを削減する。
• 転移性能を維持しつつ、tiny ViTsを事前学習および圧縮するスケーラブルなフレームワークを提案する。

提案手法

• 反復的な教師の順伝播を行わずに高速な事前学習蒸留を可能にするため、疎な教師ロジットとデータ拡張メタデータをディスクに保存する。
• 保存された教師出力から復元した疎なソフトラベルを用いた蒸留損失で小型のStudent ViTを訓練する。
• グラウンドトゥルースラベルではなくソフトな教師予測を活用するラベルなし蒸留設定を用いる。
• 大規模なシードViTを段階的に縮約して、パラメータとスループットの制約の下でTinyViTファミリを生成する。
• 窓付きアテンションとMBConvブロックを備えた階層的なSwin系アーキテクチャを採用し、精度と効率のバランスを取る。
• ImageNet-21kで事前学習し、ImageNet-1kでファインチューニングする。精度向上のために任意で高解像度のファインチューニングを追加。

実験結果

リサーチクエスチョン

• RQ1 事前学習中に大規模事前学習モデルから知識を蒸留することで、小型 Vision Transformer は競争力の性能を達成できるだろうか？
• RQ2 蒸留を如何に高速かつスケーラブルにして、大きな教師のメモリ/時間負荷を回避できるか？
• RQ3 事前学習蒸留がTiny ViTsの下流タスクへの転移性に与える影響は何か？
• RQ4 段階的なモデル縮約がTinyViTの精度と効率のトレードオフにどのような影響を与えるか？

主な発見

• TinyViT-21MはIN-21k事前学習と30エポックのIN-1kファインチューニングの後、21MパラメータでImageNet-1kのtop-1 84.8%を達成。
• より高解像度の入力で、TinyViTは86.5%のtop-1に達し、Swin-Lを僅かに上回りつつパラメータは約11%を使用。
• IN-21k上で蒸留を用いて事前学習したTinyViT-21Mは下流タスクへ良好に転移し、例えば物体検出でCOCO AP 50.2（28MパラメータのSwin-Tより2.1ポイント高い）。
• 高速な事前学習蒸留フレームワークは、疎な教師ロジットを保存しデータ拡張をエンコードすることでメモリと計算を削減し、訓練中に教師をロードせずに大規模バッチ蒸留を可能にする。
• 高品質な教師モデル（Florence、CLIP-ViT-L/14 など）を使用するとTinyViTの性能がさらに向上する一方で、オンディスクロジットのおかげで実用的な訓練コストを維持する。