論文「ACCELERATING VISION TRANSFORMERS WITH ADAPTIVE PATCH SIZES」の技術的サマリー

この論文は、Vision Transformer (ViT) の推論および学習速度を大幅に向上させるための新しい手法「Adaptive Patch Transformer (APT)」を提案しています。ViT の固定されたパッチサイズによる非効率性を解決し、画像の内容に応じた適応的なパッチ分割を実現することで、計算コストを削減しつつ精度を維持することを目的としています。

以下に、問題定義、手法、主要な貢献、実験結果、および意義について詳細にまとめます。

---

1. 背景と問題定義

Vision Transformer (ViT) の課題:
ViT は画像認識において支配的なアーキテクチャとなっていますが、そのスケーラビリティは自己注意（Self-Attention）の計算コストがシーケンス長に対して二次関数的に増加することに制限されています。従来の ViT は、入力画像を均一なサイズのパッチに分割し、それぞれをトークンとして扱います。

• 高解像度画像の問題: 高解像度画像の場合、単純な背景（均一な色など）であっても複雑な物体（顔やテクスチャなど）であっても、同じ数のパッチに分割されるため、冗長なトークンが大量に生成されます。
• 既存手法の限界: 既存のトークン削減手法（トークンのプルーニングやマージ）には以下の欠点があります。
  • 固定された削減率: 画像の複雑さに応じて動的に調整されないため、単純な画像では削減不足、複雑な画像では情報損失が起きる可能性があります。
  • 実装上のオーバーヘッド: 推論中のトークン削除はパディングや不規則な形状を生み、実際の高速化（特にバッチサイズ>1）を妨げることがあります。
  • 学習コスト: 多くの手法は学習から始めなければならず、事前学習済みモデルへの適用が困難です。

自然言語処理における Byte-Pair Encoding (BPE) や SentencePiece のように、内容に応じて可変長のトークン化を行うアプローチが、視覚タスクにも適用可能であるという洞察が本研究の起点です。

---

2. 提案手法：Adaptive Patch Transformer (APT)

APT は、画像内の領域の複雑さ（情報量）に応じて、パッチサイズを動的に変更する手法です。

2.1 パッチサイズの決定（Entropy-based Allocation）

画像を均一なグリッドではなく、クアッドツリー（Quadtree）のような階層的構造に基づいてパッチ分割を行います。

1. エントロピー計算: 画像の各領域について、ピクセル強度の分布に基づいてエントロピーを計算します。
   • 低エントロピー（均一な領域）: 空の壁や背景など、情報が少ない領域には大きなパッチ（例：64x64）を割り当てます。
   • 高エントロピー（複雑な領域）: 物体の輪郭やテクスチャなど、情報量が多い領域には小さなパッチ（例：16x16）を割り当てます。
2. 閾値による決定: 各スケールレベルでエントロピー閾値 $\tau$ を設定し、閾値以下の領域はさらに大きなパッチとしてマージされます。

2.2 異なるパッチサイズの埋め込み統合

異なるサイズのパッチを ViT の入力として統一された埋め込み次元（$d_{embed}$）に変換する必要があります。

• サブパッチの埋め込み: 大きなパッチを、標準的なパッチサイズ（例：16x16）のサブパッチに分割し、それぞれを標準的なパッチ埋め込み層 $E$ で処理します。
• 畳み込みによる集約: 分割されたサブパッチの埋め込みを、畳み込み層（Conv2d）で下采样し、元の大きなパッチのサイズに合わせて集約します。
• ゼロ初期化 MLP (Zero-initialized MLP): 集約された埋め込みと、大きなパッチをリサイズして得た埋め込みを結合します。ここで、ControlNet のアイデアを応用し、重みをゼロで初期化された MLPを使用します。
  • 効果: これにより、モデルは初期段階ではリサイズされた埋め込み（元の ViT と同等の動作）に依存し、微調整（Fine-tuning）を通じて高解像度の詳細情報を徐々に取り入れることができます。これにより、事前学習済みモデルからの転移が容易になり、1 エポック程度の微調整で収束します。

2.3 可変長の入力処理

APT では画像ごとにトークン数が異なるため、バッチ処理には**シーケンスパッキング（Sequence Packing）**を採用します。

