Zenn
🎨

Prompt Segmentation Model "Segment Anything" 要約

2023/05/01に公開
ディープラーニング
機械学習
cv
nn
tech

Alexander Kirillov, Eric Mintun, Nikhila Ravi, Hanzi Mao, Chloe Rolland, Laura Gustafson, Tete Xiao, Spencer Whitehead, Alexander C. Berg, Wan-Yen Lo, Piotr Dollár, Ross Girshick. "Segment Anything". 2023.4.5. https://arxiv.org/abs/2304.02643

11Mのフリーな画像データと10億のマスクを持つデータセット"SA-1B"で学習した革命的なゼロショット推論機能を持つSegment Anything Model(SAM)がMetaから発表されました。(NNCVの分野でSAMといえばSharpness-Aware Minimizationな気がするのですが...)

一言でまとめると、プロンプトと呼ぶ画像上の点やBBoxなどの指示に対して、リアルタイムで推論内容を修正更新するセグメンテーションモデルで、そのための大規模データ集めも自動化することに成功したという話です。

特記しない限り画像は元論文からの引用です。

概要

NLP分野では基盤言語モデルのようなLLMが巷を賑わせているが、画像分野でも"基盤画像モデル"的なものとして単語列と画像を対象学習を利用したCLIPやALIGNがあり、これらは獲得した汎化能力によりDiffusion Modelなどの上流タスクに実用されている。

この論文で提唱されるSegment Anything Modelはプロンプト(画像上の点やBBox、セグメンテーションマスク、テキスト)への応答としてセグメンテーションを行うものである。
セグメンテーションモデルにおける汎用モデルを構築するため、次の3つの問題に焦点を当てている。

  • ゼロショット汎化を可能にするタスクは何か
  • どのようなモデルアーキテクチャなら解決できるか
  • どのような学習データがモデル学習のタスクに適しているか

これらについて解決するためには次の流れが考えられた。

  1. プロンプトセグメンテーションのためのデータセット"SA-1B"を作るため、データ作成用のデータエンジンを作る
  2. プロンプトに対応するマスク出力をリアルタイムで定期的に出力させるように(入力された単語トークンごとに?)設計するため、テキストエンコーダと画像エンコーダで埋め込み表現を構成し、軽量なマスクデコーダを作る
  3. マスクがプロンプトの曖昧性に対応できるように複数のマスクを予測するように作る

pointプロンプトによる曖昧性の例

データエンジンは、まず人間の付けたアノテーションからそのサブセットを作り出し、これにより人間はそれ以外のアノテーションを行うことに注力できるようになった。この工夫でデータセットは巨大に作れるようになり、また倫理的に公正に作られている。

実際作られたSegment Anything Modelはプロンプトエンジニアリングによるゼロショット転移タスクで、エッジ検出、オブジェクト提案などのタスクで強力な性能が得られている。

SA-1Bのマスク粒度別サンプル

[CLIP] Alec Radford, Jong Wook Kim, Chris Hallacy, Aditya Ramesh, Gabriel Goh, Sandhini Agarwal, Girish Sastry, Amanda Askell, Pamela Mishkin, Jack Clark, Gretchen Krueger, Ilya Sutskever. "Learning Transferable Visual Models From Natural Language Supervision". https://arxiv.org/abs/2103.00020
[ALIGN] Chao Jia, Yinfei Yang, Ye Xia, Yi-Ting Chen, Zarana Parekh, Hieu Pham, Quoc V. Le, Yunhsuan Sung, Zhen Li, Tom Duerig. "Scaling Up Visual and Vision-Language Representation Learning With Noisy Text Supervision". https://arxiv.org/abs/2102.05918 (読んでない)

タスクの概観

Segment Anythingのプロンプトは背景前景を示すpoint、BBox、mask、自由形式textなどのセグメント情報があり得る。

プロンプトセグメンテーションというタスクでは、任意のプロンプトに対して"有効な"マスクを返すことが求められる。この有効なというのがミソで、テキストの自由度からくる曖昧性(例えば、「シャツ」というテキストに対してシャツの生地全体なのか、それともシャツの印字は含まないのか、のような問題)があるため、少なくとも妥当な1つを出力することが要件である。つまり、言語エンコーダが自由なプロンプトに対して一貫した出力をしてくれれば良い。

事前学習として、学習データのプロンプトの系列のシミュレーション(後述)からマスクを予測する。

