ジョイジョイジョイ

ジョイジョイジョイジョイジョイ

拡散モデルと最適輸送

最適輸送の理論とアルゴリズム』が重版して第 5 刷となりました。皆さまありがとうございます!

漫画家さんやイラストレーターさんが重版したときに重版感謝の描き下ろしイラストを投稿しているのを見ていいなと思ったので、僕も専門書が重版したときに重版感謝の書き下ろし専門記事を投稿します。

本稿では、最近話題の拡散モデルと最適輸送の関係を直観的に解説します。

拡散モデルは画像の生成によく用いられる生成モデルです。モデルはノイズ入りの画像を受け取り、ノイズを除去することを目指します。生成時には、完全なノイズ画像からはじめて、モデルによりノイズを除去することと、微小なノイズを加えることを繰り返して洗練させていき、自然画像を得ます。

拡散モデルの動作の図示

このように、ノイズ  x_1 から自然画像  x_0 までゆらぎながら変化する過程をブラウン橋 (Brownian bridge) と言います。ブラウン運動 (Brownian motion) は水中の花粉の運動などで皆さん耳にしたことがあると思います。ブラウン運動は基本的に初期時刻の場所や分布が固定されて、そこから自由に拡散していくのに対し、ブラウン橋は初期時刻と終了時刻両方の場所が決まっています。上記の例の場合は、ノイズ分布  p(x_1) と自然画像分布  p(x_0) の間をブラウン橋で繋いでいることになります。

最適輸送とは、ある確率分布を別の確率分布に最小コストで輸送させる方法のことです。

最適輸送の動作の図示

直観的には、寄り道なしでまっすぐ確率分布を繋ぐのが最適輸送です。

上の二つの例を見ると、拡散モデルと最適輸送が似たことをしていることが直観的に分かると思います。拡散モデルは揺らぎながら移動しているのに対し、最適輸送はまっすぐ移動しています。

実は、拡散モデル(ブラウン橋)は最適輸送の変種と密接に関わりがあります。それがエントロピー正則化付き最適輸送 (entropic regularized optimal transport) です。これは、輸送のコストを小さくしつつ、それだけではなく輸送方法のエントロピーを大きくするようにバランスを取ったものです。以下の定理が知られています。

定理 ブラウン橋の両端の同時分布  p(x_0, x_1)エントロピー正則化付き最適輸送である。

詳しくはサーベイ論文 [Léonard 2014] や最適輸送の教科書 [Peyré+ 2019, Section 7.6] をご覧ください。

本稿では数学的な詳細には立ち入りませんが、ここで重要なのは、拡散モデルは「ほぼ」最適輸送であり、最適輸送の輸送方法を少しぼやかしたものであるということが、数学的にもきちんと示せるということです。

この関係を念頭に置くと、拡散モデルを最適輸送風に解釈し直すこともできますし、最適輸送を拡散モデル風に解釈することもできます。

例えば、[Gushchin+ NeurIPS 2023] では、最適輸送を拡散モデル風に、ノイズを除去するモデルを用いて解く手法を提案しています。多くの最適輸送のアルゴリズムは、離散分布に対してしか適用できないか、連続分布について適用できてもコストの値が求まるだけで具体的な輸送方法は求まらないのに対して、ノイズを除去するモデルを用いると拡散モデルと同じように連続分布について具体的な輸送経路が求まる(サンプリングできる)というのがこの方式の良い点です。

拡散モデルと最適輸送には一つ大きな違いがあります。それは、拡散モデルは一方の分布を単純なノイズ分布に限っているということです。そうすることの利点は、データ分布  p(x_0) からノイズ分布  p(x_1) への変換(順方向)が閉じた式で書けることです。具体的には、データにノイズを加えるだけで対応が得られます。このおかげで、アルゴリズムが劇的に単純になり、訓練の安定性も向上します。

一方、最適輸送は両端の分布として任意のものが取れます。例えば、[Gushchin+ NeurIPS 2023] では、データ分布  p(x_0) から別のデータ分布  p(x_1) への変換を最適輸送で求めています。どちらも複雑な分布ですが、それでも正しく対応関係が求まります。このように、任意の分布から任意の分布への対応関係が求められるようになると、教師なし機械翻訳ドメイン変換などさまざまな応用が実現できると考えられます。

これはどちらが優れているという訳ではありません。一般性は最適輸送の方が高いですが、その分アルゴリズムは複雑になります。 例えば、[Gushchin+ NeurIPS 2023] では、移動先の分布が  p(x_1) に沿っているかを検査する敵対的な分類器を追加で用意する必要があります。 逆に、拡散モデルは最適輸送のうち、きれいに解ける部分をうまく取り出したともいえるでしょう。

一般の分布から一般の分布への最適輸送を拡散モデル的に解く方法は発展途上にあります。一般性が高いこともあり、画像生成の拡散モデルほどはうまく動作していないのが現状です。拡散モデルの発展で培われたアルゴリズム的・エンジニアリング的な知見を基にして、今後ますます成長していく分野なのではないかと思います。

おわりに

本稿で最適輸送に興味を持った方はぜひ『最適輸送の理論とアルゴリズム』をお手に取っていただけると幸いです。誤解の無いように補足しておくと、本書では拡散モデルについては直接取り扱ってはいません。関連の深いエントロピー正則化付き最適輸送や敵対的定式化などは詳しく取り扱っているので、本書で基礎を学んだのちに、本稿を改めて読んだり、[Gushchin+ NeurIPS 2023] などの論文にチャレンジすると、拡散モデルと最適輸送の関係についてより深く理解できるかと思います。

新作『グラフニューラルネットワーク』の予約も先日開始されました。発売は 4/25 の予定です。世界一わかりやすく、世界一詳しいグラフニューラルネットワークの教科書になったと自負しています。ぜひお手に取っていただけると幸いです。Chu! 宣伝まみれでごめん。

連絡先: @joisino_ / https://joisino.net