『今更ながらVAEって（以下略）』の定式化を清書する

はじめに

先日、以下の記事を書いた。

tam5917.hatenablog.com

新年なので（？）、その定式化を清書しようというわけである。長くなってしまったので、急ぐ人は「8. 考察」からどうぞ。

はじめに
定式化
おわりに

定式化

1. 問題設定：変分推論による生成モデル

我々の目的はデータセット $\mathcal{D} = \{\mathbf{x}_1, \dots, \mathbf{x}_N \}$ から学習を行い、新しいデータを生成できるような生成モデル $p(\mathbf{x})$ を構築することである。例えばMNISTデータセット上で考える場合、各 $\mathbf{x}_{i}$ は個々の手書き文字画像を表す。

データ $\mathbf{x}$ は潜在変数 $\mathbf{z}$ によって生成されると仮定する。このとき、生成モデルは以下のように表すことができる。

$\begin{align*} p(\mathbf{x}) = \int p(\mathbf{x}|\mathbf{z}) p(\mathbf{z}) d\mathbf{z} \end{align*}$

ここで $p(\mathbf{z})$ は潜在変数 $\mathbf{z}$ の事前分布であり、通常は標準正規分布 $\mathcal{N}(\mathbf{0}, \mathbf{I})$ が用いられる。 $p(\mathbf{x}|\mathbf{z})$ は、潜在変数 $\mathbf{z}$ が与えられたときのデータ $\mathbf{x}$ の条件付き確率分布であり、デコーダによってモデル化される。しかしながら、この積分は解析的に計算できないため、変分推論を用いて近似を行う。

2. 変分推論

変分推論では、真の事後分布 $p(\mathbf{z}|\mathbf{x})$ をより扱いやすい分布 $q(\mathbf{z}|\mathbf{x})$ で近似する。この分布 $q(\mathbf{z}|\mathbf{x})$ を変分事後分布と呼ぶ。このとき、対数尤度 $\log p(\mathbf{x})$ の下限（Evidence Lower Bound, ELBO）を最大化することで、変分事後分布 $q(\mathbf{z}|\mathbf{x})$ を最適化する。ELBO は以下のように表される。

$\begin{align*} \mathcal{L}_{\text{ELBO}} = \mathbb{E}_{q(\mathbf{z}|\mathbf{x})} \left[ \log p(\mathbf{x}|\mathbf{z}) \right] - D_{\text{KL}} \left( q(\mathbf{z}|\mathbf{x}) || p(\mathbf{z}) \right) \end{align*}$

第一項は対数尤度の期待値であり、デコーダの出力と元のデータとの一致度を測る尺度を表す。

第二項は分布 $q$ と分布 $p$ の間のKLダイバージェンスであり、変分事後分布 $q(\mathbf{z}|\mathbf{x})$ が事前分布 $p(\mathbf{z})$ にどれだけ近いかを表す。

3. 今回のVAEにおけるモデル

今回のVAEでは、以下の仮定とモデル化を行う。

変分事後分布 $q(\mathbf{z}|\mathbf{x})$ : ガウス分布 $\mathcal{N}(\boldsymbol{\mu}(\mathbf{x}), \boldsymbol{\sigma}^2(\mathbf{x}))$ で近似する。
- 平均 $\boldsymbol{\mu}(\mathbf{x})$ と分散 $\boldsymbol{\sigma}^2(\mathbf{x})$ は、エンコーダの出力として得られる。
デコーダ $p(\mathbf{x}|\mathbf{z})$ : 潜在変数 $\mathbf{z}$ を入力とし、元のデータ $\mathbf{x}$ を再構成するニューラルネットワークとしてモデル化する。
アフィン変換: 潜在変数 $\mathbf{z}$ の一部をアフィン変換パラメータとして扱い、デコーダが画像を生成する際に、座標グリッドにアフィン変換を適用する。

4. アフィン変換の組み込み

潜在変数 $\mathbf{z}$ を、アフィン変換パラメータ $\mathbf{z}_{a}$ とその他のパラメータ $\mathbf{z}_{r}$ に分割する。

