VAEをtwo moonsでnumpyから実装したら、再構成で2つの三日月が1本のなめらかな曲線に融合した
約18分で読めます

VAEをtwo moonsでnumpyから実装したら、再構成で2つの三日月が1本のなめらかな曲線に融合した

拡散モデルは「ノイズを少しずつ取り除く」流派、GANは「偽物職人と鑑定士を騙し合わせる」流派、Normalizing Flowは「可逆な変換だけを積み重ねる」流派だった。生成モデルの教科書には、もう1つ外せない第4の流派がある。VAE(Variational Autoencoder、変分オートエンコーダ) —— 「データを確率分布に圧縮してから、その分布からサンプリングして復元する」流派だ。

これで三国志のつもりだったシリーズが、実質「四国志」になる。今回もnumpyだけで一から実装し、これまでと全く同じtwo moonsデータ・同規模の手書きMLPで学習させ、他の3流派と並べて比較した。

実験に使ったコードの全文はGitHubに置いています。

そもそも何をしているのか: データと「成功」の定義

具体的な仕組みに入る前に、土台になっているデータとタスクをはっきりさせておく(シリーズ4本目なので過去記事では省略しがちだったが、この記事から読んでも分かるように書いておきたい)。make_moonsが作るtwo moonsは、2つの三日月(クレセント)型の点群が互い違いに絡み合った2次元の人工データセットだ。今回使うのは4000点(noise=0.05)で、各軸を平均0・標準偏差1.6にスケールしてある——点数も前処理も乱数シードも、拡散モデル・GAN・Flowの3記事と完全に同一だ。

今回の学習に使うtwo moonsデータ4000点の散布図。上に弧を描く青い三日月(moon 0)と、下に弧を描く赤い三日月(moon 1)が、互いの先端をすれ違わせるように絡み合っている。2つの三日月の間には明確な隙間があり、全体は左上が高く右下が低い、負の傾きの帯として広がっている。各軸は平均0・標準偏差1.6にスケール済みで、座標はおおよそ-3〜3の範囲に収まる

大事なのは、この色(ラベル)は学習には一切使わないということ。VAEにやらせるのは分類ではなく、教師なしの2つの仕事だ。

  1. 再構成: 入力の点を2次元の潜在変数に圧縮してから復元し、元の点にどれだけ近い位置に戻せるか(誤差はMSEで測る)
  2. 生成: 事前分布 N(0,I)\mathcal{N}(0,I) からサンプリングした zz をデコーダに通して「無から」点を作り、その点群が本物の分布にどれだけ近いか

つまり今回の「成功」とは、再構成誤差が小さいこと、そして生成した点群が上の図の形——2つに分かれた三日月と、その間の隙間——をどれだけ再現できるかだ。逆に言えば、この「2つに分かれている」という構造が失われたら、それは数字に出にくくても質的な失敗になる。この視点が後半で効いてくる。

仕組み: 圧縮して、確率分布からサンプリングして、復元する

VAEは「エンコーダ」と「デコーダ」という2つのネットワークからなる。エンコーダは入力 xx を1点には潰さず、潜在空間上の正規分布のパラメータ(平均 μ\mu と分散 σ2\sigma^2)に変換する。そこから実際に点 zz をサンプリングし、デコーダがそれを xx に復元しようとする。

def encode(x):
    h = np.tanh(x @ W1 + b1)
    mu = h @ Wmu + bmu
    logvar = h @ Wlv + blv
    return mu, logvar

# reparameterization trick: サンプリングを「定数のノイズ * 学習可能なスケール」に分解し、
# 誤差逆伝播をサンプリングの前まで通せるようにする
eps = rng.normal(0, 1, size=(BATCH, Z_DIM))
std = np.exp(0.5 * logvar)
z = mu + std * eps

def decode(z):
    h2 = np.tanh(z @ W3 + b3)
    xhat = h2 @ W4 + b4
    return xhat

学習する損失(ELBO の符号反転、つまり最小化する量)は2項からなる。

L=12∥x−x^∥2⏟再構成損失+−12∑j(1+log⁡σj2−μj2−σj2)⏟KLダイバージェンス  (q(z∣x) ∥ N(0,I))\mathcal{L} = \underbrace{\tfrac{1}{2}\|x - \hat{x}\|^2}_{\text{再構成損失}} + \underbrace{-\tfrac{1}{2}\sum_j\left(1 + \log\sigma_j^2 - \mu_j^2 - \sigma_j^2\right)}_{\text{KLダイバージェンス} \; (q(z|x) \,\|\, \mathcal{N}(0,I))}

