🔖

[実装付き]Stable Diffusionの追加学習に適する画像を、VAEを利用して選別する

2022/11/11に公開

はじめに

こんにちは。

前置きが思いつかないので、突然本題に入ります。

Stable Diffusionをはじめとする、Latent Diffusion Model(以下LDM)の追加学習手法は、その登場以来様々なやり方が提案されてきました。

例えば、学習データの表現を語に埋め込み、txt2imgの結果を直感的に操作するTextual Inverisionや、学習データによってDenoising U-Netのパラメータを変換するDNNを挿入するHyperNetworksはその代表的な手法です。

しかし、どの手法にも良い点と悪い点があり、こちらが決定的に良い、ということはなく、追加学習によるアウトカムの良し悪しを決定する最も大事なことは変わりません。
それは、適切な学習データを用意することです。

しかし、適切な学習データとは何か？という話はなかなか簡単に済ませられる話ではなく、モデルやタスクの特徴によっても千差万別です。

そこでこの記事では、Stable Diffusionにおいて適切な学習データ(ここでは画像とします)とは何か、ということにまず触れ、1つ簡単に仮説を立てた上で、学習データ候補の画像群の適・不適を選別する定量的な指標を提案します。
それから、その実装を紹介し、実際に手元の画像群を分析する手順を説明します。

Stable Diffusionにおける適切な学習データ(画像)の選別とは

前提として、本来Diffusion Modelとは、例えば1024x1024ピクセルの画像を生成する場合、1024x1024xチャンネルの大きさのノイズを徐々にデノイズして出力を得るモデルでした。しかし、これはちょっと考えられない計算コストを要するため、計算量削減の一手としてLDMが生み出され、その系譜にStable Diffusionが位置しています。

LDMでは、VAEのレイヤー構造の中間に位置する小さな潜在空間のサイズでノイズを生成し、その中でデノイズを行ってから実空間（ピクセル空間）の大きさまで伸長し、出力を得ます。これによって個人が用意できる範囲のコンピュータでも十分動作するようになったために、現在のようなブームが生み出されているわけですが、当然ながら、本来のモデル構造になかった問題点が生まれています。
それは、潜在空間のサイズでどうやっても表現できないものは上手に生成できませんし、うまく学習もできない（であろう）ということです。

例えば、めちゃくちゃ小さな文字がズラッと書いてある画像とか、異なるモノが複雑に入り組んだ画像とかを学習させようとしても、VAEで潜在空間に射影する際にかなりの情報が失われてしまい、Denoising U-Netまで重要だったはずの表現が伝わってくれません。

話が飲み込みやすいように、よく使われる説明画像を載せておきます

また、VAEによる実空間から潜在空間へのエンコード処理はもちろん、潜在空間から実空間への橋渡しについても完璧な精度で行うことは不可能なので、仮にデノイズの過程が精巧な潜在表現を作り上げたとしても、VAEが苦手な表現は上手にデコードされません。

例えば、上の説明画像をStable DiffusionのVAEで圧縮・再構成するとこんな感じになります。これでも十分すごいですが、文字は特に苦手な表現の一つなのか、だいぶグチャグチャになります

まとめると、オリジナルのDiffusion Modelや、GoogleやNVIDIAが自慢げに結果だけ見せつけてくる別形態のDiffusion Modelにとっては理解できるのに、LDMにとっては理解することが難しい画像表現とは、大体が潜在空間に落とし込むことが難しい、または潜在空間から上手くデコードできない表現と言っても過言ではないので、潜在空間と実空間を橋渡しするVAEにとって難しい画像ばかりを学習データとして使用すると、おそらくモデルの学習はうまくいきません。

つまり、Stable Diffusionの追加学習に使用する画像データの適切な選定方法として、VAEにとって難しい画像を省くことは有効だと考えられそうです。

ここまで来たらあとは単純ですが、Stable DiffusionのVAEの入出力を利用した指標を設けることで、主観に依らず適切に学習しやすい画像とそうでない画像を選り分ける手法があるのでは？と考えました。

画像の選別方法について

