スコアマッチングとスコアベース生成モデルを解説

スコアベース生成モデル(Score-based Generative Models)は、データ分布のスコア関数(対数密度の勾配)を学習し、それを用いてサンプリングを行う生成モデルの枠組みです。Song & Ermon (2019) が提案したこのアプローチは、拡散モデル(DDPM)と深い数学的なつながりを持ち、Song et al. (2021) の確率微分方程式(SDE)フレームワークによって両者が統一的に理解されるようになりました。

スコアベース生成モデルの理論的な美しさは、分配関数(正規化定数)の計算を完全に回避できる点にあります。確率密度関数そのものではなく、そのスコア(勾配)を推定するため、正規化定数が不要になるのです。本記事では、スコアマッチングの基礎理論から暗黙的スコアマッチング、デノイジングスコアマッチング、ランジュバン動力学、そしてSDEフレームワークまでを省略なく導出します。

本記事の内容

  • スコア関数 $\nabla_{\bm{x}} \log p(\bm{x})$ の定義と幾何学的意味
  • 明示的スコアマッチングの目的関数
  • 暗黙的スコアマッチング(部分積分による等価変換の導出)
  • デノイジングスコアマッチング(ノイズ付きデータでのスコア推定)
  • ランジュバン動力学によるサンプリング
  • アニーリング付きランジュバン動力学
  • SDEフレームワークによるDDPMとの統一的理解
  • Pythonで2Dデータのスコアマッチング・ランジュバンサンプリング

前提知識

この記事を読む前に、以下の記事を読んでおくと理解が深まります。

スコア関数とは

定義

データ分布 $p(\bm{x})$ のスコア関数(score function)は、対数密度の勾配として定義されます。

$$ \bm{s}(\bm{x}) = \nabla_{\bm{x}} \log p(\bm{x}) $$

ここで $\nabla_{\bm{x}}$ は $\bm{x}$ に関する勾配演算子です。$\bm{x} \in \mathbb{R}^d$ の場合、スコア関数は $d$ 次元ベクトルを返します。

$$ \bm{s}(\bm{x}) = \begin{pmatrix} \frac{\partial \log p(\bm{x})}{\partial x_1} \\ \vdots \\ \frac{\partial \log p(\bm{x})}{\partial x_d} \end{pmatrix} $$

幾何学的意味

スコア関数は、各点 $\bm{x}$ において確率密度が最も急激に増大する方向を指すベクトル場です。直感的に言えば、スコアベクトルの方向は「データがより集中している方向」を示しています。

確率密度 $p(\bm{x})$ が高い領域(データの密集地帯)に向かって、スコアベクトルが収束していくイメージです。

なぜスコア関数なのか

確率密度関数 $p(\bm{x})$ を直接モデリングする場合、正規化条件

$$ \int p(\bm{x})\,d\bm{x} = 1 $$

を満たす必要があり、この積分(分配関数 $Z$)の計算は一般に困難です。しかしスコア関数は対数の勾配であるため

$$ \nabla_{\bm{x}} \log p(\bm{x}) = \nabla_{\bm{x}} \log \frac{\tilde{p}(\bm{x})}{Z} = \nabla_{\bm{x}} \log \tilde{p}(\bm{x}) $$

となり、正規化定数 $Z$ が消えます。これがスコアベース手法の最大の利点です。

明示的スコアマッチング

目的関数

スコア関数をニューラルネットワーク $\bm{s}_\theta(\bm{x})$ で近似することを考えます。最も素直な目的関数は、真のスコア関数とモデルスコアの二乗誤差です。

$$ J_{\text{ESM}}(\theta) = \frac{1}{2}\mathbb{E}_{p(\bm{x})}\left[\|\bm{s}_\theta(\bm{x}) – \nabla_{\bm{x}} \log p(\bm{x})\|^2\right] $$

しかし、この目的関数は真のスコア $\nabla_{\bm{x}} \log p(\bm{x})$ を必要とします。データの密度関数 $p(\bm{x})$ が未知である以上、これは直接計算できません。

暗黙的スコアマッチング(Implicit Score Matching)

Hyvarinen (2005) の定理

