Differentiable Sampling of Categorical Distributions
Using the CatLog-Derivative Trick

Lennert De Smet 1 Emanuele Sansone 1 Pedro Zuidberg Dos Martires 2
fashion. Unfortunately, inference in probabilistic models
with categorical latent variables is usually computationally intractable due to its combinatorial nature. This intractability often leads to the use of sampling-based approximate inference techniques, which in turn poses problems
to gradient-based learning as sampling is an inherently nondifferentiable process.

Abstract
Categorical random variables can faithfully represent the discrete and uncertain aspects of data aspart of a discrete latent variable model. Learningin such models necessitates taking gradients withrespect 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. This trick forms
the basis of the well-known REINFORCE gradient estimator and its many extensions. While
the Log-Derivative trick allows us to differentiate
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 LogDerivative 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 the estimatesof IndeCateR outperform the state of the art forthe same number of samples.

In order to bypass this non-differentiability, two mainclasses of gradient estimators have been developed. On the
one hand, there is a range of unbiased estimators based on
the Log-Derivative trick and the subsequent REINFORCE
gradient estimator (Williams, 1992). On the other hand, we
have biased estimators that use continuous relaxations towhich the reparametrisation trick (Ruiz et al., 2016) canbe applied, such as the Gumbel-Softmax trick (Jang et al.,2017; Maddison et al., 2017).
A clear advantage of the REINFORCE estimator overrelaxation-based estimators is its unbiased nature. However,
REINFORCE tends to be sample-inefficient and its gradient estimates exhibit high variance in practice. To resolvethese issues, methods have been proposed that modify REINFORCE by, for instance, adding control variates (Richteret al., 2020; Titsias & Shi, 2022). These modified estimators
have been shown to deliver more robust gradient estimates
than standard REINFORCE.
Instead of modifying REINFORCE, we take a different approach and modify the Log-Derivative trick by explicitly
taking into account that we are working with multivariate
categorical distributions. We call this first contribution the
CatLog-Derivative trick. Interestingly, we show that the
CatLog-Derivative trick leads to Rao-Blackwellised estimators (Casella & Robert, 1996), immediately giving us a guaranteed reduction in variance. A similar result was shown by
Tokui & Sato (2017) using the Gumbel-Max reparametrisation for categorical distributions, which we discuss further
in the related work (Section 5). As a second contribution, we
propose IndeCateR (read as ‘indicator’), a gradient estimator for the special case of independent categorical randomvariables. IndeCateR is a hyperparameter-free estimator thatcan be implemented efficiently on modern AI accelerators.
Thirdly, we empirically show that IndeCateR is competitive
with comparable state-of-the-art gradient estimators.

1. Introduction
Categorical random variables naturally emerge in many domains in AI, such as language modelling, reinforcement
learning and neural-symbolic AI (De Smet et al., 2023).
They are compelling because they can faithfully represent
the discrete concepts present in data in a sound probabilistic1
Department of Computer Science, KU Leuven, Belgium 2 Center for Applied Autonomous Systems, Örebro University, Sweden.
Correspondence to: Lennert De Smet<lennert.desmet@kuleuven.be>.

Published at the Differentiable Almost Everything Workshop of the
40 th International Conference on Machine Learning, Honolulu,
Hawaii, USA. July 2023. Copyright 2023 by the author(s).

1The CatLog-Derivative Trickin Equation (4). A growing body of research has beentackling this problem by proposing variance reduction techniques (Grathwohl et al., 2018; Richter et al., 2020; Titsias& Shi, 2022; Tucker et al., 2017). In what follows we will
focus on estimating this second term and drop the first term,since it can be assumed to be unproblematic.

2. Notation and Preliminaries
Throughout this paper, we consider expectations with respect to multivariate categorical probability distributions,which we write asP(1)EX∼p(X) [f (X)] = x∈Ω(X) p(x)f (x),where we assume this expectation to be finite. The symbol
X denotes a random vector (X1 , . . . , XD ) of D categoricalrandom variables while p(X) denotes a multivariate probability distribution. The expression X ∼ p(X) indicates thatthe random vector X is distributed according to p(X). On
the right-hand side of Equation (1) we write the expectation
as an explicit sum over Ω(X), the finite sample space of therandom vector X, using x = (x1 , . . . , xD ) for the specificassignments of the random vector (X1 , . . . , XD ).

