Stable Diffusion アルゴリズムの技術解説

図：Stable Diffusionによる生成画像

AI画像生成がまたバズりだしましたね。2016あたりのGANブームから何度か流行っているものの、今度のブームでは"確率拡散モデルによる高画質化" "テキスト条件付けによる生成コンテンツの制御" などいくつか技術的にも目新しいアイディアが入っていて、興味深いです。
しかしバズった割に技術的内容に踏み込んだ解説があまりないので、勉強を兼ねてまとめておくことにしました。
内容の正しさなどに関しては保証などはありませんのであしからず。

準備: AIは画像をどう扱うのか

AIは関数です。ここでいう"関数"というのは数学でいうような関数とプログラミング的な意味での関数のダブルミーニングになっていて、プログラミングするときにAIをPythonの関数 (入出力の管理があるので実際はクラスにしますが) としてパッケージ化するし、 AIの振る舞いを数学者が数学関数とAIを見なして研究したりします。
最近のAIが普通のプログラミングの関数と異なるのは、入出力の関連付けがハードコードされる（プログラムされたとおりになる）のではなく、学習データにあうように調整して作られることです。このようにデータを学習させることを通して関数を構成する技術を機械学習と呼び、 AIつくるための必須技術になりつつあります。このために画像系では特にニューラルネットを使います。

AIは関数なんだなーということに納得すると、次は入出力の仕様を決めたくなります。デジタル画像は(縦, 横, (RGB)) の3階の多次元配列に格納されているので、基本的にはこれをそのままAIに入れたり出したりできます。さらに(縦, 横, (RGB))の画像は値だけを考えれば（フラット化をすれば）[縦×横×3] 次元の1次元配列 = ベクトルと見なせます。ここでベクトルとは、高校で習う2次元や3次元のやつではなく、数値データを好きな数だけ（例えば256次元とか）並べたもののことです。そうしてしまえば、画像AIは結局ベクトルを入れるとベクトルを出すだけです。ここまでくると、AIは結局"ベクトルをベクトルに次元を増やしたり減らしたりしながら変換する関数" だというざっくりとした理解ができます。

例えば画像を入力、出力をベクトルとして、出力ベクトルが画像カテゴリと対応するように学習する（要素0が大きかったら犬、1が大きかったらネコ、……）と"画像分類"が解けるようになります (図上)。そのほか入力を適当な分布のランダムベクトル、出力を何らかの自然に見える画像、として学習すればその適当な分布の乱数と画像を関連付ける初歩的な"画像生成"のモデルになります (図下)。このように様々な画像に関係するAIタスクを定式化できるのがこの考え方のうれしいところです。

エンコーダ-デコーダモデル

さて画像生成AIで問題になるのは入力ベクトルをどうやって作ればよいか？ということです。本当にランダムなベクトルを入れてしまうと、出力される画像も変なものになってしまうかもしれません。エンコーダ-デコーダモデルはそこでAIにいうことを聞かせる代表的な方法です。
エンコーダ-デコーダでは生成AI（この文脈ではデコーダと呼ばれます）の入力ベクトルを認識AI（エンコーダ）の出力ベクトルとします。これで入力画像と比較的近い・関係のある画像が生成されやすいようになります。"画像から画像への変換"を学習することに特化している、といってもよいでしょう。

ここでエンコーダとデコーダの間に発生する中間的なベクトルを"潜在変数"と呼びます。ではこの潜在変数はいったい何を表しているのか？ナイーブな理解は、この潜在変数がなすベクトル空間が人間が心の中に持っている内面的な思考操作の作業場で、この空間の各点が"概念"に対応している、というものです。とはいえこのようなアナロジーが正しいか確かめるすべをAI研究者は持っていないので、これは学術的には単なる言い過ぎですがそれでもふわっとした理解のためにはこの気持ちを持っておくとよいのではないかなと思っています。

エンコーダ-デコーダモデルの初期の応用として、画像の画質改善があります。画質改善では、例えばノイズの乗ってしまった画像からノイズを取り除く "ノイズ除去" 、サイズが縮小されてしまった画像を高解像度にする "超解像" などがあります。学習データはきれいな画像を用意してから、それを人工的に劣化させた画像をペアにして"この汚い画像をこんなふうにきれいにしてね" というお手本にしてあげればよいので、簡単に試せるところもよいですね。これをAI画像生成のもっとも初歩的な形だった、と考えることもできそうですね。

