Vision Transformerのタスク適応型トークン剪定：レジスタ活用で高速化と高精度を両立

Vision Transformerの効率化とマルチタスク適応の必要性

Vision Transformer (ViT) は、自然言語処理分野での成功を背景に、画像認識タスクにおいても目覚ましい成果を上げています。しかし、その計算コストの高さは、特にエッジデバイスやリアルタイム処理が求められるアプリケーションでの導入を妨げる要因となっています。推論時の高速化は喫緊の課題であり、トークン剪定（Token Pruning）はその有効な解決策の一つです。トークン剪定は、画像パッチから生成されるトークンのうち、情報量の少ないものを削減することで、計算量を削減し、推論速度を向上させます。しかし、これまでの剪定ポリシーの多くは、単一のタスク、例えば画像分類に特化して設計されてきました。私の見方では、これはViTが持つ汎用性のポテンシャルを十分に引き出せていない状況です。ViTは、画像分類だけでなく、セマンティックセグメンテーションや物体検出といった、空間的な詳細情報がより重要となる多様なタスクで再利用されることが増えています。このようなマルチタスク環境において、単一タスク向けの剪定ポリシーでは、他のタスクでの性能が著しく低下するリスクがあります。この課題を解決するためには、タスクの特性に柔軟に対応できる、より洗練された剪定メカニズムが不可欠だと私は考えています。

異なるタスクにおける剪定ポリシーの挙動解析

私は、異なる画像認識タスクがトークン剪定にどのように反応するかを詳細に調査しました。具体的には、画像分類（ImageNet-1K）、セマンティックセグメンテーション（ADE20K）、物体検出（COCO）の3つの主要タスクを対象としました。プローブと呼ばれる手法を用いて、既存の事前学習済みViTモデルのパラメータを固定したまま、様々なパラメータフリーのトークン削減基準（例：アテンションスコアに基づくもの、特徴量のL2ノルムに基づくものなど）を各層に適用し、その性能への影響を評価しました。この厳密な分析から、3つの重要な発見がありました。

第一に、セグメンテーションと検出といった高密度な予測を必要とするタスクでは、トークン剪定基準の優先順位が画像分類とは異なり、かつ両者間でも異なる傾向を示すことが分かりました。これは、各タスクが画像の異なる局所的・大局的特徴に依存しているためです。第二に、画像分類タスクは、モデルの初期層でのアテンションベースの剪定に対して特に敏感であることが判明しました。初期層で重要なトークンを誤って削除すると、その後の特徴抽出に深刻な影響を与え、最終的な分類精度が大きく低下します。第三に、高密度タスクでは、剪定後に失われた情報を回復させるための「回復エンドポイント」が、タスクによって逆の傾向を好むことが明らかになりました。これは、セグメンテーションや検出がピクセルレベル、またはオブジェクトレベルの詳細な情報を必要とするため、どの段階で情報を回復させるかが極めて重要であることを示唆しています。私の経験上、このようなタスク特性の理解なしに闇雲に剪定を行うことは、性能低下を招く最も一般的な失敗パターンです。

Task-Adaptive Pruning (TAP)の設計思想とメカニズム

これらの知見に基づき、私は「Task-Adaptive Pruning (TAP)」という新しいフレームワークを考案しました。TAPの核心は、既存のViTアーキテクチャに存在する「レジスタートークン」の概念を拡張し、タスク適応性を持たせる点にあります。従来のレジスタートークンは、タスクに依存しない汎用的な特徴を保存する役割を担うことが一般的でしたが、TAPではこれをタスク固有の情報保持に活用します。具体的には、タスクごとに1つの専用「タスクレジスタ」を導入し、ViTが特定のタスクを実行する際には、そのタスクに対応するタスクレジスタのみをアクティブにします。このタスクレジスタは、単に情報を保存するだけでなく、以下の3つの重要な機能を果たします。

1. トークンの重要度ランク付け: 各タスクレジスタは、現在のタスクにとってどのトークンが最も重要であるかを動的に評価し、ランク付けします。これにより、タスクの目標に合致したトークンを優先的に保持し、重要度の低いトークンを剪定対象とします。
2. 深さ方向の削除予算配分: モデルの層（深さ）ごとに、どれだけのトークンを削除するかという予算を正確に配分します。これは、初期層では慎重に、後期層ではより大胆に剪定するといった、タスクと層の特性に応じた柔軟な戦略を可能にします。
3. 高密度特徴の回復スケール設定: セグメンテーションや検出のような高密度な予測を必要とするタスクでは、剪定によって失われた空間的詳細情報を適切に回復させる必要があります。タスクレジスタは、この回復プロセスにおけるスケールや方法をタスクに応じて最適に設定し、性能低下を防ぎます。私の経験上、このような「ぐるぐる回す」フィードバックループの設計が、プロダクトの性能を最大限に引き出す鍵となります。

TAP-Jによるマルチタスク性能と効率の飛躍的向上

TAPの有効性を検証するため、私は共同適応モデル「TAP-J」を開発し、複数のタスクでその性能を評価しました。最終的なトークン保持率をρ=0.5（つまり、半分のトークンを剪定）というアグレッシブな設定にもかかわらず、TAP-Jは驚くべき結果を示しました。

セマンティックセグメンテーションタスクでは、ADE20Kデータセットにおいて47.0 mIoUという高精度を達成しつつ、ViTエンコーダのスループットを1.30倍に向上させることができました。これは、同等の精度を維持しながら、処理速度が30%向上したことを意味します。物体検出タスクでは、COCOデータセットにおいて53.7 box APを達成し、エンコーダのスループットは1.32倍に高速化しました。これもまた、精度を犠牲にすることなく、大幅な高速化を実現した証拠です。さらに、画像分類のベンチマークであるImageNet-1Kにおいても、TAP-Jは既存の強力なモデルと同等かそれ以上の競争力のある性能を維持しました。この結果は、TAPが単一タスクに特化することなく、複数の異なるタスクにおいてViTの効率と性能を同時に最大化できる、非常に強力で汎用性の高いアプローチであることを明確に示しています。私の見方では、このようなマルチタスクでの効率化は、今後のAIプロダクト開発において「最速」で市場投入し、競争優位性を確立するための必須要件です。設計者は、単なる精度追求だけでなく、このような実行効率も常に考慮に入れるべきだと断言します。