question about gumbel softmax

[XiaoyuShi97]

Hi, nice work! I am a bit confused about gumbel softmax. You mention in your paper that, during traininig, gumbel softmax is used. I wonder if it can be replaced by pure softmax (i.e. torch.softmax)? Could you please give more explanation on this design choice? Thx!

[LongguangWang]

Hi @btwbtm, thanks for your interest in our work. Softmax is also used in several network quantization or pruning methods to soften one-hot distributions. In my opinion, softmax may also works in our SMSR but I have not tried it. In our experiments, gumbel softmax is adopted since it is theorically identical to one-hot distribution while softmax is not.

[wangqiim]

Hi @btwbtm, thanks for your interest in our work. Softmax is also used in several network quantization or pruning methods to soften one-hot distributions. In my opinion, softmax may also works in our SMSR but I have not tried it. In our experiments, gumbel softmax is adopted since it is theorically identical to one-hot distribution while softmax is not.

SMSR/model/smsr.py

Lines 12 to 21 in daac49c

	def gumbel_softmax(x, dim, tau):
	gumbels = torch.rand_like(x)
	while bool((gumbels == 0).sum() > 0):
	gumbels = torch.rand_like(x)
	gumbels = -(-gumbels.log()).log()
	gumbels = (x + gumbels) / tau
	x = gumbels.softmax(dim)
	return x

I found the implement of gumbel softmax in your code is different from original paper("Categorical reparameterization with gumbel-softmax"), why do you modify this? which is better?