Vision Transformer（ViT）にレジスタトークンが必要な理由

 (openreview.net)
1 ポイント 投稿者 GN⁺ 2024-05-13 | 1件のコメント | WhatsAppで共有
  • ViTの特徴マップに現れる高ノルムトークンは、主に情報量の少ない背景領域に現れ、パッチトークンが内部計算用の保存領域のように再利用されている現象として解釈できる
  • 解決策は、画像パッチとは別のレジスタトークンを入力シーケンスに追加し、内部計算の役割を専用トークンに分離する方式である
  • 教師あり学習・自己教師あり学習のViTの両方でアーティファクトが減少し、特徴マップとattention mapが滑らかになって、downstreamの視覚タスクに有利になる
  • object discoveryとdense predictionの改善が報告された一方、OpenCLIPではobject localization性能がわずかに低下し、「すべてのモデルで改善される」という表現は修正された
  • 4〜16個のレジスタトークンを入れると2〜6%の追加計算量が生じ、トークン数を増やしたからといってdense taskの性能が常に向上するわけではない

ViT特徴マップで明らかになったアーティファクト

  • Vision Transformerは視覚表現学習で強力なツールとして使われているが、教師あり学習および自己教師あり学習のViTの特徴マップでアーティファクトが確認されている
  • アーティファクトは推論中に現れる高ノルムトークンとして観察される
    • 主に画像内の情報量が少ない背景領域で発生する
    • ネットワークがそれらのトークンを内部計算に再利用しているように見える
  • これらのトークンは周囲のパッチと似た背景パッチから生じるが、局所的なパッチ情報は少なく、グローバルな画像情報をより多く含むと評価される
  • DINOv2や他のViTでは、attention mapの空間解像度や精度の低下が、downstreamのobject discoveryやlocalization性能に影響する可能性がある

レジスタトークン方式

  • 提案された解決法は、Vision Transformerの入力シーケンスに画像パッチとは別のレジスタトークンを追加することである
  • レジスタはCLS tokenのように画像パッチトークンとは分離されるが、loss計算には直接使われないトークンとして整理される
  • 目的は、既存のパッチトークンが内部計算用の「保存領域」のように使われないよう、その役割を専用トークンに担わせることにある
  • この方式は既存のMemory Transformerと関連があるが、特徴アーティファクトを分析し、それを修正するために新しいトークンを入れるというアプローチはこの論文の分析に特化している

特徴マップとdownstreamタスクの変化

  • レジスタを追加すると、教師あり学習モデルと自己教師あり学習モデルの両方で、該当するアーティファクトが除去されると報告されている
  • 特徴マップとattention mapがより滑らかになり、downstreamのvisual processingに有利になる
  • 自己教師あり視覚モデルのdense visual prediction taskで新たなstate of the artを達成したと整理されている
  • より大きなモデルでもobject discovery手法を使えるようになった
  • DINOv2とDeiT-IIIではobject discovery性能が向上し、画像レベルtaskの性能は概ね維持される

モデル別の例外と追加分析

  • OpenCLIPではレジスタを追加してもobject localizationの結果は改善せず、表の数値ではわずかな性能低下が見られる
    • 著者の回答では、「すべてのモデルでobject discovery性能が改善される」という表現は数値で裏づけられていないことを認めている
    • 該当内容はSection 3.3で修正された
  • OpenCLIPはFigure 7でノルムの減少が見られる一方、それがTable 3のobject localization改善にはつながっていないという不一致が議論されている
    • LOST評価では、モデルごとにkeys、queries、valuesの中で最も高性能なembeddingを選択する
    • CLIPではvaluesが最良で、artifactsはkeysやqueriesでは見られるがvaluesでは見られないため、Table 3および定性的分析と整合している
  • MAEでは同じ「outlier patches」が現れないように見える
    • MAEはlocal lossのみで学習されるため、グローバル情報を集める必要がない、あるいは少ないためかもしれない
    • MAE-ViT-Largeのlinear probingによるImageNet分類精度は75%で、他のSSL手法より低い

性能、コスト、限界

  • レジスタはシンプルな方法だが、追加トークンのため計算量が増える
    • 4〜16個のレジスタトークンを追加した場合、**2〜6%**の計算量増加が報告されている
  • レジスタ数の増加が常にdense taskの改善につながるわけではない
    • 0個から1個のレジスタに増やしたときにアーティファクトが除去され、segmentationとdepth predictionが改善する
    • 1個以降にさらに多くのレジスタを追加しても、feature mapはすでにクリーンになっているため、segmentationとdepth predictionの性能は大きく向上しない
    • classification性能はレジスタが多いほど追加改善する傾向があったが、その理由について明確な直感はまだない
  • NYU depth datasetでは、8個から16個のレジスタに増やしたときRMSEがわずかに増加する
    • シーケンス長の増加によって最適なhyperparameterが変わった可能性がある
    • 8個から16個に増やしたときの0.03 RMSE差は、0個から1個に増やしたときの0.1 RMSE差に比べて小さく、単なるノイズの可能性もある