前項の繰り返しになりますが、ここでの内容は、Stable DiffusionのVAEが上手く扱えない画像は学習も難しいだろう、という仮説に基づいています。

もう少ししっかりと言い直します。
VAEが上手く扱えない画像とは、おそらく等価で上手に再構成できない画像であり、それは学習プロセスの中でStable Diffusionに入力されたとき、潜在空間に上手く落とし込まれることがないか、上手く落とし込まれたとしても、それやそれに近い潜在表現を上手にデコードすることができない画像と言い換えることができます。

ゆえに、それを学習に使用したところで、結局は潜在空間内の学習可能パラメータに対して表現を正確に伝えることができないか、それを抜きにしてもVAEのデコーダーがよく知らない表現なので出力できないか、のどちらかの理由で上手に扱うことができないため、ただ学習の妨げになってしまうのではないか？ということです。

そして、その仮説に基づくと、学習対象の候補となる画像と、それをVAEに入力して再構成された画像の間で誤差(再構成誤差)を計ることで、その画像がどの程度VAEの表現力の範疇になく、ひいてはStable Diffusionの追加学習に適していないのかを一側面から見ることができるはずですし、誤差の値が小さいものは、その側面において、学習データとして適していると言うことができるのではないでしょうか。

ただ一点、単純に再構成誤差をMSE等で測った場合、当然ですが画像全体の情報量が少ないものほど誤差が小さく、反対に情報量が多いものほど大きくなってしまうという問題があります。

この問題を考慮せず、画像全体の平均的な誤差を指標として画像を選別してしまうと、全体として誤差の値が小さくても局所的に大きな誤差がある画像を弾くことができません。
その大きな誤差が出ている箇所こそが潜在空間内の処理によって注目され、強く学習される対象だった場合、モデルが間違った方向に学習を進めてしまうこともあり得ます。

以下は誤差を安直に算出した例です。Stable DiffusionのVAEが学習時に使用する損失関数で再構成された画像の誤差を評価したところ、2枚目の画像は顔部分（特に目元）がだいぶ崩れているのに、誤差の値は小さく算出されてしまいました。

この問題を解決するために、画像を縦横で小さく分割してそれぞれのチャンクでMSEを算出し、その中の最大値を誤差の値とするようにしたところ、だいぶ直感に沿う評価を下すことができるようになりました。

そして、それを実装したGithubのリポジトリがこちらです。

このリポジトリをクローンし、中にあるノートブックで画像を処理・分析することで、自分が学習したい画像の中で相対的にどれが学習データとして好ましくないのかを定量的に確認できます。

ここまでで大体伝えたいことは伝えたので、後は聞かなくても分かるという方はこの先を読まなくても大丈夫です。誤差値を参考にしつつ学習に使用する画像を選定してみてください。

次の項からは、ノートブックが出来ることと、使い方を解説します。

ノートブックが出来ること・導入と使い方

出来ること

このノートブックが出来ることは以下の2つです。

与えられた画像データセットをVAEに入力出来るように最低限前処理する
データセット内の1枚1枚の画像がStable Diffusionの学習データとしてどの程度適していないのかの指標を誤差値で与え、相対的に適しているものを示唆する

やれそうでできないことは以下です。後々やれるようにしたいなぁとは思っています。

stable-diffusion-webuiの前処理機能のように、与えられたデータセットに対してもうそのまま学習にも突っ込めるぐらいのしっかりとした前処理（データオーギュメンテーション、キャプション生成など）を行う
算出した誤差値を元に、自動で学習に適する画像だけを選別して画像データセットを作り直す

2はすぐにでもできるようにしたいですが、どの程度の誤差値で適・不適を仕分けるのが良いのかについてまだまだ理解出来ていない部分があるので、今のところ手作業で選別を行う想定です。

導入

まずは、リポジトリをクローンしてディレクトリの中に入ります。

https://github.com/discus0434/evaluate-images-to-feed-diffusion.git
cd evaluate-images-to-feed-diffusion

次に環境構築ですが、3つの分岐があります。

1. 最新のstable-diffusion-webuiが動く環境が既にある場合

