重参数 (Reparameterization)

基本概念

• 重参数 (Reparameterization) 实际上是处理如下期望形式的目标函数的一种技巧：
  上面的期望式可能在如下情形中出现 (e.g. VAE)：假设我们在模型的前向传播过程中得到了随机变量 $Z$ 的概率分布 $p_{\theta }(z)$，其中 $\theta$ 为模型参数，然后需要根据 $p_{\theta }(z)$ 对随机变量 $Z$ 进行采样，再根据采样得到的值 $z$ 完成后续的前向传播过程，例如计算训练损失 $f(z)$. 此时，训练损失 $L_{\theta }$ 即可写为上述期望的形式。然而，这里存在一个很大的问题，就是采样操作是不可导的，虽然我们可以完成模型的前向传播，但反向传播时却无法计算出梯度 $\partial L_{\theta }/\partial \theta$，也就无法进行模型的训练。而 Reparameterization 则是提供了一种变换，使得我们可以直接从 $p_{\theta }(z)$ 中采样，并且保留 $\theta$ 的梯度，也就是将采样操作由不可导变为可导
• 重参数假设从分布 $p_{\theta }(z)$ 中采样可以分解为两个步骤：(1) 从无参数分布 $q(\epsilon )$ 中采样一个 $\epsilon$；(2) 通过变换 $z=g_{\theta }(\epsilon )$ 生成 $z$。那么，上述期望就变成了
  这时候被采样的分布就没有任何参数了，全部被转移到 $f$ 内部了，因此可以采样若干个点，当成普通的 loss 那样写下来了 (上述重参数过程假定 $p_{\theta }(z)={\int }_{g_{\theta }(\epsilon )=z}q(\epsilon )d\epsilon =\int \delta (z-g_{\theta }(\epsilon ))q(\epsilon )d\epsilon$，$\delta (\cdot )$ 是狄拉克函数，因此有 $L_{\theta }={\mathbb{E}}_{z\sim p_{\theta }(z)}[f(z)]=\iint q(\epsilon )\delta (z-g_{\theta }(\epsilon ))f(z)d\epsilon dz=\int q(\epsilon )f(g_{\theta }(\epsilon ))d\epsilon ={\mathbb{E}}_{\epsilon \sim q(\epsilon )}[f(g_{\theta }(\epsilon ))]$)

连续情形

• 简单起见，我们先考虑 $z$ 为连续随机变量的情形：
  在 VAE 中常见的是正态分布 $p_{\theta }(z)=\mathcal{N}\left(z;{\mu }_{\theta },{\sigma }_{\theta }^2\right)$
• 总的来说，连续情形的重参数还是比较简单的。从数学本质来看，重参数是一种积分变换，即原来是关于 $z$ 积分，通过 $z=g_{\theta }(\epsilon )$ 变换之后得到新的积分形式。一个最简单的例子就是正态分布：对于正态分布来说，重参数就是 “从 $N(z;{\mu }_{\theta },{\sigma }_{\theta }^2)$ 中采样一个 $z$” 变成 “从 $N(\epsilon ;0,1)$ 中采样一个 $\epsilon$，然后计算 $\epsilon \times {\sigma }_{\theta }+{\mu }_{\theta }$”，所以
  

离散情形

• 为了突出 “离散”，我们将随机变量 $z$ 换成 $y$，即对于离散情形要面对的目标函数是
  此时，$p_{\theta }(y)$ 是一个 $k$ 分类模型：
  
• 看到上述期望项中的求和，第一反应可能是 “求和？那就求呗，又不是求不了”。的确，对于离散的随机变量，其期望只不过是有限项求和，理论上确实可以直接完成求和再去梯度下降。但是，如果 $k$ 特别大呢？举个例子，假设 $y$ 是一个 100 维的向量，每个元素不是 0 就是 1，那么所有不同的 $y$ 的总数目就是 $2^{100}$，要对这样的 $2^{100}$ 个单项进行求和，计算量是难以接受的 (每一项都需要计算前向传播过程 $f(y)$)。所以，还是需要回到采样上去，如果能够采样若干个点就能得到期望的有效估计，并且还不损失梯度信息，那自然是最好了

Gumbel Max

• 为此，需要先引入 Gumbel Max。假设每个类别的概率是 $p_1,p_2,\dots ,p_k$，那么 Gumbel Max 提供了一种依概率采样类别的方案：
  也就是说，先算出各个概率的对数 $\log p_i$，然后从均匀分布 $U[0,1]$ 中采样 $k$ 个随机数 ${\epsilon }_1,\dots ,{\epsilon }_k$，把 $g_i=-\log (-\log {\epsilon }_i)\sim \text{Gumbel(0,1)}$ 加到 $\log p_i$ 上去，最后把最大值对应的类别抽取出来就行了。由于现在的随机性已经转移到 $U[0,1]$ 上去了，并且 $U[0,1]$ 不带有未知参数，因此 Gumbel Max 就是离散分布的一个重参数过程
