JA EN
体系生成モデル
·無料·13分で読めます

VAEを1から — 「圧縮して戻す」に確率を混ぜると生成になる

圧縮して元に戻すだけのオートエンコーダは、なぜ新しいデータを作れないのか。そこに確率を一滴混ぜると生成モデルに変わる理屈を、ELBO・再パラメータ化・潜在空間の補間まで前提知識ゼロから解説します。

対象textタスクgeneration

荷造りの比喩から始める

引っ越しで、部屋の中身をスーツケースに詰める人と、引っ越し先で開けて元通りに並べ直す人がいるとします。鞄が小さいほど詰める側は取捨選択を迫られ、「これは捨てても、開ける人が推測して置き直せるはずだ」という判断が要る。逆に言えば、小さい鞄で復元できたなら、その中身は部屋の本質だけが残った要約です。

これがオートエンコーダです。詰める側をエンコーダ、戻す側をデコーダ、途中の小さい鞄を潜在変数 zz と呼びます。学習の目標は「元に戻せたか」だけで、たとえば画素ごとの二乗誤差を最小化します。正解ラベルが入力そのものなので、人手のラベル付けは要りません。

肝は鞄の細さです。部屋と同じ大きさの鞄なら全部そのまま入れればよく、何も学びません。通り道を細くしたことが、意味のある要約を強制するわけです。

オートエンコーダでは「作れない」

学習が終わったオートエンコーダから、デコーダだけを取り出して適当な zz を入れてみます。新しい画像が出るでしょうか。出ません。だいたいはノイズか、どこかで見たような崩れた何かです。

この訓練が zz の置き場所について何も約束していないからです。訓練データ1万枚をエンコードすれば潜在空間に1万個の点が打たれ、デコーダはその点の上でだけ正しく動くよう鍛えられました。しかし点と点のあいだの土地には誰も何も書き込んでいない。街の位置だけが載っていて、街と街のあいだが平野なのか崖なのかは無言の地図です。

そのうえ、その1万点がどのあたりに散らばっているのかも分かりません。どこから zz を引けばいいのかが決まっていない以上、サンプリングのしようがない。生成モデルにするには、潜在空間に穴を空けないことと、引いてくる場所を最初から決めておくことを、同時に解く必要があります。

確率を一滴混ぜる

VAE(変分オートエンコーダ)の発想は拍子抜けするほど単純で、エンコーダに「1点」ではなく「もや」を出させるというものです。出力を zz そのものではなく中心 μ\mu広がり σ\sigma の2つにし、実際の zz はその中心と広がりで決まるガウス分布から毎回ランダムに引く。同じ画像を10回入れれば、少しずつ違う zz が10個出てきます。

これで最初の問題が解けます。デコーダから見ると、同じ画像に対して毎回ずれた zz が渡ってくる。ずれた場所からでも同じ画像に戻せなければ損をするので、μ\mu の周りの一帯がまとめて「その画像の領域」として塗り潰されていく。点だった訓練データが面積を持つのです。

ただし抜け道があります。σ\sigma を0に近づければ、もやは点に戻ってしまう。そこで2つ目の圧力をかけます。すべての入力について、エンコーダが出すもやを、あらかじめ決めた共通の分布に近づけろという罰則です。普通は平均0・分散1の標準ガウス分布 N(0,I)\mathcal{N}(0, I) を使い、これを事前分布と呼びます。

これで σ\sigma を潰す逃げ道が塞がれ、同時にもう一つの問題も片付きます。全部のもやが標準ガウスに寄っているなら、生成したいときは標準ガウスから zz を1本引いてデコーダに通せばいい。「圧縮して戻す」に確率を混ぜると生成になる、とはこの意味です。2013年に Kingma と Welling が提案し、ほぼ同時期に Rezende らが同等の考え方を独立に示しました。

仕組み: ELBO という1本の式

本当に最大化したいのは、手元のデータ xx がモデルから出てくる確率 p(x)p(x) です。ところがこれは p(x)=p(xz)p(z)dzp(x) = \int p(x \mid z)\, p(z)\, dz という積分で、ありうる zz を全部なめる必要があり、zz が数十次元あれば計算不能です。

そこでVAEは p(x)p(x) を直接いじるのを諦め、確実にそれ以下だと分かっている量を持ち上げます。下から支える板を上げれば、上に乗っている本体も上がる。この下限が ELBO(Evidence Lower BOund、変分下限)です。

