技術概要：高解像度セマンティックセグメンテーションのための STEP（SuperToken および Early-Pruning）

1. 問題提起

ビジョン・トランスフォーマー（ViT）は、長距離依存関係を捉える能力により、セマンティックセグメンテーションにおいて最先端の性能を達成しています。しかし、その実用的な展開は、主に画像解像度に対する自己注意メカニズムの二次的な複雑さに起因する、高い計算コストとメモリ消費によって阻害されています。入力解像度が増加する（例えば 1024×1024 へ）につれて、計算コストとメモリ消費は著しく増大します。

既存のトークン削減戦略は、密な予測タスクにおいて特定の限界に直面しています：

• パッチレベルの手法（例：CTS）は、しばしば固定されたハイパーパラメータ（例：マージされるパッチの固定数）に依存しており、複雑な画像では過剰マージ、単純な画像では過少マージを引き起こし、最適ではない性能をもたらします。
• トークンレベルのプルーニング（ハードプルーニング）は、トークンを永続的に削除するため、ピクセルごとの出力のために空間情報を保持する必要がある密な予測にはしばしば不適切です。
• ハイブリッドアプローチは、しばしば変化する画像コンテンツへの適応性が欠如しているか、高解像度入力の特定の拡張性課題に対処できていません。

本論文は、標準的なグリッドベースの処理が非効率である高解像度シナリオにおいて、計算効率と精度保持のバランスを確立する必要があると主張しています。

2. 手法：STEP フレームワーク

著者らは、精度を大幅に損なうことなく ViT の効率を向上させるために設計されたハイブリッド・トークン削減フレームワークであるSTEP（SuperToken and Early-Pruning）を提案します。STEP は、2 つの相補的なレベルで動作します：

2.1 動的 CTS（dCTS）：適応的パッチマージ

静的なマージ戦略とは異なり、STEP は EfficientNetLite0 に基づく軽量でクラス非依存の方針ネットワークであるdCTSを導入します。

• メカニズム：このネットワークは、$2\times2$、$4\times4$、$8\times8$、および$16\times16$のウィンドウといった隣接パッチのグループを評価し、類似度スコア $S \in [0,1]$ を予測します。
• 決定ロジック：類似度スコアが事前に定義された閾値 $\tau$ を超える場合、パッチはスーパーパッチにマージされます。このプロセスは、ネストされたグループ間の競合を防ぐために、粗いものから細かいものへの階層に従います。
• 出力：結果として、非一様なサイズを持つスーパーパッチのグリッドが得られます。これらは双線形補間を介して標準解像度（16×16）にリサイズされ、スーパートークンとして埋め込み空間に投影されます。
• 利点：これにより、トークン数を画像コンテンツのセマンティックな均質性に動的に適応させることが可能になります（例：大きな背景領域をマージしつつ、物体境界では解像度を保持する）。

2.2 早期プルーニング機構

STEP は、「容易な」トークンの処理を停止させる DToP（Dynamic Token Pruning）に着想を得た早期退出戦略を統合します。

• アーキテクチャ：ViT エンコーダは $S$ ステージに分割されます。各ステージには、各トークンの信頼度スコアを計算する補助デコーディングヘッド（Attention-to-Mask モジュールを使用）が含まれます。
• プルーニングロジック：信頼度スコアが閾値 $\tau$ を超えるトークンはマスクされ、後続の層への伝播が停止されます。低い信頼度（困難な）トークンのみが、より深い層を通過し続けます。
• 統合：この戦略は dCTS と相補的です。dCTS が初期シーケンス長を削減する一方、早期プルーニングはネットワークが深くなるにつれてアクティブなトークンのセットを漸進的に短縮します。

3. 主要な貢献

本論文は、以下の具体的な貢献を概説しています：

