第9回：拡散と凝縮 ～熱力学との融合～

注意事項

• 標準的な拡散モデル（DDPM等）はユークリッド空間 ${\mathbb{R}}^d$ 上で定義される。最終状態はガウス分布。
• 「球面上の拡散」は別の定式化であり、標準手法とは異なる。
• SDE/ODEの式は概念的な表現であり、実装では離散化・近似が入る。
• スコア関数の推定は、実際にはノイズ予測ネットワークを介して行われる。理論と実装の対応に注意。
• 本回で扱う「時間」は、物理的な時間ではなく、拡散・逆拡散プロセスの「進行度」を指す。

導入：霧から星が凝縮する

前回、「時間の発見」として、生成プロセスを軌道として捉える視点を導入した。Neural ODEは、離散的な層の積み重ねを連続時間の流れとして再解釈する枠組みを与えた。

本回では、この視点をさらに発展させる。鍵となるのは、確率過程の導入である。

拡散モデル（Diffusion Models）は、近年の生成AIの中核技術となっている。画像生成（Stable Diffusion, DALL-E）、動画生成（Sora）、音声合成など、あらゆるモダリティで成功を収めている。

その直感は驚くほどシンプルである。

インクを水に垂らすと拡散する。最初は明確な形を持っていたインクが、時間とともに広がり、最終的には一様に薄まる。これが Forward Process（順過程） である。

では、この過程を「逆再生」できたらどうなるか。一様に拡散したインクが、徐々に凝縮し、元の形を取り戻す。これが Reverse Process（逆過程） であり、拡散モデルの生成プロセスそのものである。

NOTE

熱力学との類似（ただし比喩である）： この描像は、熱力学における「エントロピー増大」と「時間の矢」を想起させる。拡散は自然に起こるが、その逆は自然には起こらない。ただし、拡散モデルの逆過程は物理法則を時間反転しているわけではない。Forward Processとは別の確率過程を学習して構成している。熱力学的な比喩は直感を助けるが、物理的可逆性とは異なる概念である。

本回では、この直感を数学的に定式化し、幾何学的な視点から再解釈する。SDE（確率微分方程式）とODE（常微分方程式）の対応、スコア関数の意味、そして「ベクトル場に沿った積分」という視点が、生成モデルの本質をどう照らすかを見ていく。

標準的な拡散モデル： ${\mathbb{R}}^d$ 上の定式化

Forward Process：データからノイズへ

拡散モデルの出発点は、データを徐々にノイズで汚していく Forward Process である。

標準的な定式化（VP-SDE: Variance Preserving SDE）では、以下のSDEで記述される：

$$\begin{gathered} d{\mathbf{x}}_t=-\frac{{\beta }_t}{2}{\mathbf{x}}_t \\ ,dt+\sqrt{{\beta }_t} \\ ,d{\mathbf{W}}_t \end{gathered}$$

ここで：

• ${\mathbf{x}}_t\in {\mathbb{R}}^d$ ：時刻 $t$ での状態
• ${\beta }_t$ ：ノイズスケジュール（時刻に依存するスカラー関数）
• $d{\mathbf{W}}_t$ ：ブラウン運動（標準ウィーナー過程の微小増分）

項	役割	物理的解釈
$$\begin{gathered} -\frac{{\beta }_t}{2}{\mathbf{x}}_t \\ ,dt \end{gathered}$$	ドリフト項（決定論的）	原点に向かう収縮
$$\begin{gathered} \sqrt{{\beta }_t} \\ ,d{\mathbf{W}}_t \end{gathered}$$	拡散項（確率論的）	ランダムな揺らぎ

解の閉形式：任意時刻への直接ジャンプ

このSDEの重要な性質は、解が閉形式で書けることである。

$$\begin{gathered} {\mathbf{x}}_t=\sqrt{{\overline{\alpha }}_t} \\ ,{\mathbf{x}}_0+\sqrt{1-{\overline{\alpha }}_t} \\ ,\mathit{ϵ},\mathit{ϵ}\sim \mathcal{N}(\mathbf{0},\mathbf{I}) \end{gathered}$$

ここで $$\begin{gathered} {\overline{\alpha }}_t=\exp (-{\int }_0^t{\beta }_s \\ ,ds) \end{gathered}$$ である。

この式は、任意の時刻 $t$ における ${\mathbf{x}}_t$ を、 ${\mathbf{x}}_0$ から直接計算できることを意味する。学習時には、 $t=0$ から $t=T$ まで逐次シミュレーションする必要がない。

$$\begin{gathered} \text{データ }{\mathbf{x}}_0⟶{\mathbf{x}}_t=\sqrt{{\overline{\alpha }}_t} \\ ,{\mathbf{x}}_0+\sqrt{1-{\overline{\alpha }}_t} \\ ,\mathit{ϵ}(\text{任意の }t\text{ へ直接ジャンプ可能}) \end{gathered}$$

最終状態：ガウス分布への収束

$t\to \mathrm{\infty }$ で ${\overline{\alpha }}_t\to 0$ となるため、Forward Processの最終状態は：

$${\mathbf{x}}_{\mathrm{\infty }}\sim \mathcal{N}(\mathbf{0},\mathbf{I})$$

すなわち、標準ガウス分布に収束する。

IMPORTANT

ガウス分布と球面の違い： ガウス分布 $\mathcal{N}(\mathbf{0},\mathbf{I})$ は ${\mathbb{R}}^d$ 全体に広がる分布であり、球面 $S^{d-1}$ 上の一様分布とは全く異なる。高次元では、ガウス分布のサンプルは「原点から $\sqrt{d}$ 程度の距離」に集中するため、「薄い球殻」に近い振る舞いを見せるが、これは球面上の分布ではない。

Reverse Process：ノイズからデータへ

Forward Processを時間反転した Reverse Process は、以下のSDEで記述される：

$$\begin{gathered} d{\mathbf{x}}_t=[-\frac{{\beta }_t}{2}{\mathbf{x}}_t-{\beta }_t{\mathrm{\nabla }}_{\mathbf{x}}\log p_t({\mathbf{x}}_t)]dt+\sqrt{{\beta }_t} \\ ,d{\overline{\mathbf{W}}}_t \end{gathered}$$

ここで $d{\overline{\mathbf{W}}}_t$ は逆時間のブラウン運動、 ${\mathrm{\nabla }}_{\mathbf{x}}\log p_t(\mathbf{x})$ はスコア関数である。

IMPORTANT

時間方向の注意： 上記の式は「サンプリング方向」（ $t:T\to 0$ ）の記法で書いている。文献によっては、逆時間変数 $\tau =T-t$ を導入して $d\tau >0$ で書く流儀もあり、その場合は符号が異なって見える。実装時には、dt の符号とスコア項の符号の対応に注意が必要である。

具体的には、実装では for t in reversed(timesteps) のように t を減らしながら更新する。この資料の式は「 $dt>0$ の形で書くが、計算は $t$ を減らして行う」という対応になる。

項	役割
$-\frac{{\beta }_t}{2}{\mathbf{x}}_t$	元のドリフト項
$-{\beta }_t{\mathrm{\nabla }}_{\mathbf{x}}\log p_t({\mathbf{x}}_t)$	スコア関数による補正
$$\begin{gathered} \sqrt{{\beta }_t} \\ ,d{\overline{\mathbf{W}}}_t \end{gathered}$$	逆時間のノイズ

スコア関数 ${\mathrm{\nabla }}_{\mathbf{x}}\log p_t(\mathbf{x})$ は、確率密度の対数の勾配であり、「確率が高い方向」を指し示す。この項がなければ、逆時間のSDEは元のForward Processを再現しない。

NOTE