Hyvarinen (2005) は、$J_{\text{ESM}}$ を部分積分によって、真のスコアを含まない等価な目的関数に変換できることを示しました。

$J_{\text{ESM}}$ を展開します。

$$ \begin{align} J_{\text{ESM}}(\theta) &= \frac{1}{2}\mathbb{E}_{p}\left[\|\bm{s}_\theta(\bm{x})\|^2 – 2\bm{s}_\theta(\bm{x})^T \nabla_{\bm{x}} \log p(\bm{x}) + \|\nabla_{\bm{x}} \log p(\bm{x})\|^2\right] \end{align} $$

第3項は $\theta$ に依存しないので最適化に影響しません。問題は第2項(クロスターム)です。

$$ \mathbb{E}_{p}\left[\bm{s}_\theta(\bm{x})^T \nabla_{\bm{x}} \log p(\bm{x})\right] $$

部分積分による変換

このクロスタームを成分ごとに書くと

$$ \mathbb{E}_{p}\left[\bm{s}_\theta(\bm{x})^T \nabla_{\bm{x}} \log p(\bm{x})\right] = \sum_{i=1}^{d} \int p(\bm{x})\, s_{\theta,i}(\bm{x})\, \frac{\partial \log p(\bm{x})}{\partial x_i}\,d\bm{x} $$

ここで $\frac{\partial \log p(\bm{x})}{\partial x_i} = \frac{1}{p(\bm{x})} \frac{\partial p(\bm{x})}{\partial x_i}$ を用いると

$$ \int p(\bm{x})\, s_{\theta,i}(\bm{x})\, \frac{1}{p(\bm{x})}\frac{\partial p(\bm{x})}{\partial x_i}\,d\bm{x} = \int s_{\theta,i}(\bm{x})\, \frac{\partial p(\bm{x})}{\partial x_i}\,d\bm{x} $$

$x_i$ について部分積分を適用します(他の変数は固定)。

$$ \int s_{\theta,i}(\bm{x})\, \frac{\partial p(\bm{x})}{\partial x_i}\,dx_i = \left[s_{\theta,i}(\bm{x})\, p(\bm{x})\right]_{-\infty}^{\infty} – \int \frac{\partial s_{\theta,i}(\bm{x})}{\partial x_i}\, p(\bm{x})\,dx_i $$

境界条件として、$|x_i| \to \infty$ で $p(\bm{x}) \to 0$ が十分速いと仮定すると、境界項は消えます。

$$ \int s_{\theta,i}(\bm{x})\, \frac{\partial p(\bm{x})}{\partial x_i}\,dx_i = -\int \frac{\partial s_{\theta,i}(\bm{x})}{\partial x_i}\, p(\bm{x})\,dx_i = -\mathbb{E}_p\left[\frac{\partial s_{\theta,i}(\bm{x})}{\partial x_i}\right] $$

全次元について合算すると

$$ \mathbb{E}_{p}\left[\bm{s}_\theta(\bm{x})^T \nabla_{\bm{x}} \log p(\bm{x})\right] = -\mathbb{E}_p\left[\text{tr}\left(\nabla_{\bm{x}} \bm{s}_\theta(\bm{x})\right)\right] $$

ここで $\text{tr}(\nabla_{\bm{x}} \bm{s}_\theta(\bm{x})) = \sum_{i=1}^d \frac{\partial s_{\theta,i}}{\partial x_i}$ はスコアモデルのヤコビアンのトレース(ダイバージェンス)です。

暗黙的スコアマッチングの最終形

以上をまとめると、定数項を除いた暗黙的スコアマッチングの目的関数は

$$ \boxed{J_{\text{ISM}}(\theta) = \mathbb{E}_{p(\bm{x})}\left[\text{tr}\left(\nabla_{\bm{x}} \bm{s}_\theta(\bm{x})\right) + \frac{1}{2}\|\bm{s}_\theta(\bm{x})\|^2\right]} $$

この目的関数は真のスコア $\nabla_{\bm{x}} \log p(\bm{x})$ を一切含みません。データのサンプルのみから計算できます。

ただし、ヤコビアンのトレースの計算は $d$ 回の逆伝播が必要で、高次元データでは計算コストが高くなります。

