Reparameterization Trick

之前在学蒸馏的时候接触了gumbel-softmax，顺势了解了一下重参数技巧，还是很有意思的一个东西

引入

重参数技巧主要是尝试对这样形式的一个东西求梯度
$$L_{\theta }=E_{z\sim p_{\theta }(z)}[f_{\theta }(z)](1)$$
其中$z\sim p_{\theta }(z)$表示随机变量$z$服从概率密度函数$p_{\theta }(z)$,显然这个密度函数是跟模型参数$\theta$有关的；$f_{\theta }(z)$一般可以表示模型某一层关于变量$z$的输出,显然它也跟模型参数$\theta$有关

不妨先来想想这个式子要如何处理。一个非常naive的思路：采样估计。但是如果直接采样的话，每次采样我们只能获得${\nabla }_{\theta }{f}_{\theta }(z)$,而不同样本之间的信息是无法共用的，我们也就无从得到${\nabla }_{\theta }{L}_{\theta }$。所以我们想想看，有没有什么好的处理方法，能在估计出$(1)$式的同时还能保留梯度信息

不妨先来做一个简化，我们先假设$p_{\theta }(z)$是一个跟$\theta$无关的概率密度函数，简记为$p(z)$,我们很快注意到现在是可以采样估计梯度了：
$$\begin{gathered} {\nabla }_{\theta }{L}_{\theta }={\nabla }_{\theta }{E}_{z\sim p(z)}[f_{\theta }(z)]={\nabla }_{\theta }[{\int }_{z}p(z)f_{\theta }(z)dz] \\ ={\int }_{z}p(z){\nabla }_{\theta }{f}_{\theta }(z)dz \\ =E_{z\sim p(z)}[{\nabla }_{\theta }{f}_{\theta }(z)] \end{gathered}$$
从而
$${\nabla }_{\theta }{L}_{\theta }\approx \frac{1}{n}\sum _{i=1}^n{\nabla }_{\theta }{f}_{\theta }(z_i),z_i\sim p(z)$$
这是因为求梯度的操作成功转移到了$f_{\theta }(z)$上面
上述过程可以用一句话来总结：期望的梯度等于梯度的期望

那我们回到$p_{\theta }(z)$,并尝试类似的步骤：
$$\begin{gathered} {\nabla }_{\theta }{L}_{\theta }={\nabla }_{\theta }{E}_{z\sim p_{\theta }(z)}[f_{\theta }(z)]={\nabla }_{\theta }[{\int }_z{p}_{\theta }(z)f_{\theta }(z)dz] \\ ={\int }_z{p}_{\theta }(z){\nabla }_{\theta }{f}_{\theta }(z)dz+{\int }_z{\nabla }_{\theta }{p}_{\theta }(z)f_{\theta }(z)dz \\ =E_{z\sim p_{\theta }(z)}[{\nabla }_{\theta }{f}_{\theta }(z)]+\underset{???}{\underbrace{{\int }_z{\nabla }_{\theta }{p}_{\theta }(z)f_{\theta }(z)dz}} \end{gathered}$$
前面一块还是可以仿照之前的处理的，但是后者就显得比较诡异了，求梯度操作转移到$p_{\theta }(z)$上面去，也就意味着我们无法将其整理成正常的关于某个东西的期望的形式。或许我们可以将${\nabla }_{\theta }{p}_{\theta }(z)$求出来，但在大部分情况下这是不现实的。

此时就可以引入重参数技巧了

重参数

顾名思义，我们需要引入新的参数来处理上述问题：
考虑一个新的无参数分布
$$ϵ\sim q(ϵ)$$
以及变换
$$z=g_{\theta }(ϵ)$$
保证变换之后得到的$z$服从$p_{\theta }$
那么对$(1)$式求梯度可以变成：
$$\begin{gathered} {\nabla }_{\theta }{L}_{\theta }={\nabla }_{\theta }{E}_{z\sim p_{\theta }(z)}[f_{\theta }(z)] \\ =E_{ϵ\sim q(ϵ)}[f_{\theta }(g_{\theta }(ϵ))](a) \\ =E_{ϵ\sim q(ϵ)}[{\nabla }_{\theta }{f}_{\theta }(g_{\theta }(ϵ))](b) \end{gathered}$$
从而
$${\nabla }_{\theta }{L}_{\theta }\approx \frac{1}{n}\sum _{i=1}^n{\nabla }_{\theta }{f}_{\theta }(g_{\theta }({ϵ}_i)),{ϵ}_i\sim q(ϵ)$$