そしてゼロショットタスク転移についても、基盤言語モデルのそれのようにプロンプトの入力を設計することにより制御できる(例えば上流タスクとして、あるオブジェクトのBBoxを出力するモデルがあるとき、そのオブジェクトのinstance segmentationタスクはBBoxをプロンプトとして入力することで達成できる)。

以上の方法を踏まえ、プロンプトとそれを構成要素としたモデルを用いれば色々なタスクに転用可能であると考えられる。CLIPがDiffusion ModelのVision-Languageアラインメントになっているのと同様に、他の基盤モデル風の合成が行えるかもしれない。

モデルの概観

モデルの概要

プロンプト可能なセグメンテーションのためのモデル"Segment Anything Model"を設計する。

  • Image Encoder: MAE(Masked Autoencoders)でSSLされたViTにViTDetの方法(FPNを用いず、純粋なスケールのViTからFeature Pyramidを作る手法)を適用した
  • Prompt Encoder: 疎なプロンプト(point / BBox / text)と密なプロンプト(mask)に対して、エンコードされた画像のトークンと対応した埋め込みが行われる
    • pointとBBoxは位置エンコーディング(ViTのstemで行われるPositional Encoding表現?)にエンコード
    • textはCLIPの埋め込みにエンコード
    • maskは畳み込みで画像埋め込みの要素ごとに合計(?)
  • Mask Decoder: 画像の埋め込みとプロンプトの埋め込みと出力されたトークンからマスクを生成する、プロンプトのSelf Attentionとプロンプト↔画像埋め込みのCross Attentionを用いたTransformer Decoder Blockを用いる

前述の曖昧さを扱うため、モデルは複数の有効なマスクを出力する。
だいたい3つ(whole, part, subpart)のマスクがあれば足りるらしい。
マスクに対するMinimum Lossのみを誤差逆伝播に伝え、各マスクの信頼度スコア(IoU)をもちいてランク付けする。
予測マスクに対してはFocal LossとDice Lossの和で教師あり学習する。
1マスクあたり11roundsのプロンプトをランダムにサンプリングし、インタラクティブな入出力をシミュレーションすることで、モデルとデータエンジンをうまいこと統合している。

[MAE] Kaiming He, Xinlei Chen, Saining Xie, Yanghao Li, Piotr Dollár, Ross Girshick. "Masked Autoencoders Are Scalable Vision Learners". https://arxiv.org/abs/2111.06377
[ViTDet] Yanghao Li, Hanzi Mao, Ross Girshick, Kaiming He. "Exploring Plain Vision Transformer Backbones for Object Detection". https://arxiv.org/abs/2203.16527
[Minimum Loss] Abner Guzman-Rivera, Dhruv Batra, and Pushmeet Kohli. Multiple choice learning: Learning to produce multiple structured outputs. NeurIPS, 2012. (読んでない)

Image Encoder

使用されているViT-H/16 (14×14 Window Attention、equally-spaced Global Attention(?)などを持つ)は、よしなにスケールされた1024×1024pxを64×64×256の埋め込み表現に変換する(ViT-Hの出力の大きなChannel数は最後に畳み込みで削って256にしている)。
画像ごとに1回行われるため軽いFLOPSで計算できる、つまりプロンプトを対応的に処理しても、それ単位で画像の計算は行わなくてよい。

使用するエンコーダのViTは、ViT-B、ViT-Lの向上は限定的だが、ViT-Hでは大幅に良いということがアブスタレーションスタディで示されている。

Prompt Encoder

プロンプトを256次元の埋め込み表現(ベクトル)に変換する。

  • point → positional encodingと、背景か前景かを示すlearned embeddings(?)のうち1つ、これらの和
  • BBox → 左上のアンカーのpositional encodingとlearned embeddingsの和と、右下のそれ
  • text → CLIPによる埋め込み表現
  • Mask(入力画像の1/4解像度) → 単純なCNNでさらに1/4にダウンスケールし、256次元ベクトルにする(Maskをプロンプトとして持たない場合は"no mask"を示すlearned embeddingsを追加する)

Lightweight Mask Decoder