デノイジングスコアマッチング(DSM)

動機

暗黙的スコアマッチングのヤコビアン計算は高次元で困難です。Vincent (2011) は、ノイズを加えたデータのスコアを推定することで、この問題を回避できることを示しました。

ノイズ付きデータの分布

データ $\bm{x}$ にガウスノイズを加えた分布を考えます。

$$ q_\sigma(\tilde{\bm{x}} | \bm{x}) = \mathcal{N}(\tilde{\bm{x}};\; \bm{x},\; \sigma^2 \bm{I}) $$

ノイズ付きデータの周辺分布は

$$ q_\sigma(\tilde{\bm{x}}) = \int q_\sigma(\tilde{\bm{x}} | \bm{x})\, p(\bm{x})\,d\bm{x} $$

DSMの目的関数

デノイジングスコアマッチングの目的関数は

$$ J_{\text{DSM}}(\theta) = \frac{1}{2}\mathbb{E}_{q_\sigma(\tilde{\bm{x}}, \bm{x})}\left[\|\bm{s}_\theta(\tilde{\bm{x}}) – \nabla_{\tilde{\bm{x}}} \log q_\sigma(\tilde{\bm{x}} | \bm{x})\|^2\right] $$

条件付きスコアの計算

$q_\sigma(\tilde{\bm{x}}|\bm{x}) = \mathcal{N}(\tilde{\bm{x}}; \bm{x}, \sigma^2\bm{I})$ のスコアは解析的に計算できます。

$$ \log q_\sigma(\tilde{\bm{x}}|\bm{x}) = -\frac{\|\tilde{\bm{x}} – \bm{x}\|^2}{2\sigma^2} + \text{const} $$

$$ \nabla_{\tilde{\bm{x}}} \log q_\sigma(\tilde{\bm{x}}|\bm{x}) = -\frac{\tilde{\bm{x}} – \bm{x}}{\sigma^2} $$

したがってDSMの目的関数は

$$ \boxed{J_{\text{DSM}}(\theta) = \frac{1}{2}\mathbb{E}_{p(\bm{x})}\mathbb{E}_{q_\sigma(\tilde{\bm{x}}|\bm{x})}\left[\left\|\bm{s}_\theta(\tilde{\bm{x}}) + \frac{\tilde{\bm{x}} – \bm{x}}{\sigma^2}\right\|^2\right]} $$

DSMとISMの等価性

Vincent (2011) は次の定理を証明しました。

定理: $J_{\text{DSM}}(\theta)$ と、ノイズ付き分布 $q_\sigma(\tilde{\bm{x}})$ に対する暗黙的スコアマッチング $J_{\text{ISM}}^{q_\sigma}(\theta)$ は、$\theta$ に依存しない定数の差を除いて等しい。

$$ J_{\text{DSM}}(\theta) = J_{\text{ISM}}^{q_\sigma}(\theta) + \text{const} $$

これは、DSMで $\bm{s}_\theta(\tilde{\bm{x}}) \approx \nabla_{\tilde{\bm{x}}} \log q_\sigma(\tilde{\bm{x}})$ を学習できることを意味します。つまり、ノイズ付きデータの条件付きスコア(既知)を目標にするだけで、ノイズ付きデータの周辺スコアを学習できるのです。

証明の概要: 明示的スコアマッチングの目的関数を $q_\sigma(\tilde{\bm{x}})$ に適用し、部分積分(暗黙的スコアマッチングの変換)を行い、$q_\sigma(\tilde{\bm{x}}) = \int q_\sigma(\tilde{\bm{x}}|\bm{x})p(\bm{x})d\bm{x}$ を代入して整理すると、DSMの目的関数と一致します。

$$ \begin{align} J_{\text{ESM}}^{q_\sigma}(\theta) &= \frac{1}{2}\mathbb{E}_{q_\sigma(\tilde{\bm{x}})}\left[\|\bm{s}_\theta(\tilde{\bm{x}}) – \nabla_{\tilde{\bm{x}}}\log q_\sigma(\tilde{\bm{x}})\|^2\right] \\ &= \frac{1}{2}\mathbb{E}_{q_\sigma(\tilde{\bm{x}})}\left[\|\bm{s}_\theta(\tilde{\bm{x}})\|^2\right] – \mathbb{E}_{q_\sigma(\tilde{\bm{x}})}\left[\bm{s}_\theta(\tilde{\bm{x}})^T \nabla_{\tilde{\bm{x}}}\log q_\sigma(\tilde{\bm{x}})\right] + C_1 \end{align} $$