実装との対応： 実際の拡散モデルでは、スコア関数を直接学習するのではなく、 ノイズ予測ネットワーク ${\mathit{ϵ}}_{\theta }({\mathbf{x}}_t,t)$ を学習する。 両者は ${\mathrm{\nabla }}_{\mathbf{x}}\log p_t({\mathbf{x}}_t)\approx -\frac{{\mathit{ϵ}}_{\theta }({\mathbf{x}}_t,t)}{\sqrt{1-{\overline{\alpha }}_t}}$ という関係で結ばれる。

スコア関数：確率の流れの方向

スコア関数の幾何学的意味

スコア関数 ${\mathrm{\nabla }}_{\mathbf{x}}\log p(\mathbf{x})$ は、確率分布の「勾配」を表す。

具体的に考えてみよう。1次元のガウス分布 $p(x)=\mathcal{N}(0,1)$ の場合：

$$\log p(x)=-\frac{x^2}{2}+\text{const}$$$${\mathrm{\nabla }}_x\log p(x)=-x$$

スコアは「 $x>0$ なら負、 $x<0$ なら正」、つまり常に原点（モードの位置）に向かう。

多峰性の分布では、スコアは概ね「近いモードに向かう」傾向を示す（ただし、これは直感的な理解であり、厳密には成分の分散・重み・次元によって複雑に変わる。特にモード間の境界付近では、単純な「最近傍モード」規則にはならない）。

p(x) = 0.5 * N(-2, 1) + 0.5 * N(+2, 1)

スコアのベクトル場（概念図）:
  ←← ← · → →→        ←← ← · → →→
     [モード1]           [モード2]

スコアマッチング：スコアの学習

スコア関数を直接学習するのは困難である。 $p_t(\mathbf{x})$ は未知であり、正規化定数の計算も難しい。

スコアマッチング（Hyvärinen, 2005）は、この問題を回避する手法である。以下の損失関数を最小化することで、スコア関数を近似するネットワーク $s_{\theta }(\mathbf{x},t)$ を学習できる：

$${\mathcal{L}}_{\text{SM}}={\mathbb{E}}_{t,{\mathbf{x}}_0,\mathit{ϵ}}\left[{\Vert s_{\theta }({\mathbf{x}}_t,t)-{\mathrm{\nabla }}_{{\mathbf{x}}_t}\log p_{t|0}({\mathbf{x}}_t\mid {\mathbf{x}}_0)\Vert }^2\right]$$

Forward Processの閉形式解を使うと、条件付きスコアは：

$${\mathrm{\nabla }}_{{\mathbf{x}}_t}\log p_{t|0}({\mathbf{x}}_t\mid {\mathbf{x}}_0)=-\frac{\mathit{ϵ}}{\sqrt{1-{\overline{\alpha }}_t}}$$

これが「ノイズ予測」と「スコア推定」が等価である理由である。

アプローチ	学習対象	損失関数
スコアマッチング	$s_{\theta }\approx {\mathrm{\nabla }}_{\mathbf{x}}\log p_t$	$|s_{\theta }-\mathrm{\nabla }\log p_{t\mid 0}{|}^2$
ノイズ予測	${\mathit{ϵ}}_{\theta }\approx \mathit{ϵ}$	$|{\mathit{ϵ}}_{\theta }-\mathit{ϵ}{|}^2$

両者は定数倍の違いを除いて等価である。

確率フローODE：決定論的な生成

SDEからODEへ

Reverse SDEは確率的であり、同じ初期ノイズから異なるサンプルが生成される。しかし、同じ周辺分布 $p_t(\mathbf{x})$ を保ちながら、決定論的に状態を遷移させる方法がある。

これが 確率フローODE（Probability Flow ODE） である：

$$\frac{d\mathbf{x}}{dt}=-\frac{{\beta }_t}{2}{\mathbf{x}}_t-\frac{{\beta }_t}{2}{\mathrm{\nabla }}_{\mathbf{x}}\log p_t({\mathbf{x}}_t)$$

SDEと比較すると、ノイズ項 $$\begin{gathered} \sqrt{{\beta }_t} \\ ,d{\overline{\mathbf{W}}}_t \end{gathered}$$ が消え、代わりにスコア項の係数が半分になっている。

NOTE

係数の由来（一般形との対応）： 一般のSDE $d\mathbf{x}=f(\mathbf{x},t)dt+g(t)d\mathbf{W}$ に対応する確率フローODEは、

$$\frac{d\mathbf{x}}{dt}=f(\mathbf{x},t)-\frac{1}{2}g(t{)}^2{\mathrm{\nabla }}_{\mathbf{x}}\log p_t(\mathbf{x})$$

という形になる。VP-SDEでは $f=-\frac{{\beta }_t}{2}\mathbf{x}$ 、 $g=\sqrt{{\beta }_t}$ なので、 $g^2={\beta }_t$ となり、スコア項の係数が $\frac{{\beta }_t}{2}$ になる。

NOTE

直感的理解： SDEでは「ランダムに揺らぎながら確率の高い方向へ」、ODEでは「揺らがずに確率の高い方向へまっすぐ」進む。両者は、ある意味で「確率的な山登り」と「決定論的な山登り」の違いである。

確率フローODEの利点

確率フローODEは、いくつかの重要な利点を持つ。

1. 決定論的サンプリング：同じ初期ノイズ ${\mathbf{x}}_T$ から、常に同じ出力 ${\mathbf{x}}_0$ が得られる。これにより、潜在空間での操作（補間、編集など）が可能になる。

2. 高速サンプリング：ODEソルバーは、SDEサンプラーより少ないステップで収束することが多い。DDIM（Song et al., 2021）やDPM-Solver（Lu et al., 2022）などの高速サンプラーは、この性質を利用している。

3. 尤度計算：ODEの軌道に沿って、対数尤度を（原理的には）計算できる。これはVAEのようなELBOではなく、真の対数尤度である。

手法	種類	ステップ数	確率性
DDPM	SDE	数百〜数千	確率的
DDIM	ODE	数十〜数百	決定論的
DPM-Solver	ODE	10〜20	決定論的

ベクトル場としての解釈

確率フローODEの右辺は、 ${\mathbb{R}}^d$ 上のベクトル場を定義する。各点 $\mathbf{x}$ に対して、「次にどの方向に動くべきか」を指示するベクトルが割り当てられている。

$$\mathbf{v}(\mathbf{x},t)=-\frac{{\beta }_t}{2}\mathbf{x}-\frac{{\beta }_t}{2}{\mathrm{\nabla }}_{\mathbf{x}}\log p_t(\mathbf{x})$$

このベクトル場に沿って積分することが、生成プロセスそのものである。

ノイズ x_T ──→ ベクトル場に沿って流れる ──→ データ x_0

    ↗ → → ↘
  ↗           ↘
↗               ↓   ← ベクトル場 v(x, t)
↑               ↓
↑             ↙
  ← ← ← ← ←

Langevin動力学：スコアによるサンプリング

スコアベースのサンプリング

スコア関数 ${\mathrm{\nabla }}_{\mathbf{x}}\log p(\mathbf{x})$ が分かれば、分布 $p(\mathbf{x})$ からサンプリングできる。これが Langevin動力学 のアイデアである。

$$\begin{gathered} {\mathbf{x}}_{k+1}={\mathbf{x}}_k+\frac{\eta }{2}{\mathrm{\nabla }}_{\mathbf{x}}\log p({\mathbf{x}}_k)+\sqrt{\eta } \\ ,{\mathit{ϵ}}_k,{\mathit{ϵ}}_k\sim \mathcal{N}(\mathbf{0},\mathbf{I}) \end{gathered}$$