3. The CatLog-Derivative Trick
The standard log-derivative trick and its corresponding gradient estimators are applicable to both discrete and continuous
probability distributions. However, this generality limits
their usefulness when it comes to purely categorical random variables. For example, the REINFORCE gradientestimator suffers from high variance when applied to problems involving high-dimensional multivariate categoricalrandom variables. In such a setting there are exponentiallymany possible states to be sampled, which makes it increasingly unlikely that a specific state gets sampled. We now
introduce the CatLog-Derivative trick that reduces the exponential number of states arising in a multivariate categorical
distribution by exploiting the distribution’s factorisation. All
subsequent statements are proven in the appendix.
Theorem 3.1 (CatLog-Derivative Trick). Let p(X) be a multivariate categorical probability distribution that depends on
a set of parameters Λ, then it holds that ∂λ EX∼p(X) [f (X)]is equal toGiven an order of the random variables in X, we can inducea factorisation of the joint probability distribution as follows
QD
p(X) = d=1 p(Xd | X<d ).
(2)Here, X<d denotes the ordered set of random variables(X1 , . . . , Xd−1 ). Similarly, X>d will denote the orderedset (Xd+1 , . . . , XD ) in subsequent sections.
When performing gradient-based learning, we are interested in partial derivatives of the expected value in 1, i.e.,∂λ EX∼p(X) [f (X)] . Here, we take the partial derivative of
the expectation with respect to the parameter λ and assume
that the distribution p(X) and the function f (X) depend on
a set of parameters Λ with λ ∈ Λ. For probability distributions to which the reparametrisation trick does not apply, we
can rewrite the partial derivative using the Log-Derivative
trick.

DX

EX<d ∼p(X<d ) [∂λ p(xδ | X<d )(5)

d=1 xδ ∈Ω(Xd )EX>d ∼p(X>d |xδ ,X<d ) [f (X̸=d , xδ )]].
Intuitively, the CatLog-Derivative trick decomposes the
log-derivative trick into an explicit sum of multiple LogDerivative tricks, one for each of the categorical randomvariables present in the multivariate distribution. We show
next that this decomposition is effectively equivalent toRao-Blackwellising (Casella & Robert, 1996) gradient estimators.
Definition 3.2 (The CateR gradient estimator). We define
the Categorical REINFORCE (CateR) estimator via theexpression

Theorem 2.1 (Log-Derivative Trick (Williams, 1992)). Letp(X) be a probability distribution and f (X) such that its
expectation is finite, with both functions depending on a setof parameters Λ. Then, it holds that
∂λ EX∼p(X) [f (X)] =X(3)EX∼p(X) [∂λ f (X)] + EX∼p(X) [f (X)∂λ log p(X)] .
In general, both expectations in Equation (3) are intractableand often estimated with a Monte Carlo scheme. The mostimmediate such estimation is provided by the REINFORCE
gradient estimator (Williams, 1992)∂λ EX∼p(X) [f (X)] ≈D
X∂λ EX∼p(X) [f (X)] ≈
(4)