第2項について、$\nabla_{\tilde{\bm{x}}}\log q_\sigma(\tilde{\bm{x}}) = \frac{\nabla_{\tilde{\bm{x}}} q_\sigma(\tilde{\bm{x}})}{q_\sigma(\tilde{\bm{x}})}$ を代入し、$q_\sigma(\tilde{\bm{x}}) = \int q_\sigma(\tilde{\bm{x}}|\bm{x})p(\bm{x})d\bm{x}$ を用いると

$$ \begin{align} \mathbb{E}_{q_\sigma(\tilde{\bm{x}})}\left[\bm{s}_\theta(\tilde{\bm{x}})^T \frac{\nabla_{\tilde{\bm{x}}} q_\sigma(\tilde{\bm{x}})}{q_\sigma(\tilde{\bm{x}})}\right] &= \int \bm{s}_\theta(\tilde{\bm{x}})^T \nabla_{\tilde{\bm{x}}} q_\sigma(\tilde{\bm{x}})\,d\tilde{\bm{x}} \\ &= \int \bm{s}_\theta(\tilde{\bm{x}})^T \int \nabla_{\tilde{\bm{x}}} q_\sigma(\tilde{\bm{x}}|\bm{x})p(\bm{x})\,d\bm{x}\,d\tilde{\bm{x}} \\ &= \mathbb{E}_{p(\bm{x})}\left[\int \bm{s}_\theta(\tilde{\bm{x}})^T \nabla_{\tilde{\bm{x}}} q_\sigma(\tilde{\bm{x}}|\bm{x})\,d\tilde{\bm{x}}\right] \end{align} $$

内側の積分に部分積分を適用すると($q_\sigma(\tilde{\bm{x}}|\bm{x})$ はガウスなので境界項は消える)

$$ \int \bm{s}_\theta(\tilde{\bm{x}})^T \nabla_{\tilde{\bm{x}}} q_\sigma(\tilde{\bm{x}}|\bm{x})\,d\tilde{\bm{x}} = -\int \text{tr}(\nabla_{\tilde{\bm{x}}} \bm{s}_\theta) q_\sigma(\tilde{\bm{x}}|\bm{x})\,d\tilde{\bm{x}} + \int \bm{s}_\theta^T \nabla_{\tilde{\bm{x}}} \log q_\sigma(\tilde{\bm{x}}|\bm{x})\,q_\sigma(\tilde{\bm{x}}|\bm{x})\,d\tilde{\bm{x}} $$

ここで暗黙的スコアマッチングの変換(部分積分)を逆に適用することで、最終的に

$$ J_{\text{ESM}}^{q_\sigma}(\theta) = J_{\text{DSM}}(\theta) + C_2 $$

が示されます。$\square$

ランジュバン動力学によるサンプリング

ランジュバン動力学

スコア関数 $\nabla_{\bm{x}} \log p(\bm{x})$ が得られれば、ランジュバン動力学(Langevin Dynamics)を用いてデータ分布からサンプリングできます。

$$ \boxed{\bm{x}_{k+1} = \bm{x}_k + \eta\,\nabla_{\bm{x}} \log p(\bm{x}_k) + \sqrt{2\eta}\,\bm{z}_k, \quad \bm{z}_k \sim \mathcal{N}(\bm{0}, \bm{I})} $$

ここで $\eta > 0$ はステップサイズです。

この更新式は、確率勾配によりデータの高密度領域に向かって移動しつつ、ガウスノイズにより確率的に探索を行います。$\eta \to 0$ かつ更新回数 $K \to \infty$ の極限で、$\bm{x}_K$ の分布は $p(\bm{x})$ に収束することが知られています。

導出(フォッカー・プランク方程式との関係)

ランジュバン動力学は、以下の確率微分方程式(SDE)の離散化です。

