Title: Differentiable Sampling of Categorical Distributions Using the CatLog-Derivative Trick
Venue: NeurIPS 2023
Accepted: yes
Keywords: gradient estimation

categorical random variables

probability theory

discrete distributions
Abstract: Categorical random variables can faithfully represent the discrete and uncertain aspects of data as part of a discrete latent variable 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. 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 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.

Differentiable Sampling of Categorical Distributions Using the CatLog-Derivative Trick

Abstract

Categorical random variables can faithfully represent the discrete and uncertain aspects of data as part of a discrete latent variable 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. 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 tricka 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.

1 Introduction

Categorical random variables naturally emerge in many domains in AI, such as language modelling, reinforcement learning and neural-symbolic AI [10]. They are compelling because they can faithfully represent the discrete concepts present in data in a sound probabilistic 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 non-differentiable process.

In order to bypass this non-differentiability, two main classes 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 [40]. On the other hand, we have biased estimators that use continuous relaxations to which the reparametrisation trick [33] can be applied, such as the Gumbel-Softmax trick [16, 22].

A clear advantage of the REINFORCE estimator over relaxation-based estimators is its unbiased nature. However, REINFORCE tends to be sample-inefficient and its gradient estimates exhibit high variance in practice. To resolve these issues, methods have been proposed that modify REINFORCE by, for instance, adding control variates [31, 37]. 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 our CatLogDerivative trick leads to Rao-Blackwellised estimators [4], immediately giving us a guaranteed reduction in variance. The CatLog-Derivative trick can also be seen as a generalisation of the Local Expectation Gradients (LEG) [36] capable of exploiting more structural, distributional properties. This connection to LEG will be clarified throughout the paper.

As a second contribution, we propose IndeCateR (read as 'indicator'), a gradient estimator for the special case of independent categorical random variables. IndeCateR is a hyperparameterfree estimator that can be implemented efficiently by leveraging parallelisation on modern graphical processing units (GPUs). Thirdly, we empirically show that IndeCateR is competitive with comparable state-of-the-art gradient estimators on a range of standard benchmarks from the literature.

2 Notation and Preliminaries

Throughout this paper, we consider expectations with respect to multivariate categorical probability distributions of the form

where we assume this expectation to be finite. The symbol X denotes a random vector ( X 1 , . . . , X D ) of D categorical random variables while p ( X ) denotes a multivariate probability distribution. The expression X ∼ p ( X ) indicates that the random vector X is distributed according to p ( X ) . On the right-hand side of Equation (2.1) we write the expectation as an explicit sum over Ω( X ) , the finite sample space of the random vector X , using x = ( x 1 , . . . , x D ) for the specific assignments of the random vector ( X 1 , . . . , X D ) .

Given an order of the random variables in X , we can induce a factorisation of the joint probability distribution as follows

Here, X <d denotes the ordered set of random variables ( X 1 , . . . , X d -1 ) . Similarly, X >d will denote the ordered set ( X d +1 , . . . , X D ) in subsequent sections.

When performing gradient-based learning, we are interested in partial derivatives of the expected value in (2.1), i.e., ∂ λ E X ∼ 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.

Theorem 2.1 (Log-Derivative Trick [40]) . Let p ( X ) be a probability distribution and f ( X ) such that its expectation is finite, with both functions depending on a set of parameters Λ . Then, it holds that

In general, both expectations in Equation (2.3) are intractable and often estimated with a Monte Carlo scheme. The most immediate such estimation is provided by the REINFORCE gradient estimator [40]

The superscript on x ( n ) denotes that it is the n th sample vector drawn from p ( X ) .

A well-known problem with the REINFORCE gradient estimator is the high variance stemming from the second term in Equation (2.4). A growing body of research has been tackling this problem by proposing variance reduction techniques [11, 31, 37, 39]. In what follows we will focus on estimating this second term and drop the first term, since it can be assumed to be unproblematic.

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 gradient estimator suffers from high variance when applied to problems involving high-dimensional multivariate categorical random variables. In such a setting there are exponentially many 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.

Theorem 3.1 (CatLog-Derivative Trick) . Let p ( X ) be a multivariate categorical probability distribution that depends on a set of parameters Λ and assume E X ∼ p ( X ) [ f ( X )] is finite. Then, it holds that ∂ λ E X ∼ p ( X ) [ f ( X )] is equal to

̸

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 expectation

To continue, we write out the expectation explicitly and write the sum for the random variable X d separately, resulting in

̸

̸

̸

̸

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 of these factors with ∂ λ log p ( x δ | x <d ) gives us ∂ λ p ( x δ | x <d ) . Finally, plugging ∂ λ p ( x δ | x <d ) into Equation (3.4) gives the desired expression.

Intuitively, the CatLog-Derivative trick decomposes the Log-Derivative trick into an explicit sum of multiple Log-Derivative tricks, one for each of the categorical random variables present in the multivariate distribution. Next, we use the CatLog-Derivative trick to define a novel gradient estimator that exploits the structure of a probability distribution by following the variable ordering of the distribution's factorisation and effectively Rao-Blackwellises [4] REINFORCE.

Definition 3.2 (The SCateR gradient estimator) . We define the Structured Categorical REINFORCE (SCateR) estimator via the expression

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

Proposition 3.3. The SCateR estimator Rao-Blackwellises REINFORCE.

Proof. We start from Equation (2.4) for the REINFORCE estimator, where we ignore the first term and factorize the probability distribution similar to Equation (3.3)

For notational conciseness, we drop the subscript on n d and simply use n to identify single samples. Now we compare Equation (3.5) and Equation (3.6) term-wise for N →∞

In the equation above, we see that SCateR takes the expected value for X d (left-hand side) and computes it exactly using an explicit sum over the space Ω( X d ) , whereas REINFORCE (right-hand side) uses sampled values. This means, in turn, that the left-hand side is a Rao-Blackwellised version of the right-hand side. Doing this for every d gives us a Rao-Blackwellised version for REINFORCE, i.e., the SCateR estimator.

