正規化フロー(Normalizing Flow)の理論と実装を解説

正規化フロー(Normalizing Flow)は、単純な確率分布(例: 標準正規分布)を可逆な変換の連鎖によって複雑なデータ分布に変換する生成モデルです。VAE や GAN とは異なり、正規化フローはデータの正確な対数尤度を解析的に計算でき、かつ可逆性により潜在変数とデータの間で双方向のマッピングが可能です。

正規化フローの数学的な美しさは、確率分布の変数変換公式とヤコビアン行列式という線形代数の基本概念から構築される点にあります。本記事では、変数変換公式の導出からフローの連鎖、計算効率のためのヤコビアンの構造的制約、代表的なフロー(Planar Flow, RealNVP, GLOW)の設計、そして連続時間フロー(Neural ODE / CNF)との接続までを省略なく解説し、Python で RealNVP を実装します。

本記事の内容

  • 正規化フローの基本原理(可逆変換による確率分布の変換)
  • 変数変換公式 $p_X(\bm{x}) = p_Z(\bm{f}^{-1}(\bm{x}))\,|\det(\partial \bm{f}^{-1}/\partial \bm{x})|$ の導出
  • フローの連鎖(変換の合成)と対数尤度の計算
  • 計算効率のためのヤコビアンの制約(三角行列 $\to$ $O(D)$ 計算)
  • Planar Flow の定式化
  • RealNVP: アフィンカップリング層とそのヤコビアンの導出
  • GLOW の概要(1×1 可逆畳み込み)
  • 連続時間フロー(Neural ODE / CNF)との接続
  • Python で RealNVP の実装(2D分布での密度推定・サンプリング)

前提知識

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

正規化フローの基本原理

基本アイデア

正規化フローのアイデアは次の通りです。

  1. 潜在変数 $\bm{z}$ を単純な分布 $p_Z(\bm{z})$(例: 標準正規分布 $\mathcal{N}(\bm{0}, \bm{I})$)からサンプリング
  2. 可逆変換 $\bm{f}: \mathbb{R}^d \to \mathbb{R}^d$ を適用して $\bm{x} = \bm{f}(\bm{z})$ を得る
  3. $\bm{x}$ の分布 $p_X(\bm{x})$ が複雑なデータ分布に一致するように $\bm{f}$ を学習する

可逆変換を用いることで、$p_X(\bm{x})$ を変数変換公式により解析的に計算できます。

なぜ「正規化フロー」と呼ぶのか

「正規化(Normalizing)」は、変換後の分布が正規化条件(確率の総和が1)を満たすことを指します。「フロー(Flow)」は、可逆変換の連鎖を「流れ」に見立てたものです。単純な分布が変換の流れに沿って複雑な分布へと変化していくイメージです。

変数変換公式の導出

1次元の場合

まず1次元で直感を得ましょう。確率変数 $Z$ が密度 $p_Z(z)$ に従い、$x = f(z)$ ($f$ は単調増加で微分可能)とします。

$z$ と $z + dz$ の間の確率質量は、$x$ と $x + dx$ の間の確率質量に等しいはずです。

$$ p_Z(z)\,|dz| = p_X(x)\,|dx| $$

$$ p_X(x) = p_Z(z)\left|\frac{dz}{dx}\right| = p_Z(f^{-1}(x))\left|\frac{df^{-1}(x)}{dx}\right| $$

多次元への一般化

$\bm{z} \in \mathbb{R}^d$ が密度 $p_Z(\bm{z})$ に従い、$\bm{x} = \bm{f}(\bm{z})$ ($\bm{f}$ は可逆で微分可能な写像)とします。

微小体積要素の変換を考えます。$\bm{z}$ の近傍の微小体積 $d\bm{z}$ は、ヤコビ行列

$$ \bm{J}_f = \frac{\partial \bm{f}}{\partial \bm{z}} = \begin{pmatrix} \frac{\partial f_1}{\partial z_1} & \cdots & \frac{\partial f_1}{\partial z_d} \\ \vdots & \ddots & \vdots \\ \frac{\partial f_d}{\partial z_1} & \cdots & \frac{\partial f_d}{\partial z_d} \end{pmatrix} $$

によって $d\bm{x} = |\det \bm{J}_f|\,d\bm{z}$ と変換されます。確率質量の保存より

$$ p_X(\bm{x})\,|d\bm{x}| = p_Z(\bm{z})\,|d\bm{z}| $$

$$ p_X(\bm{x})\,|\det \bm{J}_f|\,|d\bm{z}| = p_Z(\bm{z})\,|d\bm{z}| $$