• 可以证明，这样的过程精确等价于依概率 $p_1,p_2,\dots ,p_k$ 采样一个类别，换句话说，在 Gumbel Max 中，输出 $i$ 的概率正好是 $p_i$. 不失一般性，这里我们证明输出 1 的概率是 $p_1$. 注意，输出 1 意味着 $\log p_1-log(-log{\epsilon }_1)$ 是最大的，这又意味着：
  $$\begin{array}{rl}& \log p_1-\log (-\log {\epsilon }_1)>\log p_2-\log (-\log {\epsilon }_2)\\ & \log p_1-\log (-\log {\epsilon }_1)>\log p_3-\log (-\log {\epsilon }_3)\\ & ⋮\\ & \log p_1-\log (-\log {\epsilon }_1)>\log p_k-\log (-\log {\epsilon }_k)\end{array}$$不失一般性，我们只分析第一个不等式，化简后得到：
  $${\epsilon }_2<{\epsilon }_1^{p_2/p_1}\le 1$$由于 ${\epsilon }_2\sim U[0,1]$，所以 ${\epsilon }_2<{\epsilon }_1^{p_2/p_1}$ 的概率就是 ${\epsilon }_1^{p_2/p_1}$，这就是固定 ${\epsilon }_1$ 的情况下，第一个不等式成立的概率。那么，所有不等式同时成立的概率是
  $${\epsilon }_1^{p_2/p_1}{\epsilon }_1^{p_3/p_1}\dots {\epsilon }_1^{p_k/p_1}={\epsilon }_1^{(p_2+p_3+\cdots +p_k)/p_1}={\epsilon }_1^{(1/p_1)-1}$$然后对所有 ${\epsilon }_1$ 求平均，就是
  $${\int }_0^1{\epsilon }_1^{(1/p_1)-1}d{\epsilon }_1=p_1$$

Gumbel Softmax

• 我们希望重参数不丢失梯度信息，但是 Gumbel Max 做不到，因为 $\mathrm{argmax}$ 不可导，为此，需要做进一步的近似。首先，留意到在神经网络中，处理离散输入的基本方法是转化为 one hot 形式，包括 Embedding 层的本质也是 one hot 全连接，因此 $\mathrm{argmax}$ 实际上是 $\text{one\_hot}(\mathrm{argmax})$，然后，我们寻求 $\text{one\_hot}(\mathrm{argmax})$ 的光滑近似，它就是 $softmax$. 由此，我们得到 Gumbel Max 的光滑近似版本——Gumbel Softmax：
  其中参数 $\tau >0$ 称为退火参数，它越小输出结果就越接近 one hot 形式 (但同时梯度消失就越严重)。提示一个小技巧，如果 $p_i$ 是 $softmax$ 的输出，那么大可不必先算出 $p_i$ 再取对数，直接将 $\log p_i$ 替换为 $o_i$ 即可：
  
• 跟连续情形一样，Gumbel Softmax 就是用在需要求 ${\mathbb{E}}_{y\sim p_{\theta }(y)}[f(y)]$、且无法直接完成对 $y$ 求和的场景，这时候我们算出 $p_{\theta }(y)$（或者 $o_i$），然后选定一个 $\tau >0$，用 Gumbel Softmax 算出一个随机向量来 $\tilde{y}$，代入计算得到 $f(\tilde{y})$，它就是 ${\mathbb{E}}_{y\sim p_{\theta }(y)}[f(y)]$ 的一个好的近似，且保留了梯度信息
• 注意，Gumbel Softmax 不是类别采样的等价形式，Gumbel Max 才是。而 Gumbel Max 可以看成是 Gumbel Softmax 在 $\tau \to 0$ 时的极限。当 $\tau$ 比较小时，Gumbel Softmax 采样得到的样本接近 one-hot vector，也就比较接近实际的采样情况，但梯度的方差比较大；当 $\tau$ 比较大时，Gumbel Softmax 采样得到的样本比较平滑 (一个平滑的概率向量，向量的各个分量的值都差不多)，但梯度的方差比较小。所以在应用 Gumbel Softmax 时，开始可以选择较大的 $\tau$（比如 1），然后慢慢退火到一个接近于 0 的数（比如 0.01），这样才能得到比较好的结果

---

Gumbel Softmax v.s. Softmax

• Gumbel Softmax 通过 $\tau \to 0$ 的退火来逐渐逼近 one hot，相比直接用原始的 Softmax 进行退火，区别在于原始 Softmax 退火只能得到最大值位置为 1 的 one hot 向量，而 Gumbel Softmax 有概率得到非最大值位置的 one hot 向量，增加了随机性，会使得基于采样的训练更充分一些

