[論文レビュー] Compressing Vision Transformers for Low-Resource Visual Learning

Eric Youn, Sai Mitheran J.+2 more|arXiv (Cornell University)|Sep 5, 2023

Advanced Neural Network ApplicationsComputer Science被引用数 3

ひとこと要約

本稿では、NVIDIA Jetson Nano（4GB RAM）などのリソース制限のあるエッジデバイスへの効率的なデプロイを可能にするために、構造的プルーニング、知識蒸留、トレーニング後量子化を組み合わせたハイブリッド圧縮フレームワークを提案する。このアプローチにより、スイントランスフォーマー並みの精度を維持しながら、最小限の遅延とメモリ使用量で実現し、UAVを用いた被災地モニタリングのようなリソース制限のある状況でのリアルタイムセマンティックセグメンテーションを可能にする。

ABSTRACT

Vision transformer (ViT) and its variants have swept through visual learning leaderboards and offer state-of-the-art accuracy in tasks such as image classification, object detection, and semantic segmentation by attending to different parts of the visual input and capturing long-range spatial dependencies. However, these models are large and computation-heavy. For instance, the recently proposed ViT-B model has 86M parameters making it impractical for deployment on resource-constrained devices. As a result, their deployment on mobile and edge scenarios is limited. In our work, we aim to take a step toward bringing vision transformers to the edge by utilizing popular model compression techniques such as distillation, pruning, and quantization. Our chosen application environment is an unmanned aerial vehicle (UAV) that is battery-powered and memory-constrained, carrying a single-board computer on the scale of an NVIDIA Jetson Nano with 4GB of RAM. On the other hand, the UAV requires high accuracy close to that of state-of-the-art ViTs to ensure safe object avoidance in autonomous navigation, or correct localization of humans in search-and-rescue. Inference latency should also be minimized given the application requirements. Hence, our target is to enable rapid inference of a vision transformer on an NVIDIA Jetson Nano (4GB) with minimal accuracy loss. This allows us to deploy ViTs on resource-constrained devices, opening up new possibilities in surveillance, environmental monitoring, etc. Our implementation is made available at https://github.com/chensy7/efficient-vit.

研究の動機と目的

• NVIDIA Jetson Nano（4GB RAM）を搭載したUAVなどのメモリ制限があり、バッテリー駆動のエッジデバイスへのビジョントランスフォーマーのデプロイを可能にすること。
• モデルサイズと推論遅延を最小限に抑えながら、スイントランスフォーマーと同等のセグメンテーション精度を維持すること。
• エッジデバイスへのデプロイに適した、構造的プルーニング、蒸留、トレーニング後量子化を組み合わせた実用的でエンドツーエンドの圧縮パイプラインの開発。
• LPCVデータセットを用いて、厳密なメモリ制限と遅延制限のもとで、実世界の被災地シーンの解析を評価すること。

提案手法

• 計算グラフ構造に変更を加えずに、不要なアテンションヘッドと線形層内のニューロンを構造的にプルーニングすることで、パラメータ数を削減する。
• 段階的な蒸留戦略を採用：予測ヘッドでのログ確率知識移行、初期層での特徴類似、深層部での生成ベースの蒸留。
• トレーニング後量子化（PTQ）を用いてモデルサイズとメモリフットプリントを削減するが、PyTorchにおける量子化CUDA演算のサポートが限られているため、さらなる向上が制限される。
• プルーニング後に繰り返し微調整を組み合わせることで、性能の安定化と一般化性能の向上を図る。
• より強力なスイントランスフォーマーV2-Tモデルを教師モデルとして用い、知識を小型のMobileViT学生ネットワークに移す。
• 主にmIoU（交差率の平均）を指標として、LPCVデータセット上で圧縮モデルの評価を実施する。

Figure 1: Our framework for edge-oriented, rapid on-device disaster scene parsing

実験結果

リサーチクエスチョン

• RQ1 アテンションヘッドと線形層の重みに対する構造的プルーニングは、エッジデバイス上で顕著な精度低下を伴わずにモデルサイズを効果的に削減できるか？
• RQ2 ログ確率、特徴、生成の段階的監視による知識蒸留は、大規模な教師モデルからコンactな学生ViTに性能を効果的に移すことができるか？
• RQ3 Jetson Nano上で、トレーニング後量子化を用いることで、メモリ使用量と遅延をどの程度削減できるか、かつセグメンテーション精度を維持できるか？
• RQ4 クラスの不均衡や微細な視覚的区別（例：土石流 vs. 土の道）が、実世界の被災地シーンにおける圧縮モデルの性能に与える影響は何か？
• RQ5 一括プルーニングと比較して、繰り返しプルーニングと微調整を組み合わせることで、圧縮効率が向上するか？

主な発見

• 最終的な圧縮モデルは、4GBのシステムメモリというエッジデプロイ制約を満たしつつ、LPCVデータセットでスイントランスフォーマー並みの精度を達成した。
• アテンションヘッドを4から2にプルーニングしても性能への影響は最小限であり、パラメータ数を顕著に削減しながらも高いmIoUを維持した。
• トレーニング後量子化によりメモリ使用量が顕著に削減されたが、PyTorchにおける量子化CUDA演算のサポートが限られているため、さらなる向上は制限された。
• 支配的クラスでは良好な性能を示したが、土石流や土の道のような微細で頻度の低いクラスでは苦戦し、クラス不均衡の課題が浮き彫りになった。
• 繰り返しプルーニングは一括プルーニングと比較して僅かな改善しか得られず、この特定のアーキテクチャと圧縮目的では利点が限定的であることが示唆された。
• 本フレームワークにより、Jetson Nano上でリアルタイムセマンティックセグメンテーションが実現可能であることが示され、UAVを用いた被災地モニタリングへの実用性が裏付けられた。

Figure 2: Samples from the LPCV dataset (CVPR 2023 challenge)