$$ d\bm{x} = \nabla_{\bm{x}} \log p(\bm{x})\,dt + \sqrt{2}\,d\bm{w} $$

ここで $\bm{w}$ はウィーナー過程(ブラウン運動)です。この SDE に対応するフォッカー・プランク方程式は

$$ \frac{\partial p_t(\bm{x})}{\partial t} = -\nabla \cdot \left[p_t(\bm{x}) \nabla \log p(\bm{x})\right] + \Delta p_t(\bm{x}) $$

ここで $\Delta$ はラプラシアンです。第1項を展開します。

$$ \nabla \cdot [p_t \nabla \log p] = \nabla \cdot \left[p_t \frac{\nabla p}{p}\right] $$

定常分布 $p_t = p$ のとき

$$ \nabla \cdot \left[p \frac{\nabla p}{p}\right] = \nabla \cdot (\nabla p) = \Delta p $$

よって

$$ \frac{\partial p}{\partial t} = -\Delta p + \Delta p = 0 $$

となり、真のデータ分布 $p(\bm{x})$ が定常分布であることが確認できます。

アニーリング付きランジュバン動力学

問題点と解決策

通常のランジュバン動力学には課題があります。データ分布が複数のモードを持つ場合(混合分布)、低密度領域でのスコア推定が不正確になり、モード間の遷移が困難になります。

Song & Ermon (2019) は、複数のノイズレベル $\sigma_1 > \sigma_2 > \dots > \sigma_L$ を用いてスコアネットワークを条件付けし、大きなノイズレベルから小さなノイズレベルへと段階的にランジュバン動力学を実行するアニーリング付きランジュバン動力学(Annealed Langevin Dynamics)を提案しました。

アルゴリズム

ノイズ条件付きスコアマッチングの目的関数:

$$ J(\theta) = \sum_{l=1}^{L} \lambda(\sigma_l)\, \mathbb{E}_{p(\bm{x})}\mathbb{E}_{\tilde{\bm{x}} \sim \mathcal{N}(\bm{x}, \sigma_l^2 \bm{I})}\left[\left\|\bm{s}_\theta(\tilde{\bm{x}}, \sigma_l) + \frac{\tilde{\bm{x}} – \bm{x}}{\sigma_l^2}\right\|^2\right] $$

ここで $\lambda(\sigma_l) = \sigma_l^2$ と設定することが一般的です。

アニーリング付きランジュバン動力学のサンプリング:

  1. $\bm{x}_0 \sim \mathcal{N}(\bm{0}, \sigma_1^2 \bm{I})$ を初期化
  2. $l = 1, 2, \dots, L$ について: – $\eta_l = \epsilon \cdot (\sigma_l / \sigma_L)^2$(ステップサイズ) – $k = 1, \dots, K$ について:
    • $\bm{z} \sim \mathcal{N}(\bm{0}, \bm{I})$
    • $\bm{x}_k = \bm{x}_{k-1} + \eta_l\, \bm{s}_\theta(\bm{x}_{k-1}, \sigma_l) + \sqrt{2\eta_l}\,\bm{z}$
  3. $\bm{x}_K$ を返す

大きなノイズ $\sigma_1$ でのランジュバン動力学は、モード間を自由に遷移でき(分布が平滑化されているため)、ノイズを段階的に小さくすることで最終的に真の分布に近いサンプルを得ます。

SDEフレームワーク: DDPMとの統一的理解

連続時間への拡張

Song et al. (2021) は、拡散モデルとスコアベースモデルを確率微分方程式(SDE)のフレームワークで統一しました。

離散的な前方過程 $q(\bm{x}_t|\bm{x}_{t-1})$ をステップ幅 $\Delta t \to 0$ で連続化すると、以下の SDE になります。

$$ \boxed{d\bm{x} = \bm{f}(\bm{x}, t)\,dt + g(t)\,d\bm{w}} $$

ここで $\bm{f}(\bm{x}, t)$ はドリフト係数、$g(t)$ は拡散係数、$\bm{w}$ はウィーナー過程です。

DDPMの場合

DDPMの前方過程 $\bm{x}_t = \sqrt{1-\beta_t}\,\bm{x}_{t-1} + \sqrt{\beta_t}\,\bm{\epsilon}_t$ を連続時間の SDE として書くと

