kMaX-DeepLab: k-means マスク変換器

ビジョンタスクにおけるトランスフォーマーの台頭は、ネットワークバックボーンの設計を進歩させるだけでなく、オブジェクト検出やパノプティックセグメンテーションなどのエンドツーエンド画像認識に新たなページを開くこととなりました。自然言語処理（NLP）から発祥したトランスフォーマー構造は、自己注意（self-attention）とクロス注意（cross-attention）によって、シーケンス内の要素間の長距離相互作用を効果的に学習します。しかし、既存のトランスフォーマーを基盤とするビジョンモデルの多くが単にNLPからのアイデアを借用し、言語と画像との間に存在する重要な違い、特に空間的に平坦化されたピクセル特徴量の非常に長いシーケンス長を無視していることを観察しました。これにより、ピクセル特徴量とオブジェクトクエリ間のクロス注意学習が阻害されます。本論文では、ピクセルとオブジェクトクエリの関係を見直し、クロス注意学習をクラスタリングプロセスとして再定式化することを提案します。伝統的なk-meansクラスタリングアルゴリズムに着想を得て、セグメンテーションタスク向けにk-meansマスクトランスフォーマー（kMaX-DeepLab）を開発しました。この手法は最先端技術を向上させるとともに、シンプルで優雅な設計も享受しています。その結果、kMaX-DeepLabはテスト時の拡張や外部データセットを使用せずにCOCO valセットで58.0% PQ、Cityscapes valセットで68.4% PQおよび44.0% APおよび83.5% mIoU、ADE20K valセットで50.9% PQおよび55.2% mIoUという新しい最先端性能を達成しました。我々は本研究がビジョンタスク専用のトランスフォーマー設計において新たな洞察を与えることを願っています。TensorFlowコードとモデルはhttps://github.com/google-research/deeplab2 で提供されており、PyTorchによる再実装もhttps://github.com/bytedance/kmax-deeplab で利用可能です。