Gumbel Softmax in Quantizer

Thank you for your amazing work.
I just want to make sure I understand the code correctly. The Gumbel Sampling is not necessary here, the Argmin version (line 76) will be exactly the same result, correct?

momask-codes/models/vq/quantizer.py

Lines 76 to 78 in 0421a2b

    
           # code_idx = torch.argmin(distance, dim=-1) 
        
           code_idx = gumbel_sample(-distance, dim = -1, temperature = sample_codebook_temp, stochastic=True, training = self.training)

Hi, we set the sampling temperature here:

momask-codes/models/vq/model.py

Line 72 in 0421a2b

    
           x_quantized, code_idx, commit_loss, perplexity = self.quantizer(x_encoder, sample_codebook_temp=0.5)

So the Gumbel Sampling results would be different from Argmin version. They are from argmin w/ noise.

Thank you for clarification.

	# code_idx = torch.argmin(distance, dim=-1)

	code_idx = gumbel_sample(-distance, dim = -1, temperature = sample_codebook_temp, stochastic=True, training = self.training)