$$ d\bm{x} = -\frac{1}{2}\beta(t)\,\bm{x}\,dt + \sqrt{\beta(t)}\,d\bm{w} $$

これは Variance Preserving SDE(VP-SDE) と呼ばれます。ドリフト項 $\bm{f}(\bm{x},t) = -\frac{1}{2}\beta(t)\bm{x}$ は原点に向かう引力(信号の減衰)、拡散項 $g(t) = \sqrt{\beta(t)}$ はノイズの付加を表します。

逆時間SDE

Anderson (1982) の結果により、前方SDEの逆時間過程もまたSDEで書けます。

$$ \boxed{d\bm{x} = \left[\bm{f}(\bm{x}, t) – g(t)^2 \nabla_{\bm{x}} \log p_t(\bm{x})\right]dt + g(t)\,d\bar{\bm{w}}} $$

ここで $\bar{\bm{w}}$ は逆時間のウィーナー過程、$p_t(\bm{x})$ は時刻 $t$ でのデータの周辺分布です。

この結果は深遠です。逆時間SDEにおいて必要なのは、各時刻 $t$ でのスコア関数 $\nabla_{\bm{x}} \log p_t(\bm{x})$ だけなのです。

DDPMとスコアマッチングの統一

DDPMの逆過程をSDEフレームワークで見ると

  • 前方SDE: $d\bm{x} = -\frac{1}{2}\beta(t)\bm{x}\,dt + \sqrt{\beta(t)}\,d\bm{w}$
  • 逆時間SDE: $d\bm{x} = \left[-\frac{1}{2}\beta(t)\bm{x} – \beta(t)\nabla_{\bm{x}}\log p_t(\bm{x})\right]dt + \sqrt{\beta(t)}\,d\bar{\bm{w}}$

DDPMのノイズ予測 $\bm{\epsilon}_\theta(\bm{x}_t, t)$ はスコア関数と次の関係にあります。

$$ \nabla_{\bm{x}_t} \log p_t(\bm{x}_t) = -\frac{\bm{\epsilon}_\theta(\bm{x}_t, t)}{\sqrt{1 – \bar{\alpha}_t}} $$

これは、$q(\bm{x}_t|\bm{x}_0) = \mathcal{N}(\sqrt{\bar{\alpha}_t}\bm{x}_0, (1-\bar{\alpha}_t)\bm{I})$ のスコアが

$$ \nabla_{\bm{x}_t} \log q(\bm{x}_t|\bm{x}_0) = -\frac{\bm{x}_t – \sqrt{\bar{\alpha}_t}\bm{x}_0}{1 – \bar{\alpha}_t} = -\frac{\bm{\epsilon}}{\sqrt{1 – \bar{\alpha}_t}} $$

であることから直ちに導けます。つまり、DDPMのノイズ予測はスコア関数の推定と本質的に等価です。

Variance Exploding SDE(VE-SDE)

SMLD(Score Matching with Langevin Dynamics)に対応するSDEは

$$ d\bm{x} = \sqrt{\frac{d[\sigma^2(t)]}{dt}}\,d\bm{w} $$

ドリフト項がなく($\bm{f} = \bm{0}$)、拡散のみの SDE です。これはVE-SDE と呼ばれ、分散が時間とともに増大(explode)します。

Pythonでの実装

2Dガウス混合分布を対象に、デノイジングスコアマッチングとランジュバン動力学によるサンプリングを実装します。

import numpy as np
import matplotlib.pyplot as plt

np.random.seed(42)

# --- 2Dガウス混合分布からデータ生成 ---
def sample_gmm(n_samples=3000):
    """3つのガウスの混合分布からサンプリング"""
    means = [np.array([-2, 0]), np.array([2, 0]), np.array([0, 2.5])]
    cov = np.eye(2) * 0.3
    n_per = n_samples // 3
    samples = []
    for m in means:
        samples.append(np.random.multivariate_normal(m, cov, n_per))
    return np.vstack(samples)

data = sample_gmm(3000)

