<blockquote>
<p>I wonder whether should I take <code>log</code> of <code>attn_weights</code> before passing it into <code>gumbel_softmax</code>?</p>
</blockquote>
<p>If <code>attn_weights</code> are probabilities (sum to 1; e.g., output of a softmax), then yes. Otherwise, no.</p>
<blockquote>
<p>I wonder how to choose <code>tau</code> in <code>gumbel_softmax</code>, any guidelines?</p>
</blockquote>
<p>Usually, it requires tuning. The references provided in the docs can help you with that.</p>
<p>From <a href="https://arxiv.org/pdf/1611.01144.pdf" rel="nofollow noreferrer">Categorical Reparameterizaion with Gumbel-Softmax</a>:</p>
<ul>
<li><p>Figure 1, caption:</p>
<blockquote>
<p>... (a) For low temperatures (τ = 0.1, τ = 0.5), the expected value of a Gumbel-Softmax random variable approaches the expected value of a categorical random variable with the same logits. As the temperature increases (τ = 1.0, τ = 10.0), the expected value converges to a uniform distribution over the categories.</p>
</blockquote>
</li>
<li><p>Section 2.2, 2nd paragraph (emphasis mine):</p>
<blockquote>
<p>While Gumbel-Softmax samples are differentiable, they are not identical to samples from the corresponding categorical distribution for non-zero temperature. For learning, <strong>there is a tradeoff between small temperatures</strong>, where samples are close to one-hot but the variance of the gradients is large, <strong>and large temperatures</strong>, where samples are smooth but the variance of the gradients is small (Figure 1). <strong>In practice, we start at a high temperature and anneal to a small but non-zero temperature</strong>.</p>
</blockquote>
</li>
<li><p>Lastly, they remind the reader that tau can be learned:</p>
<blockquote>
<p>If τ is a learned parameter (rather than annealed via a fixed
schedule), this scheme can be interpreted as entropy regularization (Szegedy et al., 2015; Pereyra et al., 2016), where the Gumbel-Softmax distribution can adaptively adjust the &quot;confidence&quot; of proposed samples during the training process.</p>
</blockquote>
</li>
</ul>
input for torch.nn.functional.gumbel_softmax
<p>Say I have a tensor named <code>attn_weights</code> of size [1,a], entries of which indicate the attention weights between the given query and |a| keys. I want to select the largest one using <code>torch.nn.functional.gumbel_softmax</code>.</p>
<p>I find <a href="https://pytorch.org/docs/stable/nn.functional.html?highlight=gumbel%20softmax#torch.nn.functional.gumbel_softmax" rel="nofollow noreferrer">docs about this function</a> describe the parameter as <em><strong>logits</strong></em> <em>- […, num_features] unnormalized log probabilities</em>. I wonder <strong>whether should I take <code>log</code> of <code>attn_weights</code> before passing it into <code>gumbel_softmax</code>?</strong> And I find Wiki defines <code>logit=lg(p/1-p)</code>, which is different from barely logrithm. I wonder which one should I pass to the function?</p>
<p>Further, I wonder <strong>how to choose <code>tau</code></strong> in <code>gumbel_softmax</code>, any guidelines?</p>