自己蒸留でCVのBERTを目指すiBOT
3つの要点
✔️ ViTの画像Tokenizerが重要であることを示す
✔️ ViTの画像Tokenizerを自己蒸留で学習して、End-to-EndなMasked Image Modelingを実現
✔️ ImageNet-1KでSOTAを達成し、MAEと画像認識のBERTの地位を競う
iBOT: Image BERT Pre-Training with Online Tokenizer
written by Jinghao Zhou, Chen Wei, Huiyu Wang, Wei Shen, Cihang Xie, Alan Yuille, Tao Kong
(Submitted on 15 Nov 2021 (v1), last revised 9 Dec 2021 (this version, v2))
Comments: Published on arxiv.
Subjects: Computer Vision and Pattern Recognition (cs.CV)
code：
本記事で使用している画像は論文中のもの、紹介スライドのもの、またはそれを参考に作成したものを使用しております。
はじめに
自然言語処理におけるTransformer (BERT)が大活躍していることから、Vision Transformerをはじめ画像認識にTransformerを導入する研究が加速しています。特にBERTの鍵となるMasked Language Modeling (MLM)からヒントを得て、Masked Image Modeling (MIM)を画像のタスクで高精度を出したMAE(ついに画像認識のBERTが実現？！MAEについて)がつい最近発表されました。
MLMでは単語を意味のある潜在空間に射影するトークナイザー (Tokenizer)が非常に重要な技術です。同様にMIMにおいても画像のパッチを意味のある潜在空間に射影するためには、どんなTokenizerが必要なのかについて研究する必要があります。特に意味のある単語をモデリングするより、連続的かつ冗長な画像ピクセルを高次元なセマンティックを持つトークン (Token)に変換するのが難しいとされています。
本記事ではMAEの四日後にリリースされ、MIM Tokenizerの難点に取り組んだiBOTについて紹介します。
類似した先行研究BEiTでは学習済みのdVAEをTokenzierに用いたが、iBOTは自己蒸留のフレームワークを利用してOnline Tokenzierを提案し、End-to-EndなMIMを実現しています。モデルに関する詳細は後ほど紹介するので、まず結果を早速Figure 1に載せています。驚くほど綺麗なAttention Mapを見せたDINOよりも高いImageNet-Top1の成功率を出しています。
iBOT
Masked Image Modelingについて
MIMは画像Token系列
に対して比率rでランダムマスク
をサンプリングする。ここのNはTokenの個数です。マスクされた
は
に置き換えられて、画像Token系列
が得られます。MIMの目的関数はマスクされたTokenから元の画像を復元するもので、BEiTでは式(1)のように定義されています。
は入力をK次元の確率分布に変換するモデルで、
と
はそれぞれdiscrete VAEと学習したいモデルのパラメータです。ここではdiscrete VAEは画像パッチをKカテゴリに変換しています。
自己蒸留について
異なる
には異なる事前知識が入ります。例えばBEiTでは学習済みdiscrete VAEを用いますが、iBOTでは自己蒸留した
を用います。つまり、モデルの出力を教師データとして、モデル自身を学習させることです。ここでは、簡単に同じネットワークを持ちながら、パラメータが異なる２つのモデルを用意することで、自己蒸留を行います。具体的には、二つのデータ拡張によって入力画像
から
が得られました。それぞれの予測を
と
とした際に式(2)のように予測損失関数を取ります。
は同じ入力画像から得られたためお互いに予測することが可能です。ここでは生徒と教師ネットワークは同じ構造で異なるパラメータと
を有しています。また教師ネットワークのパラメータは生徒ネットワークのパラメータの指数移動平均とします。
iBOTのアーキテクチャについて
iBOTの全体図をFigure 3に示します。iBOTは自己蒸留することでトークナイザーの役割を担います。つまり自己蒸留でMIMの学習を行います。具体的には、画像xからデータ拡張によって得たuとvを教師と生徒ネットワークに入力します。それぞれのネットワークにはbackboneとパッチ予測ヘッドを持っていますが、教師ネットワークは勾配降下法で更新しません。代わりに生徒ネットワークのパラメータの指数移動平均(EMA)で更新することが、画像を連続した特徴分布に変換できることを示されています。
iBOTは二つの目的関数を持っています。一つ目は、式(2)の自己蒸留による予測損失関数です。cross-viewの分類トークンを予測します。二つ目は式(3)に示します。教師ネットワークの出力をラベルとして、生徒ネットワークがマスクされたパッチを復元します。
後ほどの実験を通じて、
のように分類トークンとパッチ復元のHead間でパラメータ共有した方が精度が良いことが分かります。またiBOTはone-hot token idではなくsoftmax後のtoken分布を教師信号とした方が精度が良いことも示されます。
実験
ImageNet-1Kでの実験結果
iBOTで学習した表現の良さを評価するには5つの指標を用いています。
Table 1にはk-NNとlinear probingの結果を示しています。k-NNはBackboneをfixして得られた特徴ベクトルをk-nearest分類に適用した際の精度を示します。linear probingはBackboneをfixして１層のlinear classifierを付けて分類を行った際の精度を示します。BackboneがViT-S/16やViT-B/16の時に、k-NNとlinear probingのいずれの指標でもDINOを上まっています(+~1.3%)。また最後の行に示したように、ImageNet-22Kのデータを用いて事前学習した結果linear probing指標においては81.6%を達しています。
Table 2にはImageNet-1KのみでFine-tuningした結果とTable 3にはImageNet-22Kで事前学習した結果を示しています。Fine-tuningではモデルの最後に１層のLinear Classifierを付けて全てのラベルデータを使ってBackbone部分の重みと分類器の重みを調整します。Table 2より、BackboneがViT-S/16の時に82.3%の精度が得られて、ViT-B/16の時にMAEの83.6%よりも高い精度83.8%が得られました。またTable 3からImageNet-22Kで事前学習することが精度の向上に貢献することが分かります。
Table 4にはSemi-supervised learningの結果を示しています。Semi-supervised learningは一部(1%, 10%)のラベルを利用してfine-tuningするため、ラベル効率を測る指標となります。Table 4からいずれの条件においてもDINO (SOTA)より精度が良いことが分かります。
さらにTable 5には教師なし学習の結果を示します。精度（ACC）、修正ランダム指数（ARI）、正規化相互情報量（NMI）、Fowlkes-Mallows指数（FMI）などの標準的な評価指標を用います。iBOTは、Accが2.0%、NMIが1.8%で従来のSOTA (DINO)を上回っています。
以上、ImageNet-1KにおけるiBOTの結果よりMIMが良い視覚的に意味のある特徴量を抽出できることが示唆されています。
Downstream Tasksでの実験結果
MIMの目的は多様なタスクにおいても精度の良い表現を学習することです。ここでは、COCOデータセットにおける物体検出とインスタンス分割および、ADE20Kにおけるセマンティック分割タスクでの実験結果をTable 6に示します。
iBOTはMAEを含むいずれの比較手法より精度が良い(+ 0.8%~3.8%)ことが分かります。また、Supervised Learningの結果より良いことはMIMをはじめとするSelf-Supervised Learningの手法を超えており、実用レベルまで達したことが言えます。
最後に多様なデータセットにTransfer Learningする際の精度をTable 7に示しており、これだけ多くのデータセットで良い結果を出したことでiBOTが優れていることが言えるだろう。
iBOTの性質について
MIMでのPatch Tokenがどのような特徴表現を学習したのか？この問いに答えることで、意味のあるトークナイザーが学習できているのかに繋がります。ここではFigure 4に可視化の一例を示しますが、論文中には豊富な分析実験を行なっており、大変興味深い考察がなされています。興味ある方は論文の4.3節およびAppendixを確認してみてください。
Figure 4にはImageNet-1Kの評価データのPatch Tokenの確率分布を求めた上、一部の中心となるパターンを可視化しています。左の２つがライトと犬の耳となっており、同じクラスが近くなっていることが分かります。右の２つは柄が同じようなパッチが揃っていることから、iBOTにはテクスチャに関する情報が学習できたことが推測できます。
まとめ
本論文ではVision TransformerにおけるMasked Image Modeling (MIM)なモデルiBOTを提案し、画像のセマンティックを捉えられるトークナイザーの重要性にフォーカスしました。BEiTのdiscrete VAEとは異なって、iBOTは自己蒸留でトークナイザーを学習するフレームワークを提案して、その有効性を大量な実験を通じて示しました。
iBOTとMetaFormer (恥ずかしいほどシンプルなVision Transformer)をシンプルに組み合わせるだけで強力かつ軽量な画像認識モデルが得られるのではないかと考えると、ますます今後の発展に期待してしまいます。
この記事に関するカテゴリー
