TF笔记:小trick之gumbel softmax

Posted 墨客无言

tags:

篇首语:本文由小常识网(cha138.com)小编为大家整理,主要介绍了TF笔记:小trick之gumbel softmax相关的知识,希望对你有一定的参考价值。

0. 引言

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

import tensorflow as tf

def 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))进行描述,即:

a r g m a x ( v ⃗ ) = ∑ i n i ∗ δ ( i − u ) argmax(\\vec{v}) = \\sum_{i}^{n}{i * \\delta(i-u)} argmax(v )=iniδ(iu)

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

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

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

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

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

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

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

σ ′ ( x ) = e x i / δ ∑ j e x j / δ \\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 tf

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

emmm,简单过头了……

嘛,那啥,simple is best!

3. 参考链接

  1. 漫谈重参数:从正态分布到Gumbel Softmax

以上是关于TF笔记:小trick之gumbel softmax的主要内容,如果未能解决你的问题,请参考以下文章

小trick之mklink

PHP源码分析之parse_url()的2个小trick

在 Matlab 中从 Gumbel 分布中绘制随机数

Tensorflow2.0笔记

Gumbel softmax在可微NAS的作用是什么?

深度学习刷SOTA的一堆trick