Stop gradients (for ST gumbel softmax)
Hi, Given the soft implementation of Gumbel Softmax: <SCODE>from torch.autograd import Variable
import torch.nn.functional as F

def sample_gumbel(input):
    noise = torch.rand(input.size())
    eps = 1e-20
    noise.add_(eps).log_().neg_()
    noise.add_(eps).log_().neg_()
    return Variable(noise)

def gumbel_softmax_sample(input):
    temperature = 1
    noise = sample_gumbel(input)
    x = (input + noise) / temperature
    x = F.log_softmax(x)
    return x.view_as(input)
<ECODE>
eliabruni
I did stop_gradient(x) in PyTorch as Variable(x.data). However, I’m looking for a better alternative.
Ilya_Kostrikov
apaszke
Thank you both, I’ll try it out!
eliabruni
<SCODE>def gumbel_softmax_sample(self, input):
        temperature = 1
        noise = self.sample_gumbel(input)
        x = (input + noise) / temperature
        x = F.softmax(x)

        if hard == True:
            max_val, _ = torch.max(x, x.dim()-1)
            x_hard = x == max_val.expand_as(x)
            tmp = (x_hard.float() - x)
            tmp2 = tmp.clone()
            tmp2.detach_()
            x = tmp2 + x

        return x.view_as(input)
<ECODE> But, I tried to modify it such that only one ‘1’ exists at each row in x_hard as below. <SCODE>def gumbel_softmax_sample(self, input):
    temperature = 1
    noise = self.sample_gumbel(input)
    x = (input + noise) / temperature
    x = F.softmax(x)

    if hard == True:
        _, max_inx = torch.max(x, x.dim()-1)
        x_hard = torch.cuda.FloatTensor(x.size()).zero_().scatter_(x.dim()-1, max_inx.data, 1.0)
        x2 = x.clone()
        tmp = Variable(x_hard-x2.data)
        tmp.detach_()

        x = tmp + x
        
    return x.view_as(input)
<ECODE> These two scripts work without error, but I am not sure they are equivalent to the original ST gumbel softmax. I would greatly appreciate if someone reviews them.
Seungyoung_Park
Is there anyone helps me to confirm that these two scripts are equivalent to the original tensorflow implementation ?
Seungyoung_Park
ethancaballero
Because I do not know how to make the parameters same as those of the original, I have not yet compared them.
Seungyoung_Park
Qi_Zhang
<SCODE>y = tf.stop_gradient(y_hard - y) + y
<ECODE> y is pure one-hot, in terms of value (since we add the soft y, and then subtract it again the gradients are those of soft y (since all the other terms in this expression have their gradient stripped)
hughperkins
hughperkins