$$ \boxed{p_X(\bm{x}) = p_Z(\bm{f}^{-1}(\bm{x}))\,\left|\det\left(\frac{\partial \bm{f}^{-1}}{\partial \bm{x}}\right)\right| = p_Z(\bm{z})\,\left|\det\left(\frac{\partial \bm{f}}{\partial \bm{z}}\right)\right|^{-1}} $$

対数を取ると

$$ \log p_X(\bm{x}) = \log p_Z(\bm{z}) – \log\left|\det\left(\frac{\partial \bm{f}}{\partial \bm{z}}\right)\right| $$

ここで $\bm{z} = \bm{f}^{-1}(\bm{x})$ です。第2項はヤコビアンの行列式の対数であり、変換による確率密度の体積変化を補正する役割を持ちます。

フローの連鎖と対数尤度

変換の合成

1つの変換では表現力が不足する場合、$K$ 個の可逆変換を連鎖させます。

$$ \bm{x} = \bm{f}_K \circ \bm{f}_{K-1} \circ \cdots \circ \bm{f}_1(\bm{z}_0), \quad \bm{z}_0 \sim p_Z(\bm{z}_0) $$

中間変数を $\bm{z}_k = \bm{f}_k(\bm{z}_{k-1})$($k = 1, \dots, K$)とおくと、$\bm{z}_K = \bm{x}$ です。

合成関数のヤコビアンは連鎖律により

$$ \frac{\partial \bm{x}}{\partial \bm{z}_0} = \prod_{k=1}^{K} \frac{\partial \bm{f}_k}{\partial \bm{z}_{k-1}} $$

行列式の乗法性 $\det(\bm{AB}) = \det(\bm{A})\det(\bm{B})$ より

$$ \det\left(\frac{\partial \bm{x}}{\partial \bm{z}_0}\right) = \prod_{k=1}^{K} \det\left(\frac{\partial \bm{f}_k}{\partial \bm{z}_{k-1}}\right) $$

したがって、対数尤度は

$$ \boxed{\log p_X(\bm{x}) = \log p_Z(\bm{z}_0) – \sum_{k=1}^{K} \log\left|\det\left(\frac{\partial \bm{f}_k}{\partial \bm{z}_{k-1}}\right)\right|} $$

学習

学習は、データの負の対数尤度を最小化する最尤推定です。

$$ \mathcal{L}(\theta) = -\frac{1}{N}\sum_{n=1}^{N} \log p_X(\bm{x}^{(n)}) = -\frac{1}{N}\sum_{n=1}^{N}\left[\log p_Z(\bm{z}_0^{(n)}) – \sum_{k=1}^{K}\log\left|\det\bm{J}_k^{(n)}\right|\right] $$

ここで $\bm{z}_0^{(n)} = \bm{f}_1^{-1} \circ \cdots \circ \bm{f}_K^{-1}(\bm{x}^{(n)})$ です。

計算効率のためのヤコビアンの制約

ヤコビアンの行列式の計算量

一般の $d \times d$ 行列の行列式の計算は $O(d^3)$ であり、高次元データ(例: 画像の $d = 256 \times 256 \times 3$)では計算が困難です。

しかし、ヤコビ行列が三角行列(上三角または下三角)であれば、行列式は対角要素の積になります。

$$ \det(\bm{J}) = \prod_{i=1}^{d} J_{ii} $$

この場合の計算量は $O(d)$ です。フローの設計において、この三角構造を持つヤコビアンを実現することが重要な設計原理です。

Planar Flow

定式化

Rezende & Mohamed (2015) が提案した Planar Flow は、以下の形の変換です。

$$ \bm{f}(\bm{z}) = \bm{z} + \bm{u}\,h(\bm{w}^T \bm{z} + b) $$

ここで $\bm{u}, \bm{w} \in \mathbb{R}^d$、$b \in \mathbb{R}$ はパラメータ、$h$ は活性化関数(例: $\tanh$)です。

ヤコビアンの計算

ヤコビ行列は

$$ \frac{\partial \bm{f}}{\partial \bm{z}} = \bm{I} + \bm{u}\, h'(\bm{w}^T\bm{z} + b)\,\bm{w}^T $$