# --- スコアネットワーク(MLP) ---
class ScoreNetwork:
    """スコア関数を近似するMLP(入力: x + sigma)"""
    def __init__(self, hidden=128, lr=1e-3):
        # 入力: x(2) + sigma(1) = 3次元
        self.W1 = np.random.randn(3, hidden) * 0.1
        self.b1 = np.zeros(hidden)
        self.W2 = np.random.randn(hidden, hidden) * 0.1
        self.b2 = np.zeros(hidden)
        self.W3 = np.random.randn(hidden, 2) * 0.1
        self.b3 = np.zeros(2)
        self.lr = lr
        self.params = [self.W1, self.b1, self.W2, self.b2, self.W3, self.b3]
        self.m = [np.zeros_like(p) for p in self.params]
        self.v = [np.zeros_like(p) for p in self.params]
        self.step = 0

    def forward(self, x, sigma):
        """順伝播: x (N,2), sigma (N,1) -> score (N,2)"""
        inp = np.concatenate([x, sigma.reshape(-1, 1)], axis=1)
        self.inp = inp
        self.h1_pre = inp @ self.W1 + self.b1
        self.h1 = np.maximum(0, self.h1_pre)
        self.h2_pre = self.h1 @ self.W2 + self.b2
        self.h2 = np.maximum(0, self.h2_pre)
        out = self.h2 @ self.W3 + self.b3
        return out

    def backward(self, target, pred):
        """MSE損失の逆伝播 + Adam更新"""
        N = target.shape[0]
        d_out = 2.0 / N * (pred - target)
        dW3 = self.h2.T @ d_out
        db3 = d_out.sum(axis=0)
        d_h2 = d_out @ self.W3.T
        d_h2 *= (self.h2_pre > 0)
        dW2 = self.h1.T @ d_h2
        db2 = d_h2.sum(axis=0)
        d_h1 = d_h2 @ self.W2.T
        d_h1 *= (self.h1_pre > 0)
        dW1 = self.inp.T @ d_h1
        db1 = d_h1.sum(axis=0)
        grads = [dW1, db1, dW2, db2, dW3, db3]
        self.step += 1
        beta1, beta2, eps_adam = 0.9, 0.999, 1e-8
        for i, g in enumerate(grads):
            self.m[i] = beta1 * self.m[i] + (1 - beta1) * g
            self.v[i] = beta2 * self.v[i] + (1 - beta2) * g**2
            m_hat = self.m[i] / (1 - beta1**self.step)
            v_hat = self.v[i] / (1 - beta2**self.step)
            self.params[i] -= self.lr * m_hat / (np.sqrt(v_hat) + eps_adam)
        self.W1, self.b1, self.W2, self.b2, self.W3, self.b3 = self.params

# --- デノイジングスコアマッチングで学習 ---
model = ScoreNetwork(hidden=128, lr=1e-3)
sigmas = np.array([2.0, 1.0, 0.5, 0.3, 0.1])  # 複数のノイズレベル
n_epochs = 500
batch_size = 256
losses = []

for epoch in range(n_epochs):
    idx = np.random.choice(len(data), batch_size)
    x0 = data[idx]
    # ランダムなノイズレベルを選択
    sigma_idx = np.random.randint(0, len(sigmas), batch_size)
    sigma = sigmas[sigma_idx]
    # ノイズ付加
    noise = np.random.randn(*x0.shape)
    x_noisy = x0 + sigma.reshape(-1, 1) * noise
    # 条件付きスコア(ターゲット): -(x_noisy - x0) / sigma^2
    target_score = -(x_noisy - x0) / (sigma**2).reshape(-1, 1)
    # スコア予測
    pred_score = model.forward(x_noisy, sigma)
    loss = np.mean((target_score - pred_score)**2)
    losses.append(loss)
    model.backward(target_score, pred_score)

# --- アニーリング付きランジュバン動力学でサンプリング ---
def annealed_langevin(model, sigmas, n_samples=1000, K=100, eps=0.01):
    """アニーリング付きランジュバン動力学"""
    x = np.random.randn(n_samples, 2) * sigmas[0]  # 初期値
    trajectory = [x.copy()]
    for sigma in sigmas:
        eta = eps * (sigma / sigmas[-1])**2  # ステップサイズ
        sigma_arr = np.full(n_samples, sigma)
        for k in range(K):
            score = model.forward(x, sigma_arr)
            z = np.random.randn(n_samples, 2)
            x = x + eta * score + np.sqrt(2 * eta) * z
        trajectory.append(x.copy())
    return x, trajectory

