拡散モデルとは：3ステップで理解する

はじめに

この記事は、現代当たり前に浸透した絵の生成モデルの代表「拡散モデル」」の基礎的な考え方を、具体的な実装を通して体得してもらうために作成しました。難易度を3段階に分けて説明しています。必要な理解度に応じて途中まで読むのがオススメです。それではどうぞよろしくお願いします。

Google Colaboratory

Google Colaboratory（略してColab）は、Googleが提供している基本無料のサービスで、Webブラウザ上で直接Pythonのプログラムを実行できます。とにかく環境構築がいらないのが最大のメリットで、PCやスマホで実行することができます。

今回の実験のプログラムをColaboratoryノートブックとして共有します。ぜひ自身の手で動いているところを確認しながら、拡散モデルを知ってください。

Google Colab

colab.research.google.com

Google Colab

colab.research.google.com

導入

ざっくり理解する

このステップでは、拡散モデルの全体像を、ざっくりと直感的に理解することを目指します。

「彫刻」に似たAIの絵の描き方

拡散モデルが絵を描く様子は、大きな石のブロックから彫刻を掘り出す作業によく似ています。

彫刻家は、いきなり完成形を作るわけではありません。まず石のブロックの不要な部分を大まかに削り、だんだんと細かい形を整え、最終的に滑らかに磨き上げて作品を完成させます。

拡散モデルもこれと同じです。AIはまず、「ただのノイズ（砂嵐のようなザラザラ画像）」という名の石のブロックからスタートします。そして、「この画像（石）を汚しているノイズ（不要な部分）はどれだろう？」と予測し、それを少しずつ取り除く（削る）作業を何百回も繰り返します。

この「ノイズを少しずつ取り除く」作業を続けると、ノイズの中から徐々に意味のある形が浮かび上がり、最終的に一枚の絵が完成します。

AIの学習方法：ひたすら「ノイズ除去」の練習

では、AIはどうやって「ノイズを取り除く」方法を学ぶのでしょうか？
それは、「絵をノイズに変える過程」を逆再生する訓練を繰り返すことで実現します。

まず、AIにたくさんの綺麗な絵（完成形のお手本）を見せます。

次に、その絵に少しずつノイズを混ぜて、だんだん汚していく過程を見せます。

そしてAIに課せられるのは、その逆の作業です。ノイズで汚れた画像を見せられ、「この画像を汚しているノイズはどれ？」と予測し、それを取り除く訓練をひたすら行います。

この訓練によって、AIはどんなノイズ画像からでも、それに含まれるノイズ成分を正確に見つけ出す「ノイズ除去の達人」になります。その達人芸を使って、全くのランダムノイズから新しい絵を「掘り出す」のが、拡散モデルの仕組みです。

それなりに理解する

このステップでは、拡散モデルを支える重要な概念である「マルコフ連鎖」や、具体的な学習目標について解説します。

拡散過程と逆過程：マルコフ連鎖として捉える

拡散モデルのコアとなる考え方は、「画像がノイズになる過程」と「ノイズから画像を復元する過程」の両方をマルコフ連鎖として扱うことです。マルコフ連鎖とは、「次の状態は、現在の状態にのみ依存して決まる」という性質を持つプロセスのことです。

拡散過程（Forward Process）

元の画像 $x_0$ からスタートし、各ステップで直前の画像 $x_{t-1}$ にのみ基づいてノイズを加え、次の画像 $x_t$ を生成します。これを $T$ 回繰り返すと、最終的に完全なノイズ $x_T$ になります。

逆過程（Reverse Process）

拡散過程の逆をたどります。現在のノイズ画像 $x_t$ の情報だけを使い、1ステップ前の、よりノイズの少ない画像 $x_{t-1}$ を予測します。これを $T$ 回繰り返すことで、ノイズ $x_T$ から元の画像に近い $x_0$ を復元します。

　

この「現在の状態のみに依存する」という単純な仮定が、複雑な生成プロセスを扱いやすい数式でモデル化することを可能にしています。

AIの学習目標：「ノイズの誤差」を最小にする

逆過程で、AIはノイズ画像 $x_t$ から1ステップ前の画像 $x_{t-1}$ を予測する必要があります。しかし、画像そのものを予測するのは非常に難しいタスクです。

