[論文レビュー] Vision Transformers Need Registers

Timothée Darcet, Maxime Oquab+2 more|arXiv (Cornell University)|Sep 28, 2023

Domain Adaptation and Few-Shot Learning参考文献 32被引用数 48

ひとこと要約

本論文は Vision Transformer の特徴マップにアーティファクトを引き起こす高ノルムのアウトライヤー・トークンを同定し、入力列へ学習可能なレジスタ・トークンを追加することでこれらのアーティファクトを除去し、監視付き・テキスト監視付き・自己監督型の ViT において密な予測と物体発見を改善することを示す。

ABSTRACT

Transformers have recently emerged as a powerful tool for learning visual representations. In this paper, we identify and characterize artifacts in feature maps of both supervised and self-supervised ViT networks. The artifacts correspond to high-norm tokens appearing during inference primarily in low-informative background areas of images, that are repurposed for internal computations. We propose a simple yet effective solution based on providing additional tokens to the input sequence of the Vision Transformer to fill that role. We show that this solution fixes that problem entirely for both supervised and self-supervised models, sets a new state of the art for self-supervised visual models on dense visual prediction tasks, enables object discovery methods with larger models, and most importantly leads to smoother feature maps and attention maps for downstream visual processing.

研究の動機と目的

• 監視付きおよび自己监督モデルを横断する ViT の特徴マップにおけるアーティファクトを特定する。
• 高ノルムアウトライヤー・トークンの性質と出現位置を特徴づける。
• アーティファクトを緩和するために register tokens を用いた単純なアーキテクチャ的対策を提案する。
• 密な予測性能の改善と滑らかな特徴/アテンションマップを実証する。

提案手法

• アテンションマップのアーティファクトを分析し、高ノルムのトークンをアウトライヤーとして定義する（ノルム > 150）。
• トレーニング中およびモデルサイズ別に、アウトライヤーがいつどこに現れるかを調べる。
• パッチ埋め込みに対して線形モデルを用いて高ノルム・トークンの情報量を検証する。
• パッチ埋め込みの後に追加される N 個の学習可能な入力トークン（registers）を導入してアウトライヤーを吸収する。
• register の有無で線形プロービング、ゼロショット分類、物体発見などの下流タスクへの影響を評価する。
• レジスタの数をアブレーションして性能とアーティファクト低減を検討する。）

実験結果

リサーチクエスチョン

• RQ1 監督タイプを横断した ViT のアテンションマップにアーティファクトが生じる原因は何か？
• RQ2 高ノルムのアウトライヤー・トークンは局所情報を持つのか、それとも全体情報を持つのか、またそれが密予測タスクにどう影響するのか？
• RQ3 トレーニングパラダイムを横断して、レジスタ・トークンを追加してアーティファクトを除去できるか、下流の性能を害さずに？
• RQ4 レジスタの数はアーティファクト抑制と下流タスクの性能にどう影響するか？

主な発見

• アーティファクトは、ノルムが非常に大きく結果的に全体の約2%程度のトークンに対応し、大規模 ViTs の長い学習後に中間層に現れる。
• 高ノルム・トークンはパッチ位置やピクセルに関する局所情報をあまり含まないが、画像全体に関するより多くのグローバル情報を保持する。
• パッチ埋め込み後に学習可能な register tokens を追加すると高ノルム・アウトライヤーが排除され、特徴マップ/アテンションマップが平滑になる。
• register を用いて訓練されたモデルは ImageNet 線形プロービング、ADE20k セグメンテーション、NYUd 深度推定で性能の低下を示さず、場合によってはわずかな改善を示す。
• register ベースのモデルは DeiT-III, OpenCLIP, DINOv2 のバックボーンでの教師なし物体発見（例: LOST）を改善し、1 つのレジスタでアーティファクトを除去でき、より多くのレジスタは密タスクに有利となる可能性がある。