SlowGuess/NeurIPS_2023_OCR huggingface dataset:
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.
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 (DVAE) experiment is optimised on three standard datasets. Lastly, we study the behaviour of IndeCateR on a standard problem taken from the neural-symbolic literature (Section 6.4).
We mainly compare IndeCateR to REINFORCE leave-one-out (RLOO) [17, 34] and to the Gumbel-Softmax 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. 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.
1Our code is publicly available at https://github.com/ML-KULeuven/catlog.
(a)  $D = 12$  and  $|\Omega (X_d)| = 3$
(b)  $D = 6$  and  $\left|\Omega(X_d)\right| = 10$
Figure 1: We report the empirical bias and variance for the different estimators and distributions in comparison to the exact gradient. The left y-axis in each plot indicates the bias while the right one shows the variance. Bias and variance results are averages taken over 1000 runs.
(c)  $D = 3$  and  $|\Omega (X_d)| = 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
$$
\partial_ {\theta} \mathbb {E} _ {\mathbf {X} \sim p (\mathbf {X})} \left[ \sum_ {d = 1} ^ {D} | X _ {d} - b _ {d} | \right] \tag {6.1}
$$
to its exact value. Here,  $\theta$  are the logits that directly parametrise a categorical distribution  $p(\mathbf{X}) = \prod_{d=1}^{D} p(X_d)$  and  $b_d$  denotes an arbitrarily chosen element of  $\Omega(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 \cdot 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 RAM with 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
$$
\mathbb {E} _ {\mathbf {X} \sim p (\mathbf {X})} \left[ \frac {1}{D} \sum_ {i = 1} ^ {D} (X _ {i} - 0. 4 9 9) ^ {2} \right], \tag {6.2}
$$
and  $p(\mathbf{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 \cdot K \cdot N$  function evaluations with  $K = \max_d |\Omega(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 \cdot K \cdot 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.
Figure 2: 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 number of 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.
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 Gumbel-Softmax 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  $\approx 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  $10D + 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
Figure 3: The top row shows negated training ELBO, negated test ELBO and gradient variance for the DVAE on MNIST, each plotted against time. Rows 2 and 3 show the same plots for F-MNIST (middle) and Omniglot (bottom).
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  $\sum_{d=1}^{D}$  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].
Figure 4: 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 the different estimators for 4, 8, 12, and 16 MNIST digits. We compare IndeCateR with and without sampling per variables to RLOO-F and RLOO-S. As the function  $f$  has a zero derivative almost everywhere in this case, we can not compare to the Gumbel-Softmax trick. The dashed, grey line at the top of each plot represents a hypothetical sum-classifier that has access to a  $99\%$  accurate MNIST digit classifier. The  $x$ -axis is in log-scale.
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 \cdot 10$  samples. The number of function evaluations of both methods stays the same at  $D \cdot 10 \cdot 10$ . In comparison, RLOO-F uses  $D \cdot 10 \cdot 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.
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.
This research received funding from the Flemish Government (AI Research Program), from the Flanders Research Foundation (FWO) under project G097720N and under EOS project No. 30992574, from the KU Leuven Research Fund (C14/18/062) and TAILOR, a project from the EU Horizon 2020 research and innovation programme under GA No. 952215. It is also supported by the Wallenberg AI, Autonomous Systems and Software Program (WASP) funded by the Knut and Alice Wallenberg-Foundation.
[18] Wouter Kool, Herke van Hoof, and Max Welling. Estimating gradients for discrete random variables by sampling without replacement. *ICLR*, 2020.
[19] Richard E Ladner and Michael J Fischer. Parallel prefix computation. Journal of the ACM (JACM), 27(4):831-838, 1980.
[20] Brenden M Lake, Ruslan Salakhutdinov, and Joshua B Tenenbaum. Human-level concept learning through probabilistic program induction. Science, 350(6266):1332-1338, 2015.
[21] Yann LeCun. The mnist database of handwritten digits. http://yann.lecun.com/exdb/mnist/, 1998.
[22] Chris J Maddison, Andriy Mnih, and Yee Whye Teh. The concrete distribution: A continuous relaxation of discrete random variables. *ICLR*, 2017.
[23] Robin Manhaeve, Sebastian Dumancic, Angelika Kimmig, Thomas Demeester, and Luc De Raedt. Deepproblog: Neural probabilistic logic programming. advances in neural information processing systems, 31, 2018.
[24] Robin Manhaeve, Giuseppe Marra, and Luc De Raedt. Approximate inference for neural probabilistic logic programming. KR, 2021.
[25] Andriy Mnih and Karol Gregor. Neural variational inference and learning in belief networks. ICML, 2014.
[26] Andriy Mniih and Danilo Rezende. Variational inference for monte carlo objectives. International Conference on Machine Learning, 2016.
[27] Mathias Niepert, Pasquale Minervini, and Luca Franceschi. Implicit mle: backpropagating through discrete exponential family distributions. NeurIPS, 2021.
[28] John Paisley, David M Blei, and Michael I Jordan. Variational bayesian inference with stochastic search. ICML, 2012.
[29] C Radhakrishna Rao. Information and accuracy attainable in the estimation of statistical parameters. Bulletin of the Calcutta Mathematical Society, 1945.
[30] Rajesh Ranganath, Sean Gerrish, and David Blei. Black box variational inference. AISTATS, pages 814-822, 2014.
[31] Lorenz Richter, Ayman Boustati, Nikolas Nusken, Francisco Ruiz, and Omer Deniz Akyildiz. Vargrad: a low-variance gradient estimator for variational inference. NeurIPS, 2020.
[32] Jason Tyler Rolfe. Discrete variational autoencoders. ICLR, 2017.
[33] Francisco J. R. Ruiz, Michalis K. Titsias, and David M. Blei. The generalized reparameterization gradient. NeurIPS, 2016.
[34] 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.
[35] Jiaxin Shi, Yuhao Zhou, Jessica Hwang, Michalis Titsias, and Lester Mackey. Gradient estimation with discrete stein operators. NeurIPS, 2022.
[36] Titsias, Michalis and Lázaro-Gredilla, Miguel. Local Expectation Gradients for Black Box Variational Inference. NeurIPS, 2015.
[37] Titsias, Michalis and Shi, Jiaxin. Double control variates for gradient estimation in discrete latent variable models. AISTATS, 2022.
[38] Seiya Tokui and Issei Sato. Evaluating the variance of likelihood-ratio gradient estimators. ICML, 2017.
[39] 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.
[40] Ronald J. Williams. Simple statistical gradient-following algorithms for connectionist reinforcement learning. Reinforcement learning, 1992.
[41] 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.
[42] Han Xiao, Kashif Rasul, and Roland Vollgraf. Fashion-mnist: a novel image dataset for benchmarking machine learning algorithms. arXiv preprint arXiv:1708.07747, 2017.
[43] Mingzhang Yin and Mingyuan Zhou. Arm: Augment-reinforce-merge gradient for stochastic binary networks. ICLR, 2019.
[44] Mingzhang Yin, Yuguang Yue, and Mingyuan Zhou. Arsm: Augment-reinforce-swap-merge estimator for gradient backpropagation through categorical variables. ICML, 2019.
[45] Mingzhang Yin, Nhat Ho, Bowei Yan, Xiaoning Qian, and Mingyuan Zhou. Probabilistic best subset selection via gradient-based optimization. arXiv, 2020.