Corollary 3.4 (Bias and Variance) . The SCateR estimator is unbiased and its variance is upperbounded by the variance of REINFORCE.

Proof. This follows immediately from Proposition 3.3, the law of total expectation and the law of total variance [3, 29].

Computational complexity. Consider Equation (3.5) and observe that none of the random variables has a sample space larger than K = max d ( | Ω( X d ) | ) . 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 the parallel implementation of prefix sums on GPUs [19], a time complexity of O (log D + log K + log N ) can be obtained, which allows for the deployment of SCateR in modern deep architectures.

Function evaluations. The main limitation of SCateR is the number of function evaluations it requires to provide its estimates. Indeed, each term in the three nested sums of Equation (3.5) involves a different function evaluation, meaning an overall number of at most D · max d ( | Ω( X d ) | ) · N function evaluations is necessary for every estimate. Fortunately, these evaluations can often be parallelised in modern deep architectures, leading to a positive trade-off between function evaluations and performance in our experiments (Section 6).

Connection to Local Expectation Gradients (LEG). LEG [36] was proposed as a gradient estimator tailored for variational proposal distributions. Similar to the CatLog-Derivative trick, it also singles out variables from a target expectation by summing over the variables' domains. However, it does so one variable at a time by depending on the Markov blanket of those variables. This dependency has the disadvantage of inducing a weighted expression, as variables downstream of a singled-out variable X i are sampled independently of the values assigned to X i . In this sense, LEG can intuitively be related to Gibbs sampling where a variable is marginalised given possible assignments to its Markov blanket instead of being sampled.

In contrast, the CatLog-Derivative trick is applicable to any probability distribution with a known factorisation, even in the case where different factors share parameters. It sums out variables by following the variable ordering of this factorisation rather than depending on the Markov blanket. Following the topology induced by the distribution's factorisation naturally avoids the disconnect between summed-out and downstream variables for which LEG required a weighting term and makes the CatLog-Derivative trick more related to ancestral sampling. Moreover, computing these LEG weights adds additional computational complexity. A more technical and formal comparison between LEG and the CatLog-Derivative trick can be found in the appendix.

4 The IndeCateR Gradient Estimator

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

Proposition 4.1 (IndeCateR) . Let p ( X ) be a multivariate categorical probability distribution that depends on a set of parameters Λ and factorises as p ( X ) = ∏ D d =1 p d ( X d ) , then the gradient of a finite expectation E X ∼ p ( X ) [ f ( X )] can be estimated with

̸

̸

where x ( n d ) = d are samples drawn from p ( X = d ) .

̸

Proof. We start by looking at the expression in Equation (3.1). Using the fact that we have a set of independent random variables, we can simplify p ( x δ | X <d ) to p d ( x δ ) . As a result, the gradient of the expected value can be rewritten as

̸

̸

̸

̸

̸

Drawing N samples for the D -1 independent random variables X = d and for each term in the sum over d then gives us the estimate stated in the proposition.

Example 4.2 (Independent Factorisation) . Let us consider a multivariate distribution involving three 3-ary independent categorical random variables. Concretely, this gives us

where X 1 , X 2 and X 3 can take values from the set Ω( X ) = { 1 , 2 , 3 } . Taking this specific distribution and plugging it into Equation (4.1) for the IndeCateR gradient estimator now gives us

̸

In order to understand the difference between the Log-Derivative trick and the CatLog-Derivative trick, we are going to to look at the term for d = 2 and consider the single-sample estimate

̸

where x 1 and x 3 are sampled values for the random variables X 1 and X 3 . These samples could be different ones for d = 2 . The corresponding single sample estimate using REINFORCE instead of IndeCateR would be ∂ λ p 2 ( x 2 ) f ( x 1 , x 2 , x 3 ) .

We see that for REINFORCE we sample all the variables whereas for IndeCateR we perform the explicit sum for each of the random variables in turn and only sample the remaining variables.

̸

In Section 6, we demonstrate that performing the explicit sum, as shown in Example 4.2, leads to practical low-variance gradient estimates. Note how, in the case of D = 1 , Equation (4.1) reduces to the exact gradient. With this in mind, we can interpret IndeCateR as computing exact gradients for each single random variable X d with respect to an approximation of the function E X = d ∼ p ( X = d ) [ f ( X = d , X d )] .

̸

̸

LEG in the independent case. The gap between the LEG and CatLog-Derivative trick gradient estimates grows smaller in the case the underlying distributions factorises into independent factors. The additional weighting function for LEG collapses to 1 , resulting in similar expressions. However, two main differences remain. First, LEG does not support factors with shared parameters, although this problem can be mitigated by applying the chain rule. Second, and more importantly, LEG does not allow for different samples to be drawn for different variables. While this difference seems small, we will provide clear evidence that it can have a significant impact on performance (Section 6.4).

5 Related Work

Apart from LEG, another closely related work is the RAM estimator [38]. RAM starts by reparametrising the probability distribution using the Gumbel-Max trick followed by a marginalisation. We show in Section 3 that this reparametrisation step is unnecessary, allowing SCateR to explicitly retain the conditional dependency structure. Here we also see the different objectives of our work and Tokui and Sato's. While they proposed RAM in order to theoretically study the quality of control variate techniques, we are interested in building practical estimators that can exploit the structure of specific distributions.

Moreover, just like LEG, Tokui and Sato [38] did not study the setting of a shared parameter space Λ between distributions, although this being the most common setting in modern deep-discrete and neural-symbolic architectures. Hence, an experimental evaluation for this rather common setting is missing in Tokui and Sato [38]. By providing such an evaluation, we show that efficiently implemented Rao-Blackwellised gradient estimators for categorical random variables are a viable option in practice when compared to variance reduction schemes based on control variates.

Such variance reduction methods for REINFORCE aim to reduce the variance by subtracting a mean-zero term, called the baseline, from the estimate [2, 28]. Progress in this area is mainly driven by multi-sample, i.e., sample-dependent, baselines [11, 12, 13, 25, 30, 37, 39] and leave-one-out baselines [17, 18, 26, 31, 34, 35]. Variance can further be reduced by coupling multiple samples and exploiting their dependencies [6, 7, 8, 9, 43, 44, 45]. These latter methods usually introduce bias, which has to be resolved by, for instance, an importance weighing step [8].

