Tsubatoの発信記録

主に機械学習やデータサイエンス関連で学んだことを書いています。

MetaのコンピュータビジョンモデルDINOv2はどのような自己教師あり学習を行っているか？(DINOv2: Learning Robust Visual Features without Supervision)

論文 機械学習

DINOv2とは

DINOv2は2023年4月にMetaより公開された、自己教師あり学習によるコンピュータビジョン(CV)モデルです。その名の通りDINOという既存のモデルの改良版です。
自己教師あり学習は入力データ(画像)以外のラベルなしに学習をする手法ですが、下記のtweetでは犬の部位毎の特徴を鮮明に取得できている様子が示されています。

それ以外にもセマンティックセグメンテーションなどあらゆるCVのタスクにおいてfine-tuningなしで高い性能を出すことが示されています。
こちらのデモサイトで自分の画像を使って試すこともできます。

この記事の内容

何か凄そうなモデルなので論文を読んでみましたが、どのようなタスクで自己教師あり学習をしているのかわからなかったため、関連論文も読みつつ明らかにしていきます。
DINOv2ではモデル構造や学習プロセスに多くの改良が加えられていますが、全ての詳細には触れませんのでご了承ください。

自己教師あり学習については以下のようにDINO,iBOTを組み合わせていると書かれているため、まずはそれぞれの論文を軽く見ていき、その後DINOv2の論文へ戻ってこようと思います。

We learn our features with a discriminative self-supervised method that can be seen as a combination of DINO and iBOT losses with the centering of SwAV.

DINO

Emerging Properties in Self-Supervised Vision Transformers arxiv.org

DINOという名前はSelf-distillation with no labelsから来ている通り、知識蒸留による自己教師あり学習です。 知識蒸留はモデルの軽量化のために生み出された手法で、パラメータ数が多く性能の良いteacherモデルの出力を元に、軽量なstudentモデルを学習するというものです。(画像はMetaブログより引用)

学習の流れは以下の通りです。通常の知識蒸留とは異なり、teacherモデルはstudentモデルを元に作られているというのがややこしい所です。

  1. 画像に対し異なるaugumentation(切り出し)を行い、それぞれをstudent,teacherモデルに入力する。
  2. studentモデルの出力に対してsharpening、teacherモデルの出力に対しcentering->sharpening処理を加える。sharpeningはモデル出力の均一化を、centeringは出力の特定の次元が突出することを避けるための処理です。
  3. teacher出力に対するstudent出力のcross entropy loss(CE)を算出し、studentモデルのパラメータを更新する。teacherモデルの更新はbackpropでは行わず、studentモデルのパラメータの指数移動平均で求められます。

画像の切り出しはteacherに対しては元画像の50%以上を含むglobalなものを、studentに対してはそれに加え50%以下のlocalなものを入力します。これにより、studentモデルはlocalとglobalの対応関係を学ぶように促されます。

iBOT

iBOT: Image BERT Pre-Training with Online Tokenizer arxiv.org

iBOTは自然言語処理の自己教師あり学習にも用いられる、maskを用いた学習を行います。 ただ基本的な枠組みはDINOと同じで以下のような知識蒸留モデルです。

DINOと主に異なる点は以下の通りです。

  • studentモデルへの入力について、一部のパッチをmaskする。
  • lossは以下の2項からなる。基本的にはmaskされた部分をstudentモデルが再構成できるようになることを目的としています。
    ①異なるaugumentationを加えた画像をteacher、studentモデルへそれぞれ入力した際のCLSトークン(画像全体の特徴ベクトル)のCE。これは入力がmaskされていることを除けば、DINOにおけるlossを同じものです。
    ② 同じaugumentationを加えた画像をteacher、studentモデルへ入力した際のpatchトークン(画像の一部の特徴ベクトル)について、マスクされた部分のCEの和をとったもの。

DINOv2

DINOv2: Learning Robust Visual Features without Supervision arxiv.org 先行研究の学習方法についてざっくり理解したところでDINOv2の論文に戻ります。

自己教師あり学習

論文にはDINOとiBOTの組み合わせとありましたが、iBOTがDINOの拡張のようなものなので、実質的にDINOv2はiBOTの進化版のようなものだと理解しました。
CLSトークンのCEは画像全体の、patchトークンのCEは局所的な特徴の取得に貢献すると思われ、実際に後者のlossを加えることでsegmentationの性能が向上していることが示されています。

これ以外にも改良点が数多く示されていますが、割愛させていただきます。

データキュレーション

DINOv2の改良点は学習方法だけでなく、データに関しても存在します。
データソースはImageNetなど質の高いものと、webページから収集した大量のデータの計1.2B枚の画像です。
これを元に以下の手順でデータを厳選し、142M枚の画像データセットを作成しています。

  1. 画像を学習済みViTに入力して特徴量を得る。
  2. 未厳選データ間で重複しているデータを削除、同様にベンチマークで使用する画像と重複するデータも削除。
  3. コサイン類似度を距離として、未厳選データをk-meansでクラスタリングする。その後、厳選データをクエリとして類似している未厳選データを学習用データに採用する。

性能評価

性能は既存の自己教師あり学習や弱教師あり学習モデル(CLIPなど)と比較されており、特にsegmentation,深度推定,類似画像検索タスクにおいて大きく上回る性能を示しています。
CLIPのような画像とキャプションを組み合わせたモデルは、キャプションに書かれていない画像内の詳細についての理解が困難である一方、DINOv2はこの問題を克服しているようです。

以下はDINOv2のpatch出力に主成分分析をかけ、最初の３つの主成分を色分けして可視化したものです。例えば一番左のコラムでは鳥と飛行機のような異なる物体の画像に対して、胴体、主翼、尾翼の各パーツが同じ色で表現されています。
驚くことに、これはアノテーションを加えていないのに、パーツの情報まで習得できていることを示しています。

感想

昨今ChatGPT(GPT4)を初めとしたLLMが猛威を奮っているため、CVについてもLLMの力を借りたVision-Languageモデルが主流になっていくのかな、と思っていた矢先に画像だけの自己教師あり学習でこれだけの性能を持つモデルが発表されて驚きました。

当然このモデルのコードもオープンになっていますが、まだドキュメントが整備されていないようなので、整備され次第触ってみたいと思います。 github.com