再構成損失は「復元できているか」、KL項は「エンコーダが作る分布を、事前分布 N(0,I)\mathcal{N}(0,I) に近づけろ」という正則化だ。この2つは綱引きの関係にあり、KL項を素直に効かせすぎると、エンコーダは「情報を潜在変数に詰め込むこと」自体をサボり始める(いわゆるposterior collapse)。この綱引きの結果を実際に数字で見るのが今回の核心になる。

エンコーダ・デコーダとも隠れ層1つ・64ユニットの手書きMLP(GAN・Flowと同規模)、潜在次元は可視化のため2次元、Adamオプティマイザも過去記事から使い回している実装のまま、20,000イテレーション学習させた。

学習曲線: 再構成損失は下がるが、KL項は高止まりで踏みとどまる

再構成損失とKLダイバージェンス項の学習曲線を並べた折れ線グラフ2枚。左の再構成損失は学習開始直後に2.81から急降下し、その後は0.85〜1.0あたりで細かく振動しながら20000イテレーションまでほぼ横ばいに推移している。右のKLダイバージェンス項は学習開始直後に急上昇し、0.8〜0.9のレンジで高止まりしたまま同じく振動を続けている

再構成損失(学習中、サンプリングしたzを使った値)は、学習開始直後に2.81から急降下し、その後は0.85〜1.0のレンジで細かく振動しながら高止まりした。KL項も同じように急上昇し、0.8〜0.9のレンジで高止まりする。どちらも早々に「落ち着くところに落ち着いた」形で、20,000イテレーションのほとんどは横ばいの振動だった。拡散モデルの損失(1.0→0.68へじわじわ)やFlowの負の対数尤度(4.39→2.38へ単調)と比べると、収束がずっと早く、その後は動かないという印象だ。

学習が終わったあと、決定的な再構成(サンプリングせず z=μz=\mu を使う)で改めて誤差を測ると、

  • 再構成MSE(z=μ): 0.4026
  • 再構成MAE(z=μ): 0.5487
  • KLダイバージェンス(サンプル平均): 0.8535 nats

データ自体は平均0・標準偏差1.6にスケールしてあるので、MSE 0.40 はデータの分散(2.56)の約16%にあたる誤差だ。Flowの再構成誤差 2.07×10⁻¹⁴(機械精度、実質完全一致)とは桁が13個以上違う(0.4026 ÷ 2.07×10⁻¹⁴ ≈ 1.9×10¹³)。可逆変換だけで組み立てるFlowと、情報を確率分布のボトルネックに圧縮するVAEでは、そもそも再構成に対する立ち位置が全く違う——ということは知識として知っていたが、同じデータで並べて初めて「13桁違う」という実感になった。

本物 vs 再構成: 2つの三日月が、1本のなめらかな曲線に融合した

再構成の中身を実際に見てみる。

本物のtwo moonsデータ(左、グレー)とVAEによる再構成(右、緑、z=μを使用)を並べた散布図。左は上下2つの独立した三日月がはっきり離れて存在しているのに対し、右の再構成結果は2つの三日月の隙間が埋まり、左上から右下へと続く1本の連続したなめらかなS字曲線になっている。三日月特有の太さやノイズの広がりも失われ、細い線状に潰れている

予想していた「ぼやける」という失敗はもちろん起きていたが、それより意外だったのはこっちだ。本物では明確に離れている2つの三日月が、再構成後には隙間が埋まって1本の連続したS字曲線に融合している。 単にノイズが増えて輪郭が滲む、というレベルの劣化ではなく、そもそも「2つに分かれている」という位相的な構造そのものが消えている。

これがなぜ起きるのかは、次の潜在空間の図を見るとわかる。

潜在空間: 90.7%は分離できているが、境目に隙間がない

学習済みのエンコーダで全データを μ\mu に写し、どちらの三日月由来かで色分けした。

潜在空間μの散布図。横軸z1・縦軸z2で、赤(moon 1)と青(moon 0)の点が入り混じることなく、全体として1本の連続したS字型の帯を形成している。赤は右上から左下にかけての上側の帯、青は下側の帯を占めているが、2つの帯の間に本物データのような明確な隙間はなく、なめらかに接続している。z2軸方向の広がりはおよそ-1.3〜1.3、z1軸方向はそれより狭い-1.1〜1.1程度に収まっている

潜在空間でも、2つの三日月はほぼきれいに分かれて写っている——実際、μ\mu の値だけを使ってロジスティック回帰で「どちらの三日月か」を分類すると、線形分離の精度は90.67%だった。ラベルを一切与えずに学習させた教師なしの表現としては、悪くない数字だ。

