もし生物情報科学専攻の大学院生が "StableDiffusion" を理解しようとしたら 10 ~追加学習と学習阻害~

前回

cake-by-the-river.hatenablog.jp

今回は、主にStable Diffusion自身ではなく、それに関連した話題を扱う予定です。汎用の画像生成モデルを特殊な用途に転用する追加学習、およびStable Diffusion等のLatent Diffusion Modelの学習を阻害する仕組みについて説明していきます。

(一介のしがない)イラストレーターとしては「敵(追加学習)を知り己(学習妨害)を知れば百戦危うからず」という感じですね。

追加学習

fine-tuning

今までに解説した画像生成は、特定の画風や二次創作といった特殊な目的で用いるには、モデルにとって未知の概念を含むことが多く、難しいものでした。一般に、学習データには含まれないデータ・概念を学習済みのモデルに導入することは、追加学習や、fine-tuning などと呼ばれています。Stable Diffusionを代表とする画像生成モデルでも、これら追加学習により特殊な目的の画像生成を可能とすることが出来ます。


拡散モデル、特にStable Diffusionの追加学習については、2022年後半あたりから急速に研究されているようです。主な方法としては、追加学習したい固有概念のテキスト埋め込みベクトルを習得するタイプ、および追加学習により少し変化するモデルパラメータの重みを残差として低次元近似するタイプの2種類があります。

方法1:概念の埋め込み

Textual Inversion

arxiv.org

Textual Inversionは名前の通り、学習したい概念に対応する言語の埋め込みベクトルを逆算する方法です。代表語を用意し(論文ではS*)、その語を条件に含んだ画像生成を行います。その結果と入力画像の誤差を減らすように代表語の埋め込みベクトルを学習します。

学習方法は通常のDDPMの学習とほとんど同じ形を取ります。DDPMの学習では、誤差逆伝播を用いて計算グラフを遡りモデルパラメータの勾配を計算しますが、ここではパラメータの代わりに埋め込みベクトルの勾配を計算するようにしています。各学習ステップで代表語の埋め込みベクトルのみ最適化することで、モデルを変えずに代表語の埋め込みベクトルに概念を押し込むことが出来るわけです。

Textual Inversionの原理。論文より引用


Textual Inversionで習得した埋め込みベクトルは一般の単語と同じように扱うことが可能であり、文章中に組み込むことで入力画像と異なる状況下での画像生成も可能となります。特にTextual Inversionは、"マグカップ"のような実体のみならず、特定の画風や雰囲気を学習させることも可能です。例えば、EasyNegativeと呼ばれる埋め込みベクトルは、画像生成時に生じる”好ましくない概念”をあえて習得させたものであり、これをNegative Promptに採用することで、生成する画像に好ましくない概念を登場させにくくする効果を持たせることが出来ます。

Textual Inversionによる画風の習得と適用。論文より引用。
DreamBooth

arxiv.org

DreamBoothは、GoogleがImagen(Stable Diffusion同様の画像生成の拡散モデル)用に開発した方法であり、スタンダードな追加学習法だと言えます。

Textual Inversionと同様に習得したい概念を表す代表語 identifier 、およびその概念に最も近い一般的な用語 class(例えば "dog" や "car" など)を用意します。一般的な機械学習の追加学習と同様に、identifier と class を同時に用いたテキストを条件とした画像生成に対し、入力画像との誤差を最小化するようにモデルのパラメータを更新します。なお、identifierには3文字の無意味な文字列(xks, yrqなど)を用いることが推奨されています(意味がある文字列を取り入れた場合、元の意味が学習に影響するため?)。

しかし、このままではlanguage drift、すなわち入力画像を学習することに専念した結果、元のモデルが持つ表現力や汎化性能を失う可能性が出てきます。そこで、class のみを条件に用いた画像生成の結果(正則化画像群)を用意し、この結果とidentifierを入れた場合の結果がかけ離れすぎないように正則化する項も目的関数に加えます。これにより、classで表現される概念の範囲内に収めつつ、入力画像があらわす特殊な状況を代表語に取り入れたモデルを構築することが出来るわけです。

