# How do I sample from a discrete (categorical) distribution in log space?

Suppose I have a discrete distribution defined by the vector $\theta_0, \theta_1, ..., \theta_N$ such that category $0$ will be drawn with probability $\theta_0$ and so on. I then discover that some of the values in distribution are so small that they underflow my computer's floating point number representation, so, to compensate, I do all my calculations in log-space. Now I have a log-space vector $log(\theta_0), log(\theta_1), ..., log(\theta_N)$.

Is it possible to sample from the distribution such that the original probabilities hold (category $i$ is drawn with probability $\theta_i$) but without ever leaving log-space? In other words, how do I sample from this distribution without underflows?

It is possible to sample from categorical distribution given log-probabilities without leaving log space using the Gumbel-max trick. The idea is that if you are given unnormalized log-probabilities $\alpha_1,\dots,\alpha_k$, that can be translated to proper probabilities using softmax function

$$p_i = \frac{\exp(\alpha_i)}{\sum_j \exp(\alpha_j)}$$

then to sample from such distribution you can use the fact that if $g_1,\dots,g_k \sim \mathcal{G}(0)$ are independent samples taken from standard Gumbel distribution parametrized by location $m$,

$$F(G \le g) = \exp(-\exp(-g+m))$$

then it can be shown (see references below) that

\DeclareMathOperator*{\argmax}{arg\,max} \begin{align} \argmax_i \,\{\, g_i + \alpha_i \,\} &\sim \frac{\exp(\alpha_i)}{\sum_j \exp(\alpha_j)} \\ \max_i\,\{\, g_i + \alpha_i \,\} &\sim \mathcal{G}(\; \log\sum_i\exp\{\alpha_i\}\;) \end{align}

and we can take

$$z = \argmax_i \,\{\, g_i + \alpha_i \,\}$$

as a sample from categorical distribution parametrized by $p_1,\dots,p_k$ probabilities. This approach was described in greater detail in blog entries by Ryan Adams and Laurent Dinh, moreover Chris J. Maddison, Daniel Tarlow and Tom Minka gave a speach (slides) on Neural Information Processing Systems conference (2014) and wrote a paper titled A* Sampling that generalized those ideas (see also Maddison, 2016; Maddison, Mnih and Teh, 2016; Jang and Poole, 2016), who refer to Yellott (1977) mentioning his as the one among those who first described this property.

It is pretty easy to implement it using inverse transform sampling by taking $g_i=-\log(-\log u_i)$ where $u_i$ are draws from uniform distribution on $(0,1)$. It is certainly not the most time-efficient algorithms for sampling from categorical distribution, but it let's you to stay in log-space what may be an advantage in some scenarios.

Maddison, C. J., Tarlow, D., & Minka, T. (2014). A* sampling. [In:] Advances in Neural Information Processing Systems (pp. 3086-3094).

Yellott, J. I. (1977). The relationship between Luce's choice axiom, Thurstone's theory of comparative judgment, and the double exponential distribution. Journal of Mathematical Psychology, 15(2), 109-144.

Maddison, C. J., Mnih, A., & Teh, Y. W. (2016). The Concrete Distribution: A Continuous Relaxation of Discrete Random Variables. arXiv preprint arXiv:1611.00712.

Jang, E., Gu, S., & Poole, B. (2016). Categorical Reparameterization with Gumbel-Softmax. arXiv preprint arXiv:1611.01144.

Maddison, C. J. (2016). A Poisson process model for Monte Carlo. arXiv preprint arXiv:1602.05986.

Here is one common way to avoid underflow/overflow.

Let $m = \max_i \log(\theta_i)$.

Let $\theta_i' = \exp( \log(\theta_i) - m )$.

You can sample from $\theta' = [\theta_1' , \theta_2',...]$.

• This works as long as the difference between any one value and the max isn't too great---when that happens, the exp can lose precision, leading to distributions like [1.0, 3.45e-66, 0.0, 7.54e-121]. I'd like to hold out for some answer that is robust even in that case. But for now I'm upvoting your answer. Jul 16, 2013 at 22:56