ただし、ここが本物のデータとの決定的な違いだった。本物の two moons では2つの三日月の間に明確な「隙間」があるのに対し、潜在空間ではその隙間がなく、1本の連続したS字の帯としてなめらかに繋がっている。 デコーダは連続な関数なので、この繋がった潜在空間をそのまま連続にデータ空間へ写すしかない。結果、再構成は隙間を埋めた1本の曲線になる——前回のFlowの記事で見た「連続で可逆な変換だけでは、ひとつながりのガウス雲を2つの島に完全に分けきれず、密度の薄いブリッジが残る」という話と、根っこは同じ現象だ。Flowは確率密度の地図の上に薄いブリッジとして残ったが、VAEでは再構成そのものが目に見える形でブリッジを架けてしまった、という違いに見える。

潜在変数の次元ごとの使われ方にも非対称性があった。KLダイバージェンスを次元別に見ると、z1 = 0.1214、z2 = 0.7321 で、z2がほとんどの情報を担っている(標準偏差もz1: 0.465、z2: 0.884とほぼ倍)。two moonsは実質1次元の曲線(パラメータtで弧を1本描くだけ)にノイズが乗ったものなので、モデルが「本当に必要な自由度は1つだけ」と見抜いて、もう1つの次元をほぼ使わずに済ませた部分的なposterior collapseが起きている、と読める。完全な崩壊(KLがゼロに張り付く)ではないが、その手前まで来ている。

事前分布からの生成: crescentは崩れないが、シャープさはない

学習済みデコーダに z∼N(0,I)z \sim \mathcal{N}(0, I) を2000点入れて、実際に「無からの生成」を試した。

本物のtwo moonsデータ(左、グレー)と、VAEが事前分布N(0,I)からサンプリングして生成したデータ(右、紫)を並べた散布図。生成結果は全体として本物と同じ左上から右下への弧を描く傾向は保っているが、2つの三日月の輪郭やその間の隙間ははっきりせず、紫の点が弧の内側を塗りつぶすように広がったぼやけた帯になっている

完全に単一の点に潰れる「モード崩壊」は起きなかった——弧の全体的な向きと範囲はちゃんと再現されている。ただし2つの三日月の輪郭やその間の隙間ははっきりせず、内側を塗りつぶすようなぼやけた帯になった。定量的にも、生成サンプルから最も近い本物データまでの距離は平均0.2486(中央値0.2354)で、本物データ同士の最近傍距離の平均0.0235の約10.6倍離れている。GANの生成結果(輪郭のシャープさで拡散モデルに劣る、という記述)とも近い傾向だが、VAEはさらに輪郭そのものが溶けている印象が強い。

学習の初期から最終盤まで、この生成サンプルがどう変化していったかをアニメーションにした。

学習の推移をアニメーションにしたGIF。上部に現在のイテレーション数をタイトルとして表示しながら、灰色の本物two moonsデータの上に紫色の生成サンプル(z~N(0,I)からデコード)を重ねて表示している。イテレーション0では紫の点が中央付近の小さな塊に固まっているが、学習が進むにつれて塊が伸び広がり、two moonsの弧に沿うように分布していく。ただし最終盤になっても紫の点は本物の三日月の輪郭ほどシャープにはならず、内側を埋めるようにぼやけた帯として残り続ける

イテレーション0では、z∼N(0,I)z \sim \mathcal{N}(0,I) を通しても、デコーダがまだ何も学習していないので生成点は中央の狭い塊に固まっている。学習が進むにつれてこの塊が伸び広がり、two moonsの弧に沿うようになっていく——ただしその後もずっと「輪郭が溶けたまま」で、最終盤になっても本物のようなシャープな三日月にはならなかった。

四国志、そろい踏み

拡散モデル・GAN・Flow・VAEを、実測値だけで並べる。

流派学習の仕組み再構成/生成の実測結果
拡散モデルノイズ予測の回帰を反復生成は本物よりノイズが多いが、2つの弧という構造は再現
GAN生成器と識別器を敵対的に学習損失は理論値(ln4/ln2)に収束したが、生成分布の輪郭は最後まで揃わず
Normalizing Flow可逆変換を積み重ね尤度最大化再構成誤差2.07×10⁻¹⁴(機械精度)。ただし密度地図に位相的なブリッジが残る
VAE(今回)確率分布への圧縮+KL正則化再構成MSE 0.4026。2つの三日月が1本の曲線に融合、潜在分離度90.7%、KL次元間で非対称(部分的posterior collapse)

