BiFormer:二段階ルーティングアテンションを備えたビジョンTransformer

視覚変換器(Vision Transformer)の核心的な構成要素であるアテンションは、長距離依存関係を捉える強力な手法である。しかし、その強力な性能は代償を伴う:すべての空間位置間でペアワイズなトークン相互作用を計算する必要があるため、膨大な計算負荷と大きなメモリ消費が生じる。この問題を軽減するために、多くの研究が手作業で設計された、コンテンツに依存しないスパース性をアテンションに導入するアプローチを提案している。具体的には、アテンション演算を局所的なウィンドウ内、軸方向のストライプ内、または拡張されたウィンドウ内に制限する手法が用いられている。これに対して、本研究では、二段階ルーティング(bi-level routing)を用いた新しい動的スパースアテンションを提案し、コンテンツに応じた柔軟な計算割り当てを可能にする。具体的には、クエリに対してまず粗い領域レベルで関係のないキー・バリュー対をフィルタリングし、残された候補領域の和集合(すなわちルーティングされた領域)内で、細粒度なトークン間アテンションを適用する。本研究で提案する二段階ルーティングアテンションは、シンプルでありながら効果的であり、スパース性を活用することで計算量とメモリ使用量の両方を削減しつつ、GPUに適した密行列積み算のみを用いる実装が可能である。この二段階ルーティングアテンションを基盤に構築された、新たな汎用的な視覚変換器「BiFormer」を提案する。BiFormerは、他の無関係なトークンの干渉を受けることなく、クエリに応じて関連するトークンの小さなサブセットに注目するため、特に密度予測タスクにおいて優れた性能と高い計算効率を両立する。画像分類、物体検出、セマンティックセグメンテーションといった複数のコンピュータビジョンタスクにおける実証的評価により、本設計の有効性が確認された。コードは以下のURLから入手可能である:\url{https://github.com/rayleizhu/BiFormer}。