论文笔记CATEGORICAL REPARAMETERIZATION WITH GUMBEL-SOFTMAX

Gumbel-Softmax分布

Gumbel-Softmax分布是一个定义在单纯形(simplex)上的连续分布。
Gumbel-Softmax分布可以近似categorical分布。
用$z$表示为服从$\pi =({\pi }_1,\dots ,{\pi }_k)$的categorical随机变量。categorical分布的样本表示为$k$维的one-hot向量，在$k-1$维的单纯形空间${△}^{k-1}$中。

Gumbel-Max trick是重参数技巧(reparametrization tricks)的一个特例，其提供了一个简单有效的从categorical分布采样的方法：
$$z=\text{one-hot}\left({\mathrm{argmax}}_i[g_i+\log {\pi }_i]\right)\text{\hspace{1em}(1)}$$其中$g_i\sim Gumbel(0,1)$。Gumbel分布用于对各种分布的多个样本的最大值（或最小值）的分布进行建模。Gumbel分布的概率密度是：
$$Gumbel(\mu ,\beta )=\frac{1}{\beta }\exp (-\frac{x-\mu }{\beta }+\exp (-\frac{x-\mu }{\beta }))$$
上面的argmax是不可导的，用可导的softmax函数去近似公式(1)中的argmax，可以得到样本$y\in {△}^{k-1}$：
$$y_i=\mathrm{softmax}[g_i+\log {\pi }_i]=\frac{\exp (\frac{\log {\pi }_i+g_i}{\tau })}{{\sum }_{j=1}^k\exp (\frac{\log {\pi }_j+g_j}{\tau })}$$$y_i$服从Gumbel-Softmax分布。Gumbel-Softmax分布的概率密度函数是：

随着$\tau$趋近于0，Gumbel-Softmax分布的样本逐渐变成one-hot的，Gumbel-Softmax分布也逐渐变成了categorical分布。如下图所示：

Gumbel-Softmax Estimator

Gumbel-Softmax分布的$\frac{\partial y}{\partial \pi }$是有定义的。
通过用Gumbel-Softmax样本替换categorical样本，我们可以使用反向传播来计算梯度。
把在训练阶段，用可导的Gumbel-Softmax样本替代不可导的categorical样本的过程称为Gumbel-Softmax Estimator。

在温度$\tau$小时，样本接近one-hot但梯度方差大，在温度$\tau$大时，样本平滑但梯度方差小。
实际中，我们从高温$\tau$开始，然后退火到一个很小但非零的温度。

Straight-Through (ST) Gumbel-Softmax Estimator

Straight-Through Estimator (STE)

首先介绍下Straight-Through Estimator (STE)。
STE是量化（quantization）中常见的求导方式。
比如有sign函数：
$$w_b=\mathrm{sign}(w)=\{\begin{array}{l}+1,\text{ if }w\ge 0\\ -1,\text{ otherwise }\end{array}$$这个sign函数在定义域范围内导数都是0。STE就是用来解决sign函数梯度无法反传的问题的。
二值网络训练过程可以是这样：模型中每个参数其实都是一个浮点型的数，每次迭代其实都是在更新这个浮点型数。但是，在前向传播的过程中，先用sign函数对浮点型参数二值化处理然后再参与到运算，而此时并没有把这个浮点型数值抛弃掉，而是暂时在内存中保存起来。前向传播完之后，网络得到一个输出，就可以接着通过反向传播算出二值参数的梯度，再直接用这个梯度来更新对应的浮点型参数。这样，前向反向就跑通了。等训练的差不多了，就最后对模型的这些浮点型参数做一次二值化处理形成最终的二值网络，此时浮点型的参数就完成了任务，可以被抛弃掉了。

Straight-Through (ST) Gumbel-Softmax Estimator

在前向的时候使用argmax离散化$y$，但在梯度反传的时候，使用连续近似${\nabla }_{\theta }z\approx {\nabla }_{\theta }y$。

参考

ICLR 2017 Categorical Reparameterization with Gumbel-Softmax
Emma Benjaminson blog
二值网络，围绕STE的那些事儿
gumbel-max-trick的数学证明