PN
1
(n)
(n)
(n)
) + f (x )∂λ log p(x ) .
n=1 ∂λ f (x
N

(6)

N

X

d=1 xδ ∈Ω(Xd )

1 X
(n )
(n )
(n )
∂λ p(xδ | x<dd )f (x<dd , xδ , x>dd ),N n =1d(n )The superscript on x(n) denotes that it is the nth sample
vector drawn from p(X).

where the sample x>dd is drawn while conditioning on(n )xδ and x<dd . The subscript on nd indicates that different
samples are drawn for every d.

A well-known problem with the REINFORCE gradient estimator is the high variance stemming from the second term

Proposition 3.3. The CateR estimator Rao-Blackwellises
REINFORCE.
2

The CatLog-Derivative Trick

Corollary 3.4 (Bias and Variance). The CateR estimator
is unbiased and its variance is upper-bounded by REINFORCE.

∂λ p2 (x2 )f (x1 , x2 , x3 ). We see that for REINFORCE we
sample all the variables whereas for IndeCateR we performthe explicit sum for each of the random variables in turn andonly sample the remaining variables.

Proof. This follows trivially from Proposition 3.3, the lawof total expectation and the law of total variance (Blackwell,1947; Radhakrishna Rao, 1945).

Note how, in the case of D = 1, Equation (7) reduces
to the exact gradient. With this in mind, we can interpret
IndeCateR as computing exact gradients for each single
random variable Xd with respect to an approximation of thefunction EX̸=d ∼p(X̸=d ) [f (X̸=d , Xd )].

4. The IndeCateR Gradient Estimator
Using the CatLog-Derivative trick derived in the previous
section we are now going to study a prominent special caseof multivariate categorical distributions. That is, we willassume that our probability Q
distribution admits the indepenDdent factorisation p(X) = d=1 pd (Xd ). Note that all D
different distributions still depend on the same set of learnable parameters Λ. Furthermore, we subscript the individual
distributions pd as they can no longer be distinguished bytheir conditioning sets. Plugging in this factorisation into
Theorem 3.1 gives us the Independent Categorical REINFORCE estimator, or IndeCateR for short.

Computational Complexity Consider Equation (7) andobserve that none of the random variables has a sample
space larger than K = maxd (|Ω(Xd )|). Computing our
gradient estimate requires performing three nested sums
with lower bound 1 and upper bounds equal to D, K and
N , respectively. These summations result in a time complexity of O(D · K · N ). Leveraging modern AI accelerators, they can be parallelised to obtain a time complexity ofO(log D+log K +log N ), which allows for the deployment
of IndeCateR in modern deep architectures.

Proposition 4.1 (IndeCateR). Let p(X) be a multivariate
categorical probability distribution that depends
QD on a setof parameters Λ and factorises as p(X) = d=1 pd (Xd ),then the gradient of the expectation EX∼p(X) [f (X)] canbe estimated withDX

5. Related Work
The work closest related to ours is the RAM estimator introduced by Tokui & Sato (2017). The general idea is to first
reparametrise the probability distributions such that theyno longer depend on any parameters and to then perform amarginalization. We show in Section 3 that this reparametrisation step is unnecessary. Avoiding reparametrisation has
the major advantage that we explicitly retain the conditional
dependency structure in the CateR estimator, which allows
us to trivially build a special purpose estimator for (conditionally) independent distributions (IndeCateR). Moreover,
Tokui & Sato (2017) did not study the setting of a sharedparameter space Λ between distributions, although this being the most common setting in modern deep-discrete andneural-symbolic architectures. We rectify this omission andshow that efficiently implemented Rao-Blackwellised gradient estimators for categorical random variables are a viableoption in practice orthogonal to variance reduction schemes
based on control variates.

NX

d=1 xδ ∈Ω(Xd )∂λ pd (xδ )

1 X
(n )
f (x̸=dd , xδ ),N n =1

(7)

d

(n )

where x̸=dd are samples drawn from p(X̸=d ).
Example 4.2 (Independent Factorisation). Let us consider amultivariate distribution involving three binary, independent
categorical random variables. Concretely, this gives usp(X1 , X2 , X3 ) = p1 (X1 )p2 (X2 )p3 (X3 ),(8)

where X1 , X2 and X3 can take values from the set Ω(X) ={0, 1}. Taking this specific distribution and plugging it into
Equation (7) for the IndeCateR gradient estimator now gives
us
P3 P
PN
(nd )
1
d=1
xδ ∈{0,1} ∂λ pd (xδ ) N
nd =1 f (x̸=i , xδ ). (9)Such variance reduction methods for REINFORCE aim to
reduce the variance by subtracting a mean-zero term, called
the baseline, from the estimate (Bengio et al., 2013; Paisley et al., 2012). Progress in this area is mainly driven bymulti-sample, i.e., sample-dependent baselines (Grathwohl
et al., 2018; Titsias & Shi, 2022; Tucker et al., 2017), andleave-one-out baselines (Kool et al., 2019; 2020; Mnih &Rezende, 2016; Richter et al., 2020). Variance can further
be reduced by coupling multiple samples and exploiting
their dependencies (Dimitriev & Zhou, 2021; Dong et al.,2021; Yin et al., 2020). A general drawback of baseline variance reduction methods is that they often involve a certain

In order to understand the difference between the LogDerivative trick and the CatLog-Derivative trick, we aregoing to look at the term for d = 2 and consider the singlesample estimate∂λ p2 (0)f (x1 , 0, x3 ) + ∂λ p2 (1)f (x1 , 1, x3 ),(10)

where x1 and x3 are sampled values for the random variables X1 and X3 . These samples would be different ones
for d ̸= 2. The corresponding complete single sample estimate using REINFORCE instead of IndeCateR would be
3

The CatLog-Derivative Trick

IndeCateR

IndeCateR-D

RLOO-F

RLOO-S

Reference

1.0

1.0

1.0

0.8

0.8

0.8

0.8

0.6

0.6

0.6

0.6

0.4

0.4

0.4

0.4

0.2

0.2

0.2

0.2

Sum Accuracy

1.0

0.0

10 1

100

Epochs

101

102

0.0 1
10

100

Epochs

101

0.0

102

100

101

Epochs

102

0.0

100

101

Epochs102

Figure 1: Test set accuracy of predicting the correct sum value versus number of epochs for the MNIST addition. From left
to right, the plots show curves for 4, 8, 12, and 16 MNIST digits. The dashed, grey line at the top of each plot represents ahypothetical sum-classifier that has access to a 99% accurate MNIST digit classifier. The x-axis is in log-scale.
amount of computational overhead, usually in the form oflearning optimal parameters or computing statistics. This
computational overhead can be justified by assuming thatthe function in the expectation is costly to evaluate. In suchcases, it might pay off to perform extra computations if thatmeans the expensive function is evaluated fewer times.

In the field of neural-symbolic AI such problems are either
solved exactly (Manhaeve et al., 2018) or by simplifying
the underlying combinatorial structure (Huang et al., 2021;
Manhaeve et al., 2021). Exact methods scale very poorly
while simplifying the combinatorics introduces problematic
biases. In contrast, we will study the MNIST addition problem using sampling and unbiased gradient estimators. The
general architecture is as follows. Each of the D different
MNIST images is passed through a neural classifier, which
gives probabilities for each class. These probabilities areused to sample a number between 0 and 9 for each image.
The numbers are summed up and compared to the label
using a binary cross-entropy loss.

Another popular, yet diametrically opposite, approach to low
variance gradient estimates for categorical random variables
is the concrete distribution (Maddison et al., 2017). The
concrete distribution is a continuous relaxation of the categorical distribution using the Gumbel-Softmax trick (Janget al., 2017). Its main drawback is that it results in biased
gradient estimates. Even though this bias can be controlled
by a temperature parameter, the tuning of this parameter isa highly non-trivial matter in practice.

Using IndeCateR in a neural-symbolic setting we achieve
two things. On the one hand, we use the sampling in IndeCateR as an unbiased search, replacing the usual symbolic
search. On the other, we render this stochastic search differentiable by estimating gradients instead of performing thecostly exact computation.

Both approaches can additionally benefit from RaoBlackwellisation, as illustrated by Liu et al. (2019) for REINFORCE and Paulus et al. (2021) for relaxed methods.
However, these works still only focus on the univariate case.
In contrast, we perform Rao-Blackwellisation by exploiting
multivariate dependencies. In general, reducing variance by
exploiting dependencies is orthogonal to much of the existing literature. Hence, it might prove interesting to examine
whether it is possible to combine existing techniques with
our work for a further reduction in variance.

We compare IndeCateR to REINFORCE leave-one-out(RLOO) (Kool et al., 2019; Salimans & Knowles, 2014)as it is a strong representative of methods reducing variancefor REINFORCE (Figure 1). IndeCateR is given 10 samples
leading to D · 10 · 10 function evaluations. Two instances of
RLOO are used, RLOO-F and RLOO-S. The former uses
D · 10 · 10 samples and function evaluations while the latter
only uses 10. IndeCateR-D is an augmented version of IndeCateR that takes 10 new samples for each of the D terms
in Equation (7), which would not be possible following thework of Tokui & Sato (2017). We see that IndeCateR-D isthe only method capable of scaling and solving the MNIST
addition for 16 digits.

6. Experiments
We study the behaviour of IndeCateR on a problem fromthe neural-symbolic literature, the addition of MNIST digits (Manhaeve et al., 2018). Given a set of D MNIST digits,
the task is to predict the sum of the digits. The only provided supervision is the correct sum and no direct label ofany digit. The difficulty of the problem scales exponentially,as there are 10D states in the sample space. There are only
10D +1 possible labels, resulting in very sparse supervision.
4

The CatLog-Derivative Trick

LeCun, Y. The mnist database of handwritten digits.
http://yann. lecun. com/exdb/mnist/, 1998.

Acknowledgements
This research is funded by TAILOR, a project from the EU
Horizon 2020 research. It was also supported by the Wallenberg AI, Autonomous Systems and Software Program(WASP) funded by the Knut and Alice Wallenberg Foundation. We also have to acknowledge support from Flanders
AI, FWO and the KU Leuven Research Fund. Emanuele
Sansone is partially funded by the KU Leuven Research
Fund (C14/18/062) and the Flemish Government (AI Research Program).

Liu, R., Regier, J., Tripuraneni, N., Jordan, M., andMcauliffe, J. Rao-blackwellized stochastic gradients fordiscrete distributions. ICML, 2019.
Maddison, C. J., Mnih, A., and Teh, Y. W. The concrete
distribution: A continuous relaxation of discrete randomvariables. ICLR, 2017.
Manhaeve, R., Dumancic, S., Kimmig, A., Demeester, T.,
and De Raedt, L. Deepproblog: Neural probabilistic logic
programming. advances in neural information processing
systems, 31, 2018.

R EFERENCES
Bengio, Y., Léonard, N., and Courville, A. Estimatingor propagating gradients through stochastic neurons forconditional computation. arXiv, 2013.

Manhaeve, R., Marra, G., and De Raedt, L. Approximate
inference for neural probabilistic logic programming. KR,
2021.

Blackwell, D. Conditional expectation and unbiased sequential estimation. The Annals of Mathematical Statistics,
1947.

Mnih, A. and Rezende, D. Variational inference for montecarlo objectives. International Conference on Machine
Learning, 2016.

Casella, G. and Robert, C. P. Rao-blackwellisation of sampling schemes. Biometrika, 1996.

Niepert, M., Minervini, P., and Franceschi, L. Implicit
mle: backpropagating through discrete exponential family
distributions. NeurIPS, 2021.

De Smet, L., Martires, P. Z. D., Manhaeve, R., Marra, G.,
Kimmig, A., and De Readt, L. Neural probabilistic logic
programming in discrete-continuous domains. UAI, 2023.

Paisley, J., Blei, D. M., and Jordan, M. I. Variational
bayesian inference with stochastic search. ICML, 2012.

Dimitriev, A. and Zhou, M. Carms: Categorical-antitheticreinforce multi-sample gradient estimator. NeurIPS,
2021.
Dong, Z., Mnih, A., and Tucker, G. Coupled gradientestimators for discrete latent variables. NeurIPS, 2021.

Paulus, M. B., Maddison, C. J., and Krause, A. Raoblackwellizing the straight-through gumbel-softmax gradient estimator. ICLR, 2021.

Grathwohl, W., Choi, D., Wu, Y., Roeder, G., and Duvenaud,
D. Backpropagation through the void: Optimizing control
variates for black-box gradient estimation. ICLR, 2018.

Radhakrishna Rao, C. Information and accuracy attainable
in the estimation of statistical parameters. Bulletin of the
Calcutta Mathematical Society, 1945.

Huang, J., Li, Z., Chen, B., Samel, K., Naik, M., Song,L., and Si, X. Scallop: From probabilistic deductive
databases to scalable differentiable reasoning. NeurIPS,
2021.

Richter, L., Boustati, A., Nüsken, N., Ruiz, F., and Akyildiz,
O. D. Vargrad: a low-variance gradient estimator forvariational inference. NeurIPS, 2020.
Rolfe, J. T. Discrete variational autoencoders. ICLR, 2017.

Jang, E., Gu, S., and Poole, B. Categorical reparameterization with gumbel-softmax. ICLR, 2017.

Ruiz, F. J. R., Titsias, M. K., and Blei, D. M. The generalized reparameterization gradient. NeurIPS, 2016.

Kool, W., van Hoof, H., and Welling, M. Buy 4 reinforce
samples, get a baseline for free! ICLR Deep RL Meets
Structured Prediction Workshop, 2019.

Salimans, T. and Knowles, D. A. On using control variateswith stochastic approximation for variational bayes andits connection to stochastic linear regression. arXiv, 2014.

Kool, W., van Hoof, H., and Welling, M. Estimating gradients for discrete random variables by sampling without
replacement. ICLR, 2020.

Titsias, M. and Shi, J. Double control variates for gradientestimation in discrete latent variable models. AISTATS,
2022.

Lake, B. M., Salakhutdinov, R., and Tenenbaum, J. B.
Human-level concept learning through probabilistic program induction. Science, 350(6266):1332–1338, 2015.

Tokui, S. and Sato, I. Evaluating the variance of likelihoodratio gradient estimators. ICML, 2017.
5

The CatLog-Derivative Trick

Tucker, G., Mnih, A., Maddison, C. J., Lawson, J., and SohlDickstein, J. Rebar: Low-variance, unbiased gradientestimates for discrete latent variable models. NeurIPS,
2017.
Williams, R. J. Simple statistical gradient-following algorithms for connectionist reinforcement learning. Reinforcement learning, 1992.
Wu, J., Hu, W., Xiong, H., Huan, J., Braverman, V., andZhu, Z. On the noisy gradient descent that generalizes assgd. In ICML. PMLR, 2020.
Xiao, H., Rasul, K., and Vollgraf, R. Fashion-mnist: a
novel image dataset for benchmarking machine learningalgorithms. arXiv preprint arXiv:1708.07747, 2017.
Yin, M., Ho, N., Yan, B., Qian, X., and Zhou, M. Probabilistic best subset selection via gradient-based optimization.
arXiv, 2020.

6

The CatLog-Derivative Trick

A. Proofs of Theorems
A.1. Proof of Theorem 3.1
Theorem 3.1 (CatLog-Derivative Trick). Let p(X) be a multivariate categorical probability distribution that depends on aset of parameters Λ, then it holds that ∂λ EX∼p(X) [f (X)] is equal to
D
X

X

(5)

EX<d ∼p(X<d ) [∂λ p(xδ | X<d )d=1 xδ ∈Ω(Xd )EX>d ∼p(X>d |xδ ,X<d ) [f (X̸=d , xδ )]].
Proof. We start by applying the standard log-derivative trick and fill in the product form of the categorical distribution
followed by pulling this product out of the logarithm and the expectationhi
QDEX∼p(X) [f (X)∂λ log p(X)] = EX∼p(X) f (X)∂λ log d=1 p(Xd | X<d ) ,(11)PD= d=1 EX∼p(X) [f (X)∂λ log p(Xd | X<d )] .
(12)
To continue, we write out the expectation explicitly and write the sum for the random variable Xd separately, resulting in
PD P
P
(13)
d=1
x∈Ω(X̸=d )xδ ∈Ω(Xd ) p(x̸=d , xδ )f (x̸=d , xδ )∂λ log p(xδ | x<d ).
Next, we factorize the joint probability p(x̸=d , xδ ) as p(x>d | xδ , x<d )p(xδ | x<d )p(x<d ). Multiplying the second ofthese factors with ∂λ log p(xδ | x<d ) gives us ∂λ p(xδ | x<d ). Finally, plugging ∂λ p(xδ | x<d ) into Equation (13) gives thestatement for the theorem.
A.2. Proof of Proposition 3.3
Proposition A.1. The CateR estimator Rao-Blackwellises REINFORCE.
Proof. We start from Equation (4) for the REINFORCE estimator, where we ignore the first term and factorize the probability
distribution similar to Equation (12)DN∂λ EX∼p(X) [f (X)] ≈

N

X 1 X
1 X
(n)
f (x(n) )∂λ log p(x(n) ) ≈f (x(n) )∂λ log p(xd ).
N n=1
N n=1

(14)

d=1

For notational conciseness, we dropped the subscript on nd and simply use n to identify single samples. Now we compare
Equation (6) and Equation (14) term-wise for N → ∞
N
X

n=1

N
h
i X
(n)
(n)
(n)
(n)
(n)
(n)
(n)
f (x<d , xd , x>d )∂λ log p(xd ).
EX ∼p(X |x(n) ) f (x<d , Xd , x>d )∂λ log p(Xd |x<d ) =
d

d

<d

n=1

In the equation above, we see that for CateR we take the expected value for Xd (left-hand side) and compute it exactly using
an explicit sum over the space Ω(Xd ), whereas for REINFORCE (right-hand side) we use sampled values. This means,in turn, that on the left we have the Rao-Blackwellised version of the right-hand side. Doing this for every d gives us a
Rao-Blackwellised version for REINFORCE, i.e., the CateR estimator.
A.3. Proof of Proposition 4.1
Proposition A.2 (IndeCateR). Let p(X) be a multivariate categorical probability distribution that depends on a set of
QDparameters Λ and factorises as p(X) = d=1 pd (Xd ), then the gradient of the expectation EX∼p(X) [f (X)] can be
estimated with
D
X

N

X

∂λ pd (xδ )d=1 xδ ∈Ω(Xd )

1 X
(n )
f (x̸=dd , xδ ),N n =1d(n )where x̸=dd are samples drawn from p(X̸=d ).
7(7)The CatLog-Derivative Trick
0.5

0.5
Bias

0.4

Variance

0.5

10

1

10

1

10

2

10

2

10

0.4

4

10

0.3

4

0.3
10

8

10

8

0.2

0.2

0.1

0.1

0.1

10
IndeCateR

REINFORCE

RLOO

16

10

0.0

GS

1

10

2

10

4

10

8

10

16

0.3

0.2

0.0

10

0.4

IndeCateR(a) D = 12 and |Ω(Xd )| = 3

REINFORCE

RLOO

GS

16

0.0
IndeCateR(b) D = 6 and |Ω(Xd )| = 10

REINFORCE

RLOO

GS

(c) D = 3 and |Ω(Xd )| = 100

Figure 2: We report the (empirical) bias and variance for the different estimators and distributions in comparison to the exact
gradient. Bias and variance were computed using sample means and averaging them over 1000 runs.
Proof. We start by looking at the expression in Equation (5). Using the fact that we have a set of independent randomvariables, we can simplify p(xδ | X<d ) to pd (xδ ). As a result, the gradient of the expected value can be rewritten as∂λ EX∼p(X) [f (X)] =

D
X

X



EX<d ∼p(X<d ) ∂λ pd (xδ )EX>d ∼p(X>d ) [f (X̸=d , xδ )]

(15)

D
X

X

∂λ pd (xδ )EX̸=d ∼p(X̸=d ) [f (X̸=d , xδ )](16)

d=1 xδ ∈Ω(Xd )=d=1 xδ ∈Ω(Xd )Drawing N samples for the D − 1 independent random variables X̸=d and for each term in the sum over d then gives us theestimate stated in the proposition.

B. Additional Experiments
Apart from the provided neural-symbolic experiment, we performed two more traditional experiments from the gradientestimation literature. A synthetic study of the quality of the IndeCateR gradients and the optimisation of a discrete variational
auto-encoder (DVAE).
B.1. Synthetic: Exact Gradient Comparison
For small enough problems the gradient for multivariate categorical random variables can be computed exactly via explicit
enumeration. Inspired by Niepert et al. (2021), we compare the estimates ofhPiD∂θ EX∼p(X)(17)
d=1 |Xd − bd |QD
to its exact value. Here, θ are the logits that directly parametrise a categorical distribution p(X) = d=1 p(Xd ) and bd
denotes an arbitrarily chosen element of Ω(Xd ). We compare the gradient estimates from IndeCateR, REINFORCE, RLOO,
and Gumbel-Softmax (GS) by varying the number of distributions D and the cardinality of the distributions.
In Figure 2 we show the (empirical) bias and variance for the different estimators. Each estimator was given 1000 samples,while IndeCateR was only given a single one. Hence, IndeCateR has the fewest function evaluations as D · K is smaller
than 1000 for each configuration. IndeCateR offers gradient estimates close to the exact ones with orders of magnitude
lower variance for all three settings. RLOO exhibits the smallest difference in bias, yet it can not compete in terms ofvariance. Furthermore, the computation times were of the same order of magnitude for all methods. This is in stark contrast
to the estimator presented by Tokui & Sato (2017), where a two-fold increase in computation time of RAM with respect to
REINFORCE is reported.
B.2. Synthetic: Optimisation
We now study an optimization setting (Titsias & Shi, 2022), where the goal is to maximise the expected value
h P
i
D
1
2
EX∼p(X) D,
i=1 (Xi − 0.499)
8(18)

The CatLog-Derivative Trick

IndeCateR

RLOO-F

RLOO-S

GS-F
10 8

0.25080.250810 9Function Value

Gradient Variance0.2510Function Value

0.25100.250610 100.25060.250410 110.25040.25020.25020.25000.2500
0

250

500

750

1000

1250

Iterations

1500

1750

2000

10 12
10 13

0

1

2

3

Time (s)

4

5

6

0

250

500

750

1000

1250

Iterations

1500

1750

2000

Figure 3: We plot the function value for different estimators against iterations (left) and time (middle). On the right, we plot
the variance of the gradients against iterations. Statistics were obtained by taking the average and standard error over 10
runs. IndeCateR and RLOO-S both use 2 samples, while RLOO-F and Gumbel-Softmax (GS) use 800 samples. The numberof function evaluations is equal for IndeCateR, RLOO-F, and GS-F. We performed a hyperparameter search for the learning
rate and the temperature of GS-F. Parameters were optimised using RMSProp.

where p(X) factorizes into D independent binary random variables. The true maximum is given by p(Xd = 1) = 1 forall d. This task is challenging because of the small impact of the individual values for each Xd on the expected value
for higher values of D. We set D to 200 and report the results in Figure 3, where we compare IndeCateR to RLOO andGumbel-SoftMax.
In Figure 3 and subsequent figures we use the notation RLOO-F and RLOO-S, which we define as follows. If IndeCateR
takes N samples, then it performs D · K · N functional evaluations with K = maxd |Ω(Xd )|. As such, we define RLOO-S
as drawing the same number of samples as IndeCateR, which translates to N function evaluations. For RLOO-F we matchthe number of function evaluations, which means that it takes D · K · N samples. We give an analogous meaning to GS-S
and GS-F for the Gumbel-SoftMax gradient estimator.
IndeCateR distinguishes itself by having both the lowest variance and quickest convergence across all methods, evencompared to RLOO-F. Additionally, the time to compute all gradient estimates does not differ significantly for the different
methods and leads to the same conclusions. It is striking to see that the Gumbel-Softmax estimator struggles in this task,
which is likely due to its bias in combination with the sensitive loss function.
B.3. Discrete Variational Auto-Encoder
As a third experiment we analyse the ELBO optimisation behaviour of a discrete variational auto-encoder (DVAE) (Rolfe,2017). We optimise the DVAE on the three main datasets from the literature, being MNIST (LeCun, 1998), F-MNIST (Xiaoet al., 2017) and Omniglot (Lake et al., 2015). The encoder component of the network has three dense hidden layers ofsizes 384 and 256 ending in a latent 200-dimensional Bernoulli variable. The decoder takes samples from this variable
followed by hidden layers of size 256, 384 and 784. IndeCateR again uses two samples, hence we can compare to the same
configurations of RLOO and Gumbel-Softmax as in Section B.2, i.e., equal samples (GS-S and RLOO-S) and equal function
evaluations (GS-F and RLOO-F).
As evaluation metrics, we show the negated training and test set ELBO in combination with the variance of the gradients
throughout training. We opted to report all metrics in terms of computation time (Figure 4), but similar results in terms of
iterations are given in the appendix.
A first observation is that IndeCateR performs remarkably well in terms of convergence speed. It beats all other methods on
both MNIST and F-MNIST in terms of training ELBO, only having to yield to the Gumbel-Softmax trick with equal function
evaluations on Omniglot. However, we can observe a disadvantage of the quick convergence in terms of generalisation
performance when looking at the test set ELBO. RLOO-F and IndeCateR exhibit overfitting on the training data for MNIST,
resulting in an overall higher negative test set ELBO. We speculate that the relaxation for the Gumbel-Softmax or the higher
variance (Wu et al., 2020) of other methods act as a regulariser for the network. As one would ideally like to separate
9The CatLog-Derivative Trick

IndeCateR

RLOO-SGS-S

RLOO-FGS-F

102

9 × 101

Gradient Variance

Test ELBO

Training ELBO

10 1

10 2

102

9 × 101

10 3
8 × 101

0 1000 2000 3000 4000 5000 6000 7000 8000

Time(s)0 1000 2000 3000 4000 5000 6000 7000 8000

2.65 × 102

2.65 × 102

2.6 × 102

2.6 × 102

Test ELBO

2.7 × 102

Training ELBO

2.7 × 102

2.55 × 102

Time(s)10 1

2.55 × 102

2.5 × 1022.45 × 102

2.4 × 102

2.4 × 102

2000

4000

Time(s)6000

2.35 × 102

8000

1.6 × 102

1.5 × 102

1.5 × 102

1.4 × 102

1.4 × 102

Training ELBO

1.6 × 102

2000

4000

Time(s)

1.3 × 102

1.2 × 102

1.1 × 102

102

102
0

500

1000 1500 2000 2500 3000 3500

Time(s)

8000

0

2000

4000

Time(s)

6000

8000

10 2

1.2 × 102

1.1 × 102

6000

10 1

Test ELBO

1.3 × 102

10 3
0Gradient Variance

0

10 2

2.5 × 1022.45 × 102

2.35 × 102

0 1000 2000 3000 4000 5000 6000 7000 8000

Time(s)Gradient Variance

8 × 101

10 3
0

500

1000 1500 2000 2500 3000 3500

Time(s)05001000 1500 2000 2500 3000 3500

Time(s)Figure 4: The top row shows training ELBO, test ELBO and gradient variance for the DVAE on MNIST – each plottedagainst time. Rows 2 and 3 show the same plots for F-MNIST (middle) and Omniglot (bottom).
discrete sampling layers from network regularisation, we see this overfitting as a feature and not a drawback of IndeCateR.
IndeCateR additionally is competitive in terms of gradient variance, especially considering the number of samples that are
drawn. Across all datasets, the gradient variance of IndeCateR is comparable to that of RLOO-F, even though IndeCateR
takes only 2 samples compared to the 800 of RLOO-F. The only method with consistently lower gradient variance is the
Gumbel-Softmax trick, but only for equal function evaluations. However, in that case the Gumbel-Softmax trick takes ≈ 1.6
times longer than IndeCateR to compute a its gradient estimates.

10