技術要約: CAViT - 動的な特徴融合のためのチャネル認識型ビジョン・トランスフォーマー

問題提起

Vision Transformer (ViT) は、自己注意機構（Self-Attention）を利用して長距離の空間的依存関係をモデル化することで、コンピュータビジョンの分野で強力なモデルとしての地位を確立してきました。しかし、著者らは標準的なViTにおける根本的なアーキテクチャの非対称性を指摘しています。すなわち、空間的な混合（Spatial Mixing）はマルチヘッド自己注意（MHSA）を通じて動的かつデータ依存的である一方、チャネル混合（Channel Mixing）は静的であるという点です。従来のViTにおいて、チャネル間の相互作用は固定された多層パーセプトロン（MLP）によって処理されますが、これは入力内容を考慮せずに特徴量を線形結合するに過ぎません。この静的なアプローチは、多様な視覚構造への適応能力を制限する可能性があります。なぜなら、特徴次元は意味のあるセマンティック・フィルタ（テクスチャやパーツなど）を表しており、それらはグローバルな画像コンテキストに基づいて動的に再校正されるべきだからです。SE-Netや動的MLPといった既存の試みは、多くの場合、追加のモジュールや構造的なオーバーヘッドを導入します。

手法: CAViT アーキテクチャ

著者らは、標準的なTransformerブロック内の静的なMLP層を、特徴量相互作用のための動的なアテンションベースのメカニズムに置き換えた、デュアルアテンション・アーキテクチャである CAViT (Channel-Aware Vision Transformer) を提案しています。核心となる革新は、コンテンツ認識型のチャネル混合を実現するために、特徴チャネルをトークンとして扱うことにあります。

主要なアーキテクチャ構成要素

1. デュアルアテンション・メカニズム: 各CAViTブロックは、2つの連続するアテンション段階を実行します。
   • 空間自己注意 (Spatial Self-Attention): 標準的なMHSAが、通常通り空間トークンに対して動作します。
   • チャネル自己注意 (Channel Self-Attention): 空間ステージに続き、入力テンソルは次元の入れ替え（Dimension Swap）が行われます。空間トークン ($N$) とチャネル次元 ($C$) が転置され、実質的にチャネルをシーケンス・トークンとして扱います。
2. 次元の入れ替えとCLSの取り扱い:
   • 入力テンソル $x \in \mathbb{R}^{B \times (N+1) \times C}$ は、分類トークン ($x_{cls}$) と空間トークン ($x_{spat}$) に分割されます。
   • 空間トークンは $\mathbb{R}^{B \times N \times C}$ から $\mathbb{R}^{B \times C \times N}$ へ転置されます。
   • $x_{cls}$ は新しいシーケンス軸に合わせて整形され、連結されます。これにより、テンソルの形状は $\mathbb{R}^{B \times (C+1) \times N}$ となります。
   • 極めて重要な点として、$x_{cls}$ は次元の入れ替えにおいて、そのグローバルな集約ロールを損なうことなく、チャネル・トークンとして振る舞うように明示的に管理されます。これにより、セマンティックな整合性を維持したまま、グローバルなチャネル間相互作用に参加させることができます。
3. シングルヘッド・チャネルアテンション (SHSA): 空間ステージがマルチヘッド自己注意を使用するのに対し、チャネルステージはシングルヘッドメカニズムを採用しています。著者らは、転置された空間における各チャネル・トークンはすでにグローバルな画像コンテキストをエンコードしているため、アテンションを複数のヘッドに分割すると、この抽象化が希薄化する可能性があると主張しています。SHSAは過度な断片化を避け、各チャネル・トークン内に完全な空間コンテキストを保持します。
4. 反転 (Reversal): チャネルアテンションのステージ後、次元は再び入れ替えられ、$x_{cls}$ が再挿入されます。これにより、次のブロックは空間およびチャネルレベルの両方の相互作用によって強化された表現を受け取ることができます。