DreamBoothの原理。論文より引用。


DreamBoothは、一つの実体概念を習得することに長けており、その精度も高いものとなっていますが、Textual Inversionと異なりモデル全体を更新するため、モデルの持ち運びが悪かったり、複数の概念を組み合わせて利用できなかったりする欠点があります。

DreamArtist

arxiv.org


他にも似たような方法は色々と出ているようですが、一例としてDreamArtistを紹介します。DreamAartistは、名前はDreamBoothに似ているものの、内容的にはTextual Inversionに近いものとなっています。基本的には、Texual Inversionが単なる代表語のプロンプトに対する埋め込みベクトルの逆算であったのに対し、Negative Promptに対応する代表語も用意し、両方を同時に学習させる方法です。

DreamArtistの仕組み。論文より引用。


シンプルな改良ですが、Negative Promptの効果によって(Positive)Promptの出力対象に対する効果がはっきりとし、上記の2手法に比べて文章による制御性が増したようです。その結果、入力画像が一つの場合であっても、高精度に対象を学習することが出来たと主張しています。

Textual InversionとDreamArtistのAttentionマップの違い。赤い領域が各テキストが示すAttentionの高い領域に対応する。


個人的にこの論文はちょっとはっちゃけすぎている気がする

方法2:重みの微修正

LoRA

arxiv.org

LoRA(Low-Rank Adaptation)は、Microsoftが大規模言語モデル用に開発した方法です。基本的なアイデアは、追加学習を重み行列に対する摂動として捉え、それを特異値分解のように低ランク行列で近似する、ということです。


追加学習は、事前学習に少量のサンプルを加えて学習を行うことでした。この場合、学習後の重みは(膨大なサンプルからなる)事前学習のものからそこまで多くは変化しないものと考えられます。つまり、

 W = W_0 + \Delta W

のような形式で書けると考えます。LoRAは、この  \Delta W を低ランクの行列  A, B を使って

 \Delta W = BA

として近似します。

LoRAの仕組み。論文より引用。

学習時も、事前学習の重みは変更せず、低ランク行列  A, B のみを学習させます。こうした学習は全体を再学習せずに済むため効率が良く、Parameter Efficient Fine-Tuning (PEFT) と呼ばれることがあります。

また、DreamBoothなどの手法と異なり、摂動  \Delta W は加算が可能なため、複数個組み合わせることができます。つまり、「りんご」と「青空」のそれぞれのLoRAを足し合わせるような計算が可能であり、出力結果もその組み合わせからなるように出来る、という訳です。


また、LoRAにTextual Inversionを組み合わせるPivotal Tuningというテクニックも知られています。つまり、新しい概念に対応する埋め込みベクトルと摂動に対応する重みの更新を同時に行う、ということです。

LyCORIS

LoRAは非常に単純な手法であるため、様々な派生手法があります。LoConは LoRA for Convolution Layer の略称で、LoRAが基本的にTransformer等の重み行列にしか対応していなかったのに対し、畳み込み演算などにも適用できるように拡張したものです。

一方、LoHaやLoKrなどの手法では、単に1つの低ランク行列のペアのみならず、複数のペアや非線形的な演算を加えることで、より表現力をあげようと考えたものです。これらの手法はLyCORISという統合フレームワークの形で提供されています。

個人的にはDeepでポン(≒謎テクニック)の極みのようなキモさがあって好きじゃないです

LoHaやLoKrの概念図。LyCORIS論文より引用。

openreview.net

学習阻害

画像生成AIは、特に特定の対象(人物/画風/画家自体)を「狙い撃つ」追加学習により、多くの懸念点を生み出しています。日本では、これらについて法整備等の法的な取り組みが進行している最中であるとはいえ、すぐに対処できることではないと考えられます。僕は理系であり、「どこまでが権利的に主張できるか?」といった曖昧な議論にはめっぽう弱いので、今すぐに考えられる面白い取り組みである「学習阻害」について取り上げようと思います。