多分その環境を流用すれば大体動くので、プロジェクトディレクトリの中で以下のコマンドだけ実行してください。

wget -O ./models/waifu-diffusion-v1-4/kl-f8-anime2.ckpt \
    https://huggingface.co/hakurei/waifu-diffusion-v1-4/resolve/main/vae/kl-f8-anime2.ckpt

もしそれでノートブックが動かない場合、2の手順を行う必要があります。

2. Python仮想環境を作成する場合

condaで仮想環境を作成し、その環境に入ります。

conda env create -f=environment.yaml
conda activate sd

次に、同じディレクトリ階層に存在するsetup.shを実行します。

chmod +x ./setup.sh
./setup.sh

3. Dockerを使う場合

ビルドして、

docker build . -t eval_images

実行します。

docker run -it --runtime=nvidia --gpus all -d --restart=always eval_images:latest bash

コンテナの中に入ると、condaの仮想環境がsdという名前で作成されているはずです。

使い方

まずはプロジェクトディレクトリ下のimagesに選別対象の画像群を放り込む必要があります。

画像の拡張子やサイズが統一されている必要はありません。pngとかjpg、jpegのように一般的な画像なら大体大丈夫です。

画像を放り込んだら、eval_images.ipynbを開き、ノートブックのセルを全実行すればとりあえず勝手にゴチャゴチャやり始めるはずです。

ノートブックの設定項目はノート上部のConstantsにまとめてあります。
ほとんどの項目は弄る必要がありませんが、以下の項目は画像の前処理後のサイズを変更するため、Stable Diffusionの追加学習に使用する予定の画像サイズに合わせると良さそうです。

# image size, after preprocessing
IMAGE_WIDTH = 512
IMAGE_HEIGHT = 512

また、処理が終了するとノートブックの章立てで言うVisualize Results以下に画像それぞれの損失の値が可視化されているので、それを参照しながら画像の選定を行うことができます。

参考までに例を挙げると、人物が写ったフリー画像や、いらすとやで入手した画像からなる27枚の画像データセットを処理した結果、誤差の値が大きい5つ、小さい5つの画像は以下のように出力されました。

誤差が大きい5枚の画像（上段）は、ぱっと見情報量が少なく、ちゃんと再構成できているように見えるものも含まれていますが、大体は人物の顔や紙幣のような情報量の多いモノが引きで写っていて、そのような局所的なオブジェクトがVAEに上手く扱えなかったっぽいことが分かります（拡大して見るとだいぶおかしい部分があります）。

次に誤差が小さいと見なされた下段の5枚を見ると、どこを切り取っても情報量が少なく、複雑ではないいらすとやの画像が並んでいます。順当に高い精度で再構成されているため、これらはVAEに扱いやすく、潜在空間らへんにある学習可能なパラメータにも正確な表現が伝わりそうに見えます。それなりに直感に沿う誤差値が算出されているのではないでしょうか？

また、トップ5、ワースト5の結果とは別に、個別の結果もノートブックの最後に誤差値の降順で表示されるようになっています。

おわりに

Stable Diffusionの追加学習に使用するデータセットの質を上げるには、画像の仕分け、キャプションの正確な記述などありますが、このあたりはまだまだ自動化手法が追いついておらず、人力で丁寧に行ったほうが断然良い結果が出やすい作業だと考えています。

その中でも画像の仕分けはかなり勘や感覚に頼る部分が強いと認識していますが、この記事で提案する手法を使えば、仕分けに際してStable Diffusionの気持ちがある程度定量的に取り入れられるのではないかと期待しています。

また、ノートブックを少し弄れば、大量の候補画像を仕分けるときそれなりの精度を発揮する自動化ツールとしても使えるはずなので、データセットを人力で作成するにせよ自動で作成するにせよ役に立つのではないかと思い、リポジトリを公開しました。

とはいえ、昨日ふと思いついて1人で書いた実装なので、何かしら問題があるかもしれません。もしここを直した方がいいよ、みたいな話があればIssuesにどんどん投げていただけると嬉しいです。