そこで拡散モデルは、問題をより簡単なものに置き換えます。それは、「画像 $x_{t-1}$ に加えられたノイズを予測する」というタスクです。なぜなら、「現在のノイズ画像 $x_t$」と「それに加えられたノイズ」さえ分かれば、1ステップ前の画像 $x_{t-1}$ は計算で復元できるからです。

この考えに基づき、AIの学習目標は「AIが予測したノイズ」と「実際に加えられた真のノイズ」の差（二乗誤差）を最小化することに設定されます。

具体的には、学習データから画像 $x_0$ を取り出し、それにランダムなノイズ $\varepsilon$ を加えてノイズ画像 $x_t$ を作ります。そして、AIに $x_t$ を見せてノイズ $\varepsilon$ を予測させ、その予測精度を高めるように学習を進めていくのです。

がっつり理解する

このステップでは、拡散モデルの理論的背景を数式レベルで解き明かし、関連技術についても触れます。

なぜ「ノイズ予測」に帰着できるのか？

現代の拡散モデルの成功の核心には、画像の生成を「ノイズの予測」だけに帰着させることで、とてもシンプルで計算簡単な目標を用意できたことにあります。数式も見ながら、理論的な立て付けを理解していきましょう。まずは、拡散過程と逆過程を式で表すとどう表現できるのかを見ましょう。

拡散過程（Forward Process）

ざっくりいえば、1つ前の画像 $x_{t-1}$ を持ってきて、それをほんの少しだけぼんやりさせて（平均）、そこに $\beta_t$ という強さの正規分布ノイズをちょい足し（分散）します。そうやって出来上がるのが、次の画像 $x_t$ ですよ。って感じです。

逆過程（Reverse Process）

左辺 $p_\theta(x_{t-1} \mid x_t)$ は、「ノイズが乗った画像 $x_t$ が与えられたとき、その1ステップ前の、よりノイズが少なかった画像 $x_{t-1}$ が従う確率分布」を意味します。$\theta$ は、この分布を表現するモデル（ニューラルネットワーク）のパラメータをイメージしてください。右辺が具体的にどんな分布かを示しています。正規分布（$\mathcal{N}(x; \mu, \sigma^2)$）にしたがうわけですが、ここで重要なのは平均です。多くの論文で分散は固定として扱われているので、詳細は割愛します。平均、つまり$\mu_\theta(x_t, t)$は、現在のノイズ画像 $x_t$ と現在の時刻 $t$ を入力とした、1ステップ前の画像 $x_{t-1}$ の「最もありえそうな姿」の予測にあたります。

最適化の目標

　

$\mu_q(x_t, x_0)$は、いわば教師データです。学習のとき、私たちは元のキレイな画像 $x_0$ を持っています。なので、「もし $x_0$ から $x_t$ が作られたなら、1つ前の画像 $x_{t-1}$ の平均は理論上こうなるはずだ」という完璧な答えを計算できます。一方で、$\mu_\theta(x_t, t)$はニューラルネットワークの予測した、1つ前の画像です。このとき、ニューラルネットワークは $x_0$ を知らないので、 $x_t$ と $t$ だけをヒントに予測します。

この式が言いたいことをまとめると、モデル $\mu_\theta$ が予測した「1つ前の画像の姿」が、数学的に計算できる「正解 $\mu_q$」にどれだけ近いかを測っており、このズレが小さいほどモデルの成績が良いといえるってことです。

最適化の目標の簡単化

ノイズを含む画像の一歩前を予測するためには、画像の複雑な構造を理解する必要があり、難易度が高いです。そこで、先ほどの式を上記の式にまで簡略化します。単に、一歩前に加えられたノイズを予測するだけで済むようにします。これによって、モデルは画像の複雑な構造を考慮せず、ノイズっぽいなと思う部分をハイライトするだけで済むようになります。

しかし、この簡略化はなぜ機能するのでしょうか。それは、先ほどの式に以下を代入することで導けます。

この式は、『現在のノイズ付き画像 $x_t$ から、予測したノイズ $\epsilon_\theta$ を（適切な量だけ）取り除けば、元のきれいな画像 $x_0$ が推定できる。そして、その推定した $x_0$ と現在の $x_t$ があれば、理論式に当てはめて一歩手前の画像 $x_{t-1}$ の平均を計算できる。』という考えに基づいています。つまり、ノイズを予測することは、間接的に一歩前の画像を予測することと等価なのです。

　