学習阻害は、AI側からすると「アタック」として扱われます(正義の反対は正義、のような話に近くて面白いですね)。そして、学習阻害に対抗する仕組みも「ディフェンス」として扱われます。これらのアタック&ディフェンスについての総説である以下の論文を参考に、いくらかピックアップしていきます。

arxiv.org

敵対的攻撃

DreamBoothやLoRAのような手法は、少量のサンプルに対して追加学習を行います。こうした少数のサンプルによる追加学習は一般的なAIの学習メカニズムと異なり、汎用性ではなく特殊な用途への狙い撃ちを目的とする分、「過学習」の状態になりやすいです。過学習されたモデルは、次に説明する敵対的攻撃メンバーシップ推定といったアタックに脆い状態となります。


敵対的攻撃は、Stable Diffusion等がでるよりもずっと昔から研究されてきた分野です。例えば、CNNからできた画像分類問題に対し、ちょっとしたノイズを加えることですぐに謝った判定をしてしまう、といったことが知られています。

パンダの画像に若干の敵対的ノイズを加えることで、テナガザルだと勘違いするAI。以下の論文より引用。

arxiv.org

これの延長線上として、Stable Diffusion等の拡散モデルに対する敵対的ノイズを画像へと加えるGlazeやMistといったツールが知られています。

Style Cloak (Glaze)

Glazeシカゴ大学のチームにより開発された敵対的攻撃の手法であり、「画家特有の画風」を覆い隠すようなノイズを加えることを考えます。より正確には、「まるで他の画風の絵のように見せる」ということをします。

arxiv.org

Glazeの仕組み。論文より引用。

より詳細に説明しましょう。Glazeでは、Stable Diffusion等の公開されたAIについて、特徴抽出器  \Phi が公開されていることを利用し、これに対する攻撃を考えます。入力画像  x に対し、自分の画風とは異なる他の画風  T へと変換するAI  \Omega を用いて画風を変換した画像を考えます。これを、入力画像へノイズ  \delta_x を追加したものと比較し、その特徴抽出の結果が近くなるようにします。Stable Diffusionの場合は、 \Phi, \Omega ともにStable Diffusion自体を利用することが可能です。

 \min || \Phi(\Omega(x,T)) - \Phi(x+\delta_x)||

つまり、特徴抽出器にとってノイズを加えた後の画像が他の画風の画像と類似するようにするわけです。これを用いて画像生成を行えば、特徴抽出器は元の画風ではなく他の画風に類似するようになり、全体として似つかない画像を出力するようになる、ということを狙っているわけです。

ただし、このままではめちゃくちゃなノイズを加えたほうが有利ということになってしまい、元の画像自体の見た目にも影響を及ぼしてしまいます。そこで、GlazeではLPIPSと呼ばれる画像の見た目のズレを評価するスコアを用い、LPIPSのスコアが一定以下、すなわち人の目ではそこまで変化しないようなノイズの範囲で調べる(正則化する)ように限定します。これにより、見た目はそこまでズレずとも画風を隠す効果を残そうとします。

Mist

arxiv.org

Mistは、より追加学習の方に着目した敵対的攻撃の手法です。基本的には各手法における目的関数をハッキングするようなノイズを設計する手法であり、DreamBoothやTextual Inversionに対応した目的関数を設定することで、それを最大にする(すなわち学習を非効率化させる)ノイズを加えます。

Mistによるノイズと生成の阻害。論文より引用。

こうした手法は、DreamBoothやTextual Inversionに対する設計は出来ていますが、LoRA等の手法に対応するものは(自分の知る限りは)まだないようです。また、Stable Diffusion等の「相手がよく分かっている」場合に効果的である一方、新しい画像生成手法が現れると、使ったノイズが必ずしも役に立つとは限らない点が挙げられます。