Ageneral drawback of baseline variance reduction methods is that they often involve a certain amount of computational overhead, usually in the form of learning optimal parameters or computing statistics. This computational overhead can be justified by assuming that the function in the expectation is costly to evaluate. In such cases, it might pay off to perform extra computations if that means the expensive function is evaluated fewer times.

Another popular, yet diametrically opposite, approach to low-variance gradient estimates for categorical random variables is the concrete distribution [22]. The concrete distribution is a continuous relaxation of the categorical distribution using the Gumbel-Softmax trick [16]. 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 is highly non-trivial in practice.

6 Experiments

In this section we evaluate the performance of IndeCateR on a set of standard benchmarks from the literature. Firstly, two synthetic experiments (Section 6.1 and Section 6.2) will be discussed. In Section 6.3, a discrete variational auto-encoder (DV AE) experiment is optimised on three standard datasets. Lastly, we study the behaviour of IndeCateR on a standard problem taken from the neuralsymbolic literature (Section 6.4).

We mainly compare IndeCateR to REINFORCE leave-one-out (RLOO) [17, 34] and to the GumbelSoftmax trick. The former is a strong representative of methods reducing variance for REINFORCE while the latter is a popular and widely adopted choice in deep-discrete architectures [15].

We implemented IndeCateR and RLOO in TensorFlow [1], from where we also took the implementation of the concrete distribution. 1 Furthermore, all gradient estimators were JIT compiled. The different methods were benchmarked using discrete GPUs and every method had access to the same computational resources. A detailed account of the experimental setup, including hyperparameters and hardware, is given in the appendix.

In the analysis of the experiments we are interested in four main quantities: 1) the time it takes to perform the gradient estimation, 2) the number of drawn samples, 3) the number of function evaluations and 4) the value of a performance metric such as the ELBO. The number of function calls matters in as far that they might be expensive. Note, however, that evaluating the function f is comparatively cheap and trivially parallelisable in the standard benchmarks.

IndeCateR

REINFORCE

RLOO

GS

0.0

0.1

0.2

0.3

0.4

0.5

10

16

10

8

10

4

10

2

10

1

(a)

D

= 12

and

|

Ω(

X

d

)

|

= 3

0.5

0.4

0.3

0.2

0.1

0.0

IndeCateR

(b)

D

REINFORCE

= 6

and

RLOO

|

Ω(

X

d

)

|

10

10

10

10

GS

10

= 10

0.5

0.4

0.3

0.2

0.1

0.0

IndeCateR

(c)

D

REINFORCE

= 3

and

|

RLOO

Ω(

X

d

)

|

1

2

4

8

10

10

10

10

10

16

GS

= 100

6.1 Synthetic: Exact Gradient Comparison

For small enough problems the exact gradient for multivariate categorical random variables can be computed via explicit enumeration. Inspired by Niepert et al. [27], we compare the estimates of

to its exact value. Here, θ are the logits that directly parametrise a categorical distribution p ( X ) = ∏ D d =1 p ( X d ) and b d denotes an arbitrarily chosen element of Ω( X d ) . We compare the gradient estimates from IndeCateR, REINFORCE, RLOO, and Gumbel-Softmax (GS) by varying the number of distributions D and their cardinality.

In Figure 1 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 of variance. 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 and Sato [38], where a two-fold increase in computation time of RAMwith respect to REINFORCE is reported.

6.2 Synthetic: Optimisation

We now study an optimisation setting [37], where the goal is to maximise the expected value

and p ( X ) factorizes into D independent binary random variables. The true maximum is given by p ( X d = 1) = 1 for all d . This task is challenging because of the small impact of the individual values for each X d on the expected value for higher values of D . We set D = 200 and report the results in Figure 2, where we compare IndeCateR to RLOO and Gumbel-SoftMax.

In Figure 2 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 function evaluations with K = max d | Ω( X d ) | . 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 match the 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, even compared 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 insensitive loss function.

1

2

4

8

16

0

250

500

750

1000

1250

1500

1750

2000

Iterations

0.2500

0.2502

0.2504

0.2506

0.2508

0.2510

Function Value

0

1

2

3

4

5

6

Time (s)

0.2500

0.2502

0.2504

0.2506

0.2508

0.2510

Function Value

0

250

500

750

1000

1250

1500

1750

2000

Iterations

10

13

10

12

10

11

10

10

10

9

10

8

Gradient Variance

IndeCateR

RLOO-F

RLOO-S

GS-F

6.3 Discrete Variational Auto-Encoder

As a third experiment we analyse the ELBO optimisation behaviour of a discrete variational autoencoder (DVAE) [32]. We optimise the DVAE on the three main datasets from the literature, being MNIST [21], F-MNIST [42] and Omniglot [20]. The encoder component of the network has two dense hidden layers of sizes 384 and 256 ending in a latent 200-dimensional Bernoulli variable. The decoder takes samples from this variable followed by layers of size 256, 384 and 784. IndeCateR again uses two samples, hence we can compare to the same configurations of RLOO and GumbelSoftmax as in Section 6.2. That is, 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 3), 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 as it beats all other methods on all datasets in terms of training ELBO. 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 both exhibit overfitting on the training data for MNIST and F-MNIST, resulting in an overall higher negative test set ELBO compared to the other methods. We speculate that the relaxation for the Gumbel-Softmax or the higher variance [41] of RLOO-S act as a regulariser for the network. As one would ideally like to separate 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 its gradient estimates.

6.4 Neural-Symbolic Optimisation

A standard experiment in the neural-symbolic literature is the addition of MNIST digits [23]. We will examine an alternative version of this experiment, that is, given a set of D MNIST digits, predict the sum of those digits. During learning, the only supervision provided is the sum and not any direct label of the digits. The difficulty of the problem scales exponentially, as there are 10 D states in the sample space. Note that there are only 10 D +1 possible labels, resulting in only very sparse supervision.