ここで $\eta$ はステップサイズ（学習率のようなハイパーパラメータ）である。 $k\to \mathrm{\infty }$ で、 ${\mathbf{x}}_k$ の分布は $p(\mathbf{x})$ に収束する（適切な条件下で）。

NOTE

Langevin式のバリエーション： 文献によっては、ステップサイズを $ϵ$ と書き、ノイズ項を $\sqrt{2ϵ}$ とする場合や、スコア項に係数をつけない定義など、バリエーションが無数にある。本資料の式は一般的な形の一つであり、実装時には使用する文献の定義を確認すること。

拡散モデルとの関係

Langevin動力学は、拡散モデルの Reverse Process（離散化版） と構造的に類似している。

$$\begin{gathered} {\mathbf{x}}_{t-\mathrm{\Delta }t}={\mathbf{x}}_t+{\beta }_t{\mathrm{\nabla }}_{\mathbf{x}}\log p_t({\mathbf{x}}_t)\mathrm{\Delta }t+\sqrt{2{\beta }_t\mathrm{\Delta }t} \\ ,\mathit{ϵ} \end{gathered}$$

NOTE

係数の注意： 上記の Langevin 式と拡散モデルの Reverse 離散化は「同型」だが、係数（ $1/2$ がどこに入るか、拡散係数の定義、温度、 $dt$ の符号）は定義や文献によって異なる。両者を厳密に対応させるには、各項の定義を揃える必要がある。

違いは、拡散モデルでは時間依存のスコア ${\mathrm{\nabla }}_{\mathbf{x}}\log p_t(\mathbf{x})$ を使うことである。これにより、「どの時刻のノイズレベルか」に応じた適切なデノイジングが可能になる。

NOTE

歴史的文脈： Langevin動力学は、統計物理学に起源を持つ古典的な手法である。エネルギーベースモデル（EBM）のサンプリングにも使われていたが、計算コストの問題で実用が難しかった。拡散モデルは、これを「複数のノイズレベルでの段階的なデノイジング」として再構成することで、実用化に成功した。

エネルギーベースモデルの復活

EBMとは

エネルギーベースモデル（Energy-Based Model, EBM） は、確率分布を以下の形で表現する：

$$p(\mathbf{x})=\frac{\exp (-E(\mathbf{x}))}{Z}$$

ここで $E(\mathbf{x})$ はエネルギー関数、 $Z=\int \exp (-E(\mathbf{x}))d\mathbf{x}$ は分配関数（正規化定数）である。

EBMの問題は、 $Z$ の計算が高次元では一般に困難なことである。

スコアマッチングによる突破

スコア関数は、 $Z$ に依存しない：

$${\mathrm{\nabla }}_{\mathbf{x}}\log p(\mathbf{x})=-{\mathrm{\nabla }}_{\mathbf{x}}E(\mathbf{x})$$

これが、スコアマッチングの威力である。 $Z$ を計算せずに、スコア（= 負のエネルギー勾配）を学習できる。

拡散モデルは、この洞察を発展させたものと見なせる。「複数のノイズレベルでスコアを学習し、それを使って段階的にサンプリングする」という戦略により、EBMの計算困難を回避している。

手法	分配関数 $Z$	サンプリング
従来のEBM	計算困難	Langevin動力学（遅い）
拡散モデル	不要（スコアのみ）	段階的デノイジング（速い）

幾何学的視点の価値

スコアは「確率の流れの方向」

スコア関数 ${\mathrm{\nabla }}_{\mathbf{x}}\log p_t(\mathbf{x})$ は、幾何学的には以下のように解釈できる：

• 確率密度の「登り坂」：スコアの方向に進むと、確率密度が高くなる
• ベクトル場としての時間発展：確率分布の時間発展を、ベクトル場の流れとして捉える
• 接ベクトル：各点での「次にどこへ向かうべきか」の指示

この視点は、確率分布を「静的な関数」としてではなく、「ダイナミクスの到達点」として理解することを可能にする。

word2vecの再評価：ベクトル演算の復権（解釈の一案）

ここで、一見無関係に思える話題を接続しよう。

word2vecの「王 - 男 + 女 = 女王」というベクトル演算は、かつては「おもちゃのような例題（toy problem）」と見なされることもあった。「たまたま上手くいく例」に過ぎないのではないか、と。

確率フローODEの視点から見ると、これをベクトル場に沿った積分という枠組みで再解釈することができる。

$$\text{king}\stackrel{-\text{male direction}+\text{female direction}}{\to }\text{queen}$$

ここでの king/male/female/queen は、それぞれ「王/男/女/女王」に対応する。

この操作は、「意味空間における接ベクトルに沿った移動」と解釈できる。

NOTE

解釈の限界： この見方は一つの解釈であり、「word2vecの線形演算が厳密に多様体上の測地線や接ベクトルの積分と同一である」という数学的証明があるわけではない。word2vecの埋め込み空間が実際にどのような幾何構造を持つかは、依然として研究対象である。ただし、「昔の直線的な演算が、実は高次元空間での意味ある操作だった」という直感を得る枠組みとしては有用である。

確率分布の時間発展

拡散モデルの本質は、「確率分布の時間発展」を制御することである。

Forward Processでは、データ分布 $p_{\text{data}}$ が徐々に広がり、最終的にガウス分布 $\mathcal{N}(\mathbf{0},\mathbf{I})$ に収束する。Reverse Processでは、この流れを逆転させ、ガウス分布からデータ分布を「凝縮」させる。

%%{init: {"themeVariables": {"fontSize": "12px"}, "flowchart": {"nodeSpacing": 18, "rankSpacing": 20}} }%%
graph TD
    subgraph F["Forward（t: 0 -> T）"]
        F0["p_data<br/>集中"] --> F1["p_1<br/>やや拡散"]
        F1 --> F2["p_2<br/>拡散"]
        F2 --> Fd["..."]
        Fd --> FT["p_T ≈ N(0, I)<br/>高ノイズ"]
    end

    subgraph R["Reverse（t: T -> 0）"]
        RT["p_T ≈ N(0, I)<br/>高ノイズ"] --> Rd["..."]
        Rd --> R2["p_2<br/>凝縮"]
        R2 --> R1["p_1<br/>さらに凝縮"]
        R1 --> R0["p_data<br/>集中"]
    end

    NoteF["Forward: 情報を壊して広げる"]:::state
    NoteR["Reverse: 情報を回復して凝縮する"]:::state
    F --> NoteF
    R --> NoteR

    classDef state fill:#F8FAFC,stroke:#475569,stroke-width:1.3px,color:#0F172A;
    class F0,F1,F2,Fd,FT,RT,Rd,R2,R1,R0,NoteF,NoteR state;
    style F fill:#F8FAFC,stroke:#64748B,stroke-width:1.2px,color:#0F172A;
    style R fill:#F8FAFC,stroke:#64748B,stroke-width:1.2px,color:#0F172A;
    linkStyle 0,1,2,3 stroke:#C96A1B,stroke-width:2px;
    linkStyle 4,5,6,7 stroke:#2E8B57,stroke-width:2px;
    linkStyle 8,9 stroke:#6B7280,stroke-width:1.4px;

$$\text{Forward: }{p}_0=p_{\text{data}}\to p_t\to p_T\approx \mathcal{N}(\mathbf{0},\mathbf{I})$$$$\text{Reverse: }{p}_T\approx \mathcal{N}(\mathbf{0},\mathbf{I})\to p_t\to p_0=p_{\text{data}}$$

この「確率分布の時間発展」という視点は、空間が ${\mathbb{R}}^d$ であれ球面であれ、本質的に同じである。空間の選択は、分布の「形」や「収束先」に影響するが、「時間発展を制御する」という基本原理は変わらない。

