Differentiable Sampling of Categorical Distributions using the CatLog-Derivative Trick

Categorical random variables can faithfully represent the discrete and uncertain aspects of data as part of a discrete latent variable model. For example, the discrete actions in a reinforcement learning application or the tokens in a language model. Learning in such models necessitates taking gradients with respect to the parameters of the categorical probability distributions, which is often intractable due to their combinatorial nature. A popular technique to estimate these otherwise intractable gradients is the Log-Derivative trick that forms the basis of the well-known REINFORCE gradient estimator. While the Log-Derivative trick allows differentiation through samples drawn from categorical distributions, it does not take into account the discrete nature of the distribution itself. Our first contribution addresses this shortcoming by introducing the CatLog- Derivative trick – a variation of the Log-Derivative trick tailored towards categorical distributions. Secondly, we use the CatLog-Derivative trick to introduce IndeCateR, a novel and unbiased gradient estimator for the important case of products of independent categorical distributions with provably lower variance than REINFORCE. Thirdly, we empirically show that IndeCateR can be efficiently implemented and that its gradient estimates have significantly lower bias and variance for the same number of samples compared to the state of the art.