OpenReviewでの採択とレビューの争点

  • OpenReviewでの判定はAccept (oral) である
  • メタレビューでは4人のレビュアー全員が8点を付けており、論文の技術的貢献が認められている
  • レビュアーたちは、性能の不一致、dataset biasとSSLモデルの関係、最適化設定の選択、gradient clippingの影響などを懸念していた
  • 著者らは反論プロセスで追加実験とappendix分析を加えた
  • その後の議論では、レジスタがoutlierの振る舞いを引き継ぐのか、image tokenがよりlocalな性質に戻るのか、CLS token normがどう変化するのかといった内部メカニズムの分析が補強された

関連記事

1件のコメント

 
GN⁺ 2024-05-13
Hacker News のコメント

  • 論文によると、register は Vision Transformer モデルの学習中に入力シーケンスへ追加される、追加の学習可能トークンです。
    パッチ埋め込み層の後に [CLS] トークンのような学習可能な値として追加され、Vision Transformer の最後では register トークンを捨てたうえで、[CLS] トークンとパッチトークンだけを画像表現として使います。
    このトークンは、順伝播の途中でモデルがグローバル情報を保存・処理・参照するための空間を提供し、パッチトークンをその用途に無理に使わなくて済むようにします。
    その結果、学習済み Vision Transformer の特徴マップに現れるアーティファクトや高ノルムの外れ値トークンが消え、特徴マップがより滑らかになり、dense prediction の性能や教師なし物体発見も向上します。
    推論コストは 2% しか増えないのに ViT の性能をかなり改善するので、ほとんど「ただ飯」に近い、すっきりした結果に見えます。

    • トークンをあちこちに足したら魔法のようにうまくいく、というこの流れはかなり怪しく見えるし、確率論まで絡んでくると、離散数学系の人たちは軽く流すのが難しそうです。
      確率的正規表現もその一例ですが、これは数学的な放埒さがまったく別のレベルです。
      それでもこう動くというのは本当に驚きです。
    • BERT に [CLS] トークンを複数追加しようとする試みもありましたが、結果はそれほど印象的ではありませんでした: https://arxiv.org/pdf/2210.05043
    • 一般的なテキスト Transformer にも、ここから学べることがあるのか気になります。
      attention sink や pause token に似て聞こえます。
    • こうしたトークンを使うには、最終層の結果をいったんトークンに戻し、次のパスで再び埋め込む必要があります。
      あるパスの生の出力をそのまま次の入力へ渡す方式を調べた人がいるのか気になります。

  • 関連していそうです: "Let's Think Dot by Dot: Hidden Computation in Transformer Language Models" https://arxiv.org/abs/2404.15758
    この論文は、言語モデルの chain-of-thought 性能向上が人間のようなタスク分解によるものなのか、それとも追加トークンが与える計算量によるものなのかは不明だと見ています。
    Transformer が意味のない埋め草トークン、たとえば ...... だけでも、中間トークンなしで答えると解けなかった難しいアルゴリズム課題を解けることを示し、ただしそうした埋め草トークンの使用を学習させるには、具体的で密な教師信号が必要だとしています。
    なので、テキスト LLM にも特別な計算/register トークンを入れると利点があるのか気になります。
    追加の議論: https://news.ycombinator.com/item?id=40182695, https://www.reddit.com/r/LocalLLaMA/comments/1cf2w5a/transfo...

  • ここ数年、ビジョンとマルチモーダル認識の分野で論文やアイデアを何十本、もしかすると 100 本以上試してきましたが、これは珍しくすべてがうまく機能した事例でした。
    たとえばこのモデルは、4 個の register トークンと Matryoshka 式の損失を併用して学習し、CLIP の代替を探している人に役立つかもしれない非常に小さな64 次元埋め込みを生成します: https://huggingface.co/unum-cloud/uform3-image-text-english-...

  • ICLR で見ましたが、今年最高の論文の一つで、ポスターセッションでもそれは明らかでした。
    著者の皆さん、おめでとうございます。

  • 数か月前から DinoV2 を使っていて、[CLS] + パッチトークンに加えて register トークン 4 個があるモデルも試しました。
    既存の register なしモデルから得た複数の埋め込みトークンがすでにソリューションの一部になっているため、0 register モデルと 4 register モデルの間で [CLS] トークンがそろわず、新しい register モデルは採用しませんでした。
    2 つのモデル間で [CLS] とパッチトークンが何らかの形でそろうとよいのですが。

  • ビジョンモデル向けの attention sink のように見えます