Flow Matching：統一的な視点

Flow Matchingとは

近年、Flow Matching（Lipman et al., 2023）という枠組みが注目を集めている。これは、拡散モデルと連続正規化フロー（Continuous Normalizing Flows）を統一的に扱う視点を提供する。

基本的なアイデアは、確率フローODEのベクトル場を直接学習することである：

$$\frac{d\mathbf{x}}{dt}={\mathbf{v}}_{\theta }(\mathbf{x},t)$$

学習は、条件付きベクトル場 ${\mathbf{u}}_t(\mathbf{x}\mid {\mathbf{x}}_1)$ （データ点 ${\mathbf{x}}_1$ に向かうベクトル場）を使って行う：

$${\mathcal{L}}_{\text{FM}}={\mathbb{E}}_{t,{\mathbf{x}}_0,{\mathbf{x}}_1}\left[{\Vert {\mathbf{v}}_{\theta }({\mathbf{x}}_t,t)-{\mathbf{u}}_t({\mathbf{x}}_t\mid {\mathbf{x}}_1)\Vert }^2\right]$$

拡散モデルとの関係

Flow Matchingは、拡散モデルの一般化と見なせる。

観点	拡散モデル	Flow Matching
Forward Process	SDEで定義	任意のパスを選択可能
学習対象	スコア（ノイズ予測）	ベクトル場
パスの形	ノイズスケジュールに依存	直線パスなど自由に設計可能

Optimal Transport（最適輸送） の視点から見ると、Flow Matchingは「ノイズ分布からデータ分布への輸送」を最も効率的なパスで行うことを目指している。

NOTE

直線パスの単純さ： Flow Matchingの典型的な選択は、 ${\mathbf{x}}_0$ （ノイズ）と ${\mathbf{x}}_1$ （データ）を直線で結ぶパスである。これは、拡散モデルの複雑なノイズスケジュールと比べて、概念的に単純である。

まとめ

概念	定義	本回での役割
Forward Process	データ→ノイズのSDE	拡散モデルの「汚す」側
Reverse Process	ノイズ→データのSDE	拡散モデルの「生成」側
スコア関数	${\mathrm{\nabla }}_{\mathbf{x}}\log p_t(\mathbf{x})$	確率の流れの方向
確率フローODE	SDEの決定論的版	高速サンプリングの基礎
Langevin動力学	スコアによるサンプリング	EBMとの接続

本回のポイント

拡散モデルは、「確率分布の時間発展」を明示的に扱う枠組みである。

• Forward Processは、データを徐々にノイズで汚し、最終的にガウス分布に収束させる。このプロセスは閉形式で表現でき、任意の時刻の状態を直接計算できる。
• Reverse Processは、Forward Processを時間反転したもので、スコア関数（確率密度の対数の勾配）を使って構成される。スコア関数は「確率が高い方向」を指し示すベクトル場であり、これに沿って進むことでデータ分布を「凝縮」させる。
• 確率フローODEは、SDEの決定論的版であり、同じ周辺分布を保ちながらノイズ項を除去したものである。DDIMやDPM-Solverなどの高速サンプラーは、この性質を利用している。
• Langevin動力学との接続は、拡散モデルが エネルギーベースモデル（EBM） の現代的な実現であることを示している。スコアマッチングにより、分配関数の計算を回避しつつスコアを学習できる。
• word2vecの「ベクトル演算」は、確率フローODEの視点から再解釈できる。「意味空間における接ベクトルに沿った移動」という枠組みで捉えることで、線形演算に新たな直感を与えることができる（ただし、これは解釈の一案であり、厳密な数学的同一性の証明ではない）。

拡散は「確率分布の時間発展」だが、空間選択（ ${\mathbb{R}}^d$ vs 球面）は別問題。本質は「流れの制御」にある。

以下は「発展的トピック」の節で述べる。

• 球面上の拡散は別の定式化であり、最終状態は球面上の一様分布である。標準手法とは異なるが、方向データや正規化された表現との親和性を持つ。
• 学習の熱力学は、生成過程（ $\mathbf{x}$ 空間）とは別に、学習過程（ $\mathit{\theta }$ 空間）を不可逆輸送として捉える視点を与える。

発展的トピック: 学習の熱力学：不可逆輸送と認識論的コスト

TIP

ここは読み飛ばしてもよい。

ここまでは、データ生成プロセス（ $\mathbf{x}$ 空間での拡散と凝縮）を熱力学的に捉えてきた。ここで視点を一段階上げ、「ニューラルネットワークが学習するプロセスそのもの」 を熱力学的に記述する最新(2026.01)の理論を紹介する。

Daisuke Okanohara, "A Thermodynamic Theory of Learning I" [arXiv:2601.17607]

これは、学習を「モデルのパラメータ分布（アンサンブル）の有限時間輸送」として定式化し、学習に伴うコストを熱力学的な「仕事」と「散逸」に対応づける野心的な枠組みである。

学習とは「パラメータ分布の輸送」である

拡散モデルの生成過程は「データ分布の変形」だったが、学習過程は 「パラメータ分布（または仮説空間）の変形」 と見なせる。

• 初期状態: 学習前。パラメータはランダムであり、エントロピー（仮説の不確実性）が大きい状態。
• 学習: データに適合するようにパラメータを更新する。分布は特定の領域に収束し、不確実性が減少する（知識を得た状態）。

この「パラメータ分布の輸送」を有限時間で行うと、必然的に理想的な準静的過程からの乖離が生じ、エントロピー生成（散逸） が発生する。この理論では、我々が最小化している損失関数（負の対数尤度など）を、単なる統計的な誤差ではなく、「分布を強制的に変形させるために必要な熱力学的仕事」 の上界として再解釈する。

認識論的コスト (Epistemic Costs)

この理論の重要な帰結は、認識論的コスト (Epistemic Costs) の概念である。

情報理論におけるランダウアーの原理（情報の消去には物理的なエネルギーコストが伴う）を学習に応用すると、モデルが外部データから情報を獲得し、内部状態（重み）を更新して不確実性を減らすこと自体に、物理的なコストの理論的な下限が存在することが示唆される。

• 学習のコスト: 重みの更新 $\approx$ 情報の書き込みと古い情報の消去（忘却）
• 不可逆性: 一度学習した（収束した）状態から、自然に元のランダムな初期状態に戻ることはない（学習における時間の矢）。

拡散モデルが「ノイズを除去して画像を凝縮させる」のと同様に、学習アルゴリズム（SGD等）は「パラメータ空間のノイズ（不確実性）を除去して知識を凝縮させる」熱機関として機能しているという物理的な解釈が可能になる（※これは情報処理の不可逆性に関する理論的な話であり、GPUの消費電力を直接表すものではない）。

IMPORTANT

生成と学習の区別:

• 拡散モデル（生成）: $\mathbf{x}$ 空間（画像空間）での分布変形。ノイズ $\to$ データ。
• 学習の熱力学: $\mathit{\theta }$ 空間（パラメータ空間）での分布変形。初期重み $\to$ 学習済み重み。

両者は数学的に似た構造（Fokker-Planck方程式など）で記述できるが、対象としている空間が異なる点に注意が必要である。

発展的トピック: 球面上の拡散モデル

標準手法との違い

ここまで述べた標準的な拡散モデルは、 ${\mathbb{R}}^d$ 上で定義され、最終状態はガウス分布である。

しかし、球面 $S^{d-1}$ 上の拡散という別の定式化も研究されている。この場合：

• 状態空間：単位球面 $S^{d-1}\subset {\mathbb{R}}^d$
• Forward Process：球面上のブラウン運動（heat kernel）
• 最終状態：球面上の一様分布（ガウス分布ではない）

