Gumbel Softmax数学证明及其应用解析

论文地址：

https://arxiv.org/pdf/1611.01144.pdf%20http://arxiv.org/abs/1611.01144.pdf

问题描述：

考虑离散变量 $x$ ，如果已知其分布向量 $\pi =\{{\pi }_1,\dots ,{\pi }_N\}$ ，则得到 $x$ 的采样的一个简单采样方法是：

$$x_{\pi }=one\_hot(\underset{i}{\mathrm{argmax}}({\pi }_i))$$

根据 $softmax$ 函数，其中 $\arg \max$ 这一操作取得 ${\pi }_n$ 的概率为：

$$P({\pi }_n)=\frac{e^{{\pi }_n}}{{\sum }_{j=1}^N{e}^{{\pi }_j}}$$

当我们在BP神经网络中，需要让采样结果可导的时候，这样简单的采样就行不通了。

原因在于 $\arg \max$ 这一操作是不可导的，即并没有一个表达式可以映射 $\pi$ 到 $z$ 上。

该技巧应用甚广，如深度学习中的各种 GAN、强化学习中的 A2C 和 MADDPG 算法等等。

只要涉及在离散分布上运用重参数技巧时（re-parameterization），都可以试试 Gumbel-Softmax Trick。

具体实现：

一般来说，对于 $N$ 维概率向量 $\pi$，我们可以通过添加随机 Gumbel 噪声 $G_i$ 再取样：

$$x_{\pi }=\underset{i}{\mathrm{argmax}}\left(\ln \left({\pi }_i\right)+G_i\right)$$

其中 $G_i$ 是独立同分布的标准 Gumbel 分布的随机变量。

我们重新看一下 Gumbel 分布，Gumbel 分布是一种极值型分布，它的概率密度函数（PDF）为：

$$f(x;\mu ,\beta )=e^{-z-e^{-z}},z=\frac{x-\mu }{\beta }$$

公式中，$\mu$ 是位置系数，$\beta$ 是尺度系数，标准 Gumbel 分布中有：$\mu =0$，$\beta =1$。

相应的，Gmubel 分布的累积密度函数（CDF）为：

$$F(x;\mu ,\beta )=e^{-e^{-(x-\mu )/\beta }}$$

并且我们易得它的反函数：

$$F^{-1}(y;\mu ,\beta )=\mu -\beta \ln (-\ln (y))$$

这样我们就可以通过从均匀分布中求逆得到 $G_i$：

$$G_i=-\ln \left(-\ln \left(U_i\right)\right),U_i\sim U(0,1)$$

这就是 Gmubel-Max trick。

由于上述算法中 $\arg \max$ 这一操作仍是不可导的，因此我们可以用两种方式来让该操作可导。

一种方式是使用Straight-Through Estimator 思想（例如 VQ-VAE 中使用的），重新设计采样为：

$$t_{\pi }=softmax(\ln \left({\pi }_i\right)+G_i)$$

$$z_{\pi }=one\_hot(\arg \max \left(t_{\pi }\right))$$

$$x_{\pi }=t_{\pi }+sg[z_{\pi }-t_{\pi }]$$

其中 $sg$ 是 stop gradient 的意思，这样向前传播的时候使用的是采样的 $z_{\pi }$，而反向传播则是使用 $t_{\pi }$。

但是 Gumbel Softmax 的作者指出并证明，直接用 $softmax$ 函数来代替量化过程也是可行的，即：

$$x_{\pi }=softmax\left(\ln \left({\pi }_i\right)+G_i\right)$$

具体操作为：

1. 对于 $N$ 维概率向量 $\pi$，我们生成 $N$ 个服从均匀分布 $U(0,1)$ 的独立样本 $U_1,\dots ,U_N$；
2. 通过 $G_i=-\ln (-\ln U_i)$ 计算得到 $G_i$；
3. 对应相加得到新的向量 $z$，其中 $z_i={\pi }_i+G_i$；
4. 通过 $softmax$ 函数计算 $x_{\pi }$，其中：

$$x_i=\frac{e^{z_i/\tau }}{{\sum }_{j=1}^N{e}^{z_j/\tau }}$$

前三步的目标是让新的随机变量 $z$ 与原随机变量 $\pi$ 相同，只需要证明取到 $z_n$ 的概率跟取到 ${\pi }_n$ 的概率相同，第四步则是使用温度参数 $\tau$ 来控制采样结果的分布倾向：

1. 当 $\tau$ 越小时，结果 $x_{\pi }$ 就越接近于 one-hot 分布；
2. 当 $\tau$ 越大时，结果 $x_{\pi }$ 就越接近于均匀分布；

