Gumbel-Softmax trick

Gumbel-Softmax 是一种可导采样技巧，它可以用来对离散分布进行重参数采样。与softmax区别是，gumbel-softmax得到得是样本（即采样,对于神经网络，离散输入常转化为one-hot形式，采样出来的样本是经soft后one-hot向量$[x_1,x_2,...,x_n]$），而softmax得到是样本得均值（即各分类概率,$[p_1,p_2,...,p_n]$）。
假设每个类别得概率是$p_1$, $p_2$,…,$p_k$, 那么下述过程提供了一种依概率采样类别的方案，称为Gumbel Max:
$${\mathrm{argmax}}_i[log{p}_1-log(-log{ϵ}_1),log{p}_2-log(-log{ϵ}_2),...,log{p}_k-log(-log{ϵ}_k)],{ϵ}_i\sim U[0,1]$$
但因为式子中的$argmax$不可导，因此对其进行光滑近似（即soft）。注意到在神经网络中，处理离散输入的基本方法是转化为one-hot形式，即寻求one-hot(arg max)的光滑近似，即softmax.具体可参考苏剑林老师《函数光滑化杂谈：不可导函数的可导逼近》。自此得到Gumbel Max的光滑近似版本–Gumbel Softmax:
$$\begin{gathered} softmax((log{p}_i-log(-log{ϵ}_i))/\tau {)}_{i=1}^k \\ =[\frac{e^{log{p}_1-log(-log{ϵ}_1))/\tau }}{{\sum }_{i=1}^k{e}^{log{p}_i-log(-log{ϵ}_i))/\tau }},\frac{e^{log{p}_2-log(-log{ϵ}_2))/\tau }}{{\sum }_{i=1}^k{e}^{log{p}_i-log(-log{ϵ}_i))/\tau }},...,\frac{e^{log{p}_k-log(-log{ϵ}_k))/\tau }}{{\sum }_{i=1}^k{e}^{log{p}_i-log(-log{ϵ}_i))/\tau }}] \end{gathered}$$
其中$ϵ\sim U[0,1]$，参数τ>0称为退火参数，它越小输出结果就越接近one hot形式（但同时梯度消失就越严重）。