1. 新規ハイブリッドフレームワーク：セマンティックセグメンテーションにおける効率性に対処するために、動的パッチマージ（dCTS）と早期トークンプルーニングを組み合わせる STEP の提案。
2. 高解像度評価：ViT-Base および ViT-Large バックボーンを使用した、高解像度画像（最大 1024×1024）における広範な実験。著者らは、高解像度セマンティックセグメンテーションの文脈において、トークンプルーニング戦略を体系的に評価した最初の試みであると主張しています。
3. 包括的な分析：マージの閾値パラメータ、補助プルーニングヘッドの最適な配置、計算複雑性（GFLOPs）、推論速度（FPS）、および精度（mIoU）の間のトレードオフに関する詳細なアブレーション研究。
4. スケーラビリティの実証：画像解像度とモデルサイズが増加するにつれて、STEP が競争力のある精度を維持しつつ、計算複雑性と推論時間を大幅に削減していることを示す証拠。

4. 実験結果

実験は、NVIDIA A100 GPU を使用して COCOStuff10k、ADE20K、および Cityscapes データセットで行われました。

4.1 標準解像度（512×512）における性能

• dCTS のみ：ViT-Large に適用した場合、dCTS は標準的な 16×16 パッチングと比較してトークン数を 2.5 倍削減します。これにより、精度の低下を最小限に抑えながら、計算コストを 2.6 倍削減し、スループットを 3.4 倍向上させます。
• 完全な STEP：dCTS と早期プルーニングの両方を適用することで、最大4 倍の計算複雑性の削減と1.7 倍の推論速度の向上を達成し、精度の低下は最大でも 2.0% 以内に抑えられます。
• 構成：ViT-Large の 16 層目または 18 層目に単一のプルーニングヘッドを配置することが、マルチヘッド構成と比較して効率性とレイテンシの面で有利なトレードオフを提供しました。

4.2 高解像度（768×768 および 1024×1024）における性能

• スケーラビリティ：解像度が増加するにつれて、標準的な SegViT は GFLOPs の劇的な増加と FPS の低下を招きます。STEP 構成は安定したトレードオフを維持します。
• 1024×1024 における ViT-Large：dCTS バリアントのみで、GFLOPs を 2086 から 802 に削減（2.6 倍の複雑性低下）し、FPS を 12 から 41 に増加（3.4 倍の高速化）させ、mIoU の低下は僅かでした。
• トークンプルーニングのダイナミクス：最大40% のトークンを最終エンコーダ層に到達する前に自信を持って予測し、停止させることができます。ViT-Large を用いた高解像度設定では、平均して 48% のトークンが早期に停止されます。

4.3 スループット対 FLOPs に関する観察

著者らは、重要な発見を指摘しています：スループット（FPS）は FLOP の削減と線形に比例してスケーリングしません。

• GFLOPs が大幅に減少する一方で、推論速度の向上は時として制限されます（例：GFLOP の 4 倍削減が常に FPS の 4 倍向上をもたらすわけではない）。
• 原因：早期プルーニングに必要なマスク操作は、不規則な制御フロー、動的なメモリアクセスパターン、テンソル形状の変動性によりボトルネックを引き起こします。本論文は、不要な GPU から CPU へのメモリ転送などのハードウェアの非効率性が、期待される性能向上を減衰させている可能性を指摘しています。

5. 意義と主張

本論文は、STEP が高解像度の密な予測タスクに対する ViT の実用化に向けた重要な一歩を表すと主張しています。

• 類例なき初：著者らは、セマンティックセグメンテーションにおいて異なる解像度レベルにわたってトークンプルーニングを体系的に評価した最初の研究であると主張しています。
• 効率と精度のバランス：STEP は、精度の低下を許容可能な範囲（<2.0%）に抑えながら、大幅な計算節約（最大 4 倍）と速度向上を達成可能であることを実証しています。
• 適応性：dCTS モジュールは特に堅牢であり、固定されたハイパーパラメータを必要とせずに画像の複雑さにマージ率を適応させるため、多様な現実世界のシーンに適しています。

著者らは、STEP が効率性とスケーラビリティにおいて明確な利点を提供する一方で、FLOP 削減によって予測される潜在的なスループット向上を完全に実現するためには、プルーニングメカニズム自体のハードウェアレベルの非効率性に対処する今後の研究が必要であると結論付けています。