Two-stage Vision Transformers and Hard Masking offer Robust Object Representations

この論文「Two-stage Vision Transformers and Hard Masking offer Robust Object Representations（2 段階のビジョン・トランスフォーマーとハード・マスキングは頑健な物体表現を提供する）」は、深層学習モデルが物体認識において背景などの文脈情報に依存しすぎることによる「偽の相関（spurious correlations）」の問題を解決し、分布外（OOD: Out-of-Distribution）の背景に対しても頑健なモデルを構築するための新しいアーキテクチャを提案しています。

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

1. 問題定義：文脈バイアスと既存手法の限界

深層学習モデル、特に物体認識タスクでは、背景の文脈（例：犬は屋外、猫は屋内）を利用して分類を行う傾向があります。これは訓練データでは有効ですが、推論時に背景と物体の相関が崩れる（例：屋内にいる犬）と、モデルは誤った予測を行います。

既存の「注意機構（Attention Mechanism）」を用いたアプローチには以下の 2 つの重大な欠点があることが指摘されています。

1. 遅延マスキング（Late Masking）の問題: 既存手法は、画像全体から抽出された高レベルな特徴量（Feature Tensor）に対してマスキングを適用します。しかし、現代のアーキテクチャ（ViT や CNN）は受容野（Receptive Field）が広いため、マスキングを適用する時点ですでに背景情報が特徴量に混入（汚染）してしまっています。
2. ソフトマスキングの漏洩: 多くの手法が「ソフト（連続値）」の注意マスクを使用します。これにより、無関係な領域にも 0 ではない重みが割り当てられ、情報が完全に遮断されず、残りの漏洩（Residual Leakage）が発生します。

2. 提案手法：iFAM (Inherently Faithful Attention Maps)

著者らは、**「早期マスキング（Early Masking）」と「2 段階アーキテクチャ」**を導入することで、モデルが背景情報にアクセスできないように物理的に制限するアプローチを提案しました。

2.1 2 段階フレームワーク

モデルは 2 つの段階で構成され、これらはjointly（共同）で訓練されます。

• Stage 1（Selector / 選択器）:
  • 入力画像全体を処理し、タスクに関連する前景領域（物体の一部）を特定します。
  • PDiscoFormer [2] をベースにし、画像を $K+1$ 個の領域（$K$ 個の前景部分＋背景）に分割する「パート発見（Part Discovery）」を行います。
  • 出力は、どのトークン（画像パッチ）が前景に属するかを示す離散的なバイナリマスクです。
• Stage 2（Predictor / 予測器）:
  • Stage 1 で選択された前景トークンのみを入力として受け取ります。
  • Vision Transformer (ViT) の自己注意（Self-Attention）機構を修正し、マスクされたトークン（背景）との情報交換を物理的に遮断します。
  • これにより、予測器の受容野は厳密に選択された前景領域に限定され、背景からの情報漏洩が不可能になります。

2.2 技術的詳細

• ハードマスキングの実装: ViT のアテンション計算において、マスクされたトークン（$s_i=0$）と他のトークン間のアテンションスコアを $-\infty$ に設定し、softmax 後の重みを 0 に強制します（式 2, 3）。
• 離散化と勾配: パートマップを離散的なバイナリマスクに変換するために、Gumbel Softmax を使用し、フォワードパスではハードマスク、バックワードパスではソフトマスクを使用する「Straight-Through Gradient Trick」を採用しています。
• レジスター・トークン: 背景トークンを削除しても学習が不安定にならないよう、ViT に「レジスター・トークン」を追加し、これらが前景トークンのみに注意を向けるように制限しています。

2.3 監査可能なテスト時介入（Auditable Test-Time Interventions）

モデルの透明性を活かし、推論時に人間が介入して頑健性をさらに高める 2 つの戦略を提案しています。

1. パート除去介入（Part-Removal Intervention）: 学習されたパートの中に偽の相関（例：背景の特定の要素）を捉えているものがある場合、ユーザーが手動または自動（LOO 解析）でそのパートを除外し、Stage 2 に渡さないようにします。
2. トークン除去介入（Token-Removal Intervention）: 分布外の物体が誤って前景として検出された場合、そのパートの原型（Prototype）からの距離が閾値を超えたトークンを「背景」として再分類・除去します。

3. 主要な貢献

1. 2 段階アーキテクチャの提案: 全画像処理から選択的処理へ移行し、偽の相関を軽減する新しい設計。
2. 早期マスキングによる頑健性の保証: 受容野をタスク関連領域に物理的に制限することで、背景からの情報漏洩を根本的に防ぎ、忠実性（Faithfulness）を保証する。
3. 監査可能な介入機能: 学習済みモデルに対して、追加の訓練なしで人間が介入し、頑健性を向上させるユニークな機能の実装。
4. SOTA 性能の達成: 複数のベンチマーク（MetaShift, Waterbirds, ImageNet-9, SIIM-ACR など）において、既存のデバイアス手法や大規模モデルを上回る OOD 頑健性を達成。

4. 実験結果

多様なベンチマークでの評価により、提案手法の有効性が確認されました。

• MetaShift & Waterbirds:
  • MetaShift（犬 vs 猫）において、WGA（最悪グループ精度）を 88.6%（既存の最良手法は 83% 前後）に向上。
  • Waterbirds（水鳥 vs 陸鳥）において、WGA を 97.0%（既存手法は 93.9% 前後）に達成。
  • 後段マスキング（Late Masking）やソフトマスクを使用する手法は、背景バイアスに依然として脆弱であることを示しました。
• ImageNet-9 Backgrounds Challenge:
  • 背景への依存度を測る BG-GAP メトリックにおいて、iFAM は 2.4 という最低値（最も良い値）を記録。ViT-L や ResNet-152 などの大規模モデルよりも優れた結果を示し、性能向上がモデル容量ではなくアーキテクチャによるものであることを証明しました。
• 医療画像（SIIM-ACR）:
  • 偽の相関（胸チューブなど）に依存しやすいタスクにおいて、WGA AUC を 65.9%（ベースラインは 48.1%）まで向上。介入を組み合わせることで 69.0% まで改善しました。
• アブレーション研究:
  • 2 段階構造の除去、ソフトマスクの使用、パート形状損失の除去などを行うと、OOD 性能が大幅に低下することが確認され、提案手法の各要素が不可欠であることが示されました。

5. 意義と結論

この論文は、**「モデルが背景情報にアクセスできないように物理的に制限すること」**が、偽の相関に対する頑健性を高めるための鍵であることを示しました。

• 理論的意義: 従来の「説明可能性（Explainability）」や「注意機構」が単なる可視化ツールや後処理に留まっていたのに対し、アーキテクチャレベルで「忠実な（Faithful）」推論を強制するアプローチの重要性を浮き彫りにしました。
• 実用的意義: 医療画像診断や自律走行など、背景バイアスが致命的なエラーにつながる分野において、人間がモデルの判断根拠を監査し、必要に応じて介入できる「人間-AI 協調」の新しいパラダイムを提供します。
• 限界と展望: 2 段階処理による計算コストの増加は課題ですが、推論時のスパース性（前景のみ処理）を利用することで、実用上のオーバーヘッドは許容範囲内であることが示唆されています。

総じて、この研究は Vision Transformer の頑健性を飛躍的に向上させるだけでなく、AI の意思決定プロセスを透明化・制御可能にする重要なステップとなっています。