画像とプロンプトの埋め込み表現のセットをマッピングして、マスクとして出力するため、DETRのデコーダを修正し、ViTのCLSトークンのようにプロンプトを与える。
図のようなTransformer Blockを2つ使い、プロンプト↔画像埋め込み(ここでは642個の256次元ベクトル)のCross Attentionで対応を計算、工夫によりAttentionが位置エンコーディングとプロンプト埋め込みから画像埋め込みを更新しており、プロンプトの幾何学的な空間での対応などを計算できるようになる。
Transformerにより処理された表現は、CNNで元の解像度に拡大され、マスク毎の出力トークンとMLPで合わせてセグメンテーションマスクとしてデコードされる。またIoU scoresとして出力トークンを用いる。

効率化のために、64×64の画像埋め込みが計算されるAttention部分は256次元から128次元にChannel削減して行っている。

[DETR] Nicolas Carion, Francisco Massa, Gabriel Synnaeve, Nicolas Usunier, Alexander Kirillov, Sergey Zagoruyko. "End-to-End Object Detection with Transformers". https://arxiv.org/abs/2005.12872

曖昧さへの対処

プロンプトが複数の有効なマスクに対応するという曖昧さ(シャツの例)があり、モデルはマスクの平均を学習することになる。これは1つのマスクを学習するのではなく、少数の出力トークンを以って複数のマスクを予測することで解決する。
対象をwhole、part、subpartの3つのマスクで表現し、正解のマスクとの損失が最も小さいものを誤差逆伝播することで学習する事ができる。
応用では、追加の出力トークンから各マスク間のIoUを計算するヘッドを用いることもできる。

プロンプトで曖昧さを削減することもできる。複数のプロンプトを受け取った場合3つのマスクは殆ど同じになるため、学習段階で1つの曖昧でないマスクのみを予測させるようにすることができる。これは複数プロンプト入力時にのみ作動する予測用の4番目の出力トークンを用いることで実現できる。

(このセクションはちょっとよく理解できなかった...)

学習

プロンプトによる対話型セグメンテーションをシミュレーションすることで、学習を行う。

まず、pointとBBoxのいずれかが正解マスクから確率的に選ばれる。BBoxは通常のObject Detectionデータセットのそれよりゆるい(ユーザーが適当に囲む可能性があるため)ものを混ぜる。

この処理の後、すでに得られたIoUの最も高い予測マスクと正解マスクの誤差領域から、 次のプロンプトとして予測マスクそれ自身と、誤差の種類(false negativeかfalse positive)に対してそれぞれ前景・後景のpointを与える。
この予測マスクは、情報量を削らないために閾値で2値化せずそのままロジットとして与える。
サンプリングされるpointは8個が最効率であるらしい。

この手順をさらに2回反復し、モデルがマスクから恩恵を受けるようにする。
まとめると、教師データから1つの初期入力プロンプトをサンプリングし、8つのpointを繰り返しサンプリング、外部情報を与えることなく2回繰り返してモデル自身がマスク予測を改善するように学習する。
これはがMask Decoderの軽量性による計算量の軽さが肝になっている。

学習設定として次のものを使っている。

  • AdamW β1 = 0.9, β2 = 0.999
  • 線形warm up (250 iterations, rl ~0.0008)
  • step-wise learning rate decay schedule
  • 90K iterations (2epoch) うち60000でlr10、86666でlr1/10
  • Batch size 256
  • 正則化 (weight decay 0.1, drop path 0.4, layer-wise learning rate decay 0.8)
  • データエンジンステージ1,2のデータのみで学習するとき、large-scale jitterでaugmentation

データエンジンとデータセット

SA-1Bを作成するにあたり、3つのステップでデータセットを作成するデータエンジンを用いる。

  • Assisted-manual stage: 一般的なセグメンテーションデータセットで事前学習したのち、アノテータが前景後景の対象のpoint、mask、オブジェクトのラベルを目立つものから順に付けていき、エンコーダをViT-BからViT-Hに拡大して繰り返し6回行う
  • Semi-automatic stage: マスクの多様性を高めて目立たないものもセグメントできるようにするため、モデルの提案マスクで埋め尽くせていないところをアノテータがラベル付けし、提案されたマスクは確度の高いオブジェクトとしてBBoxを与えFaster R-CNNを学習させることを、繰り返し5回行う
  • Fully automatic stage: 工夫(省略)により、十分な精度が出るようになったらアノテーション作業を全自動する

この方法で110M枚の画像からなる巨大なデータセットが作られた(非公開なのが悲しい)。
Fully automatic stageでの収集が99%であるため、確実性のために無作為抽出で5万枚のマスクについてアノテータに修正を行わせており、これにより作られた完成形のデータセットは他と比較しても精度に遜色ないことが示された。また、RAI(Responsible AI)分析により、公平性も保たれていることが示されている。