In the field of neural-symbolic AI such problems are either solved exactly [23] or by simplifying the underlying combinatorial structure [14, 24]. Exact methods scale very poorly while simplifying the

0

1000 2000 3000 4000 5000 6000 7000 8000

Time(s)

10 2

8×10 1

9×10 1

Training ELBO

0

1000 2000 3000 4000 5000 6000 7000 8000

Time(s)

10 2

8×10 1

9×10 1

Test ELBO

0

1000 2000 3000 4000 5000 6000 7000 8000

Time(s)

10

3

10

2

10

1

Gradient Variance

IndeCateR

RLOO-S

GS-S

RLOO-F

GS-F

0

2000

4000

6000

8000

Time(s)

2.35 × 10 2

2.4 × 10 2

2.45 × 10 2

2.5 × 10 2

2.55 × 10 2

2.6 × 10 2

2.65 × 10 2

2.7 × 10 2

Training ELBO

0

2000

4000

6000

8000

Time(s)

2.35 × 10 2

2.4 × 10 2

2.45 × 10 2

2.5 × 10 2

2.55 × 10 2

2.6 × 10 2

2.65 × 10 2

2.7 × 10 2

Test ELBO

0

2000

4000

6000

8000

Time(s)

10

3

10

2

10

1

Gradient Variance

0

250

500

750

1000

1250

1500

1750

Time(s)

10 2

1.1 × 10 2

1.2 × 10 2

1.3 × 10 2

1.4 × 10 2

1.5 × 10 2

1.6 × 10 2

Training ELBO

0

250

500

750

1000

1250

1500

1750

Time(s)

10 2

1.1 × 10 2

1.2 × 10 2

1.3 × 10 2

1.4 × 10 2

1.5 × 10 2

1.6 × 10 2

Test ELBO

0

250

500

750

1000

1250

1500

1750

Time(s)

10

3

10

2

10

1

Gradient Variance

combinatorics introduces problematic biases. In contrast, we will study neural-symbolic inference and learning using sampling and unbiased gradient estimators. The overall architecture is as follows, imitating inference in a general neural-symbolic system. Each of the D different MNIST images is passed through a neural classifier, which gives probabilities for each class. These probabilities are used 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, as logical supervision is either true or false.

The hard part of this problem is the sparse supervision. In neural-symbolic AI, this sparsity is resolved by relying on classic search algorithms. However, these are not amenable to be run on parallel computing hardware, such as GPUs. In turn, training has to either occur entirely on the CPU or a hefty time penalty has to be paid for transferring data from and to the GPU.

Using IndeCateR in a neural-symbolic setting we achieve two things. On the one hand, we use sampling as a stochastic yet unbiased search, replacing the usual symbolic search. On the other, we render this stochastic search differentiable by estimating gradients instead of performing the costly exact computation. To make this work, we exploit a feature of IndeCateR that we have ignored so far.

In Equation (4.1) we can draw different samples for each of the D different terms in the sum ∑ D d =1 corresponding to the D different variables. While drawing new samples for every variable does increase the total number of samples, the number of function evaluations remains identical. We indicate the difference by explicitly mentioning which version of IndeCateR draws new samples per variable. Note that this modification of the estimator would not have been possible following the work of Tokui and Sato [38] or Titsias, Michalis and Lázaro-Gredilla, Miguel [36].

10

1

10 0

10 1

10 2

Epochs

0.0

0.2

0.4

0.6

0.8

1.0

10

1

10 0

10 1

10 2

Epochs

0.0

0.2

0.4

0.6

0.8

1.0

10 0

10 1

10 2

Epochs

0.0

0.2

0.4

0.6

0.8

1.0

10 0

10 1

10 2

Epochs

0.0

0.2

0.4

0.6

0.8

1.0

IndeCateR

IndeCateR (sampling per variable)

RLOO-F

RLOO-S

Reference

In Figure 4 we show the empirical results. We let IndeCateR without sampling per variable draw 10 samples such that IndeCateR with sampling per variables draws D · 10 samples. The number of function evaluations of both methods stays the same at D · 10 · 10 . In comparison, RLOO-F uses D · 10 · 10 samples and function evaluations. We see that IndeCateR is the only method capable of scaling and consistently solving the MNIST addition for 16 digits when taking new samples for every variable. In particular, these results show how the CatLog-Derivative trick estimates can significantly outperform the LEG and RAM estimates by drawing different samples for different variables, even in the independent case.

7 Conclusion

We derived a Rao-Blackwellisation scheme for the case of gradient estimation with multivariate categorical random variables using the CatLog-Derivative trick. This scheme led straightforwardly to a practical gradient estimator for independent categorical random variables - the IndeCateR estimator. Our experimental evaluation showed that IndeCateR produces faithful gradient estimates, especially when the number of samples is the main concern, and can be run efficiently on modern GPUs. Furthermore, IndeCateR constitutes the only hyperparameter-free, low-variance estimator for (conditionally) independent categorical random variables.

Three main axes can be identified for future work. First and foremost, we aim to analyse the more general case of a non-independently factorising probability distribution by applying SCateR to more structured problems, such as differentiable planning and reinforcement learning. Secondly, we plan to further study both IndeCateR and SCateR in the neural-symbolic setting and develop estimators for multivariate discrete-continuous distributions containing symbolic structures [5]. Finally, the combination of control variates with the CatLog-Derivative trick seems to be fertile ground for further improvements in variance reduction. Considering the Gumbel-Softmax Trick generally still exhibits lower variance for an equal amount of function evaluations, such a combination could put REINFORCE-based methods on equal footing with those based on the Gumbel-Softmax Trick.

References