Straight-Through Gumbel-Softmax Estimator

• 由 Gumbel Softmax 得到的采样样本是实际采样样本的一个近似，它甚至都不在离散变量的取值范围之内，即使 $\tau$ 比较小，Gumbel Softmax 采样得到的样本也只是接近 one-hot vector，而非真正离散化的 one-hot vector. 但总存在那么一些场景，我们只想采样离散值而非连续值 (e.g. RL 中从离散的动作空间中采样)
• 假设 Gumbel Softmax 输出的采样向量为 $y$，为了利用 Gumbel Softmax 采样离散值，我们可以在前向传播时使用 $z=\text{one\_hot}(\mathrm{argmax}y)$ 得到离散的采样值，在反向传播时利用 ${\nabla }_{\theta }z\approx {\nabla }_{\theta }y$，对 ${\nabla }_{\theta }y$ 进行梯度回传：
  $$z=y+sg(\text{one\_hot}(\mathrm{argmax}y)-y)$$其中，$sg$ 为 stop gradient 操作

背后的故事: 梯度估计 (gradient estimator)

• 重参数就这样介绍完了吗？远远没有，重参数的背后，实际上是一个称为 “梯度估计”的 大家族，而重参数只不过是这个大家族中的一员。每年的 ICLR、ICML 等顶会上搜索gradient estimator、REINFORCE 等关键词，可以搜索到不少文章，说明这是个大家还在钻研的课题。要想说清重参数的来龙去脉，也要说些梯度估计的故事

SF 估计 (Score Function Estimator)

• 前面我们分别讲了连续型和离散型的重参数，都是在 “loss 层面” 讲述的，也就是说都是想办法把 loss 显式地定义好，剩下的交给框架自动求导、自动优化就是了。而事实上，就算不能显式地写出 loss 函数，也不妨碍我们对它求导，自然也不妨碍我们去用梯度下降了。比如 Score Function Estimator：
  这是对原来损失函数的最朴素的估计，在强化学习中 $z$ 代表着策略，那么上式就是一个最基本的策略梯度，所以有时候也直接称上述估计为叫 REINFORCE。现在我们可以直接从 $p_{\theta }(z)$ 中采样若干个点来估算 $\partial L_{\theta }/\partial \theta$ 的值了，不用担心会不会没梯度
• 同时注意到，重参数技巧要求 $f$ 可导，但是在诸如强化学习的场景下，$f(z)$ 对应着奖励函数，很难做到光滑可导，此时就必须使用 SF 估计

梯度方差

• SF 估计看上去很美好，得到了一个连续和离散变量都适用的估计式，那为什么还需要重参数呢？主要的原因是：SF 估计的方差太大。SF 估计是函数 $f(z)\frac{\partial }{\partial \theta }\log p_{\theta }(z)$ 在分布 $p_{\theta }(z)$ 下的期望，我们要采样几个点来算 (理想情况下，希望只采样一个点)，换句话说，我们想用下面的近似
  于是问题就来了：这样的梯度估计方差很大，这导致了我们用梯度下降优化的时候相当不稳定，非常容易崩

降方差

• 重参数就是一种降方差技巧，为此，我们写出重参数后的梯度表达式：
  对比 SF 估计，我们可以直观感知为什么上式方差更小了 (只是一般情况下，并不是绝对成立)：(1) SF 估计中包含了 $\log p_{\theta }(z)$，我们知道，作为一个合理的概率分布，一般都在无穷远处 (即 $\Vert z\Vert \to \infty$）都会有 $p_{\theta }(z)\to 0$，取了 $\log$ 之后反而会趋于负无穷，换句话说，$\log p_{\theta }(z)$ 这一项实际上放大了无穷远处的波动，从而一定程度上增加了方差；(2) SF 估计中包含的是 $f$ 而重参数之后变成了 $\partial f/\partial g$，$f$ 一般是神经网络，而通常我们定义的神经网络模型其实都是 $\mathcal{O}(z)$ 级别的模型，从而我们可以预期它的梯度是 $\mathcal{O}(1)$ 级别的（不严格成立，只能说在平均意义下基本成立），所以相对情况下更平稳一些，因此 $f$ 的方差也比 $\partial f/\partial g$ 的方差要大

References

• 苏剑林. (Jun. 10, 2019). 《漫谈重参数：从正态分布到 Gumbel Softmax 》[Blog post]. Retrieved from https://kexue.fm/archives/6705
• Gumbel Softmax paper: Jang, Eric, et al. “Categorical Reparameterization with Gumbel-Softmax.” 5th International Conference on Learning Representations, ICLR 2017