# Sampling from a categorical distribution

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?

• What is a "vector of logits"? Commented Jul 27, 2018 at 20:00
• A vector where each element is a logit: en.m.wikipedia.org/wiki/Logit Commented Jul 27, 2018 at 20:17

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.

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