[論文レビュー] PPT: Token Pruning and Pooling for Efficient Vision Transformers

WU Xin-jian, Fanhu Zeng+2 more|arXiv (Cornell University)|Oct 3, 2023

Advanced Neural Network Applications被引用数 6

ひとこと要約

本論文では、学習可能なパラメータを一切持たない新しいフレームワークPPTを提案する。PPTは、ビジョントランスフォーマーにおける計算の冗長性を低減するために、アテンションの重みを動的に選択的にプルーニングまたはプーリングすることで、トークンのプルーニングとプーリングを適応的に統合する。層の深さや入力固有のトークン分布に基づいて、動的にプルーニングまたはプーリングを選択することで、DeiT-SにおいてFLOPsを37%以上削減し、精度に影響を与えることなくスループットを45%向上する。

ABSTRACT

Vision Transformers (ViTs) have emerged as powerful models in the field of computer vision, delivering superior performance across various vision tasks. However, the high computational complexity poses a significant barrier to their practical applications in real-world scenarios. Motivated by the fact that not all tokens contribute equally to the final predictions and fewer tokens bring less computational cost, reducing redundant tokens has become a prevailing paradigm for accelerating vision transformers. However, we argue that it is not optimal to either only reduce inattentive redundancy by token pruning, or only reduce duplicative redundancy by token merging. To this end, in this paper we propose a novel acceleration framework, namely token Pruning & Pooling Transformers (PPT), to adaptively tackle these two types of redundancy in different layers. By heuristically integrating both token pruning and token pooling techniques in ViTs without additional trainable parameters, PPT effectively reduces the model complexity while maintaining its predictive accuracy. For example, PPT reduces over 37% FLOPs and improves the throughput by over 45% for DeiT-S without any accuracy drop on the ImageNet dataset. The code is available at https://github.com/xjwu1024/PPT and https://github.com/mindspore-lab/models/

研究の動機と目的

• すべてのトークンにわたる冗長なアテンションによって引き起こされるビジョントランスフォーマー（ViTs）の計算非効率性を解消すること。
• 現在の手法が「無関心なトークン」または「重複するトークン」のいずれかの冗長性にのみ焦点を当てているという限界を克服し、両方の冗長性を同時に処理すること。
• 層や入力インスタンスに応じて、トークンプルーニングとプーリングを適応的に統合できる、軽量でパラメータフリーのメカニズムを開発すること。
• 微調整やアーキテクチャの大幅な変更なしに、ViTsの効率的な推論と高速な学習を可能にすること。
• 最先端のトークン圧縮手法と比較して、より優れた精度-FLOPsのトレードオフを達成すること。

提案手法

• 各層におけるトークン重要度スコアの分布に基づき、プルーニングまたはプーリングのどちらを適用するかを動的に選択する適応的トークン圧縮モジュールを提案する。
• トークンスコアの分散に基づくヒューリスティックな意思決定基準を用い、分散が大きい（重要度が明確に異なる）場合はプルーニング、分散が小さい（類似性が高い）場合はプーリングを適用する。
• トレーニング可能なパラメータを追加せず、標準的なトランスフォーマーブロックに統合することで、事前学習済みモデルとの互換性を確保する。
• 層ごとに異なる戦略を適用する：浅い層では類似性が高いためプーリングを優遇し、深い層では重要度の差が顕著なためプルーニングを優先する。
• 入力インスタンスごとに局所的なトークンスコア分布に基づいてポリシーを動的に調整し、インスタンスに応じた圧縮を実現する。
• プルーニングとプーリングの優先度を制御する学習可能な閾値τを設定し、微調整により最適化するか、オフザシェルの状況では経験的に設定する。

実験結果

リサーチクエスチョン

• RQ1 ビジョントランスフォーマーにおいて、トークンプルーニングとプーリングを併用することで、単独で使用する場合よりも優れた効率性-精度のトレードオフが達成できるか？
• RQ2 トークンの冗長性に層依存のパターンが存在するか？ すなわち、浅い層と深い層で異なる圧縮戦略（プルーニング対プーリング）を採用することが正当化されるか？
• RQ3 トレーニング不要のインスタンスに応じたポリシー選択メカニズムが、多様な入力において無関心な冗長性と重複する冗長性の両方を効果的にバランスできるか？
• RQ4 意思決定閾値τの選択が、適応的圧縮フレームワークの性能とロバストネスにどのように影響するか？
• RQ5 提案手法は、微調整なしで標準的なベンチマーク（例：ImageNet）に容易に統合可能であり、精度の低下を最小限に抑えられるか？

主な発見

• DeiT-Sにおいて、ImageNet上でのFLOPsを37%以上削減し、精度に影響を与えることなくスループットを45%以上向上。
• 固定ルールに基づくポリシー（例：浅い層でプーリング、深い層でプルーニング）やランダム選択と比較して、特にFLOPが低い領域で優れた性能を示す。
• 本フレームワーク内では、プルーニング単体よりもプーリング単体が優れているが、適応的選択による併用が全体として最良のパフォーマンスを達成する。
• 微調整状況では最適な閾値τは6×10⁻⁵、オフザシェル状況では7×10⁻⁵であると特定され、プルーニングとプーリングの優先度を適切にバランスさせる。
• 異なるViTアーキテクチャにわたり強力な一般化性能を示し、ImageNet上での最先端の効率性-精度トレードオフを達成。
• 推論の高速化に有効であり、追加のパラメータなしでも、精度の低下を最小限に抑えつつ、高速な学習を可能にする。