定式化	状態空間	最終分布	主な用途
標準（VP-SDE等）	${\mathbb{R}}^d$	ガウス $\mathcal{N}(\mathbf{0},\mathbf{I})$	画像、音声等
球面拡散	$S^{d-1}$	球面上の一様分布	方向データ、分子構造

球面拡散の動機

球面拡散が有用な場面として、以下が挙げられる：

方向データ：風向、分子の結合角など、本質的に「方向」であるデータ。

正規化された表現：nGPT（第6回）のように、すべての表現を単位球面上に制約するアーキテクチャとの親和性。

vMF分布との接続：球面上の「ガウス分布」であるvon Mises-Fisher分布（第3回）を、拡散過程の終端として自然に扱える。

CAUTION

研究段階の注意： 球面上の拡散モデルは、標準手法ほど成熟していない。実装も複雑で、主流のアプリケーションでは依然として ${\mathbb{R}}^d$ 上の拡散が使われている。

次回予告

第10回「思考の連鎖」では、推論過程を幾何学的な軌跡として扱う。

拡散モデルでは「ノイズからデータへの軌道」を学習した。同様の視点で、「問題から解答への軌道」を考えることはできないか。Chain of Thought（CoT）は、推論の中間過程を言語として明示化する。これを「意味空間における軌跡」として解釈すると、何が見えてくるか。推論の幾何学を探求する。

実装ノート

標準的な実装

標準的な拡散モデル（DDPM, DDIM）は ${\mathbb{R}}^d$ 上で実装される。

主要なコンポーネント：

1. ノイズ予測ネットワーク ${\mathit{ϵ}}_{\theta }({\mathbf{x}}_t,t)$ ：通常はU-Netアーキテクチャ
2. ノイズスケジュール ${\beta }_t$ ：線形、コサイン、シグモイドなど
3. サンプラー：DDPM（SDE）、DDIM（ODE）、DPM-Solverなど

DDPMの学習ループ

コード例: 09_ddpm_training.py

import math

import torch
import torch.nn as nn
import torch.nn.functional as F


def linear_beta_schedule(timesteps, beta_start=1e-4, beta_end=0.02):
    """線形ノイズスケジュール

    Args:
        timesteps: 総ステップ数
        beta_start: 初期β
        beta_end: 最終β

    Returns:
        betas: [timesteps] のテンソル
    """
    return torch.linspace(beta_start, beta_end, timesteps)


def cosine_beta_schedule(timesteps, s=0.008):
    """コサインノイズスケジュール（Improved DDPM）

    Args:
        timesteps: 総ステップ数
        s: オフセット

    Returns:
        betas: [timesteps] のテンソル
    """
    steps = timesteps + 1
    x = torch.linspace(0, timesteps, steps)
    alphas_cumprod = torch.cos(((x / timesteps) + s) / (1 + s) * math.pi / 2) ** 2
    alphas_cumprod = alphas_cumprod / alphas_cumprod[0]
    betas = 1 - (alphas_cumprod[1:] / alphas_cumprod[:-1])
    return torch.clamp(betas, 0.0001, 0.9999)


class DiffusionSchedule:
    """拡散スケジュールの管理"""

    def __init__(self, timesteps=1000, schedule_type="linear"):
        self.timesteps = timesteps

        if schedule_type == "linear":
            betas = linear_beta_schedule(timesteps)
        elif schedule_type == "cosine":
            betas = cosine_beta_schedule(timesteps)
        else:
            raise ValueError(f"Unknown schedule: {schedule_type}")

        self.betas = betas
        self.alphas = 1.0 - betas
        self.alphas_cumprod = torch.cumprod(self.alphas, dim=0)
        self.sqrt_alphas_cumprod = torch.sqrt(self.alphas_cumprod)
        self.sqrt_one_minus_alphas_cumprod = torch.sqrt(1.0 - self.alphas_cumprod)

    def q_sample(self, x_0, t, noise=None):
        """Forward process: x_0 から x_t をサンプリング

        x_t = sqrt(α̅_t) * x_0 + sqrt(1 - α̅_t) * ε

        Args:
            x_0: 元データ [batch, ...]
            t: タイムステップ [batch]
            noise: ノイズ（Noneなら生成）

        Returns:
            x_t: ノイズが加わったデータ
        """
        if noise is None:
            noise = torch.randn_like(x_0)

        sqrt_alpha = self.sqrt_alphas_cumprod[t]
        sqrt_one_minus_alpha = self.sqrt_one_minus_alphas_cumprod[t]

        # 形状を合わせる
        while sqrt_alpha.dim() < x_0.dim():
            sqrt_alpha = sqrt_alpha.unsqueeze(-1)
            sqrt_one_minus_alpha = sqrt_one_minus_alpha.unsqueeze(-1)

        return sqrt_alpha * x_0 + sqrt_one_minus_alpha * noise


def ddpm_loss(model, x_0, schedule, t=None):
    """DDPMの損失関数（ノイズ予測）

    L = E_{t, x_0, ε}[||ε - ε_θ(x_t, t)||²]

    Args:
        model: ノイズ予測ネットワーク
        x_0: 元データ [batch, ...]
        schedule: DiffusionSchedule
        t: タイムステップ（Noneならランダム）

    Returns:
        loss: スカラー
    """
    batch_size = x_0.shape[0]
    device = x_0.device

    # ランダムなタイムステップ
    if t is None:
        t = torch.randint(0, schedule.timesteps, (batch_size,), device=device)

    # ノイズを生成
    noise = torch.randn_like(x_0)

    # x_t を計算
    x_t = schedule.q_sample(x_0, t, noise)

    # ノイズを予測
    predicted_noise = model(x_t, t)

    # MSE損失
    loss = F.mse_loss(predicted_noise, noise)

    return loss


# 簡単なノイズ予測ネットワーク（教育目的）
class SimpleNoisePredictor(nn.Module):
    """簡単なノイズ予測ネットワーク

    実際の実装ではU-Netを使用
    """

    def __init__(self, dim, hidden_dim=256, time_emb_dim=64):
        super().__init__()
        self.time_emb = nn.Sequential(
            nn.Linear(1, time_emb_dim),
            nn.SiLU(),
            nn.Linear(time_emb_dim, time_emb_dim),
        )

        self.net = nn.Sequential(
            nn.Linear(dim + time_emb_dim, hidden_dim),
            nn.SiLU(),
            nn.Linear(hidden_dim, hidden_dim),
            nn.SiLU(),
            nn.Linear(hidden_dim, dim),
        )

    def forward(self, x, t):
        """
        Args:
            x: 入力 [batch, dim]
            t: タイムステップ [batch]

        Returns:
            予測されたノイズ [batch, dim]
        """
        # 時刻を正規化して埋め込み
        t_emb = self.time_emb(t.float().unsqueeze(-1) / 1000.0)

        # 連結して予測
        x_t = torch.cat([x, t_emb], dim=-1)
        return self.net(x_t)


# 使用例
if __name__ == "__main__":
    dim = 64
    batch_size = 32
    timesteps = 1000

    schedule = DiffusionSchedule(timesteps, schedule_type="cosine")
    model = SimpleNoisePredictor(dim)

    # ダミーデータ
    x_0 = torch.randn(batch_size, dim)

    # 損失計算
    loss = ddpm_loss(model, x_0, schedule)
    print(f"DDPM Loss: {loss.item():.4f}")

    # Forward processの確認
    t = torch.tensor([0, 250, 500, 750, 999])
    for ti in t:
        x_t = schedule.q_sample(x_0[:1], ti.unsqueeze(0))
        print(f"t={ti.item():4d}: x_t norm = {x_t.norm().item():.4f}")

