水を運ぶように確率を変形せよ ── 正規化フローと「連続の式」
前回の第一話では、画像生成AIの真の舞台が、気の遠くなるような**「超高次元の幾何学宇宙」**であることを見てきました。
フルHD画像なら、画素の値を並べただけでも数百万次元に達します。
その暗黒の空間に、1枚の画像は「たった1つの点」として存在しています。
そして生成AIとは、その巨大な空間の中に存在する「本物の画像たちの確率分布」を学習し、そこから新しい点を生み出す仕組みだと考えることができます。
前回、私はそのイメージを、こんな風景として描きました。
ぽつんと浮かんだ1つの点が、AIの計算した「一貫した風(速度ベクトル場)」に流され、目的地である「本物のデータの山(多様体)」へと漂流していく。
これは、生成AIを「確率分布からのサンプリング」として捉えたときの、ひとつの幾何学的な見取り図です。
では、この壮大な「知の螺旋(Uターン)」の旅路において、最初の走者を務めたのは誰だったのでしょうか。
それこそが、**「正規化フロー(Normalizing Flows)」**です。
正規化フローは、確率分布を水や空気のような「流体」に見立て、ある確率分布を、逆変換可能な写像によって別の確率分布へと変形していく数理的な枠組みです。
今回は、正規化フローがどのようにして確率分布を変形させ、そこで「ヤコビアン」がなぜ重要な役割を果たすのかを見ていきます。
さらに、そのヤコビアンをめぐる計算上の問題を乗り越えるために、なぜ「連続時間」という発想が登場したのかを考えます。
そして最後には、ここからさらに別の道へとつながっていきます。
その数学的な流れを追ってみましょう。
確率分布を「水」のように流す
正規化フローの基本的な考え方は、驚くほどシンプルです。
「扱いやすい単純な確率分布を、変形によって複雑な確率分布へ変えてしまおう」
というものです。
たとえば、最初に用意するのは標準ガウス分布です。
標準ガウス分布は、平均0、分散1の正規分布で、数学的にもコンピュータ上でも扱いやすい確率分布です。
これを「出発地点」とします。
一方、私たちが本当に欲しいのは、現実の画像データが従っている複雑な確率分布です。
そこで、
単純な分布 → 変形 → ・・・ → 変形 → 複雑なデータ分布
という道筋を考えます。
たとえるなら、広い水槽に溜まった水を、迷路のようなチューブに流し込み、途中で何度も形を変えながら、最後には複雑な形の容器へと送り込むようなものです。
この「分布を流して形を変える」という発想が、Normalizing Flow という名前につながっています。
数学的には、まず扱いやすいベース分布に従う確率変数
$$
\mathbf{z} \sim p_0(\mathbf{z})
$$
を考えます。
そして、逆変換可能な1対1の写像
$$
\mathbf{x} = \mathbf{f}(\mathbf{z})
$$
によって、$\mathbf{z}$を目的のデータ空間の点$\mathbf{x}$へ変換します。
ここで重要なのは、この変換が可逆であることです。
つまり、
$$
\mathbf{z} = \mathbf{f}^{-1}(\mathbf{x})
$$
と逆向きにもたどれる必要があります。
なぜでしょうか。
生成するときには、
$$
\mathbf{z} \rightarrow \mathbf{x}
$$
という方向に進みます。
一方、学習時には、実際のデータ$\mathbf{x}$が与えられたとき、
$$
\mathbf{x} \rightarrow \mathbf{z}
$$
と逆向きにたどって、元の単純な分布のどこに対応するのかを知りたいからです。
そして、もう一つ大切な原則があります。
確率の総量は保存されなければなりません。
確率密度を全空間で積分した値は1です。
分布をどれだけ複雑に変形しても、この「全確率=1」という条件は壊れてはいけません。
ここで、水風船を思い浮かべてみましょう。
水風船をギュッと押し潰しても、中に入っている水の量そのものは変わりません。
しかし、水が占める空間は狭くなります。
すると、その狭い場所には、より多くの水が凝縮された状態になります。
確率分布でも、基本的には同じことが起こります。
ある領域が押し縮められれば、その領域に確率が集中し、確率密度が高くなる。
逆に、空間が引き伸ばされれば、同じ確率がより広い領域に広がるため、確率密度が低くなる。
では、この「空間の伸び縮み」と「確率密度の変化」を、数学的にどう結びつければよいのでしょうか。
そこで登場するのが、
ヤコビ行列(Jacobian matrix)
と、その行列式である
ヤコビアン(Jacobian determinant)
です。
ヤコビアンが記述する、空間の「伸び縮み」
多変数の写像
$$
\mathbf{x} = \mathbf{f}(\mathbf{z})
$$
を考えます。
このとき、写像$\mathbf{f}$が、ある点の近くで空間をどの方向に、どれくらい変形させるのかを記述するのがヤコビ行列です。
$$
J_{\mathbf{f}}(\mathbf{z})=
\frac{\partial \mathbf{f}(\mathbf{z})}{\partial \mathbf{z}}
$$
そして、その行列式
$$
\det J_{\mathbf{f}}(\mathbf{z})
$$
が、局所的な体積変化と結びつきます。
確率密度については、変数変換の公式から、
$$
p_1(\mathbf{x})=
p_0(\mathbf{z})
\left|
\det J_{\mathbf{f}}(\mathbf{z})
\right|^{-1}
\tag{1}
$$
となります。
この式は、正規化フローを理解するうえで極めて重要です。
両辺の対数を取ると、
$$
\log p_1(\mathbf{x})=
\log p_0(\mathbf{z})-
\log
\left|
\det J_{\mathbf{f}}(\mathbf{z})
\right|
\tag{2}
$$
となります。
この式を見ると、何が起こっているのかが、よりはっきりします。
もし空間が局所的に引き伸ばされれば、その体積が大きくなります。
すると、同じ確率がより広い空間に広がるため、確率密度は低下します。
逆に、空間が押し縮められれば、確率密度は高くなります。
ヤコビ行列の行列式の絶対値
$$
\left|
\det J_{\mathbf{f}}(\mathbf{z})
\right|
$$
は、この局所的な体積要素が何倍になったかを表しています。
たとえば、局所的な体積が2倍に引き伸ばされたとしましょう。
すると、確率密度は半分になります。
対数を取れば、
$$
-\log 2
$$
だけ変化します。
逆に、体積が半分に圧縮されれば、確率密度は2倍になり、対数密度は
$$
+\log 2
$$
だけ増加します。
つまり式(2)は、
「空間をどれだけ伸ばしたか、縮めたかによって、確率密度がどのように変化するか」
を正確に記述しているのです。
ここに正規化フローの大きな魅力があります。
実際のデータ$\mathbf{x}$が与えられたとき、逆写像
$$
\mathbf{z} = \mathbf{f}^{-1}(\mathbf{x})
$$
によって、それが単純なベース分布のどこから来たのかを求める。
そして、ヤコビアンの行列式を計算する。
それによって、そのデータ点$\mathbf{x}$に対する確率密度を直接評価できます。
確率分布全体を一つひとつ調べる必要はありません。
ここが、正規化フローの非常に美しいところです。
複雑な確率分布を作りながら、その確率密度そのものも計算できる。
しかし、ここで大きな問題が立ちはだかります。
その主役が、今まさに登場したヤコビアンです。
第1走者の挫折 ── 「ヤコビアンの呪い」
式(2)を計算するには、
$$
\det J_{\mathbf{f}}
$$
を計算しなければなりません。
ところが、高次元になると、これは非常に重い計算になります。
たとえば入力が$d$次元なら、ヤコビ行列は
$$
d \times d
$$
の巨大な行列になります。
画像データでは、$d$が数百万になることさえあります。
そんな巨大な行列を作り、その行列式を計算する。
これは、現実的なニューラルネットワークの学習にとって大きな負担になります。
そこで初期の正規化フローでは、**「ヤコビアンを計算しやすい構造にしてしまう」**という工夫が行われました。
代表例がNICE [3] やRealNVP [4] です。
これらでは、変換を特殊な構造に制限することで、各変換層のヤコビ行列を三角構造にできるようにします。
三角行列なら、行列式は対角成分の積として簡単に計算できます。
つまり、
表現の自由度をある程度犠牲にして、計算しやすさを手に入れる。
これが初期の正規化フローが採った戦略でした。
この工夫によって、正規化フローは大きく前進しました。
しかし、ここには一つのジレンマがあります。
変換を自由にすればするほど、複雑なデータ分布を表現しやすくなります。
ところが自由にしすぎると、ヤコビアンの計算が難しくなる。
反対に、ヤコビアンを簡単に計算できる構造に制限すると、今度は変換の自由度が下がります。
つまり、
表現力と計算効率のトレードオフ
が生まれます。
正規化フローは、この問題と戦うことになりました。
そして、ここから一つの大きな発想転換が生まれます。
それは、
「そもそも、巨大な変換を一気に計算する必要があるのだろうか?」
という問いでした。
行列式から発散へ ── 連続時間が開いた新しい道
もし変換を、一度に大きく行うのではなく、
無限に細かい小さな変化の積み重ね
として考えたらどうでしょうか。
2018年、Neural ODE(Neural Ordinary Differential Equations)[5] が登場しました。
そして、この連続時間の考え方を確率分布の変換に応用することで、**連続時間正規化フロー(Continuous Normalizing Flows: CNF)**への道が開かれます。
粘土を一度に大きく変形させるのではありません。
ほんの少しずつ、連続的に形を変えていく。
その「流れ」を速度ベクトル場で記述します。
粒子の位置を$\mathbf{x}(t)$、その場所での速度を
$$
\mathbf{v}_t(\mathbf{x})
$$
とすると、
$$
\frac{d\mathbf{x}(t)}{dt}=
\mathbf{v}_t(\mathbf{x}(t))
\tag{3}
$$
となります。
これは決定論的な常微分方程式(ODE)です。
ここでは、ある時刻の位置が決まれば、その後どこへ進むのかも速度場によって決まります。
では、この粒子たちが大量に存在していると考えてみましょう。
粒子一つひとつを追いかけるのではなく、空間に固定した視点から、
「それぞれの場所の確率密度が時間とともにどう変化しているか」
を観察します。
これは流体力学でいうオイラー記述です。
この視点から確率密度の保存を記述する方程式が、
連続の式(Continuity Equation)
です。
$$
\frac{\partial p_t(\mathbf{x})}{\partial t}
+
\nabla \cdot
\left(
p_t(\mathbf{x})\mathbf{v}_t(\mathbf{x})
\right)=
0
\tag{4}
$$
少し難しそうに見えますが、言っていることは非常にシンプルです。
第1項
$$
\frac{\partial p_t(\mathbf{x})}{\partial t}
$$
は、固定した場所で確率密度が時間とともにどれだけ変化するかを表します。
第2項
$$
\nabla\cdot
\left(
p_t(\mathbf{x})\mathbf{v}_t(\mathbf{x})
\right)
$$
は、確率密度と速度を掛け合わせた「確率流束」の発散です。
つまり、その小さな領域から確率がどれだけ外へ流れ出しているか、あるいは中へ流れ込んでいるかを表しています。
したがって式(4)は、
「その場所の確率密度の変化」と「そこからの確率の流出入」を合わせるとゼロになる
と言っています。
確率が突然消えたり、何もないところから突然生まれたりしない。
確率はただ、空間の中を移動しているだけです。
これは、流体における質量保存則と同じ構造です。
オイラーからラグランジュへ
ここで、少し視点を変えてみましょう。
今までは「空間に固定した場所」から確率の流れを眺めていました。
これはオイラー記述です。
一方、
「その流れに乗っている粒子自身を追いかける」
という見方もできます。
こちらがラグランジュ記述です。
式(4)を展開すると、
$$
\frac{\partial p_t}{\partial t}
+
\mathbf{v}_t\cdot\nabla p_t
+
p_t\nabla\cdot\mathbf{v}_t=
0
$$
となります。
ここで、
$$
\frac{\partial p_t}{\partial t}
+
\mathbf{v}_t\cdot\nabla p_t
$$
という二つの項に注目します。
これは、流れに乗って移動する粒子を追いかけたときの確率密度の時間変化、
$$
\frac{d}{dt}p_t(\mathbf{x}(t))
$$
に相当します。
したがって、
$$
\frac{d}{dt}p_t(\mathbf{x}(t))
+
p_t(\mathbf{x}(t))
\nabla\cdot\mathbf{v}_t(\mathbf{x}(t))=
0
$$
と書けます。
両辺を$p_t(\mathbf{x}(t))$で割ると、
$$
\frac{d}{dt}
\log p_t(\mathbf{x}(t))=
-\nabla\cdot\mathbf{v}_t(\mathbf{x}(t))
\tag{5}
$$
となります。
これが、CNFを理解するうえで極めて重要な式です。
旅する粒子が感じる対数確率密度の変化率は、その場所における速度場の発散のマイナスに等しい。
ここで突然、
「ヤコビアンはどこへ行ったの?」
という疑問が浮かびます。
実は、ヤコビアンは消えたわけではありません。
連続時間の極限では、その役割が**発散(divergence)**という、より局所的な量に姿を変えているのです。
行列式は、なぜ「トレース」に変わるのか
ここまで見てきたように、正規化フローでは、変換による確率密度の変化をヤコビ行列の行列式で測っていました。
では、この考え方を「連続時間」に持ち込むと、何が起こるのでしょうか。
ポイントは、ごく短い時間 $\Delta t$ の間に起こる、微小な変換を考えることです。
まず、微小な変換のヤコビ行列を求める
時刻 $t$ における点を $\mathbf{x}(t)$ とし、その点が速度場
$$
\mathbf{v}_t(\mathbf{x})
$$
に従って動いているとします。
その運動は、常微分方程式
$$
\frac{d\mathbf{x}(t)}{dt}=
\mathbf{v}_t(\mathbf{x}(t))
$$
で表されます。
時間をほんの少しだけ $\Delta t$ 進めると、
$$
\mathbf{x}(t+\Delta t)
\approx
\mathbf{x}(t)
+
\Delta t,\mathbf{v}_t(\mathbf{x}(t))
$$
となります。
ここで、この微小な時間発展そのものを1つの変換
$$
\mathbf{f}(\mathbf{x}(t))
\equiv
\mathbf{x}(t+\Delta t)
$$
とみなしてみましょう。
すると、
$$
\mathbf{f}(\mathbf{x})
\approx
\mathbf{x}
+
\Delta t\ \mathbf{v}_t(\mathbf{x})
$$
です。
つまり、いま計算しようとしているのは、この微小な変換 $\mathbf{f}$ のヤコビ行列 $J_{\mathbf f}$ です。
ヤコビ行列の定義から、
$$
J_{\mathbf f}=
\frac{\partial\mathbf f}{\partial\mathbf x}
$$
ですから、
$$
J_{\mathbf f}=
\mathbf I
+
\Delta t\ J_{\mathbf v_t}
$$
となります。
ここで、
$$
J_{\mathbf v_t}=
\frac{\partial\mathbf v_t}{\partial\mathbf x}
$$
は、速度場 $\mathbf v_t$ のヤコビ行列です。
したがって、この微小な時間発展による局所的な体積の変化は、
$$
\left|
\det J_{\mathbf f}
\right|=
\left|
\det
\left(
\mathbf I+\Delta t\ J_{\mathbf v_t}
\right)
\right|
$$
で測ることができます。
ここで、第二話の最初に登場した「ヤコビアン」が再び姿を現しました。
微小な変化では、行列式を簡単に扱える
では、
$$
\det
\left(
\mathbf I+\Delta t\ J_{\mathbf v_t}
\right)
$$
は、$\Delta t$ が十分小さいとき、どのように変化するのでしょうか。
微小量については、次の関係が成り立ちます。
$$
\log\left|
\det
\left(
\mathbf I+\Delta t\ J_{\mathbf v_t}
\right)
\right|=
\Delta t\
\operatorname{tr}
\left(
J_{\mathbf v_t}
\right)
+
O(\Delta t^2)
$$
ここで突然、**トレース($\operatorname{tr}$)**が登場しました。
トレースとは、行列の対角成分をすべて足し合わせたものです。
なぜ行列式からトレースが現れるのでしょうか。
直感的には、行列式が「空間全体の体積がどれだけ変化したか」を測るのに対して、トレースは「各方向の微小な伸び縮みを全部足し合わせたもの」だからです。
そして、微小な変換では、この「体積の変化率」がトレースによって表されます。
ここで重要なのが、速度場のヤコビ行列について
$$
\operatorname{tr}
\left(
J_{\mathbf v_t}
\right)=
\nabla\cdot\mathbf v_t
$$
という関係です。
つまり、
$$
\log\left|
\det J_{\mathbf f}
\right|
\approx
\Delta t\
\nabla\cdot\mathbf v_t
$$
となります。
離散的な正規化フローでは「行列式」で測っていた局所的な体積変化が、連続時間では「発散」によって表現できる。
これが、行列式からトレースへ、そしてトレースから発散へとつながる道筋です。
行列式からトレースへ
ここまでの流れを整理すると、
$$
J_{\mathbf f}=
\mathbf I
+
\Delta t\ J_{\mathbf v_t}
$$
から出発して、
$$
\log\left|
\det J_{\mathbf f}
\right|=
\Delta t\
\operatorname{tr}
\left(
J_{\mathbf v_t}
\right)
+
O(\Delta t^2)
$$
となり、さらに
$$
\operatorname{tr}
\left(
J_{\mathbf v_t}
\right)=
\nabla\cdot\mathbf v_t
$$
なので、
$$
\boxed{
\log\left|
\det J_{\mathbf f}
\right|=
\Delta t\
\nabla\cdot\mathbf v_t
+
O(\Delta t^2)
}
$$
となります。
つまり、
$$
\boxed{
\text{ヤコビアンの行列式}
\quad\longrightarrow\quad
\text{トレース}
\quad\longrightarrow\quad
\text{発散}
}
$$
という変化が起こるのです。
これは単なる数学上の偶然ではありません。
連続時間の正規化フローでは、空間の局所的な伸び縮みを「行列式そのもの」ではなく、その瞬間の変化率で捉えることができる。
その変化率を記述するのが、速度場のヤコビ行列のトレース、すなわち発散なのです。
そして、この結果を時間について積み重ねていくと、確率密度の変化を積分によって記述できるようになります。
それが、連続正規化フロー(CNF)の核心へとつながっていきます。
行列式という巨大な計算を、その瞬間の「局所的な広がり具合」を表す発散へと置き換える。
そして式(5)を$t=0$から$t=1$まで積分すれば、
$$
\log p_1(\mathbf{x}(1))=
\log p_0(\mathbf{x}(0))-
\int_0^1
\nabla\cdot\mathbf{v}_t(\mathbf{x}(t))
\ dt
\tag{6}
$$
を得ます。
これが、連続時間正規化フローにおける対数確率密度の変化を表す基本式です。
ここまで来ると、正規化フローの姿が大きく変わっていることが分かります。
離散的な正規化フローでは、
ヤコビアンの行列式を計算しやすいように、変換構造そのものを工夫する。
それに対してCNFでは、
変換を連続時間の流れとして表現し、ヤコビアンの行列式を発散の時間積分として扱う。
この発想によって、変換構造を三角行列などに制約する必要性から大きく解放されます。
より自由なニューラルネットワークによって速度場を表現できるようになったのです。
しかし――。
ここで、もう一つの問題が顔を出します。
自動微分の魔法と、立ちはだかる「シミュレーションの壁」
「ヤコビアンの行列式を発散に置き換えたのだから、これで計算問題は解決したのでは?」
そう思いたくなります。
ところが、そう簡単にはいきません。
まず、速度場の発散
$$
\nabla\cdot\mathbf{v}
$$
を計算する必要があります。
これは、速度場のヤコビ行列の対角成分をすべて足し合わせたもの、
$$
\nabla\cdot\mathbf{v}=
\operatorname{tr}(J_{\mathbf{v}})
$$
です。
高次元になると、この計算も無視できないコストになります。
そこでFFJORD [6] では、Hutchinson推定量を利用して、ヤコビ行列そのものを明示的に構成せずにトレースを推定する方法が使われました。
その基本となる性質は、
$$
\operatorname{tr}(J)=
\mathbb{E}_{\boldsymbol{\epsilon}}
\left[
\boldsymbol{\epsilon}^{\top}
J
\boldsymbol{\epsilon}
\right]
$$
です。
ここで$\boldsymbol{\epsilon}$は、たとえば各成分が独立な標準正規分布に従うランダムベクトルです。
この性質を使えば、巨大なヤコビ行列を丸ごと作らなくても、ランダムベクトルと自動微分を利用してトレースを推定できます。
これは大きな工夫でした。
しかし、CNFにとって、さらに大きな問題が残っていました。
それが、
時間積分です。
式(6)には、
$$
\int_0^1 \cdots\ dt
$$
という積分があります。
数学の紙の上では、積分記号を一つ書けば終わりです。
しかし、コンピュータはそうはいきません。
実際には、$t=0$から$t=1$までの時間を細かく刻みながら、ODEを数値的に解かなければなりません。
ルンゲ=クッタ法などの数値ODEソルバーを使い、粒子の軌道を少しずつ追跡していきます。
しかも、正規化フローを最尤法で学習する場合、ニューラルネットワークのパラメータを更新するたびに、このODE計算が必要になります。
さらに、誤差逆伝播によってパラメータの勾配を計算する必要があります。
随伴感度法(adjoint sensitivity method)のようなメモリを節約する手法もありますが、勾配計算にもODEに関わる数値計算が必要になります。
つまり、
「行列式の計算」という問題を解決したと思ったら、今度は「ODEを何度も解く」という問題が現れた。
これがCNFの前に立ちはだかった、もう一つの壁でした。
特に画像のような高次元データを扱おうとすると、この計算負荷は大きな問題になります。
こうして、決定論的な連続フローという美しいアイデアを持ったCNFは、計算コストという現実の壁に直面することになります。
そこで、ここから別の考え方へ目を向けてみましょう。
次なる走者へのバトン ── 全体を測るのをやめ、局所を学ぶ
ここまでの流れを、いったん整理してみましょう。
正規化フローは、
「単純な確率分布を、変形によって複雑なデータ分布へ変える」
という非常に美しい発想から出発しました。
ところが、高次元になるとヤコビアンの計算が問題になります。
そこで、NICE [3] やRealNVP [4] のように変換構造を工夫しました。
さらに、
「変換を連続時間の流れとして考えればいいのではないか」
という発想からCNFが登場します。
すると、ヤコビアンの行列式は発散の時間積分へと姿を変え、変換構造の自由度は大きく広がりました。
しかし、CNFにはODEの数値計算という新たな壁がありました。
ここで、私たちの物語はいったん別の方向へ向かいます。
もう一度、式(5)を見てみましょう。
$$
\frac{d}{dt}
\log p_t(\mathbf{x}(t))=
-\nabla\cdot\mathbf{v}_t(\mathbf{x}(t))
$$
この式は、ある粒子が流れに沿って移動するとき、その軌道と確率密度の変化が、速度場によって決まることを示しています。
ここで、別の問いを立てることができます。
確率分布全体を追跡する代わりに、その場所における確率分布の局所的な形だけを学ぶことはできないだろうか。
たとえば山を想像してください。
山全体の形を最初から完全に測量する代わりに、
「今、自分が立っている場所から、どちらが上り坂なのか」
だけを知ることができれば、山の頂上を目指して進むことができます。
確率分布にも、これと似た「局所的な坂道」があります。
それが、
スコア(Score)
です。
確率密度$p(\mathbf{x})$の対数を取り、その勾配をとった
$$
\nabla_{\mathbf{x}}\log p(\mathbf{x})
$$
がスコアです。
これは、確率密度が空間のどちらの方向へ増加しているのか、その「傾き」を表しています。
確率分布の山の斜面に立っているとすれば、
「どちらへ進めば、より確率の高い場所へ向かうのか」
を教えてくれる情報だと考えることができます。
ここから、正規化フローとは少し異なる考え方が見えてきます。
これまで私たちは、
「確率分布全体を流すための速度場」
を考えてきました。
これからは、
「確率分布の局所的な傾き」
に注目します。
これは、確率分布そのものを直接扱うのではなく、その局所構造を学習するという新しい道です。
そして、この「局所的な傾き」を利用する考え方が、やがて拡散モデルの中心概念であるスコアベース生成モデルへとつながっていきます。
しかし、ここにも一つ大きな問題があります。
スコア
$$
\nabla_{\mathbf{x}}\log p(\mathbf{x})
$$
を学習するためには、そもそも未知のデータ分布$p(\mathbf{x})$が必要です。
「分布を知らないから学習したいのに、その分布の対数勾配をどうやって計算するのか?」
ここで、さらにもう一つのブレイクスルーが登場します。
それが、
デノイジング・スコアマッチング(Denoising Score Matching: DSM)[7]
です。
複雑なデータ分布そのものを直接扱うのではなく、いったんデータにノイズを加える。
すると、元のデータを条件としたノイズ付加分布については、そのスコアを計算することができます。
この既知のスコアを教師信号として利用することで、ノイズを加えたデータ分布そのもののスコアを学習できる――。
これがデノイジング・スコアマッチングの核心です。
そして、このアイデアが、後に私たちが「拡散モデル」と呼ぶ巨大な流れへとつながっていくことになります。
第二話のまとめ
今回の話を、一本の流れとして振り返ってみましょう。
最初に登場した正規化フローは、
「単純な確率分布を、可逆な変換によって複雑なデータ分布へ変形する」
という発想でした。
そのとき、確率密度の変化を計算する鍵となったのが、
ヤコビアン
でした。
ところが、高次元ではヤコビアンの計算が重くなります。
そこで、NICE [3] やRealNVP [4] では変換構造を工夫しました。
そして次に、
「変換を連続時間の流れとして考えればいい」
という発想が生まれます。
これがCNFです。
有限の変換で現れていたヤコビアンの行列式は、連続時間の極限では、
$$
\operatorname{tr}(J_{\mathbf{v}})=
\nabla\cdot\mathbf{v}
$$
という発散へと姿を変えます。
さらに時間方向に積分することで、
$$
\log p_1(\mathbf{x}(1))=
\log p_0(\mathbf{x}(0))-
\int_0^1
\nabla\cdot\mathbf{v}_t(\mathbf{x}(t))
\ dt
$$
という、連続時間版の変数変換公式が得られました。
しかし、CNFにはODEの数値計算という新たな壁がありました。
そこで、私たちは別の考え方にも目を向けることになります。
「分布全体を流すこと」から、「分布の局所的な傾きを学ぶこと」へ。
その局所的な傾きが、
$$
\nabla_{\mathbf{x}}\log p(\mathbf{x})
$$
というスコアです。
そして次回は、このスコアをどうやって学習するのかという問題に挑みます。
「元の確率分布を知らなくても、スコアを学習できるのか?」
その答えが、
デノイジング・スコアマッチング(DSM)[7]
です。
ここから、いよいよ「拡散モデル」の物語が本格的に始まります。
(第三話へ続く)
第二話 参考文献リスト
-
岡野原大輔『生成AIのしくみ』技術評論社、2024年。
-
岡野原大輔『拡散モデル』岩波書店、2023年。
-
L. Dinh, D. Krueger, Y. Bengio, “NICE: Non-linear Independent Components Estimation,” arXiv, 2014.
-
L. Dinh, J. Sohl-Dickstein, S. Bengio, “Density estimation using Real NVP,” ICLR, 2017.
-
R. T. Q. Chen, Y. Rubanova, J. Bettencourt, D. Duvenaud, “Neural Ordinary Differential Equations,” NeurIPS, 2018.
-
W. Grathwohl, R. T. Q. Chen, J. Bettencourt, I. Sutskever, D. Duvenaud, “FFJORD: Free-form Continuous Dynamics for Scalable Reversible Generative Models,” ICLR, 2019.
-
P. Vincent, “A Connection Between Score Matching and Denoising Autoencoders,” Neural Computation, 2011.
第二話 専門用語の簡単な解説
正規化フロー(Normalizing Flows):単純な確率分布を、逆変換可能な写像によって複雑なデータ分布へ変換し、その確率密度も計算できるようにする生成モデルの枠組み。
ヤコビ行列(Jacobian matrix):多変数関数の1階偏微分を並べた行列。写像による局所的な空間の変形を記述する。
ヤコビアン(Jacobian determinant):ヤコビ行列の行列式。局所的な体積要素が何倍になるかを記述する。
対数尤度(log-likelihood):観測されたデータに対してモデルが与える確率密度の対数。最尤学習では、これを最大化することが基本的な目標となる。
トレース(trace):正方行列の対角成分をすべて足し合わせたもの。速度場のヤコビ行列のトレースは、その速度場の発散に等しい。
ニューラルODE(Neural ODE):ニューラルネットワークによって常微分方程式の右辺を定義し、連続時間の状態変化を表現する枠組み。
発散(Divergence):ベクトル場がある場所からどれだけ外向きに広がっているか、あるいは内向きに集束しているかを表す量。
連続の式(Continuity Equation):流体などの総量が保存されることを、密度と流れの関係として記述する方程式。
Hutchinson推定量(Hutchinson trace estimator):ランダムベクトルを利用して、巨大な行列を明示的に構成せずに行列のトレースを推定する方法。
VJP(Vector-Jacobian Product):ベクトルとヤコビ行列の積を、ヤコビ行列そのものを明示的に構成せず、自動微分によって効率よく計算する方法。
瞬間的変数変換方程式(Instantaneous Change of Variables Equation):CNFにおいて、流れに沿った対数確率密度の時間変化が速度場の発散のマイナスに等しいことを表す方程式。
スコア(Score):確率密度の対数の勾配。確率密度がどの方向へ増加しているかを示す局所的な情報。
デノイジング・スコアマッチング(DSM):データに既知のノイズを加え、そのノイズ付加過程を利用して、ノイズを加えたデータ分布のスコアを学習する方法。
付録:なぜ「トレース=発散」なのか
本編では、
$$
\operatorname{tr}(J_{\mathbf v})=
\nabla\cdot\mathbf v
$$
という関係が登場しました。
「なぜトレースと発散が同じになるのか?」と思った方のために、ここで確認しておきましょう。
速度場を (d) 次元のベクトル
\mathbf v(\mathbf x)=
\begin{pmatrix}
v_1(\mathbf x)\\
v_2(\mathbf x)\\
\vdots\\
v_d(\mathbf x)
\end{pmatrix}
とします。
この速度場のヤコビ行列は、
J_{\mathbf v}=
\frac{\partial\mathbf v}{\partial\mathbf x}=
\begin{pmatrix}
\displaystyle\frac{\partial v_1}{\partial x_1}
&
\displaystyle\frac{\partial v_1}{\partial x_2}
&
\cdots
&
\displaystyle\frac{\partial v_1}{\partial x_d}
\\[6pt]
\displaystyle\frac{\partial v_2}{\partial x_1}
&
\displaystyle\frac{\partial v_2}{\partial x_2}
&
\cdots
&
\displaystyle\frac{\partial v_2}{\partial x_d}
\\
\vdots
&
\vdots
&
\ddots
&
\vdots
\\
\displaystyle\frac{\partial v_d}{\partial x_1}
&
\displaystyle\frac{\partial v_d}{\partial x_2}
&
\cdots
&
\displaystyle\frac{\partial v_d}{\partial x_d}
\end{pmatrix}.
ここで、行列のトレースとは、対角成分をすべて足し合わせたものです。
したがって、
$$
\operatorname{tr}(J_{\mathbf v})=
\frac{\partial v_1}{\partial x_1}
+
\frac{\partial v_2}{\partial x_2}
+
\cdots
+
\frac{\partial v_d}{\partial x_d}.
$$
一方、ベクトル場 (\mathbf v) の発散は、定義によって
$$
\nabla\cdot\mathbf v =
\frac{\partial v_1}{\partial x_1}
+
\frac{\partial v_2}{\partial x_2}
+
\cdots
+
\frac{\partial v_d}{\partial x_d}
$$
と表されます。
両者を並べてみると、
$$
\operatorname{tr}(J_{\mathbf v})=
\frac{\partial v_1}{\partial x_1}
+
\cdots
+
\frac{\partial v_d}{\partial x_d}=
\nabla\cdot\mathbf v.
$$
したがって、
$$
\operatorname{tr}(J_{\mathbf v})=
\nabla\cdot\mathbf v
$$
となります。
つまり、これは何か難しい定理を証明しているわけではありません。
「ベクトル場のヤコビ行列の対角成分の和」を「発散」と呼び、「行列の対角成分の和」を「トレース」と呼んでいるため、この2つが一致するのです。
この関係が重要なのは、正規化フローでは「ヤコビアンの行列式」で測っていた空間の体積変化が、連続時間の極限では「ヤコビ行列のトレース」、すなわち「発散」によって表せるからです。
これによって、
$$
\text{行列式}
\quad\longrightarrow\quad
\text{トレース}
\quad\longrightarrow\quad
\text{発散}
$$
という、第二話の重要な数学的な橋渡しが完成します。
【本記事におけるAI(生成AI)の活用について】
本連載(および本記事)は、著者が約4ヶ月間にわたり積み重ねてきた専門論文の精読・数理的考察と、AIとの徹底的なディスカッションを経て構築したプロットに基づいています。
記事内の本文記述および見出しの言語化・構造化にあたっては、論理の透明性と読みやすさを高めるパートナーとしてAI(LLM)を活用・編集しています。数理的な構想や文脈の意図はすべて筆者の思索に基づくものです。