Reparameterization Trick

在 VAE 中，我们要对 ELBO 求关于 $\boldsymbol{\phi}$ 的偏导数 $\mathbb{E}_{q_{\boldsymbol{\phi}}(\bold{z}|\bold{x})} [\boldsymbol{f}(\bold{z})]$，$\bold{z}$ 是连续随机变量， $\bold{z}$ given $\bold{x}$ 的条件密度函数为 $g_{\boldsymbol{\phi}}(\bold{z|x})$。在 VAE 中 $\bold{z}$ 作为 encoder 的输出，其具有灵活性。

Change of variables

我们可以认为 $\bold{z}$ 是在 $\bold{x}$ 和 $\boldsymbol{\phi}$ 给定情况下，由另外一个确定的随机变量 $\boldsymbol{\epsilon} \sim p(\boldsymbol{\epsilon})$ 经过可微和可逆的变换 $\bold{g}$ 得到的，

$$\bold{z} = \bold{g}_{\boldsymbol{\phi}}(\boldsymbol{\epsilon}, \bold{x})$$

$\boldsymbol{\epsilon}$ 和 $\bold{x}$ 和 $\boldsymbol{\phi}$ 是无关的。

忽略 $\bold{x}$ 和 $\boldsymbol{\phi}$ ，因为它们都是给定的。则

$$\bold{z} = \bold{g}(\boldsymbol{\epsilon}) = (g_1(\epsilon_1, ..., \epsilon_n), ..., g_n(\epsilon_1, ..., \epsilon_n))$$

由于 $\bold{g}$ 可逆，因此存在逆映射

$$\boldsymbol{\epsilon}=(\epsilon_1, ..., \epsilon_n) = (h_1(z_1, ..., z_n), ..., h_n(z_1, ..., z_n))=\bold{h}(\bold{z})$$

因此我们有 $\bold{z} = \bold{g} (\bold{h}(\bold{z}))$，根据链式法则：

$$\frac{\partial \bold{z}}{\partial \bold{z}} = \frac{\partial \bold {g}}{\partial \bold{h}} \frac{\partial \bold {h}}{\partial \bold{z}} = \frac{\partial \bold {g}}{\partial \boldsymbol{\epsilon}} \frac{\partial \bold {h}}{\partial \bold{z}}$$

等式左边是单位矩阵 $\boldsymbol{I}$，我们对两边取行列式，则有

$$1 = \det (\boldsymbol{I}) = \det (\frac{\partial \bold {g}}{\partial \boldsymbol{\epsilon}} \frac{\partial \bold {h}}{\partial \bold{z}}) = \det \Big(\frac{\partial \bold {g}}{\partial \boldsymbol{\epsilon}}\Big) · \det \Big(\frac{\partial \bold {h}}{\partial \bold{z}}\Big)$$

$q_{\boldsymbol{\phi}}(\bold{z|x})$ 是随机变量函数的概率密度函数，则有

$$q_{\boldsymbol{\phi}}(\bold{z|x}) = p(\boldsymbol{\epsilon}) \Big| \det \Big( \frac{\partial \bold {h}}{\partial \bold{z}} \Big)\Big|$$

对 $d\bold{z}$ 做变量替换 $\bold{z} = \bold{g}_{\boldsymbol{\phi}}(\boldsymbol{\epsilon}, \bold{x})$，则有

$$d\bold{z} = \Big| \det \Big( \frac{\partial \bold {g}}{\partial \boldsymbol{\epsilon}}\Big) \Big| d \boldsymbol{\epsilon}$$

因此

$$q_{\boldsymbol{\phi}}(\bold{z|x}) d\bold{z} = p(\boldsymbol{\epsilon})\Big| \det \Big( \frac{\partial \bold {g}}{\partial \boldsymbol{\epsilon}}\Big) \Big|·\Big| \det \Big( \frac{\partial \bold {h}}{\partial \bold{z}} \Big)\Big| = p(\boldsymbol{\epsilon}) d \boldsymbol{\epsilon}$$

Monte Carlo estimte under change of variables

在使用了上面的变量替换之后，对于任何一个随机变量函数 $\boldsymbol{f}(\bold{z})$，我们有

$$\begin{aligned} \mathbb{E}_{q_{\boldsymbol{\phi}}(\bold{z}|\bold{x})} [\boldsymbol{f}(\bold{z})] &= \int \boldsymbol{f}(\bold{z}) q_{\boldsymbol{\phi}}(\bold{z}|\bold{x}) d \bold{z}\\ &= \int \boldsymbol{f}(\bold{g}_{\boldsymbol{\phi}}(\boldsymbol{\epsilon}, \bold{x})) p(\boldsymbol{\epsilon}) d \boldsymbol{\epsilon}\\ &= \mathbb{E}_{p(\boldsymbol{\epsilon})} [\boldsymbol{f}(\bold{z})] \end{aligned}$$

期望的 Monte Carlo estimte：

$$\begin{aligned} \mathbb{E}_{q_{\boldsymbol{\phi}}(\bold{z}|\bold{x})} [\boldsymbol{f}(\bold{z})] &= \mathbb{E}_{p(\boldsymbol{\epsilon})} [\boldsymbol{f}(\bold{z})]\\ &= \mathbb{E}_{p(\boldsymbol{\epsilon})}[\boldsymbol{f}(\bold{z})] \\ &\simeq \frac{1}{L} \sum_{l=1}^L \boldsymbol{f}(\bold{g}_{\boldsymbol{\phi}}(\boldsymbol{\epsilon}^{(l)}, \bold{x})) \end{aligned}$$

其中 $\boldsymbol{\epsilon}^{(l)} \sim p(\boldsymbol{\epsilon})$.

如上所示，利用 reparameterization 可以用来重写关于 $q_{\boldsymbol{\phi}}(\bold{z|x})$ 的期望，使得对期望的 Monte Carlo estimate 关于 $\boldsymbol{\phi}$ 可微。

Example

单变量高斯分布 $z\sim p(z|x) = \mathcal{N}(\mu, \sigma^2)$，其中 $\mu$ 和 $\sigma$ 由 $x$ 和 $\phi$ 决定，根据上述对变换 $\bold{g}$ 的要求，令 $z = \mu + \sigma \epsilon$，其中 $\epsilon \sim \mathcal{N}(0, 1)$，则

$$\begin{aligned} \mathbb{E}_{\mathcal{N}(\mu, \sigma^2)}[f(z)] &= \mathbb{E}_{\mathcal{N}(\mu, \sigma^2)}[f(\mu + \sigma \epsilon)] \\ &\simeq \sum_{l=1}^L f(\mu + \sigma \epsilon^{(l)}) \\ \end{aligned}$$

其中 $\epsilon^{(l)} \sim \mathcal{N}(0, 1)$.