TF笔记：小trick之gumbel softmax

TF笔记：小trick之gumbel softmax
- 0. 引言
- 1. gumbel softmax
- 2. tf代码实现
- 3. 参考链接

0. 引言

故事的起因在于我们在实际工作中遇到的一个小的需求，即我们在模型定义当中需要用到argmax的信息，因此，我们就快速地写下了如下一段代码：

import tensorflow as tfdef get_argmax(x):h = get_shape_list(x)[-1]y = tf.one_hot(tf.argmax(x, axis=-1), h)return y

由此，我们就可以找到tensor当中每一行的最大元素，并使用onehot向量将其表示出来。

但是，在实际的使用中，我们发现了一个问题，即这样定义的模型能够正常工作，但是其训练出来的模型特征表征却和我们的预期大相径庭。

原因相比大多数读者也都注意到了，即我们在这种函数定义当中，由于使用了argmax，使得梯度回传被中断了，这就导致了模型训练失败，无法达到预期的目标。

而要解决这里argmax导致的梯度回传中断的问题，gumbel softmax方法就是一种常用的方法，下面，我们就来对其进行一些简单的介绍。

1. gumbel softmax

gumbel softmax方法的本质在于说用一个连续可导的函数来模拟argmax函数的结果表达，使得其可以在不截断梯度回传的情况下完成argmax函数的功能。

argmax函数的函数曲线可以通过狄拉克函数(δ(x)\delta(x)δ(x))进行描述，即：

argmax(v⃗)=∑ini∗δ(i−u)argmax(\vec{v}) = \sum_{i}^{n}{i * \delta(i-u)} argmax(v)=i∑ni∗δ(i−u)

其中，uuu为向量v⃗\vec{v}v中最大元素的下标。

如果用one-hot向量进行argmax的表达的话，即有其中任一元素的值为δ(i−u)\delta(i-u)δ(i−u)。

由此，我们只需要使用一个连续可导的函数来模拟δ(x−u)\delta(x-u)δ(x−u)函数即可，而对于这个问题，gumbel softmax采用的方式是基于softmax函数进行参数调制的方式进行实现。

基础的softmax函数的表达式如下：

σ(x⃗)=exi∑jexj\sigma(\vec{x}) = \frac{e^{x_i}}{\sum_j e^{x_j}} σ(x)=∑jexjexi

而gumbel softmax函数事实上就是在softmax的基础上加上参数调制。

我们给出gumbel softmax的函数表达式如下：

σ′(x)=exi/δ∑jexj/δ\sigma'(x) = \frac{e^{x_i / \delta}}{\sum_j e^{x_j / \delta}} σ′(x)=∑jexj/δexi/δ

其中，delta为一个小量。

2. tf代码实现

基于此，我们可以比较快速地写出gumbel softmax函数的tf代码了。

import tensorflow as tfdef gumbel_softmax(x, delta=1e-3, axis=None):return tf.nn.softmax(x/delta, axis=axis)

emmm，简单过头了……

嘛，那啥，simple is best！

3. 参考链接

漫谈重参数：从正态分布到Gumbel Softmax