samples, traj = annealed_langevin(model, sigmas, n_samples=2000, K=100, eps=0.005)

# --- 可視化 ---
fig, axes = plt.subplots(1, 3, figsize=(18, 5))

# 学習損失
axes[0].plot(losses)
axes[0].set_xlabel('Iteration')
axes[0].set_ylabel('DSM Loss')
axes[0].set_title('Training Loss (Denoising Score Matching)')
axes[0].set_yscale('log')

# 元データ
axes[1].scatter(data[:, 0], data[:, 1], s=1, alpha=0.5, c='blue')
axes[1].set_title('Original Data (GMM)')
axes[1].set_xlim(-5, 5)
axes[1].set_ylim(-3, 5)
axes[1].set_aspect('equal')

# 生成サンプル
axes[2].scatter(samples[:, 0], samples[:, 1], s=1, alpha=0.5, c='red')
axes[2].set_title('Generated Samples (Annealed Langevin)')
axes[2].set_xlim(-5, 5)
axes[2].set_ylim(-3, 5)
axes[2].set_aspect('equal')

plt.tight_layout()
plt.show()

次に、学習したスコア関数のベクトル場を可視化します。

import numpy as np
import matplotlib.pyplot as plt

# --- スコア場の可視化 ---
# (上のコードブロックで model が学習済みであることを前提)

fig, axes = plt.subplots(1, 3, figsize=(18, 5))
sigma_vis = [2.0, 0.5, 0.1]

for ax, sigma in zip(axes, sigma_vis):
    # グリッド上でスコアを計算
    xx, yy = np.meshgrid(np.linspace(-5, 5, 20), np.linspace(-3, 5, 16))
    grid = np.stack([xx.ravel(), yy.ravel()], axis=1)
    sigma_arr = np.full(len(grid), sigma)
    score = model.forward(grid, sigma_arr)

    # スコアの大きさでクリップ(可視化のため)
    norm = np.sqrt((score**2).sum(axis=1, keepdims=True))
    score_clipped = score / (norm + 1e-8) * np.minimum(norm, 3.0)

    ax.quiver(grid[:, 0], grid[:, 1],
              score_clipped[:, 0], score_clipped[:, 1],
              norm.ravel(), cmap='viridis', alpha=0.8)
    ax.scatter(data[:500, 0], data[:500, 1], s=1, alpha=0.3, c='gray')
    ax.set_title(f'Score Field (sigma={sigma})')
    ax.set_xlim(-5, 5)
    ax.set_ylim(-3, 5)
    ax.set_aspect('equal')

plt.suptitle('Learned Score Vector Fields at Different Noise Levels', fontsize=14)
plt.tight_layout()
plt.show()

スコア場の可視化から、ノイズレベルが大きい($\sigma=2.0$)ときは広域的にデータの中心方向を指すベクトル場になり、ノイズレベルが小さい($\sigma=0.1$)ときは各ガウスクラスタの中心に向かう局所的なベクトル場になっていることが確認できます。

まとめ

本記事では、スコアマッチングとスコアベース生成モデルの理論を導出しました。

  • スコア関数 $\nabla_{\bm{x}} \log p(\bm{x})$ は対数密度の勾配であり、正規化定数に依存しない
  • 暗黙的スコアマッチングは、部分積分により真のスコアを使わずに学習できる
  • デノイジングスコアマッチングは、ノイズ付き条件付きスコア $-(\tilde{\bm{x}}-\bm{x})/\sigma^2$ を目標に学習する
  • ランジュバン動力学 $\bm{x}_{k+1} = \bm{x}_k + \eta\,\nabla \log p(\bm{x}_k) + \sqrt{2\eta}\,\bm{z}$ でサンプリングを実行
  • SDEフレームワークにより、DDPMのノイズ予測はスコア関数の推定と等価であることが示された

次のステップとして、以下の記事も参考にしてください。