Martín Abadi, Ashish Agarwal, Paul Barham, Eugene Brevdo, Zhifeng Chen, Craig Citro, Greg S. Corrado, Andy Davis, Jeffrey Dean, Matthieu Devin, Sanjay Ghemawat, Ian Goodfellow, Andrew Harp, Geoffrey Irving, Michael Isard, Yangqing Jia, Rafal Jozefowicz, Lukasz Kaiser, Manjunath Kudlur, Josh Levenberg, Dandelion Mané, Rajat Monga, Sherry Moore, Derek Murray, Chris Olah, Mike Schuster, Jonathon Shlens, Benoit Steiner, Ilya Sutskever, Kunal Talwar, Paul Tucker, Vincent Vanhoucke, Vijay Vasudevan, Fernanda Viégas, Oriol Vinyals, Pete Warden, Martin Wattenberg, Martin Wicke, Yuan Yu, and Xiaoqiang Zheng. TensorFlow: Large-scale machine learning on heterogeneous systems, 2015. URL https://www.tensorflow.org/ . Software available from tensorflow.org.

Yoshua Bengio, Nicholas Léonard, and Aaron Courville. Estimating or propagating gradients through stochastic neurons for conditional computation. arXiv , 2013.

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

George Casella and Christian P Robert. Rao-blackwellisation of sampling schemes. Biometrika , 1996.

Lennert De Smet, Pedro Zuidberg Dos Martires, Robin Manhaeve, Giuseppe Marra, Angelika Kimmig, and Luc De Readt. Neural probabilistic logic programming in discrete-continuous domains. UAI , 2023.

Alek Dimitriev and Mingyuan Zhou. Carms: Categorical-antithetic-reinforce multi-sample gradient estimator. NeurIPS , 2021.

Aleksandar Dimitriev and Mingyuan Zhou. Arms: Antithetic-reinforce-multi-sample gradient for binary variables. ICML , 2021.

Zhe Dong, Andriy Mnih, and George Tucker. Disarm: An antithetic gradient estimator for binary latent variables. NeurIPS , 33, 2020.

Zhe Dong, Andriy Mnih, and George Tucker. Coupled gradient estimators for discrete latent variables. NeurIPS , 2021.

Artur d'Avila Garcez, Sebastian Bader, Howard Bowman, Luis C Lamb, Leo de Penning, BV Illuminoo, Hoifung Poon, and COPPE Gerson Zaverucha. Neural-symbolic learning and reasoning: a survey and interpretation. Neuro-Symbolic Artificial Intelligence: The State of the Art , 2022.

Will Grathwohl, Dami Choi, Yuhuai Wu, Geoffrey Roeder, and David Duvenaud. Backpropagation through the void: Optimizing control variates for black-box gradient estimation. ICLR , 2018.

Karol Gregor, Ivo Danihelka, Andriy Mnih, Charles Blundell, and Daan Wierstra. Deep autoregressive networks. ICML , 2014.

Shixiang Gu, Sergey Levine, Ilya Sutskever, and Andriy Mnih. Muprop: Unbiased backpropagation for stochastic neural networks. arXiv , 2015.

Jiani Huang, Ziyang Li, Binghong Chen, Karan Samel, Mayur Naik, Le Song, and Xujie Si. Scallop: From probabilistic deductive databases to scalable differentiable reasoning. NeurIPS , 2021.

Iris AM Huijben, Wouter Kool, Max B Paulus, and Ruud JG Van Sloun. A review of the gumbel-max trick and its extensions for discrete stochasticity in machine learning. IEEE TPAMI , 2022.

Eric Jang, Shixiang Gu, and Ben Poole. Categorical reparameterization with gumbel-softmax. ICLR , 2017.

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

Wouter Kool, Herke van Hoof, and Max Welling. Estimating gradients for discrete random variables by sampling without replacement. ICLR , 2020.

Richard E Ladner and Michael J Fischer. Parallel prefix computation. Journal of the ACM (JACM) , 27(4):831-838, 1980.

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

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

Chris J Maddison, Andriy Mnih, and Yee Whye Teh. The concrete distribution: A continuous relaxation of discrete random variables. ICLR , 2017.

Robin Manhaeve, Sebastijan Dumancic, Angelika Kimmig, Thomas Demeester, and Luc De Raedt. Deepproblog: Neural probabilistic logic programming. advances in neural information processing systems , 31, 2018.

Robin Manhaeve, Giuseppe Marra, and Luc De Raedt. Approximate inference for neural probabilistic logic programming. KR , 2021.

Andriy Mnih and Karol Gregor. Neural variational inference and learning in belief networks. ICML , 2014.

Andriy Mnih and Danilo Rezende. Variational inference for monte carlo objectives. International Conference on Machine Learning , 2016.

Mathias Niepert, Pasquale Minervini, and Luca Franceschi. Implicit mle: backpropagating through discrete exponential family distributions. NeurIPS , 2021.

John Paisley, David M Blei, and Michael I Jordan. Variational bayesian inference with stochastic search. ICML , 2012.

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

Rajesh Ranganath, Sean Gerrish, and David Blei. Black box variational inference. AISTATS , pages 814-822, 2014.

Lorenz Richter, Ayman Boustati, Nikolas Nüsken, Francisco Ruiz, and Omer Deniz Akyildiz. Vargrad: a low-variance gradient estimator for variational inference. NeurIPS , 2020.

Jason Tyler Rolfe. Discrete variational autoencoders. ICLR , 2017.

Francisco J. R. Ruiz, Michalis K. Titsias, and David M. Blei. The generalized reparameterization gradient. NeurIPS , 2016.

Tim Salimans and David A Knowles. On using control variates with stochastic approximation for variational bayes and its connection to stochastic linear regression. arXiv , 2014.

Jiaxin Shi, Yuhao Zhou, Jessica Hwang, Michalis Titsias, and Lester Mackey. Gradient estimation with discrete stein operators. NeurIPS , 2022.

Titsias, Michalis and Lázaro-Gredilla, Miguel. Local Expectation Gradients for Black Box Variational Inference. NeurIPS , 2015.

Titsias, Michalis and Shi, Jiaxin. Double control variates for gradient estimation in discrete latent variable models. AISTATS , 2022.

Seiya Tokui and Issei Sato. Evaluating the variance of likelihood-ratio gradient estimators. ICML , 2017.

George Tucker, Andriy Mnih, Chris J Maddison, John Lawson, and Jascha Sohl-Dickstein. Rebar: Low-variance, unbiased gradient estimates for discrete latent variable models. NeurIPS , 2017.