DDIMサンプラー

注意： 以下は教育目的の概念実装であり、そのまま実行すると edge case でエラーになる可能性がある。実用には公式実装（diffusers 等）を参照のこと。

実装差異について： DDIMの更新式・スケジュールの取り方は実装により差がある（例：indexの取り方、 $t$ の離散化、 $\eta$ の扱い）。本コードでは sigma_t はスカラーを想定しているが、より複雑な実装ではバッチや次元ごとに異なる場合もある。

コード例: 09_ddim_sampler.py

import torch
import torch.nn as nn


class DiffusionSchedule:
    """拡散スケジュール（簡略版）"""

    def __init__(self, timesteps=1000):
        self.timesteps = timesteps
        betas = torch.linspace(1e-4, 0.02, timesteps)
        self.alphas = 1.0 - betas
        self.alphas_cumprod = torch.cumprod(self.alphas, dim=0)


class SimpleNoisePredictor(nn.Module):
    """簡易ノイズ予測器"""

    def __init__(self, dim):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(dim + 64, 256), nn.SiLU(), nn.Linear(256, 256), nn.SiLU(), nn.Linear(256, dim)
        )
        self.time_emb = nn.Linear(1, 64)

    def forward(self, x, t):
        t_emb = self.time_emb(t.float().unsqueeze(-1) / 1000.0)
        return self.net(torch.cat([x, t_emb], dim=-1))


@torch.no_grad()
def ddim_sample(model, schedule, shape, steps=50, eta=0.0, device="cpu"):
    """DDIMサンプリング（概念実装）

    Args:
        model: ノイズ予測ネットワーク
        schedule: DiffusionSchedule
        shape: 出力形状 (batch, dim)
        steps: サンプリングステップ数
        eta: 確率性の制御（0=決定論的ODE、1=確率的SDE）
        device: デバイス

    Returns:
        x_0: 生成されたサンプル
    """
    # サンプリングする時刻のリスト
    timesteps = torch.linspace(schedule.timesteps - 1, 0, steps + 1).long()

    # 純粋なノイズから開始
    x = torch.randn(shape, device=device)

    alphas_cumprod = schedule.alphas_cumprod.to(device)

    for i in range(steps):
        t = timesteps[i].item()
        t_next = timesteps[i + 1].item()

        # 現在のα
        alpha_t = alphas_cumprod[int(t)]
        # t_next < 0 の場合は alpha = 1（完全にデノイズされた状態）
        if t_next >= 0:
            alpha_t_next = alphas_cumprod[int(t_next)]
        else:
            alpha_t_next = torch.tensor(1.0, device=device, dtype=alpha_t.dtype)

        # ノイズを予測
        t_tensor = torch.full((shape[0],), t, device=device, dtype=torch.long)
        eps_pred = model(x, t_tensor)

        # x_0 を予測
        x0_pred = (x - torch.sqrt(1 - alpha_t) * eps_pred) / torch.sqrt(alpha_t)

        # ノイズの分散を計算（eta=0で決定論的）
        sigma_t = (
            eta
            * torch.sqrt((1 - alpha_t_next) / (1 - alpha_t + 1e-8))
            * torch.sqrt(1 - alpha_t / (alpha_t_next + 1e-8))
        )

        # x_{t-1} を計算
        dir_xt = torch.sqrt(torch.clamp(1 - alpha_t_next - sigma_t**2, min=0)) * eps_pred

        # sigma_t が実質ゼロかどうかを float で判定
        sigma_val = sigma_t.item() if sigma_t.numel() == 1 else sigma_t.mean().item()
        if sigma_val > 1e-8:
            noise = torch.randn_like(x)
        else:
            noise = torch.zeros_like(x)

        x = torch.sqrt(alpha_t_next) * x0_pred + dir_xt + sigma_t * noise

    return x


# スコア関数としての解釈
def noise_to_score(eps_pred, sqrt_one_minus_alpha):
    """ノイズ予測からスコアへの変換

    ∇_x log p_t(x) ≈ -ε / sqrt(1 - α̅_t)

    Args:
        eps_pred: 予測されたノイズ
        sqrt_one_minus_alpha: sqrt(1 - α̅_t)

    Returns:
        score: スコア関数の近似
    """
    return -eps_pred / sqrt_one_minus_alpha


# 使用例
if __name__ == "__main__":
    dim = 64
    batch_size = 8
    timesteps = 1000

    schedule = DiffusionSchedule(timesteps)
    model = SimpleNoisePredictor(dim)

    # DDIMサンプリング（決定論的、eta=0）
    samples_deterministic = ddim_sample(model, schedule, (batch_size, dim), steps=50, eta=0.0)

    # DDIMサンプリング（確率的、eta=1）
    samples_stochastic = ddim_sample(model, schedule, (batch_size, dim), steps=50, eta=1.0)

    print(f"Deterministic samples shape: {samples_deterministic.shape}")
    print(f"Deterministic samples norm: {samples_deterministic.norm(dim=-1).mean():.4f}")

    print(f"Stochastic samples shape: {samples_stochastic.shape}")
    print(f"Stochastic samples norm: {samples_stochastic.norm(dim=-1).mean():.4f}")

スコアマッチングの可視化

コード例: 09_score_visualization.py

import matplotlib.pyplot as plt
import numpy as np
import torch


def gaussian_score(x, mean=0.0, std=1.0):
    """ガウス分布のスコア（解析解）

    p(x) = N(mean, std²)
    ∇_x log p(x) = -(x - mean) / std²

    Args:
        x: 入力
        mean: 平均
        std: 標準偏差

    Returns:
        score: スコア関数の値
    """
    return -(x - mean) / (std**2)


def mixture_score(x, means, stds, weights):
    """混合ガウス分布のスコア

    注意: log_prob の計算では正規化定数（2π等）を一部省略している。
    スコア（勾配）自体は定数項の影響を受けないため問題ないが、
    表示される log_prob の絶対値は厳密な対数確率密度ではなく「比例」した値である。

    Args:
        x: 入力 [batch, dim]
        means: 各成分の平均 [K, dim]
        stds: 各成分の標準偏差 [K]
        weights: 混合係数 [K]

    Returns:
        score: スコア関数の値 [batch, dim]
    """
    K = len(means)

    # 各成分の確率密度
    log_probs = []
    for k in range(K):
        diff = x - means[k]
        log_prob = -0.5 * (diff**2).sum(dim=-1) / (stds[k] ** 2)
        log_prob -= x.shape[-1] * np.log(stds[k])
        log_prob += np.log(weights[k])
        log_probs.append(log_prob)

    log_probs = torch.stack(log_probs, dim=-1)  # [batch, K]

    # Softmax重み
    probs = torch.softmax(log_probs, dim=-1)  # [batch, K]

    # 各成分のスコアを重み付き平均
    score = torch.zeros_like(x)
    for k in range(K):
        score += probs[:, k : k + 1] * (-(x - means[k]) / (stds[k] ** 2))

    return score