$\mathbf{z}_{a} \in \mathbb{R}^6$ : 2次元アフィン変換のパラメータ（スケーリング、回転、平行移動）
$\mathbf{z}_{r}$ : その他の潜在変数

4.1 座標グリッドの定義

今回のVAEでは、画像上の各ピクセルに対応する座標を [-1, 1] x [-1, 1] の範囲で正規化したグリッドとして定義している。

具体的には、画像のサイズを $H \times W$ とすると、座標グリッド $\mathbf{c} \in \mathbb{R}^{H \times W \times 2}$ は以下のように定義される。

$\begin{align*} \mathbf{c}_{i,j} = \begin{bmatrix} \frac{2j}{W-1} - 1 \ \frac{2i}{H-1} - 1 \end{bmatrix} \end{align*}$

ここで、

$i$ は画像の縦方向のピクセルインデックス $(0 \leq i \leq H-1)$
$j$ は画像の横方向のピクセルインデックス $(0 \leq j \leq W-1)$

この定義により、画像の左上のピクセルは [-1, -1]、右下のピクセルは [1, 1] に対応し、画像内のピクセルが [-1, 1] x [-1, 1] の範囲に正規化される。

4.2. アフィン変換の適用

アフィン変換は上記の正規化された座標グリッド $\mathbf{c}$ に対して適用される。アフィン変換行列 $\mathbf{A} \in \mathbb{R}^{3 \times 3}$ は潜在変数 $\mathbf{z}_{a}$ から生成され、以下の変換を行う。

$\begin{align*} \mathbf{c}' = \mathbf{A} \mathbf{c} \end{align*}$

ただし $\mathbf{c}$ を斉次座標系として扱うために、 $\mathbf{c}$ に1を要素にもつ次元を足して、 $\mathbf{c}$ を $\mathbb{R}^{3}$ のベクトルとみなす。

ここで $\mathbf{c}'$ も斉次座標系で表されているので、最後の次元の要素を落とすことで、2次元の座標グリッドに戻すことができる。

$\begin{align*} \mathbf{c}' \in \mathbb{R}^{H \times W \times 2} \end{align*}$

4.3. デコーダの役割

デコーダは変換された座標グリッド $\mathbf{c}'$ を考慮しながら、潜在変数 $\mathbf{z}_{r}$ に基づいて画像を生成する。デコーダの出力は再構成された画像 $\hat{\mathbf{x}}$ であり、その形状は元の画像 $\mathbf{x}$ と同じである。

5. 損失関数

全体の損失関数は、再構成損失、アフィン変換パラメータに関するKLダイバージェンス、およびその他のパラメータに関するKLダイバージェンスの和として表される。

$\begin{align*} \mathcal{L}_{\text{total}} = \mathcal{L}_{\text{rec}} + \beta_{a} \mathcal{L}_{\text{KL}a} + \beta_{r} \mathcal{L}_{\text{KL}_{r}} \end{align*}$

ここで、

再構成損失（負の対数尤度）:

$\begin{align*} \mathcal{L}_{\text{rec}} = -\mathbb{E}_{q(\mathbf{z}|\mathbf{x})} \left[ \log p(\mathbf{x}|\mathbf{z}) \right] \end{align*}$

損失関数を最小化するため、対数尤度の期待値に負号をつけたものを再構成損失と呼ぶ。デコーダによる再構成の誤差を測る尺度と考えることができる。

アフィン変換パラメータに関するKLダイバージェンス:

$\begin{align*} \mathcal{L}_{\text{KL}a} = D_{\text{KL}} \left( q(\mathbf{z}_{a}|\mathbf{x}) || p(\mathbf{z}_{a}) \right) \end{align*}$

$q(\mathbf{z}_{a}|\mathbf{x})$ は、平均 $\boldsymbol{\mu}_{a}(\mathbf{x})$ と分散 $\boldsymbol{\sigma}_{a}^2(\mathbf{x})$ を持つガウス分布で近似される。 $p(\mathbf{z}_{a})$ は標準正規分布と仮定する。

