[論文レビュー] Token Merging: Your ViT But Faster

Daniel Bolya, Cheng-Yang Fu+4 more|arXiv (Cornell University)|Oct 17, 2022

Generative Adversarial Networks and Image Synthesis被引用数 62

ひとこと要約

Token Merging (ToMe) は、類似トークンを高速な bipartite matching プロセスで結合することにより Vision Transformers の処理能力を向上させ、およそ2x（動画は2.2x）までスループットを高め、訓練の有無にかかわらず、画像、動画、音声タスクで使用可能です。

ABSTRACT

We introduce Token Merging (ToMe), a simple method to increase the throughput of existing ViT models without needing to train. ToMe gradually combines similar tokens in a transformer using a general and light-weight matching algorithm that is as fast as pruning while being more accurate. Off-the-shelf, ToMe can 2x the throughput of state-of-the-art ViT-L @ 512 and ViT-H @ 518 models on images and 2.2x the throughput of ViT-L on video with only a 0.2-0.3% accuracy drop in each case. ToMe can also easily be applied during training, improving in practice training speed up to 2x for MAE fine-tuning on video. Training with ToMe further minimizes accuracy drop, leading to 2x the throughput of ViT-B on audio for only a 0.4% mAP drop. Qualitatively, we find that ToMe merges object parts into one token, even over multiple frames of video. Overall, ToMe's accuracy and speed are competitive with state-of-the-art on images, video, and audio.

研究の動機と目的

• 重い再訓練やアーキテクチャ変更を伴わずに、ViT 推論と訓練をより高速にすることを動機づける。
• トランスフォーマー層あたりのトークン数を減らすことでスループットを向上させる、単純なトークン結合モジュールを導入する。
• 画像・動画・音声の分野で適用可能性を確保しつつ、競争力のある精度を維持する。
• 市販のモデルに適した軽量で並列化可能な結合アルゴリズムを提供する。

提案手法

• 各トランスフォーマーブロックの attention と MLP のブランチの間にトークン結合ステップを挿入する。
• 結合候補を選択するために attention keys 之间のドット積類似性を用いてトークン類似性を定義する。
• トークンサイズで重み付けした特徴を平均化してトークンを結合する、fast bipartite soft matching を使用する。
• 複数の入力パッチを表す結合トークンを考慮して、比例的な注意を適用する。
• 結合をプーリングとして扱い、結合トークンを通じて誤差伝搬を行うことで ToMe を任意で訓練する。
• 層全体で削除される総トークン数を制御するため、一定および減衰する結合スケジュールを検討する。

実験結果

リサーチクエスチョン

• RQ1 軽量なトークン結合モジュールは、再訓練なしで精度の損失を最小限に抑えつつViTのスループットを大幅に向上させることができるか？
• RQ2 画像・動画・音声モダリティ全体で情報内容を保持するように、トークンをどのように照合・結合すべきか？
• RQ3 結合スケジュールや特徴選択が、異なるViTモデルと事前訓練方式における速度アップと精度にどのような影響を及ぼすか？
• RQ4 ToMe は推論のみでも訓練中でも、さまざまなモダリティに対して効果的か？

主な発見

• ToMe は画像上で、state-of-the-art ViT-L@512 および ViT-H@518 のスループットを概ね2倍にし、精度損失は約0.2–0.3%程度。
• 動画では、訓練なしで ViT-L のスループットを約2.2x達成し、精度損失は0.2–0.3%程度。
• ToMe で訓練した場合、動画での MAE 微調整の訓練速度向上は最大約2xを観測。
• 音声では、MAE pretraining 付きの ViT-B で約2xのスループットを達成し、mAP は約0.4%低下。
• ToMe はフレーム間で物体の部位を結合し、動画で部位追跡を行うことができる一方、背景と前景情報を結合して大きな精度劣化を招かずに済む。
• ToMe は追加の訓練テクニックや新しいパラメータを必要とせず、画像・動画・音声の各領域で最先端手法と競争力がある。