新時代の pix2pix？ ControlNet 解説

入力に対して画像で条件づけする Diffusion モデルである ControlNet の解説をします。
paper: https://t.co/hIkh8qgjcL
github: https://github.com/lllyasviel/ControlNet

([1] の Fig1 より引用)

本論文は Style2Paints 開発者の lllyasviel(@lvminzhang) が出されていることから、高精度に入力線画に着彩する Styl2Paints V5 (のDorothyモデル？）に使用されているモデルである可能性が高いです。

https://twitter.com/lvminzhang/status/1610836703765495810

Diffusion Model （特に Latent Diffusion）では入力の条件付けを忠実には考慮できていないモデルが多いため、本論文の結果はなかなかにインパクトがあります。個人的には GAN 時代の pix2pix がそのまま Diffusion モデルになってパワーアップした印象があります。勿論 Taming Transformer[2] や Palette [3] など Diffusion based な pix2pix モデルはありましたが今回のモデルは今広く使われている Stable Diffusion ベースとなっていることからも利便性が高いと思われます。（※ ControlNetは Stable Diffusion 以外にも適応可能です。）今後様々なタスクのベースとして使われる可能性があるため、自分の理解も含め簡単な解説を書きます。
以下の解説では説明のため、論文の内容と多少自分の説明用の解説が混じっています。気になる方は是非とも原著論文を読んでください。

Introduction

近年高精度な画像生成AIが話題となっていますが、プロンプトで生成した画像は我々のニーズを満たすような画像か？という疑問があります。プロンプトで生成する画像はあくまで生成したい画像のイメージを反映したものに過ぎず、欲しい構図や形状を正確に反映した画像を得ることができません。
生成能力は非常に高いのだから、この画像生成モデルを特定のタスク（ここでは既存のimg2imgのタスク）に応用できないか、ということを考えます。

この問題について調査したところ、以下の３つの知見があったようです。

• 固有のタスクを得くために必要なデータは一般的な text-image のタスクほど大きくはない。多くの特定の問題（姿勢-画像、白黒-着彩）で必要なデータは最大でも10万枚であり、 LAION の 5B の5*10^4 倍小さい。
• 画像処理タスクをデータ毎に特化したソリューションとして利用する場合、大規模な計算クラスタは常に利用できるわけではないため、限られた計算リソースで学習できる転移学習が必要となる（そう‼）。
• 画像処理問題には様々な形式の問題定義があります。例えばノイズ除去や深度からの画像生成などです。これらを全て手続き的に処理することは難しいため、多くのタスクで学習するためには end2end の学習が必須になる。

本論文では、Stable Diffusion のような大きな Diffusion モデルを個々のタスクに特化して学習が可能な ControlNet の紹介をします。この手法により Cannyエッジ、Hough line、人間の姿勢情報からの画像生成など、様々なタスクを学習することが可能となります。また、個人レベルのマシン（ex. Nvidia RTX3090Ti）で学習させた場合、大規模計算クラスタで学習したモデルにも引けをとらないモデルが学習できたようです。

手法

コアコンセプトは以下の２つです。

• モデルを、ベースとなる重みを固定したモデル（lockd）とそこからコピーしたモデル（trainable copy）したモデルに分け、trainable copy を学習することで、小さなデータセットでも過学習を避けられ、かつ収束も早い
• zero convolution という重みを 0 に初期化したブロックを trainable copy に加えることにより、安定させつつ学習を進めることができる

下図の (a) が式(1)、(b) が式(2) に対応します。


neural network に相当する部分がベースとなる stable diffusion (の UNet) に相当します。

([1] の Fig2 より引用)

zero convolution

zero convolution では畳み込み層の重みとバイアスはゼロとして初期化されるので学習の初期ステップでは式(3)のようになります。つまり式(4)のように学習初期はControlNetが存在しないようになるため、学習が安定し、通常の fine-tuning のように高速に学習が可能となります。

zero-convolution でも重みの勾配は式(6)のように0ではないため、通常の学習と同じ様に学習が進みます。

Stable Diffusion におけるモデルアーキテクチャ

Stable Diffusion における ControlNet のアーキテクチャは以下の様になります。Unet のエンコーダー部分を trainable copy として元の Unet にくっつけています。入力の前とデコーダーがあった場所に zero convolution が入っています。

([1] の Fig3 より引用)

学習

学習は通常の Diffusion モデルとほぼ同様です。テキストプロンプト（c_t）と個別の入力条件マップ（c_f）からノイズを少しづつ取り除きます。
工夫として50%の確率でテキストを空文字にしています。これにより入力条件マップの情報をより考慮する学習になるそうです。

学習の工夫

• 計算リソースが限られている場合は、locked と trainable copy の結合を減らすことで学習が早くなる。具体的には図３では SD Middle Block と SD Decoder1~4 で結合しているが、SD Midel Block の結合を解除することで学習速度が1.6倍になる
• 大規模学習の場合（例えば A100 80G と100万組のデータ）は、ControlNet のみを反復学習させたあと、lockedの部分を解除してモデル全体を学習することでより良い精度が得られる

実験

正直この後のパートは延々と結果がはられているだけなので特に書くことはないです。論文を実際に見たほうが良いでしょう。
いくつかのタスクにおけるデータセット量とマシンと学習時間についてピックアップします。

• Canny Edge
  • 3M のペア画像
  • A100 で 600GPU 時間
• 人間のポーズ
  • 80k のペア画像
  • 3090Ti で 400GPU 時間
• セグメンテーション
  • 164k のペア画像
  • 3090Ti で 400GPU 時間
• イラストの着彩
  • 1M のペア画像
  • A100 で 300GPU 時間
  • ベースモデルは Waifu Diffusion

([1] の Fig4, 10, 11, 15 より引用)

感想

感想としては、とりあえず結果が強い、です。
入力条件付けのアーキテクチャを追加する手法はいくつもありますが、元のモデルを完全に固定したまま同程度のサイズのモデルを横にくっつけることにより、過学習を防ぎつつ安定して学習させるアイディアは良いと思いました。
今流行りのSDがベースになっていることからも、この手法をベースに様々な個別のタスクに特化したモデルが出てきそう、といった所感です。
そう、ここは　 R　E　D　　 O　C　E　A　N　。

参考文献

[1] Lvmin Zhang and Maneesh Agrawala, Adding Conditional Control to Text-to-Image Diffusion Models
[2] Esser, Patrick, Robin Rombach, and Bjorn Ommer. "Taming transformers for high-resolution image synthesis." Proceedings of the IEEE/CVF conference on computer vision and pattern recognition. 2021.
[3] Saharia, Chitwan, et al. "Palette: Image-to-image diffusion models." ACM SIGGRAPH 2022 Conference Proceedings. 2022.