[論文レビュー] AdaViT: Adaptive Tokens for Efficient Vision Transformer

Hongxu Yin, Arash Vahdat+4 more|arXiv (Cornell University)|Dec 14, 2021

CCD and CMOS Imaging Sensors被引用数 16

ひとこと要約

A-ViT は、入力画像の複雑さに応じて推論中に処理対象となるトークン数を動的に削減する、ビジョントランスフォーマー向けのアダプティブトークン計算メカニズムを導入する。既存のネットワークパラメータを再利用して各トークンの停止確率を学習し、分布的プライア共正則化を適用することで、アーキテクチャの変更なしに、DeiT-Tiny で 62⁠%、DeiT-Small で 38⁠% のスループット向上を達成し、わずか 0.3⁠% の精度低下で、ハードウェアに最適化された効率的で、アーキテクチャ変更なしの推論を実現する。

ABSTRACT

We introduce A-ViT, a method that adaptively adjusts the inference cost of vision transformer (ViT) for images of different complexity. A-ViT achieves this by automatically reducing the number of tokens in vision transformers that are processed in the network as inference proceeds. We reformulate Adaptive Computation Time (ACT) for this task, extending halting to discard redundant spatial tokens. The appealing architectural properties of vision transformers enables our adaptive token reduction mechanism to speed up inference without modifying the network architecture or inference hardware. We demonstrate that A-ViT requires no extra parameters or sub-network for halting, as we base the learning of adaptive halting on the original network parameters. We further introduce distributional prior regularization that stabilizes training compared to prior ACT approaches. On the image classification task (ImageNet1K), we show that our proposed A-ViT yields high efficacy in filtering informative spatial features and cutting down on the overall compute. The proposed method improves the throughput of DeiT-Tiny by 62% and DeiT-Small by 38% with only 0.3% accuracy drop, outperforming prior art by a large margin. Project page at https://a-vit.github.io/

研究の動機と目的

• 入力に応じた推論を実現し、画像の複雑さに応じて計算コストを低減する。
• 追加のパラメータやネットワークアーキテクチャの変更なしに、情報量が少ない空間的トークンの処理を適応的に停止するメカニズムを設計する。
• 新しい分布的プライア共正則化を用いて、訓練の安定性を高め、停止深度を目標となる計算予算に適応させる。
• 標準的なハードウェア上で顕著なスループット向上を達成しつつ、モデルの精度を維持する。

提案手法

• A-ViT は、アダプティブコンピュテーションタイム（ACT）をトークンレベルに拡張し、各トークンが停止条件を満たした時点で破棄できるようにする。
• 各トークンの停止確率は、各トランスフォーマーブロックの最終全結合層から得られる単一のニューロンを用いて計算され、追加のパラメータや計算を追加せずに、既存のパラメータを再利用する。
• 期待される停止深度を目標に合わせて導くために、分布的プライア共正則化を導入し、訓練の安定性と収束性を向上させる。
• 情報量の多いトークンのみを深く処理する、空間的に適応的な動的計算を可能にし、FLOPs と推論時間を削減する。
• 停止メカニズムは、補助ネットワークやファインチューニングを必要とせず、ワンステップのポンダー損失と分布的プライア共正則化を用いてエンドツーエンドで訓練される。

実験結果

リサーチクエスチョン

• RQ1 アーキテクチャの変更なしに、トークンレベルでのアダプティブ計算を導入することで、ビジョントランスフォーマーの推論効率を向上させられるか？
• RQ2 既存のネットワークパラメータのみを用いて、効率的に停止意思決定を学習できるか？
• RQ3 分布的プライア共正則化により、訓練の安定性が向上し、目標となる計算予算に収束しやすくなるか？
• RQ4 アダプティブトークン停止は、画像の意味的コンテンツや人間の認知とどの程度相関するか？

主な発見

• A-ViT は、ImageNet1K におけるトップ1精度をわずか 0.3⁠% 低下させるだけで、DeiT-Tiny で 62⁠%、DeiT-Small で 38⁠% のスループット向上を達成した。
• 細粒度で空間的に適応的なトークンプルーニングにより、従来の ACT よりも 25⁠% の FLOPs 削減が達成された。
• 分布的プライア共正則化により、目標の停止深度への収束が早くなり、ポンダー損失のみを用いた場合に比べて 6.4⁠% の精度向上が得られた。
• 停止計算に単一のベクトル要素を用いることで、DeiT-Tiny では 0.08⁠%、DeiT-Small では 0.04⁠% の精度低下に留まり、極めて低いパラメータオーバーヘッドを実証した。
• 補助ネットワークを追加して停止学習を支援すると、精度向上はわずか 0.06⁠% に留まり、スループットに 12.6⁠% のペナルティが生じるため、実用的ではない。