logp(x)    Eq(zx) ⁣[logp(xz)]再構成項    DKL ⁣(q(zx)p(z))KL項\log p(x) \;\ge\; \underbrace{\mathbb{E}_{q(z \mid x)}\!\left[\log p(x \mid z)\right]}_{\text{再構成項}} \;-\; \underbrace{D_{\mathrm{KL}}\!\left(q(z \mid x)\,\|\,p(z)\right)}_{\text{KL項}}
(1)

記号を一つずつ。q(zx)q(z \mid x) はエンコーダが出すもや、つまり中心 μ\mu・広がり σ\sigma のガウス分布。p(xz)p(x \mid z) はデコーダで、zz からどれくらいの確からしさで xx が出るか。p(z)p(z) は事前分布 N(0,I)\mathcal{N}(0, I)E\mathbb{E} は平均、DKLD_{\mathrm{KL}} は2つの分布のズレを測る量です。

平たく言えば式(1)は、「もやから引いた zz で元に戻せた度合い」から「もやが引き出し口からズレている度合い」を引いた値が、本当に上げたい量の下限になっていると言っています。左辺は計算できないが右辺は計算できる、というのが全てで、学習では符号を反転して ELBO-\text{ELBO} を損失として下げます。

KL情報量そのものはKL情報量を1から — 2つの分布の「ズレ」を測るで扱っています。ここで要る理解は一行、同じ分布なら0、離れるほど大きくなる罰金です。そして qqp(z)p(z) もガウス分布なので、KL項は積分せずに閉じた式で書けます。

DKL ⁣(q(zx)N(0,I))=12j=1d(1+logσj2μj2σj2)D_{\mathrm{KL}}\!\left(q(z \mid x)\,\|\,\mathcal{N}(0, I)\right) = -\frac{1}{2}\sum_{j=1}^{d}\left(1 + \log \sigma_j^2 - \mu_j^2 - \sigma_j^2\right)
(2)

dd は潜在次元の数、jj はその1本1本です。要するにμ\mu が0から離れるほど罰金、σ\sigma が1から離れるほど罰金で、μj=0, σj=1\mu_j = 0,\ \sigma_j = 1 を入れると中身が 1+001=01 + 0 - 0 - 1 = 0 になります。この項をサンプリングなしで正確に計算できるので、乱数由来のブレが入るのは再構成項だけで済みます。

FIG 12変数の損失面を降りていく様子。VAEのELBOも「再構成項+KL項」を足した1枚の面で、学習率を上げれば同じように暴れ、平らな谷に入れば抜け出せなくなります

再パラメータ化トリック — 乱数を経路の外へ

ここで実装が詰まります。q(zx)q(z \mid x) から zz をサンプルする」という操作は微分できないのです。

ニューラルネットの学習は、出力の誤差を入力側へ連鎖律で遡って伝えます(誤差逆伝播を1から解説)。ところが経路の途中に乱数生成器が挟まると、そこで道が途切れる。デコーダまで勾配が来ているのに、エンコーダへ渡せません。

解決は、乱数を経路の外に追い出すことです。

z=μ(x)+σ(x)ε,εN(0,I)z = \mu(x) + \sigma(x) \odot \varepsilon, \qquad \varepsilon \sim \mathcal{N}(0, I)
(3)

\odot は要素ごとの掛け算です。標準ガウスから引いた ε\varepsilonσ\sigma 倍して μ\mu だけずらす。得られる zz の分布は中心 μ\mu・広がり σ\sigma のガウス分布そのもので、さっきまでと同じです。違うのは経路の形だけ。ε\varepsilonネットワークと無関係に外から降ってくる入力であり、zz から μ\muσ\sigma へは掛け算と足し算しか挟まっていないので、勾配は普通に流れます。サイコロを振ってから加工するか、加工の途中で振るか。分布は同じでも、前者なら加工の部分を微分できる、というだけの話です。

FIG 2重みを σ、バイアスを μ、入力を ε と読み替えると、再パラメータ化 z = μ + σε は1個のニューロンとまったく同じ形。重みとバイアスを動かすと出力が動く=勾配が届く、という当たり前がここで効いています

コードで書くVAE

学習1ステップの中心は10行ほどです。

mu, logvar = encoder(x)                  # σ ではなく log σ² を出させる
std = torch.exp(0.5 * logvar)
eps = torch.randn_like(std)              # ネットワークの外から来る乱数
z = mu + std * eps                       # 再パラメータ化

x_hat = decoder(z)
recon = F.mse_loss(x_hat, x, reduction="sum") / x.size(0)
kl = -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp()) / x.size(0)
loss = recon + beta * kl                 # beta=1 が素のVAE