ノイズ予測器としての「U-Net」

ノイズ予測器 $\varepsilon_{\theta}$ として活躍するのが、U-Netと呼ばれるニューラルネットワークアーキテクチャです。U-Netは、エンコーダ（画像圧縮）とデコーダ（画像復元）がU字型に対称的に配置され、両者がスキップ接続で結ばれているのが特徴です。

この構造が、画像の全体的な構造と局所的なディテールの両方を捉えられるようにしてくれました。

実践

ノートブックではここまでの理論を実際のコードに落とし込み、簡単な例ではありますが実際に学習をさせて拡散モデルについて知れる構成になっています。学習のプロセス含めて実行時間は10分未満です。ぜひ、そちらもご覧ください。

Google Colab

colab.research.google.com

Google Colab

colab.research.google.com

拡散モデルの進化

今回紹介した拡散モデルを基礎として、現在も様々な派生モデルが生まれています。

こういった話題もどんどん取り上げて説明していくので、ぜひこのブログをブックマークしてチェックしてください。

参考文献

U-Net: Convolutional Networks for Biomedical Image Segmentation

There is large consent that successful training of deep networks requires many thousand annotated training samples. In this paper, we present a network and training strategy that relies on the strong use of data augmentation to use the available annotated samples more efficiently. The architecture consists of a contracting path to capture context and a symmetric expanding path that enables precise localization. We show that such a network can be trained end-to-end from very few images and outperforms the prior best method (a sliding-window convolutional network) on the ISBI challenge for segmentation of neuronal structures in electron microscopic stacks. Using the same network trained on transmitted light microscopy images (phase contrast and DIC) we won the ISBI cell tracking challenge 2015 in these categories by a large margin. Moreover, the network is fast. Segmentation of a 512x512 image takes less than a second on a recent GPU. The full implementation (based on Caffe) and the trained networks are available at http://lmb.informatik.uni-freiburg.de/people/ronneber/u-net .

arxiv.org

U-Net: Convolutional Networks for Biomedical Image Segmentation

There is large consent that successful training of deep networks requires many thousand annotated training samples. In this paper, we present a network and training strategy that relies on the strong use of data augmentation to use the available annotated samples more efficiently. The architecture consists of a contracting path to capture context and a symmetric expanding path that enables precise localization. We show that such a network can be trained end-to-end from very few images and outperforms the prior best method (a sliding-window convolutional network) on the ISBI challenge for segmentation of neuronal structures in electron microscopic stacks. Using the same network trained on transmitted light microscopy images (phase contrast and DIC) we won the ISBI cell tracking challenge 2015 in these categories by a large margin. Moreover, the network is fast. Segmentation of a 512x512 image takes less than a second on a recent GPU. The full implementation (based on Caffe) and the trained networks are available at http://lmb.informatik.uni-freiburg.de/people/ronneber/u-net .

arxiv.org

おすすめ記事

ニューラルネットワークを目で理解する

現代AIの核心「ニューラルネットワーク」を、目で見て理解する入門記事。なぜAIは賢く、時に間違うのか？プログラムで成功と失敗を体験しながら、その強力さと繊細さを学びます。AI開発に不可欠な思考の土台を築きましょう。

ニューラルネットワークを目で理解する

現代AIの核心「ニューラルネットワーク」を、目で見て理解する入門記事。なぜAIは賢く、時に間違うのか？プログラムで成功と失敗を体験しながら、その強力さと繊細さを学びます。AI開発に不可欠な思考の土台を築きましょう。

AIをつくる側になるための考え方

「AIで仕事を自動化したいけど、何から学べば…？」その答えは、AIを「入力→出力」というシンプルな箱として捉えることにあります。この記事を読めば、AIの専門用語に惑わされず、その本質的な仕組みが分かります。「AIを作ってみたい」「AIで業務を効率化したい」と考えるあなたのための、第一歩となる知識がここにあります。

AIをつくる側になるための考え方

「AIで仕事を自動化したいけど、何から学べば…？」その答えは、AIを「入力→出力」というシンプルな箱として捉えることにあります。この記事を読めば、AIの専門用語に惑わされず、その本質的な仕組みが分かります。「AIを作ってみたい」「AIで業務を効率化したい」と考えるあなたのための、第一歩となる知識がここにあります。