我们就成功实现了在采样的同时保持了梯度

注意，在这个过程中最重要的一步转化就是：
$$L_{\theta }=E_{ϵ\sim q(ϵ)}[f_{\theta }(g_{\theta }(ϵ))]$$
它将随机性从参数$\theta$转移到了内部无参数的$ϵ$上面，从而可以利用我们之前讨论过的对无参数分布（或者说无可变参数）而言成立的“期望的梯度等于梯度的期望”这一性质来处理

例子

不妨就取$p_{\theta }(z)$是一个正态分布，即
$$p_{\theta }(z)=N({\mu }_{\theta },{\sigma }_{\theta }^2)$$
那么$q(ϵ)$我们就取标准正态分布
$$q(ϵ)=N(0,1)$$
那么显然有
$${\sigma }_{\theta }ϵ+{\mu }_{\theta }\sim N({\mu }_{\theta },{\sigma }_{\theta }^2)$$
所以我们就取
$$g_{\theta }(ϵ)={\sigma }_{\theta }ϵ+{\mu }_{\theta }$$
最后有
$$E_{z\sim N({\mu }_{\theta },{\sigma }_{\theta }^2)}[f_{\theta }(z)]=E_{ϵ\sim N(0,1)}[f_{\theta }({\sigma }_{\theta }ϵ+{\mu }_{\theta })]$$

离散情况的重参数处理

上述过程处理的是分布为连续密度函数的情况，但我们也经常遇到离散分布的情况，这种该如何处理？
为做区分，我们换一种写法：
$$L_{\theta }=E_{y\sim p_{\theta }(y)}[f_{\theta }(y)]=\sum _y{p}_{\theta }(y)f_{\theta }(y)(2)$$
一般来说，此时$y$是可枚举的，它在大部分情况下都对应了一个k分类问题，也就是说，$y$可以表示为
$$p_{\theta }(y)=softmax(o_1,o_2,...o_k{)}_y=\frac{1}{\sum e^{o_i}}{e}^{o_y}(3)$$
其中$o_i$一般就是模型的logits，它当然也是关于参数$\theta$的函数

还是同一个问题，$(2)$式直接用求和的形式是没法计算梯度的，我们还是得试试重参数方法。

所以现在问题就变成了：

找到一个合适的无参数分布$q(ϵ)$以及对应的变换$g_{\theta }(ϵ)$保证它服从$p_{\theta }$这个分布

事实上也确实已经有对应的成果了，它叫做

Gumbel Max

取
$$ϵ\sim U(0,1)$$
对应的$q_{\theta }(ϵ)$为：
$$argma{x}_i(log{p}_i-log(-log{ϵ}_i){)}_{i=1}^k(4)$$
这里第$p_{\theta }(i)$简记为$p_i$了
我们只需证明$(3)$式与$(4)$式是同一个分布，即$(4)$式输出数字$i$的概率为$p_i$

不失一般性地，我们考虑$(4)$式输出数字1的概率：
此时意味着$log{p}_1-log(-log{ϵ}_1)$是$1-k$中最大的，即
$$log{p}_1-log(-log{ϵ}_1)\ge log{p}_i-log(-log{ϵ}_i),\forall i\in (1,k]$$
得到
$${ϵ}_i\le {ϵ}_1^{p_i/p_1}\le 1,\forall i\in (1,k]$$
又$e_i\sim U(0,1)$,从而
$$P({ϵ}_i\le {ϵ}_1^{p_i/p_1})={ϵ}_1^{p_i/p_1},\forall i\in (1,k]$$
从而$(4)$式输出1的概率为
$$P({ϵ}_2\le {ϵ}_1^{p_2/p_1},{ϵ}_3\le {ϵ}_1^{p_3/p_1},...{ϵ}_k\le {ϵ}_1^{p_k/p_1})=\prod _{i=2}^k{ϵ}_1^{p_i/p_1}={ϵ}_1^{(1-p_1)/p_1}$$
对${ϵ}_1$的所有情况求个平均，得到
$${\int }_0^1{ϵ}_1^{(1-p_1)/p_1}d{ϵ}_1=p_1$$
这就是$(4)$式输出1的概率，它恰好为$p_1$
从而我们证明了$(4)$式与$(3)$式确实是同分布，所以我们就成功找到了合理的无参数分布$q(ϵ)$以及对应的变换$g_{\theta }(ϵ)$ $\square$
那么所有过程似乎到这里就圆满结束了。

