Gumbel-Softmax

分类: 深度学习基础

定义

Gumbel-Softmax（又称 Concrete 分布）是一种可微的离散采样近似方法，通过向 logits 加入 Gumbel 噪声后接 softmax，将离散的 categorical 采样松弛为连续分布，从而支持梯度反向传播。

$y_i = \frac{\exp((\log \pi_i + g_i) / \tau)}{\sum_{j=1}^{k} \exp((\log \pi_j + g_j) / \tau)}$

其中 $g_i \sim \text{Gumbel}(0,1) = -\log(-\log(u_i))$ ， $u_i \sim \text{Uniform}(0,1)$

$\tau$ 为温度参数： $\tau \to 0$ 时趋近 one-hot（离散）； $\tau \to \infty$ 时趋近均匀分布

解决了离散随机变量不可微的根本问题，使得包含离散选择的网络结构可以端到端训练

Straight-Through 变体（ST Gumbel-Softmax）：前向用 argmax，反向用 softmax 梯度

温度退火（temperature annealing）策略：训练过程中逐步降低 $\tau$ ，从探索到利用

在 NAS、剪枝门控、离散 VAE 等场景广泛应用

Gumbel-Sigmoid：二值版本，用于二元门控决策（如 HiAP 中的剪枝门控）

Jang et al., 2017: “Categorical Reparameterization with Gumbel-Softmax”

Maddison et al., 2017: “The Concrete Distribution”