「拡散モデル データ生成技術の数理」を読みました(すごい)
画像生成AIブームによって注目されるようになったアルゴリズムである拡散モデルのおそらく日本語では最初の解説書です*1
拡散モデルの学習、生成手法はGANなど既存の生成型モデルとはかなり異なる部分が多く、理解が難しいところがあり、またその高性能を説明する理論は未だ完全ではないです。
本書は各種関連論文に記載されている導出の計算式が順を追って丁寧に書かれているにも関わらず本文が100ページくらいしかなく非常に密度の高い本です。広範な既存研究を順を追って理解できるように書かれているので極めて有用です。
内容
拡散モデル - 岩波書店で「はじめに」が公開されています。歴史的経緯としては2015年にJascha Sohl-Dicksteinさんにより非平衡熱力学のアイデアから提唱されたのが端緒であり(本ブログの
拡散モデルについて思ったこと、統計モデリング等との関係 - xiangze's sparse blog
も参照)その後のYang SongさんによるScore matchingとの関連性や温度を変えながらサンプリングすることで問題点を解決するという一連の研究があり、2020 年に Jonathan HoさんによりDDPM(Denoising Diffusion Probabilistic Models)が提唱されました。そして3章で説明される確率微分方程式(SDE)を使った記述が(“Score-Based Generative Modeling through Stochastic Differential Equations.”
などで)2021年までに行われたこと、さらにまたtext2imgの条件付き生成法が示され、GANなどの既存手法に比べ画質等の性能が優れていることが示され始めて2022年の相次ぐ商用化、オープンソース化が画像生成AIの爆発的普及の基礎となったことが述べられています。
1章 生成モデルの一つとしての拡散モデル、拡散モデルとスコアマッチング法との関係について
機械学習における生成モデルとその一部であるエネルギーベースモデルの簡単な説明がされ、生成モデルの大分類として尤度
拡散モデルを含む他のほとんどの生成モデル(VAE,自己回帰モデル)が尤度ベースモデルという分類がされる一方GANは暗黙的生成モデル呼ばれ、尤度を直接計算するのではなく生成モデルによって生成された同時に学習される識別モデルにより評価されその結果をまた生成モデルが学習するという形になっています。現時点でGAN以外に暗黙的生成モデルがあるのか、思いつきませんが尤度ベースモデルより自由度の高い分布形状を表現できることから高い表現力を持っています。2つのニューラルネットを交互に学習させるため不安定になりがちというのがあります。尤度ベースモデルと暗黙的生成モデルはそれぞれ別のKLダイバージェンスを最小化する問題として定式化されます。
近年異なる分布のあいだのダイバージェンス(”距離”)を最小化する問題としての最適輸送問題が注目されていて情報幾何的な理解もできるのかもしれません。
Workshop OT 2023 – 最適輸送とその周辺 – 機械学習から熱力学的最適化まで
に注目です。
の性質と生成(サンプリング)過程で用いるランジュバンが説明されたあとに
学習目標の分布関数p(x)を含んだ形の明示的スコアマッチングの式は
ですがそれをp(x)を含まない形(暗黙的スコアマッチング)
が特定の条件では等価(定数を除いて一致)になることの証明が書かれています(Estimation of Non-Normalized Statistical Models by Score Matching)
さらにdenoising score matching
が明示的score matchingと等価であることが証明されます。
(参考 A Connection Between Score Matching and Denoising Autoencoders | Neural Computation | MIT Press Denoising Autoencoderとその一般化 )
2章 SBMとDDPMについて、学習対象の導出、統一的理解
SBM(Score Based Model)とDDPM(Denoising Diffusion Probabilistic Model)について説明されています。
まずSBMのサンプリング シミュレーテッドアニーリング(焼きなまし法)との違い、利点が詳しく説明されています*2。
スコア関数にも係数αiがかかっているのが違いです。これによって高温(αi大)のときは大きく動くことになります。
Generative Modeling by Estimating Gradients of the Data Distribution | Yang Song
(Score-Based Generative Modeling through Stochastic Differential Equationsでは生成過程をSMLD(Score Matching Langevin Dynamics)と呼んでいます。)
そしてデータに徐々にノイズを加えていく拡散過程とその逆の拡散モデルであるDDPMの説明、学習対象のLoss関数の導出、階層的なVAEとも言えなくもないそのELBO(evidence lower bound)の導出はAppendixと合わせて行われています。長い式変形ですが直感的な理解が豊富に取り込まれています。そのLoss関数はSBMのそれとは時刻ごとの係数以外は同じ形をしています*3。
さらにSBMとDDPMは統一的に理解できること、そしてノイズスケジュールも学習可能であることが説明されます。
prompton.io
3章 時間連続極限の場合と拡散モデルの高性能について
時間刻みを無限小にとった極限では拡散モデルSBMとDDPMは確率微分方程式(SDE)として記述できることが説明されてます。またこのSDEは推定された分布を常微分方程式(ODE)でも表現できること、それによってノイズを使わないサンプリング手法が成立しうること、それでもノイズを使ったほうが離散化誤差を改善できる(ディザ)ことなどが説明されます。また離散時間の場合とは違ったsignal-noise関係を使ったSBMとDDPMの統一的理解( Elucidating the Design Space of Diffusion-Based Generative Models のAppendixでは詳しく書かれている)が3.6.3とAppendixで成されます
拡散モデルが他のモデルに比べ優れている点(学習が安定している、生成過程を部分生成問題に分割できる、様々な条件を組み合わせられデータの対称性も組み込める、同一のスコア関数を何回も繰り返して適用できるのでパラメーターサイズの割に表現力が高い)とその直感的理由が説明されています。一方で生成のためのステップ数が多いので時間がかかるという難点もあります。なぜ拡散モデルが高い汎化性能を示すのかについてはまだわかっていないようです。
深層学習がパラメーターが多いにもかかわらず高い汎化性能を示す現象(暗黙的な正則化)に関しては
www.iwanami.co.jp
の著者今泉先生の研究
www.c.u-tokyo.ac.jp
がありますが、拡散モデルでの汎化性能の高さが深層学習一般と同じ原因によるのか、異なるものなのかはまだわかりません*4。
4章 条件付き生成、高解像度化、群不変、同変性
stablediffusion,midjourneyなどの目玉機能であるtext2imgを実現した条件付き生成(ガイダンス)についての説明であり1章で説明されたエネルギーベースモデルが構成的(別々のモデルのエネルギーを足し合わせて新たなモデルを作れる)であることから
のように条件付きと条件なしのスコア関数を任意のスケールγで足し合わせることができること、さらに条件付きの方にdropoutのような処理を行うことができることが説明されています。
Stable Diffusionでは高解像度の画像を出力するために画像にVAEをかけて圧縮したlatent variableに対して拡散モデルを適用しています。これとは別の方法として部分空間拡散モデルという時間発展に従い空間を狭めていく方法が説明されています。これは多様体仮説を用いた賢い方法で温度変化に依らず「データが希薄な領域においてscore functionがうまく推定できない」という困難を解決する方法のように見えます。画像集合などデータ集合固有の狭め方があるようにも思います。温度を変えていくと有効な次元が変わるというのは従来の数学や情報科学で出てきた多様体の概念では統一的には説明できないように思われ興味深いです(知らないだけかもしれません、persistent homologyの解像度を変えてデータの位相的性質をみていく様子は近いようにも思えます)。
著者の方が計算化学の会社の社長さんなので分子構造の持つ対称性を組み入れて生成ができるモデルとして注目され、特に分子構造の持つ回転、並進などの対称性(SE(3))に対して不変、同変性を組み込む拡散モデルも説明されていて特に興味深いです(群をなす変換に対して不変なニューラルネットに関してはGroup Equivariant CNNとM理論について - xiangze's sparse blogも参考にしてください)。
prompton.io
5章 アプリケーション
その他画像処理、動画、音声、分子化合物、敵対的摂動に対する頑健性向上、圧縮などの応用が紹介されています。Stable diffusionのパラメーター全体が4GB程度しかないのは驚くべきことです。圧縮効率の高さとtext2imgでは画像そのものではなくpromptと画像の関係を学習していることによるものなのでしょうか。
感想
式展開の説明が非常に丁寧で、記号の説明も冒頭でされているので丁寧に追っていけば学習対称の関数型、生成(サンプリング)アルゴリズムとそれらの間の関係が理解できます。しかし拡散モデルの説明を高密度で行なっているために機械学習や画像生成の理論面の知識、他の生成モデルに慣れた経験がない人にはイメージがわかず難解かもしれません。1,2,3章では既存の手法や数式の等価性が詳しく説明されていて機械学習の分野の中でも理論的研究が進んでいるけれど高性能の理由はまだはっきりとはわからないという拡散モデルの特徴が出ているような気がします。
エネルギーベースモデルについては第一章で軽く触れられているものの分配関数から物理量が計算できるという統計力学の考え方、MCMCなどモンテカルロ法によるシミュレーションや(制限)ボルツマンマシンなどについて知っているとよりとっつきやすくなります。特にscore functionが分配関数に依存しないという性質は重要でMCMCとの共通点もあります。「ベイズ深層学習」やPRMLをよんでおくと良さそうですがどちらも結構なボリュームがあります。
stable diffusionでの画像生成に必要な他の部分の説明もなく特にstable diffusionではscore関数はcross attentionが組み込まれたU-netであることを意識しておく必要があります。
逆に画像生成のみならず動画、音声、分子化合物の生成、圧縮など幅広いアプリケーションを紹介されていて拡散モデルのポテンシャルが感じられます。
他の本
Stable Diffusionで使われているアルゴリズムの解説 - 過学習ショップ - BOOTH
stable diffusionのアルゴリズム内で中心的位置を占める拡散モデルの解説としては本書と同様に式の説明は詳細です。ほかの生成モデル手法や統計モデリング手法との比較がされていて拡散モデルの位置づけがわかりやすいかもしれません。stable diffusionの他の構成要素(CLIP, VAE, Cross Attension Unet)などについても少ないながら説明があります。一方でDDIMとSBMの違いの認識が曖昧でその本書2,3章で詳しく述べられた相違点および統一的理解やDDIM,ノイズを使わないColdDiffusion法など多彩なサンプリング手法の解説などは全く成されていません。
「Stable Diffusionで使われているアルゴリズムの解説」の参考文献 - 過学習ショップ - BOOTH
本書で挙げられた参考文献のいくつかのリンクがあります。
画像生成業界の発展は非常に早く書籍の形では追いつくのが難しいです。最近では姿勢の指定に優れたControlNetや学習用画像に僅かなノイズを透かしのように入れることで特定のアーティストの模倣を防ぐGLAZEなどが注目されています*5。githubなどのソースコード、ブログその他ネット上の情報を発信してキャッチアップしていきましょう(他力本願)。
https://prompton.io ではAIイラストを描いてほしい、描きたい人を募集しています。
*1:同人誌はいくつかあるようです
*3:p44.の拡散モデルは拡散過程で(中略)途中の層をランダムに抜き出し、それぞれを独立に分けて学習する のところがちょっとわかりませんでした。Loss関数が時間ステップの和の形にできるということでしょうか
*4:上記論文の参考文献ではLangevin力学を使った [1901.04609] Tightening Mutual Information Based Bounds on Generalization Error,https://arxiv.org/pdf/1801.04295.pdfが紹介されているので拡散モデルの生成過程も同様に汎化誤差を解析できるかもしれません
*5:敵対的摂動に対する頑健性向上とは反対の立場の研究です