def visualize_score_field_2d():
    """2次元でのスコア場の可視化"""
    # 2成分混合ガウス
    means = torch.tensor([[-2.0, 0.0], [2.0, 0.0]])
    stds = torch.tensor([0.8, 0.8])
    weights = torch.tensor([0.5, 0.5])

    # グリッドを作成
    x_range = torch.linspace(-5, 5, 20)
    y_range = torch.linspace(-3, 3, 15)
    X, Y = torch.meshgrid(x_range, y_range, indexing="xy")

    points = torch.stack([X.flatten(), Y.flatten()], dim=-1)

    # スコアを計算
    scores = mixture_score(points, means, stds, weights)
    U = scores[:, 0].reshape(X.shape)
    V = scores[:, 1].reshape(X.shape)

    # 確率密度も計算（可視化用）
    def mixture_density(x, means, stds, weights):
        density = torch.zeros(x.shape[0])
        for k in range(len(means)):
            diff = x - means[k]
            density += (
                weights[k] * torch.exp(-0.5 * (diff**2).sum(dim=-1) / (stds[k] ** 2)) / (2 * np.pi * stds[k] ** 2)
            )
        return density

    Z = mixture_density(points, means, stds, weights).reshape(X.shape)

    # プロット
    fig, axes = plt.subplots(1, 2, figsize=(14, 5))

    # 左: 確率密度
    ax1 = axes[0]
    contour = ax1.contourf(X.numpy(), Y.numpy(), Z.numpy(), levels=20, cmap="viridis")
    plt.colorbar(contour, ax=ax1, label="p(x)")
    ax1.scatter(means[:, 0], means[:, 1], c="red", s=100, marker="x", label="Modes")
    ax1.set_xlabel("x₁")
    ax1.set_ylabel("x₂")
    ax1.set_title("Probability Density p(x)")
    ax1.legend()

    # 右: スコア場（ベクトル場）
    ax2 = axes[1]
    ax2.contour(X.numpy(), Y.numpy(), Z.numpy(), levels=10, colors="gray", alpha=0.5)
    ax2.quiver(
        X.numpy(),
        Y.numpy(),
        U.numpy(),
        V.numpy(),
        color="blue",
        alpha=0.7,
        scale=50,
    )
    ax2.scatter(means[:, 0], means[:, 1], c="red", s=100, marker="x", label="Modes")
    ax2.set_xlabel("x₁")
    ax2.set_ylabel("x₂")
    ax2.set_title("Score Field ∇ log p(x)")
    ax2.legend()

    plt.tight_layout()
    plt.savefig("score_field_2d.png", dpi=150)
    plt.close()

    print("Saved: score_field_2d.png")


def visualize_langevin_trajectory():
    """Langevin動力学の軌跡を可視化"""
    # 2成分混合ガウス
    means = torch.tensor([[-2.0, 0.0], [2.0, 0.0]])
    stds = torch.tensor([0.8, 0.8])
    weights = torch.tensor([0.5, 0.5])

    # Langevinサンプリング
    def langevin_sample(x_init, score_fn, steps=1000, step_size=0.01):
        x = x_init.clone()
        trajectory = [x.clone()]

        for _ in range(steps):
            score = score_fn(x)
            noise = torch.randn_like(x)
            x = x + step_size * score + np.sqrt(2 * step_size) * noise
            trajectory.append(x.clone())

        return torch.stack(trajectory)

    # 複数の軌跡をサンプリング
    n_samples = 5
    x_init = torch.randn(n_samples, 2) * 3

    def score_fn(x):
        return mixture_score(x, means, stds, weights)

    trajectories = langevin_sample(x_init, score_fn, steps=500, step_size=0.05)

    # プロット
    fig, ax = plt.subplots(figsize=(8, 6))

    # 密度の等高線
    x_range = torch.linspace(-5, 5, 50)
    y_range = torch.linspace(-4, 4, 40)
    X, Y = torch.meshgrid(x_range, y_range, indexing="xy")
    points = torch.stack([X.flatten(), Y.flatten()], dim=-1)

    def mixture_density(x, means, stds, weights):
        density = torch.zeros(x.shape[0])
        for k in range(len(means)):
            diff = x - means[k]
            density += (
                weights[k] * torch.exp(-0.5 * (diff**2).sum(dim=-1) / (stds[k] ** 2)) / (2 * np.pi * stds[k] ** 2)
            )
        return density

    Z = mixture_density(points, means, stds, weights).reshape(X.shape)
    ax.contour(X.numpy(), Y.numpy(), Z.numpy(), levels=10, colors="gray", alpha=0.5)

    # 軌跡をプロット
    colors = plt.cm.tab10(np.linspace(0, 1, n_samples))
    for i in range(n_samples):
        traj = trajectories[:, i].numpy()
        ax.plot(traj[:, 0], traj[:, 1], color=colors[i], alpha=0.7, linewidth=0.5)
        ax.scatter(traj[0, 0], traj[0, 1], color=colors[i], marker="o", s=50, label=f"Start {i + 1}")
        ax.scatter(traj[-1, 0], traj[-1, 1], color=colors[i], marker="x", s=50)

    ax.scatter(means[:, 0], means[:, 1], c="red", s=200, marker="*", zorder=5, label="Modes")
    ax.set_xlabel("x₁")
    ax.set_ylabel("x₂")
    ax.set_title("Langevin Dynamics Trajectories")
    ax.legend(loc="upper right", fontsize=8)

    plt.tight_layout()
    plt.savefig("langevin_trajectory.png", dpi=150)
    plt.close()

    print("Saved: langevin_trajectory.png")


# 実行
if __name__ == "__main__":
    visualize_score_field_2d()
    visualize_langevin_trajectory()

球面拡散の概念実装

コード例: 09_spherical_diffusion.py

"""球面上の拡散モデル（概念的なデモ）

!!!!! 重要な注意 !!!!!
このコードは教育目的の「直感的なデモ」であり、
**真の球面ブラウン運動の正しい離散化ではない**。

問題点:
1. 「接空間でノイズ→射影」は、真の球面上SDE離散化と一致しない
2. この近似は統計的バイアスを生む（特に大きな dt や高次元で顕著）
3. 正しい実装には、指数写像（exponential map）や測地線に沿った移動、
   または heat kernel の厳密な計算が必要

位置づけ:
- 射影ベースの近似も、非常に小さい dt かつ低次元であれば、
  「球面上でも拡散的な現象が起きる」という直感を得る目的には使える
- ただし、理論的な整合や定量的な精度は保証されない

研究・実用には:
- 専門文献（Riemannian Score-based Generative Models 等）を参照
- 専用ライブラリ（geomstats, geoopt 等）の使用を検討すること
"""

import matplotlib.pyplot as plt
import numpy as np
import torch
import torch.nn.functional as F


def project_to_sphere(x):
    """ベクトルを単位球面に射影

    Args:
        x: 入力ベクトル [batch, dim]

    Returns:
        x_proj: 単位球面上のベクトル [batch, dim]
    """
    return F.normalize(x, dim=-1)


def spherical_brownian_step(x, dt, temperature=1.0):
    """球面上のブラウン運動の1ステップ（非常に粗い近似）

    接空間でガウスノイズを加え、球面に射影し直す。
    ※ これは真の球面ブラウン運動の離散化ではなく、
       直感的なデモのための簡易実装である。

    Args:
        x: 現在の位置（単位球面上） [batch, dim]
        dt: 時間刻み
        temperature: ノイズの強度

    Returns:
        x_new: 次の位置（単位球面上）
    """
    # 接空間でのノイズ（xに直交する成分のみ）
    noise = torch.randn_like(x)
    # xに平行な成分を除去
    noise = noise - (noise * x).sum(dim=-1, keepdim=True) * x

    # 移動して射影
    x_new = x + np.sqrt(2 * temperature * dt) * noise
    return project_to_sphere(x_new)


def spherical_forward_process(x_0, timesteps, dt=0.01):
    """球面上のForward Process

    Args:
        x_0: 初期データ（単位球面上） [batch, dim]
        timesteps: ステップ数
        dt: 時間刻み

    Returns:
        trajectory: 軌跡 [timesteps+1, batch, dim]
    """
    x = x_0.clone()
    trajectory = [x.clone()]

    for _ in range(timesteps):
        x = spherical_brownian_step(x, dt)
        trajectory.append(x.clone())

    return torch.stack(trajectory)