Ronald J. Williams. Simple statistical gradient-following algorithms for connectionist reinforcement learning. Reinforcement learning , 1992.

Jingfeng Wu, Wenqing Hu, Haoyi Xiong, Jun Huan, Vladimir Braverman, and Zhanxing Zhu. On the noisy gradient descent that generalizes as sgd. In ICML . PMLR, 2020.

Han Xiao, Kashif Rasul, and Roland Vollgraf. Fashion-mnist: a novel image dataset for benchmarking machine learning algorithms. arXiv preprint arXiv:1708.07747 , 2017.

Mingzhang Yin and Mingyuan Zhou. Arm: Augment-reinforce-merge gradient for stochastic binary networks. ICLR , 2019.

Mingzhang Yin, Yuguang Yue, and Mingyuan Zhou. Arsm: Augment-reinforce-swap-merge estimator for gradient backpropagation through categorical variables. ICML , 2019.

Mingzhang Yin, Nhat Ho, Bowei Yan, Xiaoning Qian, and Mingyuan Zhou. Probabilistic best subset selection via gradient-based optimization. arXiv , 2020.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper proposes a variance-reduced statistic for $\\mathbb{E}_{x \\sim p(x)} [f(X)]$, which could be the gradient of a loss for example. The main idea (known as Rao-Blackwellization) is simple to understand and can be summarized thusly: the quantity can be rewritten to distinguish any dimension $d$:\n\n$E_{(x_d, x_{\\neq d}) \\sim p(x_{d} | x_{\\neq d}) p(x_{\\neq d})} [f(X)]$ = $E_{x_{\\neq d} \\sim p(x_{\\neq d})} E_{x_{d} \\sim p(x_{d} | x_{\\neq d})} [f(X)]$ \n\nNotationally, the difference in both sides is how $x_d$ is drawn. On the left hand side, $x_d$ is drawn jointly with the other dimensions on the right side: this means a one-to-one ratio of $x_d$ to $x_{\\neq d}$. On the right hand side, $x_d$ is drawn separately from the other dimensions: this allows for a many-to-one ratio of $x_d$ to $x_{\\neq d}$. In fact, when $x_d$ has finite outcomes, all of them can be used and the inner expectation (right hand side) can be computed exactly. The extra computation from the \"many-to-one\" sampling on the right hand side, is the price to pay for the variance-reduction that comes from using (averaging over) more $x_d$. And that computational price is acceptable for a categorical variable. \n\nThe authors show that the variance reduction is effective for a variety of tasks, whether it is the actual estimation of  $\\mathbb{E}_{x \\sim p(x)} [f(X)]$, or other tasks (e.g. optimization) that use the statistic for some other purpose such as stochastic optimization of a loss function.\n\nSTRENGTHS:\nThe exposition is simple and clear, and the results are varied and convincing. \n\nFigure 2 is also appreciated, as it empirically links the variance of the gradients (right panel) to the speed of the optimization (left and middle panels), which although it is common wisdom that has motivated certain algorithms (e.g. SVRG) is not a proven result for any optimizer.\n\nWEAKNESSES:\nNo major weaknesses to indicate so far.\n\nQUESTIONS:\nIt is not very clear from the text, that the statistic for  $\\mathbb{E}_{x \\sim p(x)} [f(X)]$ is sometimes the gradient of a loss (e.g. ELBO). For example, line 133 refers to a gradient without a loss without explicitly making the connection to $f(X)$. Could the authors clarify this in the text?\n\nLIMITATIONS:\nNo specific concern on negative societal impact.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-06-11"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper discusses improved gradient estimation through multivariate categorial probability distributions. A special case of independent random variables is discussed more thoroughly and experiments are conducted only for this case. To summarize for independent variables i.e. $p(X) = \\prod_{i \\in [N]} p(X_i)$ following is my understanding: \n\n1. Assume for applying REINFORCE trick one would use $N$ many $d-$dimensional samples $x^{(n)}$ drawn from $p(X)$.\n2. This paper proposes the following: \n* Create $N$ many $d-$dimensional samples similar to 1.\n* For each variable $i$ in $[D]$ enumerate all possible values (assuming $K$ many) of variable $i$ and replace $X_i$ in each  $x^{(n)}$. Thus creating $K$-times additional samples for each variable.\n\nSTRENGTHS:\n1. The idea of looking 'inside' a joint probability distribution instead of treating it as a black-box is interesting. For independently distributed variables the paper makes use of this factorization. \n2. Experiments (Fig. 2, 3, 4) indicate that proposed method _only sometimes_ produces better gradient estimates than REINFORCE on equal footing.\n\nWEAKNESSES:\n3. It would have been great to have an illustration of the case discussed in Example 4.2 to see things (and symbols) pictorially. Preferably at the start of paper. \n4. As discussed in line 139, the computational complexity of this approach is $\\mathcal{O}(D \\cdot K \\cdot N)$. This looks much expensive than what REINFORCE has which is just $\\mathcal{O}(N)$. Although for a fair comparison in experiments the authors do allow REINFORCE to have more samples which is termed as RLOO-F. Therefore in following I will only compare with such baselines.\\\na. First two rows in Figure 2 has only positive results of the proposed approach although only with slight improvements over RLOO-F.  \\\nb. In last row of Figure 2, GS-F is better (gumbel softmax on equal footing) with RLOO-F not so far behind. \\\nc. In Figure 4 the proposed method does far worse than RLOO-F! Unless it is allowed to sample $x^{(n)}$ individually for each $d$ in $[D]$. This is a serious shortcoming.  \n5. Figure 1 does not contain RLOO-F, why?\n6. The discussed benchmarks are too small and old. It would have been better to discuss a more practical and large-scale problem to solve. \n7. In Example 4.3 it will be good to not have $K = D = 3$ as it makes it more difficult to parse. Possibly make $K=2$ (i.e., binary variables).\n8. Can there by any benefits of proposed approach on cases where variables are not independent? Possibly experiments can be added to put estimator in eq. 3.5 to test in a future submission.\n\nQUESTIONS:\n9. How does time complexity change due to parallelization in line 140?\n10. Line 138 end: respectively w.r.t. what?\n11. What about training accuracy for Figure 4?\n\nLIMITATIONS:\n12. Overall the method falls short of its promise. It needs much computational complexity to create additional samples and REINFORCE trick does sometimes even better than proposed method if tested on equal footing.", "TITLE": "Official Review", "RECOMMENDATION": 3, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-07-05"}

