誰でもわかるStable Diffusion その4：U-Net
前の記事で、Stable Diffusionには「U-Net」と呼ばれる仕組みが使われると説明しました。
今回はそのU-Netについての概要です。
U-Netの構造
「U」という名前の由来はその構造の形からです。まるでUの字に見えるから(下の図はゆがんだUですが)U-Netと名付けられました。
U-Net自体はStable Diffusionの発表よりも前からありましたが、Stable DiffusionがU-Netを改造して独自に実装しました。
下の図はStable Diffusionで使われているU-Netの構造です。
とてもユニークな形をしています。
画像を左上から入れると、右に向かって次々と画像が処理されて、右上から結果が出てくる仕組みです。
色のついた四角いブロックがそれぞれ何かの処理をすると考えてください。処理が終わったら隣にあるブロックにバケツリレーのようにデータを渡します。
(ピンク線のように途中をすっとばしてデータを渡すこともあります)。
ブロックの中の数字は「画像の持つ特徴の数*1」を表しています。
図を見てみると
大まかにINパート、MIDパート、OUTパートに分かれています。
INパート
INパートはIN0からIN11まで12のブロックに分けられます。
ここでは画像をどんどん小さくしていって、画像が持ついろいろな特徴を取り出していきます。また、与えられたテキスト(プロンプト)を画像に反映させる処理(後述)も行います。
MIDパート
MIDパートはひとつのブロックです。
ここでは画像が最小(最初の画像に比べて8分の1)になります。
OUTパート
OUTパートはOUT0からOUT11まで12のブロックに分けられます。
ここでは画像をどんどん大きくしていって、最初の画像のサイズまで戻します。INパートやMIDパートで取り出してきた画像の特徴をもとに、OUTパートで画像サイズを戻しながら特徴をまとめていきますが、面白いのが「INパートで見つけた特徴を直接使う」という点です。INパートから特徴情報を直接受け取って使う機能があり、「スキップコネクション」と呼ばれます。
OUTパートから出てきた画像が、最終的に欲しい画像です。
スキップコネクション：データを遠くに渡す
図のピンク色の線のように、ブロックごとの処理に加えてINパートからOUTパートへとデータがダイレクトに渡されています。なぜこんなことをするのでしょう？
隣のブロックからバケツリレーのように渡されてきたデータにも画像のデータは入っているのですが、そのデータの内容はブロックを通るごとにどんどん内容が変わっているので、まるで「伝言ゲーム」のように最初に取り出した特徴情報が分からなくなってしまいます。そこで、その情報を「カンニング」するのです。
実際に伝言ゲームを例にとってみましょう。
10人の人が一列に並びます。この人たちが隣にいる人に聞いたことをそのままコッソリ伝えます。
誰もが知っていると思いますが、1番目の人が聞いた内容は伝言の途中でどんどん変わってしまい、10番目の人が聞くころには全然違う内容になっています。そこで、正確に情報を伝えるために「カンニング」を許すことにします。
U-Netでやっていることもまさにこれと同じ事です。
ちょっと実用的な話
つながっているブロック同士の関係は強くなります。どのブロックがペアになっているかという情報は、LoraやLyCORIS(Locon)といった追加学習の層別適用を行うときに重要になるでしょう。例えばLoraでIN0層を無効にした場合、OUT11層にも追加学習の効果が反映されにくくなります。
Stable Diffusion独自技術：テキスト取り込み
本来のU-Netはテキスト情報を読んだりはしません。Stable Diffusionではこれを改良し、「Transformer」という別の仕組みをくっつけて、テキスト(つまりプロンプト)によって画像処理を操作できるようにしました。
Transformerはテキスト処理が大得意です。が、本来は画像処理用ではありませんでした。これを改良してテキストと画像を比べられるようにしたのが「CLIP」という仕組みです。
Stable DiffusionはこのCLIPを利用し、さらにStable Diffusionに合わせて改良し、U-Netの中に取り込みました。
上の図のオレンジの部分がテキストを取り込んで画像処理を行う部分です。
結局何をしているの？
Stable DiffusionのU-Netがやっていることは、以下のことです
- 画像を受け取って、そこから画像の持っている特徴をたくさん見つける
- 見つけ出した特徴から、その画像に乗っているノイズを予測する
- 予測したノイズの画像を吐き出す
ネコの絵を描きたいとします。
ノイズだらけの画像をU-Netに入れると、U-Netはそのノイズだらけの画像が持つ特徴をたくさん見つけ、「ここはネコっぽいな、ここはネコっぽくないな」と判断します。そして、ネコっぽくないところを「ノイズ」と判断し、ネコっぽくない特徴を集めて「ノイズ画像」を作り、吐き出します。
U-Netから出てきたノイズ情報を元の画像から除去すると、「ネコっぽい特徴」が残り、よりネコの絵に近づきます。これを何回も繰り返すと、最終的にネコの絵になります。
「なぜU-Netは『ネコっぽい特徴』じゃなくて『ネコっぽくない特徴』を吐き出すの？最初から『ネコっぽい特徴』を吐き出せば、それがネコの絵になるじゃないか」と思われるかもしれませんが、確かにその通りです。ただ、U-Netは特定の絵を予測するよりノイズを予測する方が得意なのです。計算量はどちらもほとんど変わりません。
まとめ
Stable Diffusionの心臓部、U-Netの概要を説明しました。
次回は特徴を取り出すしくみ、「畳み込み」を説明する予定。
*1:この数字は「処理前」の特徴数であることに注意してください！！処理後に増えたり減ったりすることがあります