• 異なる長さのシーケンスを連結し、ブロック対角マスクを使用して、異なる画像のトークン間の注意を遮断します。
• FlashAttention や xFormers などの最適化されたアテンションバックエンドと互換性があり、オーバーヘッドを最小限に抑えます。

2.4 密な視覚タスクへの適応

物体検出やセマンティックセグメンテーションなどのピクセルレベルのタスクでは、特徴マップの形状維持が必要です。

• APT は、大きなパッチに対応するトークンを、そのパッチの面積比（$2^{2i}$ 倍）だけ反復複製することで、完全な微分可能な特徴マップを再構築します。
• ウィンドウアテンション（Window Attention）との互換性も確保しており、高解像度の密なタスクでも適用可能です。

---

3. 主要な貢献

1. APT の提案: 画像の内容に応じた適応的なパッチサイズを用いることで、ViT の学習・推論を最大 40% 高速化する新しいアーキテクチャを提案しました。
2. 事前学習モデルへの即応性: 既存の事前学習済み ViT（ImageNet などで微調整済み）に対して、ゼロ初期化 MLP を用いることで、追加の学習なし、あるいは1 エポックの微調整のみで元の精度を回復・維持できることを示しました。
3. 多様なタスクでの汎用性: 画像分類だけでなく、Visual QA (VQA)、物体検出、セマンティックセグメンテーションなど、高解像度かつ密な視覚タスクにおいても、精度を損なわずに大幅な高速化を実現することを証明しました。

---

4. 実験結果

実験は ImageNet 分類、VQA (LLaVA)、物体検出 (COCO)、セマンティックセグメンテーション (ADE20K) などで実施されました。

• ImageNet 分類:
  • ViT-L/14 (336x336): 学習速度が**61%**向上（スループット 527 img/s、GFLOPS 76.8）。精度はベースライン（88.1%）と同等（88.1%）を維持。
  • ViT-H/14 (336x336): 学習速度が**50%**向上。精度は 88.4%（ベースライン 88.5%）。
  • 高解像度・大規模モデルほど高速化の恩恵が大きくなります。
• 1 エポック微調整:
  • 事前学習済みモデルに対して、APT を適用し 1 エポックだけ学習させるだけで、ベースラインと同等の精度を達成しました。これに対し、既存手法（DynamicViT など）は 50 エポック以上の学習を必要とします。
• VQA (LLaVA-1.5):
  • 視覚トークンを削減することで、視覚バックボーンとクロスアテンション層の処理を高速化。**22-23%**の推論速度向上を実現し、複数のベンチマークでベースライン以上の精度を記録しました。
• 物体検出・セグメンテーション:
  • 高解像度（1536x1536）の物体検出タスクで**30%**の高速化を達成し、mAP 精度を維持しました。
  • セグメンテーションタスクでも、28-32% のトークン削減により高速化しつつ、ピクセルレベルの詳細を保持しました。
• オーバーヘッド分析:
  • パッチの再配置やマスク処理によるオーバーヘッドは存在しますが、トークン削減によるアテンション計算の削減効果がこれを上回り、実質的な高速化が得られています。

---

5. 意義と結論

技術的意義:

• 効率性のパラダイムシフト: 従来の「固定パッチ」や「層内でのトークン削減」ではなく、「入力段階での適応的パッチ化」によって、ViT の根本的な非効率性を解決しました。
• 実用性の高さ: 複雑な学習プロセスを必要とせず、既存のモデルを「すぐに」高速化できる点は、実社会での導入障壁を大きく下げます。
• リソース削減: 大規模モデルや高解像度処理における計算コスト（FLOPs）と時間を大幅に削減し、環境負荷の低減と推論コストの削減に寄与します。

限界と将来の展望:

• 現在の手法はエントロピーに基づくヒューリスティックに依存しており、タスク固有の最適化（例：検出タスクではエッジ情報が重要など）には手動の閾値調整が必要です。
• 現在は画像理解タスクに特化しており、画像生成タスクへの適用は今後の課題です。

総じて、APT は Vision Transformer のスケーラビリティと効率性を飛躍的に向上させる有望な手法であり、高解像度・大規模モデル時代における計算リソースの最適化に重要な貢献を果たすと考えられます。