下面我们来证明。

数学证明：

证明取到 $z_n$ 的概率跟取到 ${\pi }_n$ 的概率相同可以写为：

$$P\left(z_n\ge z_{n^{\prime }};\forall n^{\prime }\ne n\mid {\left\{{\pi }_{n^{\prime }}\right\}}_{n^{\prime }=1}^N\right)=P({\pi }_n)$$

也就是 $z_n$ 比其他所有 $z_{n^{\prime }}$ 都大的概率为 $P({\pi }_n)$。

根据条件累积概率分布函数，我们可以得到：

$$P\left(z_n\ge z_{n^{\prime }};\forall n^{\prime }\ne n\mid {\left\{{\pi }_{n^{\prime }}\right\}}_{n^{\prime }=1}^N\right)=\prod _{n^{\prime }\ne n}P(z_n\ge z_n^{\prime })$$

注意到，$z_n={\pi }_n+G_n$，并且 $G_n$ 服从 $\mu =0$，$\beta =1$ 的标准 Gumbel 分布，那么将 ${\pi }_n$ 看作常数时，$z_n$ 服从 $\mu ={\pi }_n$，$\beta =1$ 的标准 Gumbel 分布，它的 CDF 为：

$$F_{z_n}(x)=e^{-e^{-(x-{\pi }_n)}}$$

也就是：

$$F_{z_n^{\prime }}(x)=e^{-e^{-(x-{\pi }_{n^{\prime }})}}$$

那么根据 CDF 的定义，我们可得：

$$P(z_n\ge z_n^{\prime })=P(z_n^{\prime }\le z_n)=F_{z_n^{\prime }}(z_n)=e^{-e^{-(z_n-{\pi }_{n^{\prime }})}}$$

即：

$$P\left(z_n\ge z_{n^{\prime }};\forall n^{\prime }\ne n\mid {\left\{{\pi }_{n^{\prime }}\right\}}_{n^{\prime }=1}^N\right)=\prod _{n^{\prime }\ne n}{e}^{-e^{-(z_n-{\pi }_{n^{\prime }})}}$$

同时我们可得 $z_n$ 分布的 CDF 为：

$$f_{z_n}(x)=e^{-(x-{\pi }_n)-e^{-(x-{\pi }_n)}}$$

对 $z_n$ 求积分可得边缘累积概率分布函数：

$$\begin{array}{rl}P(z_n\ge z_{n^{\prime }}& ;\forall n^{\prime }\ne n\mid {\left\{{\pi }_{n^{\prime }}\right\}}_{n^{\prime }=1}^N)\\ =& \int P\left(z_n\ge z_{n^{\prime }};\forall n^{\prime }\ne n\mid {\left\{{\pi }_{n^{\prime }}\right\}}_{n^{\prime }=1}^N\right)\cdot f_{z_n}\left(z_n\right)d{z}_n\end{array}$$

带入 CDF 可得：

$$\begin{array}{rl}P(z_n\ge z_{n^{\prime }}& ;\forall n^{\prime }\ne n\mid {\left\{{\pi }_{n^{\prime }}\right\}}_{n^{\prime }=1}^N)\\ =& \int \prod _{n^{\prime }\ne n}{e}^{-e^{-(z_n-{\pi }_{n^{\prime }})}}\cdot e^{-(z_n-{\pi }_n)-e^{-(z_n-{\pi }_n)}}d{z}_n\end{array}$$

化简可得：

