[論文レビュー] SPViT: Enabling Faster Vision Transformers via Soft Token Pruning

Zhenglun Kong, Peiyan Dong+10 more|arXiv (Cornell University)|Dec 27, 2021

Advanced Neural Network Applications被引用数 27

ひとこと要約

SPViTはVision Transformerに対してレイテンシを意識したソフトトークンプルーニングフレームワークを導入し、マルチヘッドトークンセレクタとトークンパッケージングを用いて画像ごとに適応させ、エッジデバイスやFPGAでの精度低下を最小限に抑えつつ大幅なレイテンシ削減を実現する。

ABSTRACT

Recently, Vision Transformer (ViT) has continuously established new milestones in the computer vision field, while the high computation and memory cost makes its propagation in industrial production difficult. Pruning, a traditional model compression paradigm for hardware efficiency, has been widely applied in various DNN structures. Nevertheless, it stays ambiguous on how to perform exclusive pruning on the ViT structure. Considering three key points: the structural characteristics, the internal data pattern of ViTs, and the related edge device deployment, we leverage the input token sparsity and propose a computation-aware soft pruning framework, which can be set up on vanilla Transformers of both flatten and CNN-type structures, such as Pooling-based ViT (PiT). More concretely, we design a dynamic attention-based multi-head token selector, which is a lightweight module for adaptive instance-wise token selection. We further introduce a soft pruning technique, which integrates the less informative tokens generated by the selector module into a package token that will participate in subsequent calculations rather than being completely discarded. Our framework is bound to the trade-off between accuracy and computation constraints of specific edge devices through our proposed computation-aware training strategy. Experimental results show that our framework significantly reduces the computation cost of ViTs while maintaining comparable performance on image classification. Moreover, our framework can guarantee the identified model to meet resource specifications of mobile devices and FPGA, and even achieve the real-time execution of DeiT-T on mobile platforms. For example, our method reduces the latency of DeiT-T to 26 ms (26%$\sim $41% superior to existing works) on the mobile device with 0.25%$\sim $4% higher top-1 accuracy on ImageNet.

研究の動機と目的

• エッジデバイスとリアルタイム展開のためのVision Transformerの高い計算コストを動機づけ、対処する。
• 画像ごとに適応的なプルーニングを可能にするレイテンシを意識したソフトトークンプルーニングフレームワーク（SPViT）を提案する。
• 注意機構ベースのマルチヘッドトークンセレクタとトークンパッケージング技術を開発し、プルーニングされたトークンから情報を保存する。
• デバイス間のハードウェアレイテンシ制約を満たすためのレイテンシを意識したトレーニング戦略を導入する。
• モバイルデバイスとFPGA上でViTモデルのリアルタイムエッジ展開を、有意義なレイテンシと精度のトレードオフで実証する。

提案手法

• ViTブロック全体に軽量なマルチヘッドトークンセレクタを挿入し、各アテンションヘッドごとにトークンの重要性をスコア付けする。
• 情報量が少ないトークンを破棄するのではなくパッケージトークンにまとめてソフトトークンプルーニングを適用し、文脈情報を保持する。
• アテンションベースの分岐を介してヘッドごとのトークンスコアを集約し、Gumbel-Softmaxを用いて微分可能な保持/プルーニング決定を行う。
• レイテンシ-スパース性損失を導入し、ブロックごとのプルーニング率をハードウェアのレイテンシ予算に制約するためにレイテンシ-スパース性ルックアップテーブルを介して制御する。
• 挿入点とプルーニング率を決定するために層ごと・フェーズごとの順次トレーニングを用い、精度とハードウェアレイテンシのバランスを取る。
• モバイル（Samsung Galaxy S20）とFPGA（Xilinx ZCU102）上でSPViTをデプロイし、リアルタイム推論と計算-精度のトレードオフを実証する。

実験結果

リサーチクエスチョン

• RQ1 ViTにおけるトークンプルーニングをデバイス固有の制約を満たしつつ精度を維持するようレイテンシを意識させるにはどうすればよいか？
• RQ2 トークンパッケージングを伴うソフトトークンプルーニングアプローチは、エッジデバイス上でハードプルーニングや他のプルーニング戦略を上回るか？
• RQ3 異なるViTブロックにトークンセレクタを挿入することが精度とレイテンシに与える影響は何か？
• RQ4 SPViTは軽量な階層型ViT（例：Swin、PiT）およびエッジハードウェアでどのように機能するか？

主な発見

モデル	GFLOPs	Top1 精度（%）	レイテンシ（ms）
Swin-S	8.70	83.20	-
SPViT (Ours)	6.35 (26.4%↓)	82.71 (↓0.49)	-
Swin-T	4.50	81.20	-
SPViT (Ours)	3.47 (↓23.0%↓)	80.70 (↓0.50)	-
PiT-S	2.90	80.90	-
SPViT (Ours)	2.22 (↓23.3%↓)	80.32 (↓0.58)	-
PiT-XS	1.40	78.10	-
SPViT (Ours)	1.13 (↓18.7%↓)	77.86 (↓0.24)	-

• SPViTはバックボーン全体で31%–43%のViT計算を削減し、精度損失は0.1%–0.5%で収まる。
• DeiT-Tではモバイルで26 msのレイテンシを達成し、他モデルでは最大で40%–60%のレイテンシ削減を実現するが、精度損失はほとんどない。
• SPViTはDeiT-Tでモバイルデバイス上のリアルタイム推論を実現し、FPGAでは固定小数点実装で顕著なレイテンシ削減を達成する。
• トークンパッケージングはプルーニングされたトークンから情報を保存し、プルーニング率を高めつつ精度を維持するのに役立つ。
• SPViTは平坦型および階層型ViTsの精度-レイテンシトレードオフにおいて、いくつかの最先端プルーニング法より優れている。
• レイテンシを意識したデプロイ結果は、Samsung Galaxy S20およびXilinx ZCU102ハードウェアで顕著な改善を示す。