Gumbel-Softmax

Gumbel-Softmax 松弛是一种把离散采样近似成可微连续采样的方法，常用于需要对类别变量做反向传播的模型里，比如可微路由。

核心问题：从 categorical 分布采样一个 one-hot 类别：

这个采样操作不可导，梯度很难直接传回类别概率 。

Gumbel-Softmax 的做法是先用 Gumbel-Max trick 表示离散采样：

其中

然后把不可导的 arg⁡max 换成可导的 softmax：

这里的 是温度参数。

然后还有一个改进版本 Straight-Through Gumbel-Softmax。

前向传播时：把 soft 输出变成 hard one-hot，用最大的那个类别继续计算。

反向传播时：假装刚才用的是 soft 输出，让梯度从 soft 版本传回来。