但是！但是，这里还是有点问题：argmax这个运算本身也是无法求导的…
也就是说，我们将求梯度运算转移到了$argmax$运算上面，结果它还是没有办法求梯度？
不过没关系，这一步其实并不是很难处理。我们知道$argmax$其实可以扩展成$one\_hot(argmax)$,而后者的一个光滑近似就是$softmax$:对于这一点，我相信接触过蒸馏的同学肯定是很清楚的，我们只需要调整蒸馏的温度就能使得$softmax$无限趋近于$ont\_hot$
而$softmax$显然是可以求梯度的，我们就顺利解决了这个遗留的问题。
这种策略被称为

Gumbel Softmax

具体来说，我们的$g_{\theta }(ϵ)$要改成：
$$softma{x}_i((log{p}_i-log(-log{ϵ}_i))/\tau {)}_{i=1}^k(5)$$
其中$\tau$就是蒸馏的温度，当$\tau \to 0$的时候，$softmax$就可以看成$ont\_hot$,当然此时梯度消失现象也会很严重。
由此我们也可以得到训练策略：对参数$\tau$进行退火，最后得到接近于$ont\_hot$形式对应的结果。常见的一个退火策略为：
$${\tau }_p={\tau }_0({\tau }_p/{\tau }_0{)}^{p/P}$$
其中${\tau }_p$是第$p$次训练的温度，${\tau }_0$是初始温度，$P$是总轮数。

---

总结一下，对于总体的$k$个情况，我们从0到1的均匀分布中取$k$个值，利用Gumbel softmax得到一个$k$维向量$\tilde{p}$,
那么
$$\sum _y{\tilde{p}}_y{f}_{\theta }(y)$$
就是$L_{\theta }$的一个良好估计，并且它成功保留了梯度信息

需要指出的是，Gumbel Max是原式的等价形式，但是Gumbel Softmax并不是，它是Gumbel Max的一个光滑近似，当$\tau$足够小的时候，它可以近似看成Gumbel Max

顺便提一嘴这个东西为啥叫Gumbel Max/Softmax：
我们仔细观察$(5)$式：
$$softma{x}_i((log{p}_i-log(-log{ϵ}_i))/\tau {)}_{i=1}^k$$
按照原本的思路，我们可以先从均匀分布里采样$ϵ$，然后再做log运算，再做log运算，再与$log{p}_i$做差，不过实际上实际从一个$-log(-logϵ)$服从的分布里直接采样也是完全OK的，那我们就来看看这个分布长什么样子：
记
$$x=-log(-logϵ)$$
那么
$$F_X(x)=P_X(X\le x)=P_{ϵ}(-log(-logϵ)\le x)=P_{ϵ}(ϵ\le e^{-e^{-x}})=F_{ϵ}(e^{-e^{-x}})$$
从而
$$F_X(x)=exp(-exp(-x))$$
这就是这个分布的累积分布函数，它就被称为Gumbel分布。实际上Gumbel分布还带有另外两个参数
$$F_X(x,\mu ,\beta )=exp(-exp(-\frac{x-\mu }{\beta }))$$
也就是说这里是$\mu =\beta =0$的特殊情况。不过这一点不必细讲，感兴趣的读者可以再去了解一下。

最后讲一个实现细节：
在求原分布$q_{\theta }$的时候，我们需要从 { o i } \{o_i\} {oi​}出发做softmax得到 { p i } \{p_i\} {pi​},但是实际上$(5)$式可以直接替换为
$$softma{x}_i((o_i-log(-log{ϵ}_i))/\tau {)}_{i=1}^k$$
那么我们就不必去做softmax了
至于证明其实也很简单：
$$log{p}_i=log(softmax(o_i))=log(\frac{e^{o_i}}{{\sum }_j{e}^{o_j}})$$
从而
$$log{p}_i=o_i-C$$
从而
$$softmax((log{p}_i+g_i)/\tau )=\frac{e^{(log{p}_i+g_i)/\tau }}{{\sum }_j{e}^{(log{p}_j+g_j)/\tau }}=\frac{e^{(o_i-C+g_i)/\tau }}{{\sum }_j{e}^{(o_j-C+g_j)/\tau }}$$
显然可以将常数$C$对应的部分提出来
$$=\frac{e^{(o_i+g_i)/\tau }}{{\sum }_j{e}^{(o_j+g_j)/\tau }}=softmax((o_i+g_i)/\tau )$$
这里$g_i$就指之前讲的Gumbel分布

总结

以上就是重参数在连续和离散两个场景的应用了，它最初也是最多的应用应该是在VAE里面，我以后应该也会接触，到时候也许会对这篇文章加以补充。