主な貢献

• アーキテクチャの変更: 標準的なViTブロック内の静的なMLPを、次元の入れ替えによるチャネル単位のアテンションへと置き換えることで、モデルの深さを増すことなく、適応的なチャネル間相互作用を可能にするシンプルかつ汎用的な置換を実現しました。
• 統一フレームワーク: 空間およびチャネル方向のトークン混合を、完全にアテンション駆動型のフレームワークの下で統一するデュアルアテンション定式化により、アーキテクチャの単純さを維持しつつ表現力を高めました。
• 理論的根意: 動的なチャネル混合（アテンションによるもの）が、特徴次元間の情報フローと特徴の整列を改善することを実証し、静的な線形結合の限界と対比させました。
• 効率性と性能: このアプローチが、パラメータ数と計算コスト（FLOPs）を削減しながら精度を向上させることを実証的に検証しました。

実験結果

著者らは、自然画像および医療画像にわたる5つの多様なデータセットを用いて、CAViT（特にCAViTtinyバリアント）を標準的なViTtinyベースラインと比較評価しました。

• データセット: CIFAR-10 (自然), Cats vs Dogs (自然), Malaria (顕微鏡), PneumoniaMNIST (X線), BreastMNIST (超音波)。
• 性能向上: CAViTは、Top-1精度において一貫した向上を達成しました。顕著な利得は以下の通りです：
  • CIFAR-10 で +3.64%
  • BreastMNIST で +2.50%
  • PneumoniaMNIST で +1.14%
  • Malaria および Cats vs Dogs では同等の性能。
• 効率性: 提案モデルは、ベースラインと比較して、パラメータ数を約 32%（5.75M から 3.91M へ）、FLOPs を 33%（2.26G から 1.52G へ）削減しました。
• 定性的分析: DINOスタイルの可視化によって生成されたアテンションマップは、CAViTがより鮮明で意味のある活性化パターンを生み出すことを明らかにしました。エッジに集中したりノイズを示したりすることが多かったベースラインのViTとは異かり、CAViTは微細な構造、医療スキャンにおける病変領域、自然界の物体の境界を正確に強調することに成功しました。

アブレーション研究の洞察

• 空間 vs チャネル: 空間的なMHSAを完全にチャネル指向のSHSAに置き換えると、性能が大幅に低下しました。これは、空間アテンションが位置依存性のために不可欠であることを裏付けています。
• ヘッド数: チャネルブランチにマルチヘッド自己注意を使用した場合は、シングルヘッドよりもわずかに性能が悪化しました。これは、統合されたグローバル記述子がチャネル混合においてより効果的であることを示唆しています。
• CLSトークンの整合性: CLSトークンの次元を無差別に（標準的なチャネル・トークンとして）入れ替えると、性能が著しく低下しました。これは、提案された特定のハンドリング戦略の必要性を検証するものです。

重要性と主張

本論文は、CAViTが、Vision Transformerにおけるチャネル混合の静的な性質に対処することで、精度と効率性の良好なトレードオフを提供すると主張しています。著者らは、この研究を「構造化されたアテンションベースのトークン混合への第一歩」と位置づけ、複雑なモジュールを追加することなく、次元の入れ替えという最小限のアーキテクチャ変更を通じて、動的でコンテンツ適応型の特徴融合が可能であることを示しました。

その重要性は、計算オーバーヘッドを削減しながら、表現力と解釈性（より明確なアテンションマップによる）を向上させる能力にあります。著者らは、これらの利得が、アーキテクチャの深化、追加のパラメータ、あるいは強力な学習正規化なしに達成されていることを強調しており、設計の核心的な有用性はアテンション機構自体に内在していることを示唆しています。この研究は、自然画像および医療画像の双方に適用可能であり、将来的な大規模バックボーン（例：ImageNet）やマルチモーダル設定への拡張を視野に入れています。