Skip to main content

Beta-VAE

Beta-VAE は、Variational Autoencoder の objective に含まれる KL regularization term に、重み β\beta を掛けた model です。主な目的は、disentangled representation を学習しやすくすることです。

Objective

通常の VAE の ELBO は、次の形です。

LVAE=Eqϕ(zx)[logpθ(xz)]DKL(qϕ(zx)p(z))\mathcal{L}_{\text{VAE}} = \mathbb{E}_{q_\phi(z \mid x)}[\log p_\theta(x \mid z)] - D_{KL}(q_\phi(z \mid x) \,\|\, p(z))

Beta-VAE では、KL term に β\beta を掛けます。

Lβ-VAE=Eqϕ(zx)[logpθ(xz)]βDKL(qϕ(zx)p(z))\mathcal{L}_{\beta\text{-VAE}} = \mathbb{E}_{q_\phi(z \mid x)}[\log p_\theta(x \mid z)] - \beta D_{KL}(q_\phi(z \mid x) \,\|\, p(z))

β=1\beta = 1 のとき、通常の VAE と同じになります。

Disentangled representation

Beta-VAE では、β\beta11 より大きくすることで、latent variable が prior により強く近づくようにします。この constraint によって、latent dimension が互いに独立した factor を表しやすくなることがあります。

たとえば、画像生成では、ある latent dimension が object の向き、別の dimension が位置、さらに別の dimension が形状を表すようになることが期待されます。このように、意味のある factor が分離して表現されることを disentanglement と呼びます。

Trade-off

β\beta を大きくすると、latent representation はより制約されます。その結果、disentanglement は改善される可能性がありますが、reconstruction quality は悪化することがあります。

つまり、Beta-VAE には次の trade-off があります。

  • β\beta が小さい場合には、reconstruction が良くなりやすいです。
  • β\beta が大きい場合には、latent representation が disentangled になりやすいです。

数式で見る β-VAE

β-VAE は、VAE の KL 項に重み β\beta をかけます。

Lβ-VAE=Eqϕ(zx)[logpθ(xz)]βDKL(qϕ(zx)p(z))\mathcal{L}_{\beta\text{-}VAE} =\mathbb{E}_{q_\phi(z\mid x)}[\log p_\theta(x\mid z)] -\beta D_{KL}(q_\phi(z\mid x)\|p(z))

β>1\beta>1 にすると、latent posterior は prior により強く近づけられます。この式の気持ちは、「復元性能を少し犠牲にしてでも、latent space をより整理された、独立性の高い表現にしたい」というものです。

ただし、β\beta を大きくしすぎると posterior collapse が起こり、decoder が latent を使わなくなることがあります。Disentanglement と reconstruction quality は trade-off になりやすいです。

関連ページ