その他のパラメータに関するKLダイバージェンス:

$\begin{align*} \mathcal{L}_{\text{KL}_{r}} = D_{\text{KL}} \left( q(\mathbf{z}_{r}|\mathbf{x}) || p(\mathbf{z}_{r}) \right) \end{align*}$

$q(\mathbf{z}_{r}|\mathbf{x})$ は、平均 $\boldsymbol{\mu}_{r}(\mathbf{x})$ と分散 $\boldsymbol{\sigma}_{r}^2(\mathbf{x})$ を持つガウス分布で近似される。 $p(\mathbf{z}_{r})$ は標準正規分布と仮定する。

$\beta_{a}$ と $\beta_{r}$ はそれぞれのKLダイバージェンス項の重みを調整するパラメータである。冒頭の記事ではそれぞれの重みを1に設定し、2つのKL項に分離せず一体のものとして実装している（つまり $D_{\text{KL}} \left( q(\mathbf{z}|\mathbf{x}) || p(\mathbf{z}) \right)$ のまま）。

6. 学習と生成

学習: 損失関数 $\mathcal{L}_{\text{total}}$ を最小化するように、エンコーダとデコーダのパラメータを最適化する。
生成: 標準正規分布からランダムに潜在変数 $\mathbf{z}$ をサンプリングし、デコーダに入力することで、新しい画像を生成する。アフィン変換を有効にする場合は、アフィン変換パラメータ $\mathbf{z}_a$ を生成し、座標グリッドに変換を適用する。

7. ここまでのまとめ

今回のVAEの手法は、変分推論の枠組みに基づき、以下のステップで定式化できる。

問題設定: 生成モデルの構築と変分推論の導入
モデル化: エンコーダとデコーダのモデル化、アフィン変換の組み込み
損失関数: ELBOを最大化するように、再構成損失とKLダイバージェンスを定義
学習と生成: 損失関数を最小化するようにモデルを学習し、潜在空間からのサンプリングで新しい画像を生成

8. 考察：アフィン変換の有無と幾何学的変動の正規化について

8.1. アフィン変換の有無による影響

モデル訓練後、アフィン変換を適用する場合と適用しない場合で、VAEの動作がどのように変わるかを明確にする。ただし、いずれの場合もモデルの訓練時には座標グリッドに対してアフィン変換を常に適用したものとする。これは冒頭の記事における実験条件に沿っている。

アフィン変換あり：潜在変数 $\mathbf{z}$ の一部 $\mathbf{z}_{a}$ をアフィン変換パラメータとして使用し、座標グリッド $\mathbf{c}$ に変換を適用する。デコーダは、変換後の座標グリッド $\mathbf{c}'$ と残りの潜在変数 $\mathbf{z}_{r}$ に基づいて画像を再構成する。この場合、生成される画像の変動は $\mathbf{z}_{a}$ と $\mathbf{z}_{r}$ の両方に依存する。特に、 $\mathbf{z}_{a}$ は画像の「位置、向き、スケール」といった幾何学的な変動を制御する。
アフィン変換なし：アフィン変換パラメータ $\mathbf{z}_{a}$ を無視し、デコーダは元の座標グリッド $\mathbf{c}$ と潜在変数 $\mathbf{z}$ のみを使用して画像を再構成する。この場合、生成される画像の変動は $\mathbf{z}_{r}$ にのみ依存する。アフィン変換による幾何学的な変動は、ほぼ取り除かれる。

8.2. 幾何学的変動の正規化

アフィン変換を外した場合、生成される画像の幾何学的変動が「正規化される」とは、具体的に以下のことを意味する。

空間的な変動の抑制：アフィン変換パラメータ $\mathbf{z}_{a}$ が生成過程に影響を与えなくなるため、生成される画像の「位置、向き、スケール」などの幾何学的な変動が抑制される。
幾何学的変動以外の特徴の表現：生成される画像の変動は、アフィン変換パラメータ以外の潜在変数 $\mathbf{z}_{r}$ にのみ依存するようになる。これにより、画像の本質的な特徴（例えば、文字の太さや形状など、幾何学的変動以外の特徴）を表現した生成が行われる。
データセットの平均的な表現：アフィン変換を外して生成された画像は、学習データセットにおける画像の平均的な位置、向き、スケールに近いものになる。