確率拡散モデル: 高画質生成の鍵

確率拡散モデルは、上記の"ノイズ除去"技術の発展形と見なすことができます。画像にノイズを少しずつ足していくとどうなるでしょうか？ -- 何ステップも繰り返すといつかは元の画像の影も形もなくなり、ただのノイズになってしまうはずです。ではノイズにノイズ除去を繰り返したら？ -- いつか何らかの画像になるはず、これが確率拡散モデルのアイディアです。

これによってノイズから画像を生成することが学習されます。学習時は学習データとして使った画像しか生成できないわけですが、シードにするノイズを変えると新しい画像を生成する、という性質が発生します。
このやり方は

何度もノイズ除去をするので細部まできれいな画像が出る
学習データが人工的なノイズののせで自動で作れる
学習時は1ステップずつに分解できるので、比較的学習が速い（学習より生成のほうが遅い）

といったメリットがあります。

確率拡散モデルの学習もう少しの詳細は

入力: ノイズがのった画像
出力: ノイズが少しとれた画像
学習データ: 画像にちょっとずつノイズを乗せて作った画像の系列
使うニューラルネットワーク: エンコーダ-デコーダ

です。学習データの作り方で"ノイズをちょっとずつ足す"をいうのはコツが要るところで、試行錯誤により1000回ぐらいがよいのではないか、と作者はしています。また実際に画像を生成するときはそのためノイズに1000回程度モデルを適用を繰り返してノイズ除去 == 画像生成をしないといけないのですが、これは計算に時間がかかるというデメリットにも繋がっています。
これでランダムなノイズにノイズ除去を繰り返して画像が生成されるようになりました。画質改善AIと比べると「ノイズの乗った画像をきれいにする」から「ノイズだけから画像が生まれるようになった」という大きな進歩がありますね。

AIがノイズ除去や解像度上げ (Waifux2が代表例です) のような"画質操作"のような低レベル画像処理が得意なのはかなり知られてきていたので、これを画像生成にも使ってみよう、というのは個人的にはかなり筋の良いアイディアだと思いました。やっぱりAI開発者は賢いなあ（笑い）

Stable Diffusion: エンコーダ-デコーダ、注視機構、確率拡散モデルの合わせ技

ここまで来ていよいよ本題のStable Diffusionの説明に入れます。とはいっても各モジュールを説明してしまった後では、Stable Diffusionもそれらの組み合わせでしかないことがわかりますね！

Denoising step: これは上記の"ノイズ除去"のことですが、Stable Diffusionでは先ほどの説明のように画像そのものにではなく、潜在変数空間 (latent space) で適用しています。これでも同じような画像生成ができますがこちらの方が効率が良いようです。
Conditioning: 入力として塗り絵 (semantic map) や文章 (text) を入れられるようにしています。これは重要で、これのおかげで人間が好きな画像をAIに出させられるようになります
Crossattention: これは"ノイズ除去"モジュールの内部で使われているもので、日本語では"注意機構"などと呼ばれます。入力ベクトル (z_T) や条件付けに注意を払いながらノイズ除去を行うので、画像が良くなります。

これで生成した画像の例を見てみましょう。上の例ではバウンディングボックス（この四角にこの物体を配置してね、という指示のこと）をもとに画像を生成しています。電車や花瓶を置く位置が画像に反映されていますね。下の例はテキストをもとに画像を生成した例で、"さくらんぼ (cherries)"や"サッカーボール (football)"が自然な位置に発生していることがわかります。

なおWebデモだと既にアニメ調の美麗画像の生成に成功した人が現れ始めましたが、残念ながら論文ではそのような例は載っていませんでした……
もしかしたら論文用のデータ（他の研究とそろえた比較用のもの）とデモ用の学習データが違うのかも

今後の展望プロンプトプログラミング

AIモデルそのものをいじるのは大変なので、最近はAIに食わせる指示テキストをうまいことこの指示テキストをうまくしてAIの制御をする必要が出てきています。ちょうど今までPythonやC++でコンピュータへの命令をプログラミングしていたように、AI の指示テキストで出力をコントロールするので "プロンプトプログラミング" とこの技術が呼ばれ始めています。優れたプロンプトプログラマーになることが当面はAIを使いこなすために必要になりそうです。

図・データなどは明記しない限り https://github.com/CompVis/stable-diffusion 及び該当論文からの引用です。