人間がマスクの品質を確認するためのガイドラインの様子

ゼロショット転移

学習には無いデータセットとタスクについて、次の5つのタスクについて検証して汎化性能を考察している。

  1. single point to mask
  2. edge detection
  3. object proposal generation
  4. instance segmentation
  5. text to mask(object segmentatino from free-form text)

前景のpointは他プロンプトより影響が大きいらしい。
MAE ViT-Hを用いたモデルにSA-1Bを学習させたものを用いて実験された。

edge detectionの例

instance segmentationの例

潜在空間でのマスク埋め込み表現を閾値で可視化したもの。explicit semantic supervision学習していないのに、同じ画像内の類似したマスク埋め込みは、クエリのマスク埋め込みと類似している様子が伺える

  1. single point to maskは前景の1点を指定するプロンプトからセグメンテーションを求めるタスクで、最もconfidentialなマスクのみを評価する。結果、特に曖昧さを含むような多くのデータセットで対話型セグメンテーションモデルのベースラインRITMを超えるmIoUが得られたが、いくつかのデータセットでは超えられなかった。
  2. 画像のエッジ検出タスクでは、16×16の規則的な前景pointのgridをプロンプトとして、point毎に3つのマスクを求めてNMSで処理、ソーベルフィルタとその他後処理でエッジを計算した。結果、SOTAモデルの結果には劣るが、ゼロショット転移の結果としては大きく優位である。
  3. object proposalはFaster RCNNのRPNのように中間的に特徴マップから物体の位置提案を生成するタスクで、DMP(Detector Masquerading as Proposal generator; PASCAL VOCのような部分的なアノテーションでなく、ほぼ完全なアノテーションを行ってRPを評価する)相当のベースラインを用いてViTDet-Hと比較した。結果、average recal値において中型・大型・希少なオブジェクト・一般的なオブジェクトにおいてはSegment Anything Modelが優位だった。
  4. instance segmentationにおいて、ViTDet後の出力BBoxをプロンプトとして与える。結果、少し劣るもののViTDetと似た精度がでており、ViTDetより境界部分はきれいに出力できている。
  5. Text to Maskでは自由形式のテキストをCLIPで埋め込み表現にしたプロンプトを与えることでセグメンテーションマスクを得る。結果、追加修正が必要な場合も多いものの、単純な物体の名詞や、修飾されたフレーズに対しても有効に働くことが示された。

[RITM] Konstantin Sofiiuk, Ilya A Petrov, and Anton Konushin. "Reviving iterative training with mask guidance for interactive segmentation". ICIP, 2022 (読んでない)

[DMP] Neelima Chavali, Harsh Agrawal, Aroma Mahendru, Dhruv Batra. "Object-Proposal Evaluation Protocol is 'Gameable'". https://arxiv.org/abs/1505.05836

読んだ感想

GPTなどが世間で喧伝されるなか、画像においても基盤モデルを作るというのは急がれることだと思う。
論文の考察でも触れられているが、セグメンテーションタスクは画像分野のほんの1つのタスクであり、LLM成功の要である自己教師あり学習という点でまだ及んでいない。ただ、プロンプト対セグメンテーションという対照学習で得られた基盤モデル本体はCLIPのように他上流タスクを解決する良い方法になることが考えられる。
また、panoptic segmentationをプロンプトとして実現する筋のいい方法はまだ見つかっていないと書かれており、現段階で書かれていたプロンプトも少し不自然というか、GPT系列の美しいシンプルさには及んでいないように見える。

Segment Anythingはタスク開発とそれを実現するモデルの点で革命的だが、プロンプトやタスク設計自体の不自然さ(めちゃめちゃ主観的)はまだ改善の余地がありそうに思える。

(これを読んだあとになって気づいたが、自分がやりたいのはCLIPSegとOVPSで、Segmentat Anytingのような複雑なプロンプトや応答性は求めていなかった... サーベイ不足...)

[CLIPSeg] Timo Lüddecke, Alexander S. Ecker. "Image Segmentation Using Text and Image Prompts". https://arxiv.org/abs/2112.10003
[OVPS] Zheng Ding, Jieke Wang, Zhuowen Tu. "Open-Vocabulary Panoptic Segmentation with MaskCLIP". https://arxiv.org/abs/2208.08984

Discussion