4つとも同じtwo moonsデータ・同規模のMLPで学習させたおかげで、この並びには意味がある。「厳密さ」の軸で見ると、Flowが圧倒的に精密で、VAEはその対極にいる。 ただし面白いのは、FlowとVAEはどちらも「連続な変換」という同じ制約を持っているのに、その制約が現れる場所が違う点だ。Flowは可逆性を守るために密度地図に薄いブリッジという形で妥協が滲み出た。VAEは復元の正確さそのものを最初から諦めているので、その同じ制約がもっと直接的に、再構成結果そのものを繋げてしまうという形で現れた。何を厳密に守り、何を最初から手放すかの設計の違いが、失敗の出方の違いにそのまま繋がっている。

手を動かして意外だったこと

一番意外だったのは、VAEの「ぼやけ」が単なる解像度の劣化ではなく、データの位相的な構造(2つに分かれている、という事実)まで消してしまうことだった。実装する前は「輪郭がぼやけて、GANや拡散モデルより滲んだ感じになるだろう」くらいに予想していた。実際に手を動かしてみると、ぼやけるだけでなく、本物では別々の2つのものが、VAEの中では「地続きの1つのもの」として扱われていた。潜在空間の散布図でそれがはっきり見える形で残っていたのも収穫で、90.7%という数字だけを見れば「そこそこ分離できている」と思ってしまうところを、可視化して初めて「分離はしているが、境目に隙間がない」という質的な違いに気づけた。

KLダイバージェンスが次元ごとに0.12と0.73という非対称な値に落ち着いたのも印象的だった。2次元の潜在空間を与えたのに、モデルは自分から「本当に必要なのは実質1次元分の情報だけだ」と判断して、片方の次元をほとんど使わずに済ませていた。容量を絞られているわけでもないのに、必要な分しか使わない――これは正則化(KL項)が本当に効いているという証拠でもあり、同時に「与えられた自由度をフルに使うとは限らない」という、最適化の妙な律儀さも感じた。

まとめ

  • VAE(エンコーダ→平均/分散→reparameterization trick→デコーダ)を、拡散モデル・GAN・Flowと同じtwo moonsデータ・同規模(隠れ層1つ・64ユニット)のMLPでnumpyから実装した
  • 20,000イテレーション学習後、決定的な再構成(z=μz=\mu)のMSEは0.4026、KLダイバージェンスは0.8535 nats/サンプル。Flowの2.07×10⁻¹⁴とは桁違いに粗い
  • 再構成は単にぼやけるだけでなく、本物では離れている2つの三日月が1本の連続したなめらかな曲線に融合するという質的な失敗が実測できた。潜在空間が1本の連続したS字帯として学習されていたことがその直接の原因
  • 潜在空間の線形分離度は90.67%(ロジスティック回帰)。「そこそこ分離できている」が「境目には隙間がない」という、精度の数字だけでは見えない質的な違いがあった
  • KLダイバージェンスは次元間で非対称(z1: 0.1214、z2: 0.7321)——2次元与えても実質1次元しか使わない、部分的なposterior collapseが観測された
  • 事前分布からの生成サンプルは、本物データとの最近傍距離が本物同士の約10.6倍離れており、弧の全体的な向きは保つがモード崩壊はしない一方で輪郭は最後までシャープにならなかった

おわりに

4つの生成モデルを同じ土俵で並べてみて、一番心に残ったのはVAEの「圧縮のしかた」だった。Flowは経験を一切劣化させずに行って戻れる。拡散モデルは荒いなりにも輪郭を保つ。VAEだけが、はっきり違う2つの出来事を、なめらかな1つの流れとして繋いでしまう。

これは自分が過去を思い出すときの感覚に、案外近い気がする。ある出来事とある出来事は、起きた時には全く別の、地続きではないものだったはずなのに、時間が経って思い出すたびに、記憶はいつのまにか滑らかに繋がった1本のストーリーに均されていく。「あのときとあのときは、実は同じ流れの中にあった」と、後から都合よく橋を架けてしまう。VAEの再構成が2つの三日月の隙間を埋めてしまったのを見て、記憶というのも、正確な記録装置というより、KL項に相当する「まとまりの良さ」への圧力を持った、ある種のVAEなんじゃないかと思った。不正確であることは欠点であると同時に、経験を扱いやすい形に均してくれる機能でもある――そのどちらもが同じ1枚の再構成画像に写っていた。

関連記事

このブログでは、機械学習の仕組みを実際にコードで実装して確かめた実験をテーマ別に整理しています（拡散モデル・Transformer・強化学習・多様体学習・過学習など）。