見出し画像

CLIP Text Deprojectorを使って画像生成してみる

CLIP Text Deprojectorというのは、Stable Diffusionでテキストエンコーダーに使われるCLIP Text Modelの最終出力のEmbeddingを、Stable Diffusionの入力に使われるlast hidden stateに変換するAIモデルです。

最近、私が作ってLayered Diffusion Pipelineから試験的に利用可能になったので、解説を書いておこうと思います。なお、英語での解説はgithubにあげてあります。

他のStable Diffusionの関連記事

Layered Diffusion Pipelineを使うためのリンク集

CLIP Text Deprojectorとは何か?

下図はCLIP Text Modelの簡単な模式図です。

CLIP Text Modelの概要

最終出力は、図の最上部の「embeds」で、これが対応する画像に対するCLIP Image Modelの出力に一致するようにトレーニングされています。

本来ならば、これか、あるいはその下の「pooled state」をStable Diffusionの入力として使うのが自然ですが、現実にはその代わりに最終層(last hidden state)全体を使うアーキテクチャーとなっています。そのため、本来ならば可能なembeddingに対する操作が使えなくなってしまっています。

この問題には、以前の記事でも時折触れていました。

CLIP Text Deprojectorはこの問題を対症療法的に解決するアイデアで、最終出力(embedding)のみから最終層(last hidden state)を再現するAIモデルで、これを使うことで、Stable Diffusionの入力を本来の最終出力に変更することができます。

モデルのアーキテクチャー

基本的な構造は、1層のTransformerレイヤーからなります。また、最終出力をpooled stateに戻すために、逆行列を掛け算しています。

CLIP Text Deprojectorのモデル構造

逆行列は解析的に計算可能なので、トレーニングの対象になるのはTransformerレイヤーのみです。また、これに加えて次の要素が含まれます。

  • Position Embedding(位置符号化)

  • Final normalization Layer(最終正規化レイヤー)

これらは全てCLIP Text Modelのコンポーネントを再利用することで作られているため、初期値としてCLIPの事前学習ウエイトを流用することができます。

トレーニング

使用したデータセットはLAION 400Bです。CLIPの訓練データのWebImageTextの代替として選択しました。

Google Colabノートブック上でトレーニングを走らせる関係上、1万件ごとのデータに分割して、各分割データ上でearly stoppingを実施しています。

その上で、各分割データでの学習が終わった段階のモデルを使ってStable Diffusionで画像生成を行い、生成された画像をチェックして最適な学習量を選びました。

結果として、3万件のデータで学習した時点のモデルが最もよいという結論を得ました。

トレーニングに使用したColabノートブックはgithubで参照できます。

学習データ量の違いによる生成画像の変化

テストに使用したプロンプトは、次のテキストに画質調整用のタグをいくつか追加したものを使用しています。

  • cat maid (猫耳メイド)

  • 1girl red hair blue eye black skirt(赤髪 青目 黒スカート)

  • 1boy 1girl in class room(少年 少女 教室)

学習データ量の違い

どの時点のモデルも、生成された画像に明らかな崩壊は見られず、プロンプトの内容を一定程度反映していることが分かります。

3万件(赤枠)のモデルを最適と判断した理由は、プロンプトの指示がどの程度生成画像に正確に反映されているかどうかに注目して判断しました。

「猫耳メイド」の場合、「猫耳」が全てのモデルで無視されていますが、学習量が増えるにつれ「メイド」要素もあいまいになってしまっています。

「赤髪 青目 黒スカート」では、「青目」が全てのモデルで無視されていますが、モデルによって「赤髪」や「黒スカート」の色が髪やスカート以外にも影響を与えています。

「少年 少女 教室」では、Deprojectorを使っていないものも含めて「少年」が完全に無視されています。また、モデルによって背景が「教室」ではない画像が生成されています。

使い方

CLIP Text Deprojectorは、すでにLayered Diffusion Pipelineに含まれていて、モデルデータはHugging Faceから利用可能です。

Layered Diffusion Pipelineを使った画像生成スクリプト中で、次のように記述してください。

# 通常通りにLayered Diffusion pipelineを生成
# 詳しい使い方は https://note.com/tomo161382/n/n7cf1d92cd2ff

deprojector = CLIPTextDeprojector.from_pretrained("nanashi161382/clip-text-deprojector")
pipe.text_model.SetDeprojector(deprojector)
pipe.text_model.UseDeprojector()

# 通常通りに画像を生成

この記事が気に入ったらサポートをしてみませんか?