Reparameterization Trick

November 5, 2023 · 5 min read

AI, CVer, Pythoner, Half-stack Developer

Motivation

假设我们有个在参数 $\theta$ 下的正态分布 $q$ 。我们想要求解下面这样一个问题

\min_{\theta} E_{q}[f(x)]

其中 $E_{q}[f(x)]$ 的意思是求满足 $q$ 分布下的随机变量函数 $f(x)$ 的均值，而最外层的 $\min_{\theta}$ 则是求使得该均值最小时的 $\theta$

有一种做法就是直接对该期望求 $\theta$ 的导数 $\nabla_{\theta} E_{q}\left[f(x)\right]$

\begin{split} \nabla_{\theta} E_{q}\left[f(x)\right] &=\displaystyle \nabla_{\theta} \int q_{\theta}(x) f(x) d x \\ &=\displaystyle \int f(x) \nabla_{\theta} q_{\theta}(x) \frac{q_{\theta}(x)}{q_{\theta}(x)} d x \quad \text{(积分变量是$x$)} \\ &=\displaystyle \int q_{\theta}(x) \nabla_{\theta} \log q_{\theta}(x) f(x) d x\quad \text{(Log Derivative Trick)} \\ &=E_{q}\left[f(x) \nabla_{\theta} \log q_{\theta}(x)\right] \end{split}

tip

上面公式告诉我们，期望的梯度，可以转化为梯度的期望。

于是我们可以利用 $E_{q}\left[f(x) \nabla_{\theta} \log q_{\theta}(x)\right]$ 去估计梯度。

warning

这样的梯度估计式称为 $SF$ 估计，Score Function Estimator，在强化学习中 $q$ 代表着策略，那么上式就是一个基本的策略梯度，有时也叫成 Reinforce

^c5d6c1

并且上述的推导对于任意的随机变量 $x$ 不管是连续还是离散变量，都是通用的。这样我们可以直接从 $p_{\theta}(x)$ 采样若干个点来估算损失函数的梯度了。

既然上述 $SF$ 估计对于连续离散都适用，为什么我们还需要重参数化呢？

主要的原因是： $SF$ 估计的方差太大。

我们将随机变量 $x$ 视为另一个随机变量经过变换 $g_{\theta}(\epsilon)$ 得到的

\begin{split} \epsilon \sim p(\epsilon) \\ x = g_{\theta}(\epsilon) \end{split}

这样我们对于原式可以写为：

E_{q}\left[f(x)\right]=E_{p}\left[f(g_{\theta}(\epsilon))\right]

现在我们对其求梯度：

\nabla_{\theta} E_{q}\left[f(x)\right]=\nabla_{\theta} E_{p}\left[f(g_{\theta}(\epsilon))\right]=E_{p}[\frac{\partial f}{\partial g} \cdot \frac{\partial g}{\partial \theta}]

此时 $\theta$ 分布参数参与了前向过程，所以保留了 $\theta$ 的梯度，使得能够优化参数 $\theta$ 。但是这对于随机变量 $\epsilon$ 有什么要求呢？

info

既然上述都在讲梯度估计，我们自然很关心梯度估计的稳定性，我们不如求一下上面两个公式的方差

info

为了便于求解，我们取 $f(x)=x^2$ ， $q=N(\mu,1)$ 另外 $g(\epsilon) = \mu+ x$ ， $p =N(0,1)$

\begin{cases} \operatorname{Var}[f(x) \nabla_{\theta} \log q_{\theta}(x)] = \displaystyle \mu^4+14\mu^2+15 \\ \displaystyle \operatorname{Var}[\frac{\partial f}{\partial g} \cdot \frac{\partial g}{\partial \theta}] = 4 \end{cases}

所以在 一般情况下 使用从参数化后的梯度估计方差更小，更稳定。（显然你可以找出一个反例，只是说一般情况）

并且对比两个式子：

\begin{cases} \nabla_{\theta} E_{q}\left[f(x)\right] = E_{q}\left[f(x) \nabla_{\theta} \log q_{\theta}(x)\right] \\ \nabla_{\theta} E_{q}\left[f(x)\right] = E_{p}[\frac{\partial f}{\partial g} \cdot \frac{\partial g}{\partial \theta}] \end{cases}

可以看到 $SF$ 估计具有 $\log$ 。我们知道，作为一个合理的概率分布，一般都在无穷远处（即 $\left \| x \right \|\to \infty$ ，都会有 $q_{\theta}(x)\to 0$ ，而 $\log$ 将远处的扰动噪声进行了一定程度的放大，所以方差会大