8.3. 数式での表現

アフィン変換の有無による幾何学的変動の違いを数式で表すために、デコーダの出力を明示的に記述する。

アフィン変換あり：デコーダの出力 $\hat{\mathbf{x}}$ は、変換後の座標グリッド $\mathbf{c}'$ と潜在変数 $\mathbf{z}_{r}$ の関数として表せる。

$\begin{align*} \hat{\mathbf{x}} = f_{\text{dec}}(\mathbf{c}', \mathbf{z}_{r}) \end{align*}$

アフィン変換なし：デコーダの出力 $\hat{\mathbf{x}}$ は、元の座標グリッド $\mathbf{c}$ と潜在変数 $\mathbf{z}$ (この場合は $\mathbf{z}_{r}$ と同等) の関数として表せる。

$\begin{align*} \hat{\mathbf{x}} = f_{\text{dec}}(\mathbf{c}, \mathbf{z}) \end{align*}$

このときデコーダの入力に影響を与えるのは、潜在空間のうちアフィン変換パラメータに対応しない部分のみである。

$\begin{align*} \hat{\mathbf{x}} = f_{\text{dec}}(\mathbf{c}, \mathbf{z}_{r}) \end{align*}$

ここで、

$f_{\text{dec}}$ はデコーダを表す関数
$\mathbf{c}$ は正規化された座標グリッド
$\mathbf{c}'$ はアフィン変換後の座標グリッド
$\mathbf{z}$ は潜在変数
$\mathbf{z}_{r}$ はアフィン変換パラメータを除いた潜在変数

をそれぞれ表す。アフィン変換を外すことで、デコーダが座標グリッドを変換することによる影響を無視し、代わりに潜在変数 $\mathbf{z}_{r}$ を主として画像の生成に用いていることがわかる。このとき、生成された画像は、訓練データにおける位置、向き、スケールの平均に近いものになり、それらの幾何学的な変動が正規化されることになる。

8.4. 損失関数の役割

損失関数は学習時に以下の役割を果たす。

再構成損失：デコーダが元の画像を正確に再構成できるように学習を促す。
KLダイバージェンス：潜在空間が正規分布に近い分布になるように学習を促す。

KLダイバージェンスをアフィン変換パラメータとそれ以外のパラメータに分離することで、より独立した潜在空間表現の学習が期待できる。

蛇足ながら、学習開始時点からアフィン変換を適用しない場合、今回のVAEは「普通のVAE」と等価であり、入力画像を忠実に再現するように学習が進む。デコーダが参照する座標グリッド $\mathbf{c}$ は定数になるからである。結果として生成される画像は、訓練データに含まれる変動を（幾何学的変動を含めて）そのまま反映する。

9. 所感など

「アフィン変換 + オートエンコーダ」や「アフィン変換 + VAE」な先行研究は適当に調べるだけでも色々と見つかる。いずれの研究も、元画像に対して直接アフィン変換を適用する形だった。今回は元画像と並行に座標グリッドを導入し、そこにアフィン変換を適用する形が独創的なのである。確かに「その発想はなかった」と唸らされた。結果的にアフィン変換パラメータが潜在変数から分離できたのだから、大したものである。

本定式化を様々な画像変動に一般化すべく、ちょっと思索するとチャレンジングな課題がどんどん出てくる。そもそも座標グリッドに乗らない・当てはまらない変動のほうが多いのだから、すぐに定式化は破綻しそうなものである。それでも潜在変数のdisentanglementの観点から、興味深いテーマをいくつか示唆してくれるのではないかと思う。何だか考えるのが楽しくなってこないかな？（日頃から深く思索せよ、という自戒を込めて）

おわりに

こういう記事は新年にしか書けない。