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

论文地址：

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

问题描述：

考虑离散变量 $x$ ，如果已知其分布向量 π = { π 1 , … , π N } \pi=\{\pi_1,\dots,\pi_N\} π={ π1​,…,π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 ( z n ≥ z n ′ ; ∀ n ′ ≠ n ∣ { π n ′ } n ′ = 1 N ) = P ( π n ) P\left(z_{n} \geq z_{n^{\prime}} ; \forall n^{\prime} \neq n \mid\left\{\pi_{n^{\prime}}\right\}_{n^{\prime}=1}^{N}\right)=P(\pi_n) P(zn​≥zn′​;∀n′​=n∣{ πn′​}n′=1N​)=P(πn​)

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

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

P ( z n ≥ z n ′ ; ∀ n ′ ≠ n ∣ { π n ′ } n ′ = 1 N ) = ∏ n ′ ≠ n P ( z n ≥ z n ′ ) P\left(z_{n} \geq z_{n^{\prime}} ; \forall n^{\prime} \neq n \mid\left\{\pi_{n^{\prime}}\right\}_{n^{\prime}=1}^{N}\right)=\prod_{n^{\prime} \neq n}P(z_n\geq z_n^{\prime}) P(zn​≥zn′​;∀n′​=n∣{ πn′​}n′=1N​)=n′​=n∏​P(zn​≥zn′​)

注意到，$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 ( z n ≥ z n ′ ; ∀ n ′ ≠ n ∣ { π n ′ } n ′ = 1 N ) = ∏ n ′ ≠ n e − e − ( z n − π n ′ ) P\left(z_{n} \geq z_{n^{\prime}} ; \forall n^{\prime} \neq n \mid\left\{\pi_{n^{\prime}}\right\}_{n^{\prime}=1}^{N}\right)=\prod_{n^{\prime} \neq n}e^{-e^{-(z_n-\pi_{n^{\prime}})}} P(zn​≥zn′​;∀n′​=n∣{ πn′​}n′=1N​)=n′​=n∏​e−e−(zn​−πn′​)

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

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

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

P ( z n ≥ z n ′ ; ∀ n ′ ≠ n ∣ { π n ′ } n ′ = 1 N ) = ∫ P ( z n ≥ z n ′ ; ∀ n ′ ≠ n ∣ { π n ′ } n ′ = 1 N ) ⋅ f z n ( z n ) d z n \begin{aligned} P(z_{n} \geq z_{n^{\prime}} &; \forall n^{\prime} \neq n \mid\left\{\pi_{n^{\prime}}\right\}_{n^{\prime}=1}^{N})\\=& \int P\left(z_{n} \geq z_{n^{\prime}} ; \forall n^{\prime} \neq n \mid\left\{\pi_{n^{\prime}}\right\}_{n^{\prime}=1}^{N}\right) \cdot f_{z_{n}}\left(z_{n} \right) d z_{n} \end{aligned} P(zn​≥zn′​=​;∀n′​=n∣{ πn′​}n′=1N​)∫P(zn​≥zn′​;∀n′​=n∣{ πn′​}n′=1N​)⋅fzn​​(zn​)dzn​​

带入 CDF 可得：

P ( z n ≥ z n ′ ; ∀ n ′ ≠ n ∣ { π n ′ } n ′ = 1 N ) = ∫ ∏ n ′ ≠ n e − e − ( z n − π n ′ ) ⋅ e − ( z n − π n ) − e − ( z n − π n ) d z n \begin{aligned} P(z_{n} \geq z_{n^{\prime}} &; \forall n^{\prime} \neq n \mid\left\{\pi_{n^{\prime}}\right\}_{n^{\prime}=1}^{N})\\=& \int \prod_{n^{\prime} \neq n}e^{-e^{-(z_n-\pi_{n^{\prime}})}} \cdot e^{-(z_n-\pi_n)-e^{-(z_n-\pi_n)}} d z_{n} \end{aligned} P(zn​≥zn′​=​;∀n′​=n∣{ πn′​}n′=1N​)∫n′​=n∏​e−e−(zn​−πn′​)⋅e−(zn​−πn​)−e−(zn​−πn​)dzn​​

化简可得：

P ( z n ≥ z n ′ ; ∀ n ′ ≠ n ∣ { π n ′ } n ′ = 1 N ) = ∫ ∏ n ′ ≠ n e − e − ( z n − π n ′ ) ⋅ e − ( z n − π n ) − e − ( z n − π n ) d z n = ∫ e − ∑ n ′ ≠ n e − ( z n − π n ) − ( z n − π n ) − e − ( z n − π n ) d z n = ∫ e − ∑ n ′ = 1 N e − ( z n − π n ′ ) − ( z n − π n ) d z n = ∫ e − ( ∑ n ′ = 1 N e π n ′ ) e − z n − z n + π n d z n = ∫ e − e − z n + ln ⁡ ( ∑ n ′ = 1 N e π n ′ ) − z n + π n d z n = ∫ e − e − ( z n − ln ⁡ ( ∑ n ′ = 1 N e π n ′ ) ) − ( z n − ln ⁡ ( ∑ n ′ = 1 N e π n ′ ) ) − ln ⁡ ( ∑ n ′ = 1 N e π n ′ ) + π n d z n = e − ln ⁡ ( ∑ n ′ = 1 N e π n ′ ) + π n ∫ e − e − ( z n − ln ⁡ ( ∑ n ′ = 1 N e π n ′ ) ) − ( z n − ln ⁡ ( ∑ n ′ = 1 N e π n ′ ) ) d z n = e π n ∑ n ′ = 1 N e π n ′ ∫ e − e − ( z n − ln ⁡ ( ∑ n ′ = 1 N e π n ′ ) ) − ( z n − ln ⁡ ( ∑ n ′ = 1 N e π n ′ ) ) d z n = e π n ∑ n ′ = 1 N e π n ′ ∫ e − ( z n − ln ⁡ ( ∑ n ′ = 1 N e π n ′ ) ) − e − ( z n − ln ⁡ ( ∑ n ′ = 1 N e π n ′ ) ) d z n \begin{array}{l} P(z_{n} \geq z_{n^{\prime}} ; \forall n^{\prime} \neq n \mid\{\pi_{n^{\prime}}\}_{n^{\prime}=1}^{N})\\ =\int \prod_{n^{\prime} \neq 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} \neq 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} P(zn​≥zn′​;∀n′​=n∣{ πn′​}n′=1N​)=∫∏n′​=n​e−e−(zn​−πn′​)⋅e−(zn​−πn​)−e−(zn​−πn​)dzn​=∫e−∑n′​=n​e−(zn​−πn​)−(zn​−πn​)−e−(zn​−πn​)dzn​=∫e−∑n′=1N​e−(zn​−πn′​)−(zn​−πn​)dzn​=∫e−(∑n′=1N​eπn′​)e−zn​−zn​+πn​dzn​=∫e−e−zn​+ln(∑n′=1N​eπn′​)−zn​+πn​dzn​=∫e−e−(zn​−ln(∑n′=1N​eπn′​))−(zn​−ln(∑n′=1N​eπn′​))−ln(∑n′=1N​eπn′​)+πn​dzn​=e−ln(∑n′=1N​eπn′​)+πn​∫e−e−(zn​−ln(∑n′=1N​eπn′​))−(zn​−ln(∑n′=1N​eπn′​))dzn​=∑n′=1N​eπn′​eπn​​∫e−e−(zn​−ln(∑n′=1N​eπn′​))−(zn​−ln(∑n′=1N​eπn′​))dzn​=∑n′=1N​eπn′​eπn​​∫e−(zn​−ln(∑n′=1N​eπn′​))−e−(zn​−ln(∑n′=1N​eπn′​))dzn​​

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

P ( z n ≥ z n ′ ; ∀ n ′ ≠ n ∣ { π n ′ } n ′ = 1 N ) = e π n ∑ n ′ = 1 N e π n ′ = P ( π n ) P(z_{n} \geq z_{n^{\prime}} ; \forall n^{\prime} \neq 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) P(zn​≥zn′​;∀n′​=n∣{ πn′​}n′=1N​)=∑n′=1N​eπn′​eπn​​=P(π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}$$

刚好为指数分布。