def estimate_vMF_concentration(x_samples):
    """サンプルからvMF分布の集中度を推定（粗い近似）

    注意: この推定式は近似であり、次元・サンプル数・集中度によって
    精度が大きく揺れる。厳密な推定には最尤推定や Bessel 関数の逆関数が必要。

    Args:
        x_samples: サンプル [batch, dim]

    Returns:
        kappa: 推定された集中度（参考値）
        mean_dir: 推定された平均方向
    """
    mean_dir = x_samples.mean(dim=0)
    R = mean_dir.norm()
    mean_dir = F.normalize(mean_dir, dim=0)

    # 近似式（高次元での近似、精度は限定的）
    dim = x_samples.shape[-1]
    kappa = R * (dim - R**2) / (1 - R**2 + 1e-8)

    return kappa.item(), mean_dir


def visualize_spherical_diffusion_3d():
    """3次元球面上の拡散を可視化"""
    # 初期分布：北極付近に集中
    n_samples = 100
    kappa_init = 50  # 高い集中度

    # vMF分布からサンプリング（近似）
    mean_dir = torch.tensor([0.0, 0.0, 1.0])
    noise = torch.randn(n_samples, 3)
    x_0 = F.normalize(mean_dir + noise / np.sqrt(kappa_init), dim=-1)

    # Forward Process
    trajectory = spherical_forward_process(x_0, timesteps=200, dt=0.05)

    # 3Dプロット
    fig = plt.figure(figsize=(15, 5))

    # 球面を描画
    u = np.linspace(0, 2 * np.pi, 30)
    v = np.linspace(0, np.pi, 20)
    sphere_x = np.outer(np.cos(u), np.sin(v))
    sphere_y = np.outer(np.sin(u), np.sin(v))
    sphere_z = np.outer(np.ones(np.size(u)), np.cos(v))

    timesteps_to_show = [0, 50, 100, 200]
    titles = ["t=0 (Concentrated)", "t=50", "t=100", "t=200 (Diffused)"]

    for idx, (t, title) in enumerate(zip(timesteps_to_show, titles, strict=True)):
        ax = fig.add_subplot(1, 4, idx + 1, projection="3d")

        # 球面
        ax.plot_surface(sphere_x, sphere_y, sphere_z, alpha=0.1, color="gray")

        # サンプル点
        points = trajectory[t].numpy()
        ax.scatter(points[:, 0], points[:, 1], points[:, 2], c="blue", s=10, alpha=0.6)

        # 集中度を推定
        kappa, _ = estimate_vMF_concentration(trajectory[t])

        ax.set_title(f"{title}\nκ≈{kappa:.1f}")
        ax.set_xlim([-1.2, 1.2])
        ax.set_ylim([-1.2, 1.2])
        ax.set_zlim([-1.2, 1.2])
        ax.set_xlabel("x")
        ax.set_ylabel("y")
        ax.set_zlabel("z")

    plt.tight_layout()
    plt.savefig("spherical_diffusion_3d.png", dpi=150)
    plt.close()

    print("Saved: spherical_diffusion_3d.png")


def compare_euclidean_vs_spherical():
    """ユークリッド空間と球面の拡散の比較"""
    n_samples = 500
    dim = 3
    timesteps = 100

    # 初期分布（同じ点から開始）
    x_0 = torch.zeros(n_samples, dim)
    x_0[:, 2] = 1.0  # 北極

    # ユークリッド空間での拡散
    def euclidean_diffusion(x_0, timesteps, dt=0.1):
        x = x_0.clone()
        trajectory = [x.clone()]
        for _ in range(timesteps):
            x = x + np.sqrt(2 * dt) * torch.randn_like(x)
            trajectory.append(x.clone())
        return torch.stack(trajectory)

    traj_euclidean = euclidean_diffusion(x_0.clone(), timesteps)

    # 球面での拡散
    x_0_sphere = F.normalize(x_0, dim=-1)
    traj_spherical = spherical_forward_process(x_0_sphere, timesteps, dt=0.1)

    # 統計量の比較
    print("=" * 60)
    print("Euclidean vs Spherical Diffusion Comparison")
    print("=" * 60)

    for t in [0, 25, 50, 100]:
        euc_norm = traj_euclidean[t].norm(dim=-1)
        sph_norm = traj_spherical[t].norm(dim=-1)

        print(f"\nt={t}:")
        print(f"  Euclidean: norm mean={euc_norm.mean():.3f}, std={euc_norm.std():.3f}")
        print(f"  Spherical: norm mean={sph_norm.mean():.3f}, std={sph_norm.std():.3f}")

        kappa, _ = estimate_vMF_concentration(traj_spherical[t])
        print(f"  Spherical κ (concentration): {kappa:.1f}")


# 実行
if __name__ == "__main__":
    visualize_spherical_diffusion_3d()
    compare_euclidean_vs_spherical()

参考文献

拡散モデルの基礎

• Ho, J., Jain, A., & Abbeel, P. (2020). Denoising Diffusion Probabilistic Models. NeurIPS 2020. arXiv: arXiv:2006.11239.
  • DDPMの原論文。現代の拡散モデルの基礎。
• Sohl-Dickstein, J., Weiss, E., Maheswaranathan, N., & Ganguli, S. (2015). Deep Unsupervised Learning using Nonequilibrium Thermodynamics. ICML 2015. arXiv: arXiv:1503.03585.
  • 拡散モデルの初期の定式化。非平衡熱力学との接続。

スコアベース生成モデル

• Song, Y., & Ermon, S. (2019). Generative Modeling by Estimating Gradients of the Data Distribution. NeurIPS 2019. arXiv: arXiv:1907.05600.
  • スコアマッチングに基づく生成モデル。
• Song, Y., Sohl-Dickstein, J., Kingma, D. P., Kumar, A., Ermon, S., & Poole, B. (2021). Score-Based Generative Modeling through Stochastic Differential Equations. ICLR 2021. arXiv: arXiv:2011.13456.
  • SDEに基づく統一的な定式化。VP-SDE、VE-SDE、確率フローODEを導入。

高速サンプリング

• Song, J., Meng, C., & Ermon, S. (2021). Denoising Diffusion Implicit Models. ICLR 2021. arXiv: arXiv:2010.02502.
  • DDIMの原論文。確率フローODEに基づく高速サンプリング。
• Lu, C., Zhou, Y., Bao, F., Chen, J., Li, C., & Zhu, J. (2022). DPM-Solver: A Fast ODE Solver for Diffusion Probabilistic Model Sampling in Around 10 Steps. NeurIPS 2022. arXiv: arXiv:2206.00927.
  • 高速ODEソルバー。10〜20ステップでの高品質サンプリング。

Flow Matching

• Lipman, Y., Chen, R. T. Q., Ben-Hamu, H., Nickel, M., & Le, M. (2023). Flow Matching for Generative Modeling. ICLR 2023. arXiv: arXiv:2210.02747.
  • Flow Matchingの原論文。拡散モデルと連続正規化フローの統一。

スコアマッチング

• Hyvärinen, A. (2005). Estimation of Non-Normalized Statistical Models by Score Matching. JMLR, 6, 695–709.
  • スコアマッチングの原論文。

エネルギーベースモデル

• LeCun, Y., Chopra, S., Hadsell, R., Ranzato, M., & Huang, F. J. (2006). A Tutorial on Energy-Based Learning. In Predicting Structured Data. MIT Press.
  • EBMのチュートリアル。

Langevin動力学

• Welling, M., & Teh, Y. W. (2011). Bayesian Learning via Stochastic Gradient Langevin Dynamics. ICML 2011.
  • 確率的勾配Langevin動力学（SGLD）。ベイズ学習への応用。