Beyond Self-Attention: Sub-Quadratic Vision Transformers for Fast Image Captioning

技術要約：自己注意（Self-Attention）を超えて：高速な画像キャプション生成のための劣二次（Sub-Quadratic）ビジョン・トランスフォーマー

問題提起
画像キャプション生成は、視覚的な内容から一貫性があり、意味のあるテキスト記述を生成することを必要とし、コンピュータビジョンと自然言語処理の架け橋となる。Transformerベースのアーキテクチャは、自己注意メカニズムを通じて長距離の文脈依存性を捉えることで、この分野を大きく進展させてきたが、主に2つの制限に直面している。第一に、標準的な自己注意メカニズムは、すべての画像パッチ間のペアワイズな相互作用を計算するため、二次的な計算複雑度（$O(n^2)$）をもたらす。これは、高解像度の画像や大規模なアプリケーションにおいて、相当なメモリおよび計算のオーバーヘッドを生じさせる。第二に、既存のアプローチは、豊かな局所的特徴表現の欠如や計算の冗長性に苦しむことが多い。なぜなら、多くのパッチが意味的に重要な情報を欠いている場合でも、すべてのパッチを一様に処理してしまうからである。

手法
提案されたアーキテクチャは、従来の自己注意メカニズムをガウス混合モデル（GMM）に基づく確率的クラスタリングに置き換えることで、ビジョン・トランスフォーマー（ViT）の枠組みを再構築し、これらの制限に対処する。システムは主に2つのステージで動作する：

1. 画像エンコーディング（GMM-Transformer）：

   • パッチ埋め込み（Patch Embedding）： 入力画像は重なり合わないパッチに分割され、$D$次元の埋め込みに投影され、位置エンコーディングが付加される。
   • 文脈重み付きGMMクラスタリング： ペアワイズの注意スコアを計算する代わりに、モデルはパッチ埋め込みを、$K$個の多変量ガウス分布（$K \ll N$）の混合によって生成される連続的な確率変数として扱う。
   • 学習可能な文脈重み： サブレイヤーは、統計的特性（平均と分散）によって調整される学習可能な空間重みをパッチに割り当て、モデルが意味的に重要な領域を優先できるようにする。
   • EMアルゴリズム： ソフトクラスタリングを実行するために、期待値最大化（EM）アルゴリズムが採用される。Eステップでは、モデルは各パッチがどの程度クラスターに属しているかを示す「責任（responsibilities）」（ソフト割り当て確率）を計算する。Mステップでは、クラスターのパラメータ（平均、共分散、混合係数）が更新される。
   • 複雑度の削減： このプロセスにより、$N \times N$の注意行列が$N \times K$の責任行列に置き換わり、計算複雑度が二次的な$O(n^2)$から、ほぼ線形な$O(nK)$へと削減される。
   • 集約： 精緻化されたパッチ特徴は、残差接続とレイヤー正規化を伴うフィードフォワードネットワークを通じて処理される。最終的に、グローバル平均プーリング操作によって、クラスターの重心がコンパクトなグローバル特徴ベクトルへと集約される。

2. 言語デコーディング（統合されたGPTベースのデコーダー）：

   • デコーダーは、自己回帰的なGPTスタイルのTransformerを利用する。
   • 統合された注意（Unified Attention）： マスク付き自己注意とクロス注意を分離する従来のエンコーダ・デコーダモデルとは異なり、本アーキテクチャはこれらを単一のマルチヘッド注意ブロックへと統合している。
   • メカニズム： テキストトークンはクエリベクトルとして機能し、視覚的トークン（クラスター重心）とそれ以前のテキストトークンの結合されたシーケンスが、キーおよび値のベクトルとして機能する。
   • 因果的マスキング（Causal Masking）： 特殊な因果的マスクにより、テキストトークンはすべての視覚的トークンに注意を向けることができる一方で、自身より前に生成されたテキストトークンにのみ注意を向けることが保証され、自己回帰的な生成特性が維持される。
   • 生成： 出力はソフトマックス層を通じて投影され、次に来るトークンを反復的に予測する。

主な貢献

• 複雑度の削減： 自己注意をGMMベースのソフトクラスタリングに置き換えることで、コンテキスト的に一貫したキャプションを生成する能力を維持しつつ、トランスフォーマー・エンコーダの時間複雑度を二次的なものからほぼ線形（$O(nK)$）へと削減した。
• 効率的なデコーダー設計： マスク付き自己注意とクロス注意を単一の計算ブロックへと統合したカスタマイズされたGPTベースのデコーダーを採用し、マルチモーダルな相互作用を強化し、冗長性を低減した。
• 実証的検証： Flickr30kデータセットを用いた包括的な評価により、モデルが既存の最先端（SOTA）の成果物に対して競争力のある性能を発揮し、計算効率とキャプションの品質のバランスを実現していることを示した。

実験結果
モデルは、8層のトランスフォーマー、4つの注意ヘッド、および768の隠れ次元の構成を用い、Flickr30kデータセット（Karpathyスプリットを使用）で評価された。

• 性能指標： 提案モデルは、高度な意味的指標においてトップスコア（BLEU-4: 0.364、METEOR: 0.368、SPICE: 0.199）を達成した。これらの結果は、いくつかのSOTAモデル（例：TAKG、GNN-Visual、Scene Graph）と比較して、優れたフレーズ構造とシーングラフのマッピングを示している。
• ベースライン比較： モデルは、基本的な精度指標（BLEU-1: 0.733 vs 0.794、ROUGE: 0.568 vs 0.608、CIDEr: 0.723 vs 0.946）においてADAモデルに後塵を拝しているものの、TopDownやGNN-Visualなどの他のベースラインよりも優れた性能を示している。著者らは、このアーキテクチャが単純な複製よりも深い文脈を優先していると述べている。
• 効率性： モデルは約21.17 GFLOPsの計算コストで動作し、約3,280万個の学習可能パラメータを含んでいる。混同行列の分析では、高い真陽性率（45,580）と低い偽陽性・偽陰性数が示され、強力な分類能力が示された。

意義と主張
本論文は、提案されたフレームワークが標準的なビジョン・トランスフォーマーに固有の計算上の非効率性をうまく解決していると主張している。ガウス混合モデルと期待値最大化（EM）アルゴリズムを活用することで、モデルは意味のある特徴関係を捉える能力を損なうことなく、計算複雑度を二次的なものからほぼ線形へと大幅に削減することに成功した。著者らは、モデルが強力な予測能力とスケーラブルな特徴エンコーディングを示しており、キャプションの品質、計算効率、およびスケーラビリティの間の改善されたバランスを提供すると断言している。今後の展望として、高度なマルチモーダル事前学習の組み込み、ドメイン特化型データセットへの適応、およびリアルタイムのエッジアプリケーション向けの軽量モデルの展開が挙げられている。