σ\sigma ではなく logσ2\log \sigma^2 を出させるのがコツです。σ\sigma は必ず正でなければならないのに、ネットワークの出力は正負どちらも取ります。logσ2\log \sigma^2 なら全実数が許され、exp を通せば必ず正になる。数式に現れないのに実装には必ず要る種類の工夫です。

潜在空間を歩く

2枚の画像をエンコードして z1,z2z_1, z_2 を得て、あいだを少しずつ動かしながらデコードすると、VAEでは中間が滑らかに変形していく映像が得られます。素のオートエンコーダでは崩れた何かしか出なかったところです。「メガネをかけた顔」の zz の平均から「かけていない顔」の平均を引けば属性の方向を指すベクトルが取れ、任意の顔に足してデコードすればメガネがかかる。単語ベクトルの足し引きと同じ構造です。

ここに高次元特有の落とし穴があります。dd 次元の標準ガウスから引いたベクトルの長さは d\sqrt{d} 付近に強く集中し、原点付近はほとんど空っぽです。すると z1z_1z2z_2 を直線で結んだ中点は両端より明らかに短いベクトルになり、モデルが学習中に一度も通らなかった原点寄りを横切ります。補間の真ん中だけぼやけるのはたいていこれで、長さを保ったまま角度だけ回す球面線形補間(slerp)を使うのが定番です。

なぜVAEの出力はぼやけるのか

VAEには昔から「画像がぼやける」という評判があります。これは実装の下手さではなく、損失の選び方から論理的に出てくる性質です。

再構成項に二乗誤差を使うのは、「p(xz)p(x \mid z) は分散一定のガウス分布」と仮定するのと同じことです。すると、ある zz に対して正解がありうる形が複数あるとき、ガウス分布の最尤推定はその平均を返します。輪郭が右にも左にもありうるなら、両方の平均、つまり中間のぼんやりした輪郭になる。鮮明な候補のどれかを選ぶより、真ん中を取るほうが二乗誤差の期待値は小さいからです。

この「平均を返してしまう」問題への回答が、その後の生成モデル史そのものになっています。敵対的な判別器に鮮明さを判定させたのがGAN、潜在変数を離散のコードブックにしたのがVQ-VAE、ノイズ除去を何百段も重ねて一発の平均化を避けたのが拡散モデルです(拡散モデルを1から理解する)。それでもVAEは消えませんでした。Stable Diffusion のような潜在拡散では、VAEは生成の主役を降りて、画像を小さな潜在テンソルに畳む「圧縮器」として残っています。

現場ではこう使う

誰が、いつ触るか。 まず画像生成パイプラインの担当者です。Stable Diffusion 系を業務で回すと、AutoencoderKL の差し替え、潜在テンソルの扱い、VAEデコード時のVRAM見積もりは必ず通ります。次に異常検知の担当者で、正常データだけでVAEを学習し再構成誤差やELBOを異常スコアに使う構成は外観検査などで今も現役です。三つ目に表現学習で、ラベルのないログや購買履歴から低次元の特徴を作る用途があります。

実際に触るパラメータ名。

知らないと事故になる落とし穴。

  1. posterior collapse(事後分布の崩壊): KL損失が0に張り付いて動かなくなる症状。エンコーダが「どんな入力にも事前分布そのまま」を返すようになり、zz に情報が乗りません。デコーダが強力なとき(特に自己回帰型)に起きやすく、こうなるとVAEは入力を無視した無条件生成器に成り下がります。学習ログでKL項を必ず別々に記録してください。合計損失だけ見ていると気づけません
  2. 損失のスケール不一致: 再構成誤差を sum で取るか mean で取るかで、KL項との相対比が次元数の分(画像なら数千倍)変わります。論文の結果が再現できないときは、まずここを疑うのが早い
  3. logvar の発散: 学習が不安定になると logσ2\log \sigma^2 が振り切れ、expinf になります。clamp で範囲を切るのが定石
  4. 異常検知での過信: VAEは学習していない種類の入力もそこそこ再構成してしまうことがあり、再構成誤差だけの閾値では見逃しが出ます
  5. 潜在拡散で scaling_factor を掛け忘れる: 潜在の分散がノイズスケジュールの想定とずれ、出力が壊れます。自前でVAEを差し替えたときに踏みやすい

説明を求められる問い。 「なぜ再パラメータ化が必要か」には、期待値の勾配をサンプリング操作を通さずに取るため。「KL項がゼロになったら何が起きるか」には、q(zx)q(z \mid x) が入力に依らず事前分布と一致する=潜在に情報が乗っていない状態、と答えられれば十分です。

まとめ

コメント

コメントにはログインが必要です