comments: SUMMARY:
In this paper the authors provide an alternative to the log derivative trick, e.g. the REINFORCE estimator, for categorical distributions, which is unbiased and has provably lower variance than the REINFORCE estimator; the estimator is called the CatLog-Derivative trick. For the case D-dimensional multivariate categorical distributions, where each dimension is independent, the authors introduce a simplification of the CatLog-Derivative trick called IndeCateR.

STRENGTHS:
The proposed idea is elegant and simple and the authors did a great job explaining it! Actually, I'm surprised they were able to fit the proofs inside the main text, which speaks to the conciseness of the paper. I also think the authors did a good job explaining the limitations of the base CatLog trick approach and the variety of experiments were great as well.

WEAKNESSES:
The major weakness is that there is no experiment demonstrating the performance of the CatLog trick. While emphasis is placed on the IndeCateR trick, it would be great the CatLog trick in action. As a suggestion, maybe a variational hidden markov model?

QUESTIONS:
- There are a couple of typos in the paper and sentences that aren't complete, i.e. beginning of line 177
- Figure 1 missing y labels, making it hard to parse without reading
- It would be great to have a section in the appendix doing a brief overview of alternative methods
- Figure 4 seems to have problems rendering. I suggest rasterizing the figure and putting it back in the appendix.
- In general, all line plots are hard to read. I suggest increasing the thickness of the lines.

