I am confused by these lines of code:


def sample(self):
    u = tf.random_uniform(tf.shape(self.logits))
    return U.argmax(self.logits - tf.log(-tf.log(u)), axis=1)

This is supposed to sample from a categorical distribution.

You can ignore the tf that prepends the commands (these are basically tensorflow commands)

The function receives a vector of logits.

The first line takes the shape of the logits vector (self.logits) and samples a vector of independent random values from a uniform distribution on $[0,1]$.

The second line is what really confuses me. Why the $\textsf{logits} - \log(-\log(x))$ and substracting from the logits?

My first intuitive approach would have been going from logits to probability.. What is going on there?

  • 1
    $\begingroup$ What is a "vector of logits"? $\endgroup$
    – Clement C.
    Commented Jul 27, 2018 at 20:00
  • $\begingroup$ A vector where each element is a logit: en.m.wikipedia.org/wiki/Logit $\endgroup$
    – Juan Leni
    Commented Jul 27, 2018 at 20:17

1 Answer 1


Let me unravel your question to remove all the not-so-relevant fluff around it.

Given a tuple of $n$ values of the form $(p_i)_{1\leq i\leq n}$, where each $p_i\in(0,1)$, the question is to characterize the distribution of the following process:

  1. Draw independent random variables $U_1,\dots,U_n$ uniformly distributed in $[0,1]$

  2. Return $$\arg\!\max_{1\leq i\leq n} \left( \log \frac{p_i}{1-p_i} - \log\log\frac{1}{U_i}\right)$$

What we will show:

The output of this algorithm is a categorical random variable $Z$ such that $$ \forall i\in [n], \quad \mathbb{P}\{ Z = i\} \propto \frac{p_i}{1-p_i}\,. \tag{$\dagger$} $$

Detailed proof.

For every $1\leq i\leq n$, let $$X_i \stackrel{\rm def}{=} \log \frac{p_i}{1-p_i} - \log\log\frac{1}{U_i} = - \log\left(\frac{1-p_i}{p_i}\log\frac{1}{U_i}\right)\,.$$

By a standard result, we have that $\log\frac{1}{U_i}$ follows an exponential distribution with parameter $1$, and therefore (by properties of the exponential distribution) this implies that $$ Y_i \stackrel{\rm def}{=}\frac{1-p_i}{p_i}\log\frac{1}{U_i} \sim \mathrm{Exp}(\frac{p_i}{1-p_i})\,.\tag{1}$$

Now, since $$ \arg\!\max_{1\leq i\leq n} X_i = \arg\!\max_{1\leq i\leq n} (-\log Y_i) = \arg\!\min_{1\leq i\leq n} Y_i\tag{2} $$ we have that the output $Z$ of the algorithm has probability mass function $$ \forall i\in[n],\quad\mathbb{P}\{Z=i\} = \mathbb{P}\{ Y_i = \min_{1\leq j\leq n} Y_j\} = \frac{\frac{p_i}{1-p_i}}{\sum_{j=1}^n \frac{p_j}{1-p_j}} \tag{3} $$ using e.g. the result from this other question for the last equality.

In other term,

The output of the algorithm is $i$ with probability proportional to $\frac{p_i}{1-p_i}$.

Alternative. if you don't care about the derivation, here is the quick and simple explanation: this is called the Gumbel trick. If $U$ is uniform on $[0,1]$, then $-\log\log\frac{1}{U}$ has a standard Gumbel distribution. And you can use the following theorem (well-known, I reckon, to whoever knows it)

Theorem. If $G_1,\dots,G_n$ are independent standard Gumbel r.v.'s, and $\alpha_1,\dots, \alpha_n > 0$, then the random variable $$ X = \arg\!\max_{1\leq i\leq n}(\log \alpha_i + G_i) $$ takes values proportional to the $\alpha_i$: $$ \forall i\in [n], \quad \mathbb{P}\{X=i\} \propto \alpha_i\,. $$

See e.g. this for more.

  • $\begingroup$ @purpletentacle You're welcome! I did learn a few things in the process of answering. $\endgroup$
    – Clement C.
    Commented Jul 27, 2018 at 21:22
  • 1
    $\begingroup$ I think it is worth adding a link to this paper: arxiv.org/abs/1611.01144 $\endgroup$
    – Juan Leni
    Commented Jul 28, 2018 at 17:06

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.

Not the answer you're looking for? Browse other questions tagged or ask your own question.