概念消去

敵対的攻撃のディフェンスとして、「非学習」というものが挙げられています。「学習した概念を忘れさせる」の方が分かりやすいかもしれません。これは必ずしもディフェンスという意味だけでなく、セーフティ上問題があるコンテンツを生成しないようにする、といった用途でも利用することが可能です。


ESDは初めて概念消去を扱った手法と考えられており、いたってシンプルなアイデアからなります。Stable DiffusionではCFGを用いることで条件付きの画像生成を可能としましたが、この時の係数をマイナスにすることで、「その概念がない状態」を条件とした生成にする、というものです。

 \displaystyle \epsilon_{\theta, \gamma}(y_c) = \epsilon_\theta(y_{uc}) - \gamma (\epsilon_\theta(y_c) - \epsilon_\theta(y_{uc}))

ESDの概念図。真ん中の数式が上で上げたCFGの係数のマイナス化に対応している。論文より引用。

arxiv.org

一方、LoRAにおける概念消去を考える方法も提案されています。LECOはESDと同様の計算をLoRAに対して適用したものといえます。詳しくは以下の記事をご覧ください。

zenn.dev

メンバーシップ推定(MIA)

今までは入力画像にAIに対して不利な要素を組み込むアタックを考えてきました。敵対的攻撃の他の例として、そもそも事前学習段階で学習されるデータ(例えばインターネット上の画像)に「毒」を入れる、といった面白いアイデアもあります(バックドア攻撃)。しかし、別の観点として「私の画像がAIに読み込まれているか」を調べるというものがあります。これは、AIの学習データに特定の画像が含まれているかどうかを推定するメンバーシップ推定(MIA)と呼ばれています。


MIAは、どのようなAIを相手にするか、という問題で分類されます。つまり、「Stable Diffusionを相手は使っている、と仮定できる」状況なのか「よくわからんけどきっとAIが生成している」とするのか、です。また、そのモデルの学習に使ったデータがどのようなものかをどの程度知っているかによっても変わってきます。完全に情報が分かっている場合はホワイトボックス、完全に分からない(比較したい画像しかない)場合はブラックボックス、その中間はグレーボックスと呼ばれています。


これらのMIAを行うモデルは色々と挙げられています。基本的なアイデアは「過学習」の検出です。学習に用いたデータはそうでないデータと比べてノイズの推定誤差が少ない場合が多いため、これを検出することを目標とします。ホワイトボックスのモデルでは当然手元で自分の画像を与えたときの挙動が見えるため、それを参考に推定精度をあげることができます。具体的には、GAN-LeaksやPIAなど色々と挙げられています。

一般に、拡散モデルは他の手法(例えば、GANなど)と比べてこうしたメンバーシップ推定に弱いことが示されています。ただし、これらの手法はまだ追加学習、とくにLoRAなどには直接対応できていないようです。いわゆる「狙い撃ちLoRA」がされているかどうか、という問題における「DNA鑑定」のような手法の開発は、これからになるかもしれないですね(DNA鑑定ほどの精度はそもそも難しいですが)。

まとめ

現状、様々なアタックとディフェンスが研究されています。これらをすべて理解するのは難しいと思いますが、画像生成AIの仕組みを理解してきた今、これらの手法のメリット・デメリットを抑えることができれば、味方につけるでも敵にするでも役に立つ可能性が高いです。

拡散モデルに対するアタック&ディフェンス。論文より引用。


今回で画像生成AIであるStable Diffusionのシリーズ解説を終わりにしたいと思います。深層学習の基礎からアタック&ディフェンスに至るまで、多くの内容を理論的に抑えることができました。自分は実は既に画像生成AIに対する興味は薄くなってきてはいるものの、タンパク質等の分子生成といった別の分野への応用もあるため、今後も適宜情報を追っていこうと思っています。