$$\begin{array}{l}P(z_n\ge z_{n^{\prime }};\forall n^{\prime }\ne n\mid \{{\pi }_{n^{\prime }}{\}}_{n^{\prime }=1}^N)\\ =\int {\prod }_{n^{\prime }\ne n}{e}^{-e^{-(z_n-{\pi }_{n^{\prime }})}}\cdot e^{-(z_n-{\pi }_n)-e^{-(z_n-{\pi }_n)}}d{z}_n\\ =\int e^{-{\sum }_{n^{\prime }\ne n}{e}^{-(z_n-{\pi }_n)}-(z_n-{\pi }_n)-e^{-(z_n-{\pi }_n)}}d{z}_n\\ =\int e^{-{\sum }_{n^{\prime }=1}^N{e}^{-(z_n-{\pi }_{n^{\prime }})}-(z_n-{\pi }_n)}d{z}_n\\ =\int e^{-({\sum }_{n^{\prime }=1}^N{e}^{{\pi }_{n^{\prime }}})e^{-z_n}-z_n+{\pi }_n}d{z}_n\\ =\int e^{-e^{-z_n+\ln ({\sum }_{n^{\prime }=1}^N{e}^{{\pi }_{n^{\prime }}})}-z_n+{\pi }_n}d{z}_n\\ =\int e^{-e^{-(z_n-\ln ({\sum }_{n^{\prime }=1}^N{e}^{{\pi }_{n^{\prime }}}))}-(z_n-\ln ({\sum }_{n^{\prime }=1}^N{e}^{{\pi }_{n^{\prime }}}))-\ln ({\sum }_{n^{\prime }=1}^N{e}^{{\pi }_{n^{\prime }}})+{\pi }_n}d{z}_n\\ =e^{-\ln ({\sum }_{n^{\prime }=1}^N{e}^{{\pi }_{n^{\prime }}})+{\pi }_n}\int e^{-e^{-(z_n-\ln ({\sum }_{n^{\prime }=1}^N{e}^{{\pi }_{n^{\prime }}}))}-(z_n-\ln ({\sum }_{n^{\prime }=1}^N{e}^{{\pi }_{n^{\prime }}}))}d{z}_n\\ =\frac{e^{{\pi }_n}}{{\sum }_{n^{\prime }=1}^N{e}^{{\pi }_{n^{\prime }}}}\int e^{-e^{-(z_n-\ln ({\sum }_{n^{\prime }=1}^N{e}^{{\pi }_{n^{\prime }}}))}-(z_n-\ln ({\sum }_{n^{\prime }=1}^N{e}^{{\pi }_{n^{\prime }}}))}d{z}_n\\ =\frac{e^{{\pi }_n}}{{\sum }_{n^{\prime }=1}^N{e}^{{\pi }_{n^{\prime }}}}\int e^{-(z_n-\ln ({\sum }_{n^{\prime }=1}^N{e}^{{\pi }_{n^{\prime }}}))-e^{-(z_n-\ln ({\sum }_{n^{\prime }=1}^N{e}^{{\pi }_{n^{\prime }}}))}}d{z}_n\end{array}$$

注意到积分内为符合 $\mu =\ln \left({\sum }_{k^{\prime }=1}^K{e}^{x_{k^{\prime }}}\right)$ 的 Gumbel 分布，所以积分的结果为 $1$，即：

$$P(z_n\ge z_{n^{\prime }};\forall n^{\prime }\ne n\mid \{{\pi }_{n^{\prime }}{\}}_{n^{\prime }=1}^N)=\frac{e^{{\pi }_n}}{{\sum }_{n^{\prime }=1}^N{e}^{{\pi }_{n^{\prime }}}}=P({\pi }_n)$$

代码实现：

在 pytorch 中已经给出其实现：

def gumbel_softmax(logits: Tensor, tau: float = 1, hard: bool = False, eps: float = 1e-10, dim: int = -1) -> Tensor:

    if has_torch_function_unary(logits):
        return handle_torch_function(gumbel_softmax, (logits,), logits, tau=tau, hard=hard, eps=eps, dim=dim)
    if eps != 1e-10:
        warnings.warn("`eps` parameter is deprecated and has no effect.")

    gumbels = (
        -torch.empty_like(logits, memory_format=torch.legacy_contiguous_format).exponential_().log()
    )  # ~Gumbel(0,1)
    gumbels = (logits + gumbels) / tau  # ~Gumbel(logits,tau)
    y_soft = gumbels.softmax(dim)

    if hard:
        # Straight through.
        index = y_soft.max(dim, keepdim=True)[1]
        y_hard = torch.zeros_like(logits, memory_format=torch.legacy_contiguous_format).scatter_(dim, index, 1.0)
        ret = y_hard - y_soft.detach() + y_soft
    else:
        # Reparametrization trick.
        ret = y_soft
    return ret

注意这里生成的 $G_i$ （对应代码中 gumbels）并没有采用论文中的方式，而是用了其等价方式，即直接从指数分布中采样（tensor.exponential_()），然后再取负对数。

这是因为 $-\ln U_i,U\sim U(0,1)$ 的分布符合指数分布，证明如下：

设 $Y=-ln(X)$，且 $X\sim U(0,1)$，有：

$$\begin{array}{l}{F}_y(Y)=P(Y<y)=P(-\ln x<y)=P(x>e^{-y})=1-P(x\le e^{-y})=1-F_x(e^{-y})=1-e^{-y}\end{array}$$

对其求导可得，其概率密度函数为：

$$f_Y(y)=e^{-y}$$

刚好为指数分布。