これはランク1更新の形($\bm{I} + \bm{u}\bm{v}^T$、ただし $\bm{v} = h'(\bm{w}^T\bm{z}+b)\bm{w}$)です。行列式の公式(Matrix Determinant Lemma)より

$$ \det(\bm{I} + \bm{u}\bm{v}^T) = 1 + \bm{v}^T\bm{u} $$

したがって

$$ \det\left(\frac{\partial \bm{f}}{\partial \bm{z}}\right) = 1 + h'(\bm{w}^T\bm{z} + b)\,\bm{w}^T\bm{u} $$

計算量は $O(d)$ です。ただし、可逆性を保証するために $1 + h’\bm{w}^T\bm{u} > 0$ の条件が必要です。$h = \tanh$ のとき $h’ \in (0, 1]$ なので、$\bm{w}^T\bm{u} \geq -1$ であれば十分です。

Planar Flow は表現力が限定的(1層でヤコビアンがランク1の補正のみ)であるため、実用的には多数の層を重ねる必要があります。

RealNVP: アフィンカップリング層

基本アイデア

Dinh et al. (2017) が提案した RealNVP(Real-valued Non-Volume Preserving)は、アフィンカップリング層(Affine Coupling Layer)を用いた効率的かつ表現力の高いフローです。

入力 $\bm{z} \in \mathbb{R}^d$ を2つのグループに分割します。

$$ \bm{z} = (\bm{z}_{1:d’}, \bm{z}_{d’+1:d}) $$

ここで $d’ < d$ です(典型的に $d' = d/2$)。

アフィンカップリング層の変換は

$$ \begin{cases} \bm{y}_{1:d’} = \bm{z}_{1:d’} \\ \bm{y}_{d’+1:d} = \bm{z}_{d’+1:d} \odot \exp(\bm{s}(\bm{z}_{1:d’})) + \bm{t}(\bm{z}_{1:d’}) \end{cases} $$

ここで $\bm{s}$ と $\bm{t}$ はスケール関数と平行移動関数で、任意のニューラルネットワークです。$\odot$ は要素ごとの積を表します。

重要な点: $\bm{s}$ と $\bm{t}$ は可逆である必要がありません。変換全体の可逆性は構造的に保証されます。

逆変換

逆変換は解析的に計算できます。

$$ \begin{cases} \bm{z}_{1:d’} = \bm{y}_{1:d’} \\ \bm{z}_{d’+1:d} = (\bm{y}_{d’+1:d} – \bm{t}(\bm{y}_{1:d’})) \odot \exp(-\bm{s}(\bm{y}_{1:d’})) \end{cases} $$

順方向と同じ計算量で逆変換ができます。

ヤコビアンの導出

アフィンカップリング層のヤコビ行列を計算しましょう。$\bm{y} = (\bm{y}_{1:d’}, \bm{y}_{d’+1:d})$、$\bm{z} = (\bm{z}_{1:d’}, \bm{z}_{d’+1:d})$ として

$$ \frac{\partial \bm{y}}{\partial \bm{z}} = \begin{pmatrix} \frac{\partial \bm{y}_{1:d’}}{\partial \bm{z}_{1:d’}} & \frac{\partial \bm{y}_{1:d’}}{\partial \bm{z}_{d’+1:d}} \\ \frac{\partial \bm{y}_{d’+1:d}}{\partial \bm{z}_{1:d’}} & \frac{\partial \bm{y}_{d’+1:d}}{\partial \bm{z}_{d’+1:d}} \end{pmatrix} $$

各ブロックを計算します。

  • $\frac{\partial \bm{y}_{1:d’}}{\partial \bm{z}_{1:d’}} = \bm{I}_{d’}$(恒等変換)
  • $\frac{\partial \bm{y}_{1:d’}}{\partial \bm{z}_{d’+1:d}} = \bm{0}$($\bm{y}_{1:d’}$ は $\bm{z}_{d’+1:d}$ に依存しない)
  • $\frac{\partial \bm{y}_{d’+1:d}}{\partial \bm{z}_{d’+1:d}} = \text{diag}(\exp(\bm{s}(\bm{z}_{1:d’})))$(対角行列)
  • $\frac{\partial \bm{y}_{d’+1:d}}{\partial \bm{z}_{1:d’}}$ は一般に稠密行列

したがってヤコビ行列は下三角ブロック行列です。

$$ \frac{\partial \bm{y}}{\partial \bm{z}} = \begin{pmatrix} \bm{I}_{d’} & \bm{0} \\ \frac{\partial \bm{y}_{d’+1:d}}{\partial \bm{z}_{1:d’}} & \text{diag}(\exp(\bm{s}(\bm{z}_{1:d’}))) \end{pmatrix} $$

三角ブロック行列の行列式は対角ブロックの行列式の積です。

$$ \det\left(\frac{\partial \bm{y}}{\partial \bm{z}}\right) = \det(\bm{I}_{d’}) \cdot \det(\text{diag}(\exp(\bm{s}(\bm{z}_{1:d’})))) = \prod_{j=d’+1}^{d} \exp(s_j(\bm{z}_{1:d’})) $$

対数行列式は

$$ \boxed{\log\left|\det\left(\frac{\partial \bm{y}}{\partial \bm{z}}\right)\right| = \sum_{j=d’+1}^{d} s_j(\bm{z}_{1:d’})} $$

これは $\bm{s}$ ネットワークの出力の単純な和であり、$O(d)$ で計算できます。

交互マスキング

1つのカップリング層では $\bm{z}_{1:d’}$ が恒等変換されるため、$\bm{z}$ の全次元を変換するには、層ごとに分割のパターンを入れ替えます(交互マスキング)。奇数層では前半を固定、偶数層では後半を固定とします。

GLOW の概要

Kingma & Dhariwal (2018) が提案した GLOW は、RealNVP を拡張した画像生成のためのフローモデルです。主な改良点は以下の通りです。

1×1 可逆畳み込み: チャネルの置換を固定ではなく、学習可能な $c \times c$ の可逆行列($c$ はチャネル数)で行います。

$$ \bm{y}_{h,w} = \bm{W} \bm{z}_{h,w} $$

ここで $\bm{W} \in \mathbb{R}^{c \times c}$ は可逆行列です。対数行列式は

$$ \log|\det \bm{J}| = h \cdot w \cdot \log|\det \bm{W}| $$

Actnorm: バッチ正規化の代わりに、アフィン変換 $\bm{y} = \bm{s} \odot \bm{x} + \bm{b}$ をデータ依存で初期化する正規化層です。

連続時間フロー(Neural ODE / CNF)

Neural ODE

Chen et al. (2018) は、離散的なフローの層数を無限大にした極限として、常微分方程式(ODE)で変換を定義する Neural ODE を提案しました。

$$ \frac{d\bm{z}(t)}{dt} = \bm{v}_\theta(\bm{z}(t), t) $$

ここで $\bm{v}_\theta$ はニューラルネットワークで定義されるベクトル場、$t \in [0, 1]$ は連続的な「時刻」です。$\bm{z}(0) = \bm{z}_0$(潜在変数)から $\bm{z}(1) = \bm{x}$(データ)への写像がフローです。

連続的な変数変換公式

連続時間フローにおける対数密度の変化はインスタンタネオス変数変換公式で表されます。

$$ \frac{d \log p(\bm{z}(t))}{dt} = -\text{tr}\left(\frac{\partial \bm{v}_\theta}{\partial \bm{z}}\right) $$

これは $d \times d$ のヤコビアンの行列式ではなくトレース(ダイバージェンス)のみで表されるため、Hutchinson のトレース推定量を使って

$$ \text{tr}\left(\frac{\partial \bm{v}_\theta}{\partial \bm{z}}\right) \approx \bm{\epsilon}^T \frac{\partial \bm{v}_\theta}{\partial \bm{z}} \bm{\epsilon}, \quad \bm{\epsilon} \sim \mathcal{N}(\bm{0}, \bm{I}) $$

と効率的に近似できます。これにより、$O(d)$ の計算量で対数密度変化を追跡できます。

導出

インスタンタネオス変数変換公式の導出を示します。離散フロー $\bm{z}_{k+1} = \bm{z}_k + \Delta t \cdot \bm{v}_\theta(\bm{z}_k, t_k)$ の対数行列式は

$$ \log\left|\det\left(\frac{\partial \bm{z}_{k+1}}{\partial \bm{z}_k}\right)\right| = \log\left|\det\left(\bm{I} + \Delta t \frac{\partial \bm{v}_\theta}{\partial \bm{z}_k}\right)\right| $$

$\Delta t \to 0$ の極限で、$\det(\bm{I} + \epsilon \bm{A}) \approx 1 + \epsilon \,\text{tr}(\bm{A}) + O(\epsilon^2)$ を用いると

$$ \log|\det(\bm{I} + \Delta t \bm{J}_v)| \approx \log(1 + \Delta t \,\text{tr}(\bm{J}_v)) \approx \Delta t \,\text{tr}(\bm{J}_v) $$

よって

$$ \frac{d\log p(\bm{z}(t))}{dt} = -\text{tr}\left(\frac{\partial \bm{v}_\theta(\bm{z}(t), t)}{\partial \bm{z}}\right) $$

Pythonでの実装

2D分布を対象に RealNVP を NumPy でスクラッチ実装します。

import numpy as np
import matplotlib.pyplot as plt

np.random.seed(42)

# --- データ生成(2つの三日月形) ---
from sklearn.datasets import make_moons
data, _ = make_moons(n_samples=4000, noise=0.05)
data = data.astype(np.float64)

# --- RealNVPのアフィンカップリング層 ---
class AffineCouplingLayer:
    """RealNVPのアフィンカップリング層(2D用)"""
    def __init__(self, mask, hidden=64, lr=1e-3):
        self.mask = mask  # [1,0] or [0,1]
        self.lr = lr
        d_in = 1  # マスクされない部分の次元
        # sネットワーク: d_in -> hidden -> 1
        self.Ws1 = np.random.randn(d_in, hidden) * 0.1
        self.bs1 = np.zeros(hidden)
        self.Ws2 = np.random.randn(hidden, 1) * 0.1
        self.bs2 = np.zeros(1)
        # tネットワーク: d_in -> hidden -> 1
        self.Wt1 = np.random.randn(d_in, hidden) * 0.1
        self.bt1 = np.zeros(hidden)
        self.Wt2 = np.random.randn(hidden, 1) * 0.1
        self.bt2 = np.zeros(1)
        # 全パラメータリスト
        self.params = [self.Ws1, self.bs1, self.Ws2, self.bs2,
                       self.Wt1, self.bt1, self.Wt2, self.bt2]
        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 _s_net(self, z_masked):
        """スケールネットワーク"""
        self.s_h = np.maximum(0, z_masked @ self.Ws1 + self.bs1)
        s = self.s_h @ self.Ws2 + self.bs2
        return np.tanh(s) * 2.0  # スケール制限

    def _t_net(self, z_masked):
        """平行移動ネットワーク"""
        self.t_h = np.maximum(0, z_masked @ self.Wt1 + self.bt1)
        t = self.t_h @ self.Wt2 + self.bt2
        return t

    def forward(self, z):
        """順変換: z -> y"""
        z_masked = z[:, self.mask == 1].reshape(-1, 1)
        s = self._s_net(z_masked)
        t = self._t_net(z_masked)
        y = z.copy()
        idx = (self.mask == 0)
        y[:, idx] = z[:, idx] * np.exp(s) + t
        self.s_val = s
        self.z_in = z
        return y

    def inverse(self, y):
        """逆変換: y -> z"""
        y_masked = y[:, self.mask == 1].reshape(-1, 1)
        s = self._s_net(y_masked)
        t = self._t_net(y_masked)
        z = y.copy()
        idx = (self.mask == 0)
        z[:, idx] = (y[:, idx] - t) * np.exp(-s)
        return z

    def log_det_jacobian(self):
        """対数ヤコビアン行列式"""
        return self.s_val.sum(axis=1)  # sum of s values

    def get_grads_and_update(self, grad_loss_params):
        """Adam更新"""
        self.step += 1
        beta1, beta2, eps = 0.9, 0.999, 1e-8
        for i, g in enumerate(grad_loss_params):
            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)
        (self.Ws1, self.bs1, self.Ws2, self.bs2,
         self.Wt1, self.bt1, self.Wt2, self.bt2) = self.params


class RealNVP2D:
    """2D RealNVPフロー(数値微分で学習)"""
    def __init__(self, n_layers=6, hidden=64, lr=1e-3):
        self.layers = []
        for i in range(n_layers):
            mask = np.array([1, 0]) if i % 2 == 0 else np.array([0, 1])
            self.layers.append(AffineCouplingLayer(mask, hidden, lr))

    def forward(self, z):
        """z -> x(生成方向)"""
        log_det = np.zeros(len(z))
        x = z
        for layer in self.layers:
            x = layer.forward(x)
            log_det += layer.log_det_jacobian()
        return x, log_det

    def inverse(self, x):
        """x -> z(推論方向)"""
        log_det = np.zeros(len(x))
        z = x
        for layer in reversed(self.layers):
            z = layer.inverse(z)
            # 逆方向のlog_detは順方向の符号反転
        # 正確にlog_detを計算するため順変換を再実行
        _, log_det = self.forward(z)
        return z, log_det

    def log_prob(self, x):
        """データ点の対数尤度を計算"""
        z, log_det = self.inverse(x)
        # 標準正規分布の対数密度
        log_pz = -0.5 * (z**2).sum(axis=1) - np.log(2 * np.pi)
        return log_pz + log_det

    def sample(self, n_samples):
        """標準正規分布からサンプリングしてフローで変換"""
        z = np.random.randn(n_samples, 2)
        x, _ = self.forward(z)
        return x


# --- 数値微分による学習 ---
flow = RealNVP2D(n_layers=8, hidden=64, lr=5e-4)
n_epochs = 800
batch_size = 256
losses = []

for epoch in range(n_epochs):
    idx = np.random.choice(len(data), batch_size)
    x_batch = data[idx]

    # 現在の負の対数尤度
    log_p = flow.log_prob(x_batch)
    nll = -np.mean(log_p)
    losses.append(nll)

    # 数値勾配による各パラメータの更新
    eps_num = 1e-4
    for layer in flow.layers:
        grads = []
        for p_idx in range(len(layer.params)):
            grad = np.zeros_like(layer.params[p_idx])
            it = np.nditer(layer.params[p_idx], flags=['multi_index'])
            # パラメータ数が多い場合はサンプリング
            indices = []
            while not it.finished:
                indices.append(it.multi_index)
                it.iternext()
            # ランダムに一部のインデックスのみ更新(確率的座標降下法)
            n_update = min(len(indices), 20)
            sampled = [indices[j] for j in np.random.choice(len(indices), n_update, replace=False)]
            for mi in sampled:
                orig = layer.params[p_idx][mi]
                layer.params[p_idx][mi] = orig + eps_num
                (layer.Ws1, layer.bs1, layer.Ws2, layer.bs2,
                 layer.Wt1, layer.bt1, layer.Wt2, layer.bt2) = layer.params
                lp_plus = -np.mean(flow.log_prob(x_batch))
                layer.params[p_idx][mi] = orig - eps_num
                (layer.Ws1, layer.bs1, layer.Ws2, layer.bs2,
                 layer.Wt1, layer.bt1, layer.Wt2, layer.bt2) = layer.params
                lp_minus = -np.mean(flow.log_prob(x_batch))
                grad[mi] = (lp_plus - lp_minus) / (2 * eps_num)
                layer.params[p_idx][mi] = orig
                (layer.Ws1, layer.bs1, layer.Ws2, layer.bs2,
                 layer.Wt1, layer.bt1, layer.Wt2, layer.bt2) = layer.params
            grads.append(grad)
        layer.get_grads_and_update(grads)

    if (epoch + 1) % 100 == 0:
        print(f"Epoch {epoch+1}, NLL: {nll:.4f}")

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

# 学習損失
axes[0].plot(losses)
axes[0].set_xlabel('Epoch')
axes[0].set_ylabel('Negative Log-Likelihood')
axes[0].set_title('Training Loss')

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

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

# 潜在空間の可視化
z_mapped, _ = flow.inverse(data[:2000])
axes[3].scatter(z_mapped[:, 0], z_mapped[:, 1], s=1, alpha=0.5, c='green')
axes[3].set_title('Latent Space (z)')
axes[3].set_xlim(-4, 4)
axes[3].set_ylim(-4, 4)
axes[3].set_aspect('equal')

plt.tight_layout()
plt.show()

上のコードでは、RealNVP の核心であるアフィンカップリング層を実装し、2D三日月形データの密度推定とサンプリングを行いました。交互マスキングにより各次元が変換され、潜在空間では標準正規分布に近い分布が得られることを確認できます。

各種フローの比較

フロー ヤコビアン計算量 表現力 逆変換 特徴
Planar Flow $O(d)$ 低い 困難 ランク1の変換
RealNVP $O(d)$ 高い 解析的 アフィンカップリング
GLOW $O(d + c^3)$ 高い 解析的 1×1可逆畳み込み
Neural ODE $O(d)$ (推定) 高い ODE求解 連続時間、自由なアーキテクチャ

まとめ

本記事では、正規化フロー(Normalizing Flow)の理論を導出しました。

  • 変数変換公式 $p_X(\bm{x}) = p_Z(\bm{z})\,|\det \bm{J}|^{-1}$ により正確な対数尤度を計算可能
  • フローの連鎖で対数行列式が和に分解される
  • 三角ヤコビアンにより $O(d)$ での行列式計算を実現
  • RealNVP のアフィンカップリング層は解析的に可逆かつ効率的なヤコビアン計算を持つ
  • 連続時間フロー(Neural ODE)ではインスタンタネオス変数変換公式によりトレース計算のみで対数密度変化を追跡

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