LIMITATIONS:
Yes, they have addressed the limitations of the method.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-06

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe authors derive a gradient estimator for discrete random variables that sums out one dimension while keeping a sample for the other dimensions fixed. Their estimator works for generally fully factorised distributions, and they derive a variant for fully independent distributions, with which they perform experiments on the discrete VAE and a neurosymbolic task.\n\nSTRENGTHS:\nThe paper is very easy to follow and well-motivated. The estimator is quite interesting and competitive with the strong RLOO estimator. The experiments are useful to the community, and especially showing that likelihood-ratio-based methods work for neurosymbolic methods is a useful insight.\n\nWEAKNESSES:\nI'm afraid that the paper is very limited in novelty. The approach presented is almost the same as Local Expectation Gradients (LEG) [1], but LEG is not discussed in the paper. From my understanding, this is how they compare:\n\nSimilarities:\n- The CatLog-Derivative Trick (Eq 3.5) is the LEG (Eq 9) under a fully factorised distribution / autoregressive model.\n- The IndeCateR (Eq 4.1) is precisely the simplification of LEG for fully independent distributions in Eq 11\n- LEG also connects the estimator to Rao-Blackwellization\n\nDifferences:\n- LEG uses a 'pivot' sample and then sums over dimensions, while CatLog resamples for each dimension\n- LEG (like RAM) did not study the shared parameter setting, although, in my opinion, this is a trivial extension by backpropagation\n\nTherefore, I don't think the paper can claim a new trick, as (at best) it slightly modified an established work. That is not to say the results of this paper are not useful but need to be recontextualised, given that the method is not novel. Suggestions are a more thorough experimental setup or a focus on neurosymbolic tasks (which are somewhat understudied in the literature on discrete gradient estimation). \n\n[1] Titsias, M.K., & Lázaro-Gredilla, M. (2015). Local expectation gradients for black box variational inference. Advances in neural information processing systems, 28.\n\nEDIT: The author rebuttal showed that CatLog indeed has some novelty in its sampling mechanism, and I increased my score from 3 to 4. I still think this paper lacks novelty: While CatLog has some novelty, there are no experiments that use it. The only experiments are with IndeCateR, but that method is not novel (LEG and REM on independent distributions are the same as IndeCateR as the authors acknowledge in their rebuttal). IndeCateR-D has some novelty compared to LEG and REM in sampling, but this method is only tested in a single experiment and is not highlighted in the writing as the focus of the paper.\n\nQUESTIONS:\n- Experiments: Why not compare to RAM, as it is closest to the method? \n- Figure 1: What are the 2 y-scales?\n- Line 199: Do you assume one-hot encoded X_d? \n- Line 235: You claim IndeCateR uses two samples, but I do not see how that works. There are 200 dimensions, so you will get 200 samples if you resample per dimension. Or do you set N=2, and use a pivot like in LEG? \n- Did you evaluate the MNIST experiments on multiple runs? I do not see error bars here. These MNISTAdd experiments can have significant variance between runs, so I believe having at least 10 runs is necessary before claiming IndeCateR is the only method that can scale to 16 digits.\n- Clarify that the MNIST experiments are different from the multi-digit MNISTAdd experiments discussed in [1] (it's not $100x_1 + 10x_2 + x_3+100x_4+10x_5+x_6$, but rather $\\sum_i x_i$). These have quite different optimisation properties (the multi-digit MNISTAdd problem has $10^{2/D}-2$ labels rather than $10D+1$)\n\n[1] Manhaeve, R., Dumančić, S., Kimmig, A., Demeester, T., & De Raedt, L. (2021). Neural probabilistic logic programming in DeepProbLog. Artificial Intelligence, 298, 103504.\n\nTypos:\n- Line 177: Unfinished sentence\n- Line 207: RLOO-F mentioned before its introduction\n- Line 219: Comma after samples\n- Line 265: Comma instead of point\n\nLIMITATIONS:\nI would say the limitations are adequately addressed.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 5, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-13"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper develops a new efficient estimator for gradients of an expectation computed over a multivariate discrete random variable. The gradients are computed wrt the continuous parameters of the discrete distribution of interest.\n\nPrevious work consists of two major threads\n- unbiased estimators, such as REINFORCE (aka the score function trick or log-derivative trick)\n- biased estimators, such as relaxed Gumbel-softmax trick aka \"concrete\" random variables [refs 15,20]\n\nTo build on these, this paper presents an approach that is:\n* formally unbiased (like REINFORCE, unlike GS) \n* has low variance (unlike REINFORCE)\n* no free hyperparameters (other than the number of Monte Carlo samples), unlike control variate extensions of REINFORCE\n\nIn Sec 3, the general approach, called CateR, is presented.\nThe big idea is to Rao-Blackwellize REINFORCE across each dimension of the multivariate discrete vector X being sampled. See Eq 3.1 for the estimator. \n\nIn Sec 4, a straightforward specialization of the estimator to a model that assumes each dimension is *independent*, caled IndeCateR, is presented. Example 4.2 gives nice intuition. This requires nested sums over the dimensions D, the K possible values of each variable, and N sampled values of remaining variables, with runtime O(DKN).\n\nExperiments in Sec 6 assess:\n* empirical bias/variance on synthetic problems (Fig 1)\n* performance in optimization on a toy problem (Fig 2)\n* performance in training of discrete VAEs (Fig 3)\n* a \"neuro-symbolic\" model that tries to compute the sum of 10 MNIST digit images, by sample a predicted digit label (discrete value in 0-9) for each image then adding\n\nAcross the board, the experiments suggest the estimator is competitive in performance when computation cost is similar to its alternatives.\n\nSTRENGTHS:\nI found a lot to like about this paper.\n\n+ Attacks a significant problem: many models need such a gradient estimator and current methods (like REINFORCE) are known to be difficult due to high variance\n+ Elegant yet simple method: derivable from first principles and interpretable as Rao-Blackwellization\n+ Comprehensive experiments on toy data help gain intuition about when/why the method works\n+ Overall message and formal mathematics are both clearly communicated throughout\n+ Effort to make costs fair for all methods (in terms of num func evaluations and/or num samples) is appreciated\n\nThanks to the authors for their hard work\n\nWEAKNESSES:\nOverall I don't think there are show-stopping weaknesses here. I'll list some issues below, but I'd overall rate these as definitely worth addressing but minor.\n\n### Improve discussion of AI accelerators\n\nThe paper throughout refers to \"modern AI accelerators\" without much elaboration or citation, leaving unfamiliar readers in the dark. I think there could be many different kinds of hardware accelerator the authors are thinking of... \n\n* What does it mean that this approach can be implemented on a modern AI accelerator? Would this be true of alternatives? like REINFORCE or GS?\n* How would such accelerators improve the runtime of IndeCateR from O(DKN) to O(log D + log K + log N)? Aren't they just reducing constant factor runtime by moving computation from software to hardware?\n\n### Claim that temperature hyperparameter of GS methods is difficult to tune needs elaboration\n\nIn line 172-173, GS / concrete methods are criticized because \"tuning of this [temperature] hyperparameter\" is \"highly non-trivial in practice\".\n\nCan you provide some evidence for this claim?\nI don't doubt that it *could* be sensitive, but I wonder if your experiments could reveal more about this. For example, in Fig 1 or Fig 2 you could show GS with tuned hyperparameter compared to GS with a reasonable default value.\n\n### Is there recent work extending GS/concrete that's worth discussing/comparing?\n\nThe concrete distribution / GS methods [refs 15,20] were published about 6 years ago, in 2017.\n\nSeems like the current related work discussion doesn't really touch on any progress that might have been made since 2017. I'm not aware off the top of my head of such work, but I plan to dig in more during discussion period to see if there's relevant work.\n\nIf so, definitely seems worth citing a few more papers. If not, perhaps highlighting the lack of further progress there is interesting.\n\n\n### Presentation quality decent but needs refinement to catch isolated issues\n\nAt a few spots there are incomplete thoughts or awkward phrasings, such as \n* line 41\n* line 177\n\nIn further revision, please read carefully to catch such issues and spare a future reader any confusion.\n\n\n### DVAE experiments need more details/elaboration\n\n* Need to motivate why the DVAE is an interesting model\n* Does the original DVAE recommend a specific gradient estimator? If so, what?\n* For encoder arch, are there really 3 *hidden* layers, or just 2 hidden layers and one output layer?\n* Why 2 samples for IndeCateR instead of just 1 or many more? How sensitive is this choice?\n* Can you define the probabilistic model and the optimization problem? (perhaps in supp.). In particular, is your likelihood treating the pixels as unconstrained real floats, as floats in 0.0 - 1.0, or binary values?\n* Can you clarify the optimization algorithm used (SGD? ADAM?)? \n* Where methods compared using the same initial parameter values? Or similar sampling scheme for initialization?\n* Is there any regularization applied? (e.g. L2 penalty on encoder/decoder parameters)\n\nQUESTIONS:\n### Q1: How did you compute bias and variance in Fig 1?\n\nPresumably your estimated gradient is a vector, not a scalar.\nDid you just add or average each component's bias/variance to get the scalar metrics reported here?\n\n### Q2: Can you try to give insight into why IndeCateR has lowest variance in Fig 1 and 2, but GS has lower variance with DVAEs in Fig 3?\n\nWhat is the key difference here?\n\n\n### Q3: Why do we need sampling at all in Sec 6.4? And why frame as a binary classifier?\n\nI'm fine with what's presented in Sec 6.4 as a proof of concept that your approach can work on this problem.\n\nBut is there an *advantage* to formaluating the problem by requiring each per-image predictor to produce a random sample, that is then summed? Why not just have each per-image classifier produce the expected value of the categorical distribution over digits 0-9, and sum that across images?\n\nAlso, why use a binary (right/wrong) signal to supervise? Why not use a more continuous error metric, since if the true sum is 33 but I predict 32 I'd probably prefer that vastly compared to predicting 0.\n\n\n\nMinor questions\n\n* Why isn't GS-S shown in Fig 2?\n\nLIMITATIONS:\nCould say more about limitations in Sec 7. There's at least 1 paragraph worth of room for it.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-13"}