SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$$
\begin{array}{l} E \left[ \widehat {\nabla} _ {\text {R e i n M a x}, k} \right] = E _ {\boldsymbol {D} \sim \boldsymbol {\pi}} \left[ \frac {\partial f (\boldsymbol {D})}{\partial \boldsymbol {D}} \left(2 \cdot \frac {\boldsymbol {\pi} _ {k} + \delta_ {\boldsymbol {D I} _ {k}}}{2} \left(\boldsymbol {D} _ {k} - \sum_ {i} \frac {\boldsymbol {\pi} _ {i} + \delta_ {\boldsymbol {D I} _ {k}}}{2} \boldsymbol {I} _ {i}\right) - \frac {\boldsymbol {\pi} _ {k}}{2} \left(\boldsymbol {D} _ {k} - \sum_ {i} \boldsymbol {\pi} _ {i} \boldsymbol {I} _ {i}\right)\right) \right] \\ = \frac {1}{2} E _ {D \sim \pi} \left[ \frac {\partial f (D)}{\partial D} \left(\pi_ {k} \left(I _ {k} - I _ {D}\right) + \delta_ {D I _ {k}} \left(I _ {k} - \sum_ {i} \pi_ {i} I _ {i}\right)\right) \right] \tag {11} \\ \end{array}
$$
Since  $\delta_{D\pmb{I}_k}(\pmb{I}_k - \sum_i\pmb{\pi}_i\pmb{I}_i) = \delta_{D\pmb{I}_k}(\pmb{I}_D - \sum_i\pmb{\pi}_i\pmb{I}_i)$ , together with Equation 10 and 11, we have:
$$
E [ \widehat {\nabla} _ {\text {R e i n M a x}, k} ] = \widehat {\nabla} _ {2 \text {n d - o r d e r}, k}
$$
C Remark 4.1
Remark 3.1. When  $\sum_{i}\phi_{i}f(I_{i})$  is used as the baseline and  $f(I_i) - f(I_j)$  is approximated as  $\frac{\partial f(\pmb{I}_j)}{\partial\pmb{I}_j} (\pmb {I}_i - \pmb {I}_j)$ , we mark the resulting first-order approximation of  $\nabla$  as  $\widehat{\nabla}_{1st}$ -order-avg-baseline. Then, we have:
$$
E [ \frac {\phi_ {D}}{\pi_ {D}} \hat {\nabla} _ {S T} ] = \hat {\nabla} _ {I s t - o r d e r - a v g - b a s e l i n e}
$$
Proof. Using  $\sum_{i}\phi_{i}f(I_{i})$  as the baseline, we have:
$$
\nabla = \sum_ {i} (f (\boldsymbol {I} _ {i}) - \sum_ {j} \phi_ {j} f (\boldsymbol {I} _ {j})) \frac {d \boldsymbol {\pi} _ {i}}{d \boldsymbol {\theta}} = \sum_ {i} \sum_ {j} \phi_ {j} (f (\boldsymbol {I} _ {i}) - f (\boldsymbol {I} _ {j})) \frac {d \boldsymbol {\pi} _ {i}}{d \boldsymbol {\theta}}
$$
Approximating  $f(\pmb{I}_i) - f(\pmb{I}_j)$  as  $\frac{\partial f(\pmb{I}_j)}{\partial \pmb{I}_j} (\pmb{I}_i - \pmb{I}_j)$ , we have:
$$
\begin{array}{l} \widehat {\nabla} _ {1 \text {s t - o r d e r - a v g - b a s e l i n e}} = \sum_ {i} \sum_ {j} \phi_ {j} \frac {\partial f (\boldsymbol {I} _ {j})}{\partial \boldsymbol {I} _ {j}} \left(\boldsymbol {I} _ {i} - \boldsymbol {I} _ {j}\right) \frac {d \boldsymbol {\pi} _ {i}}{d \boldsymbol {\theta}} \\ = \sum_ {j} \frac {\phi_ {j}}{\pi_ {j}} \cdot \pi_ {j} \cdot \frac {\partial f (\boldsymbol {I} _ {j})}{\partial \boldsymbol {I} _ {j}} \sum_ {i} \boldsymbol {I} _ {i} \frac {d \boldsymbol {\pi} _ {i}}{d \boldsymbol {\theta}} \\ { = } { E [ \frac { \phi _ { D } } { \pi _ { D } } \widehat { \nabla } _ { \mathrm { S T } } ] } \\ \end{array}
$$
D Remark 4.2
Remark 3.2. In Equation 8, we approximate  $f(\pmb{I}_k) - f(\pmb{I}_i)$  as  $\frac{1}{2} (\frac{\partial f(\pmb{I}_i)}{\partial \pmb{I}_i} + \frac{\partial f(\pmb{I}_k)}{\partial \pmb{I}_k})(\pmb{I}_k - \pmb{I}_i)$ , and mark the resulting second-order approximation of  $\nabla_k$  as  $\widehat{\nabla}_{2nd-order-wo-baseline,k} = \pi_k \sum_i \pi_i \frac{1}{2} (\frac{\partial f(\pmb{I}_i)}{\partial \pmb{I}_i} + \frac{\partial f(\pmb{I}_k)}{\partial \pmb{I}_k})(\pmb{I}_k - \pmb{I}_i)$ . Then, we have:
$$
E [ \widehat {\nabla} _ {R e i n M a x} ] = \widehat {\nabla} _ {2 n d - o r d e r - w o - b a s e l i n e}
$$
Proof. Here, we aim to proof,  $\forall k\in [1,n]$ , we have  $E[\widehat{\nabla}_{\mathrm{ReimMax},k}] = \widehat{\nabla}_{2\mathrm{nd - order - wo - baseline},k}$ .
$$
\begin{array}{l} \widehat {\nabla} _ {2 \text {n d - o r d e r - w o - b a s e l i n e}, k} = \pi_ {k} \sum_ {i} \pi_ {i} \frac {1}{2} \left(\frac {\partial f (\boldsymbol {I} _ {i})}{\partial \boldsymbol {I} _ {i}} + \frac {\partial f (\boldsymbol {I} _ {k})}{\partial \boldsymbol {I} _ {k}}\right) (\boldsymbol {I} _ {k} - \boldsymbol {I} _ {i}) \\ = \pi_ {k} \sum_ {i} \pi_ {i} \frac {1}{2} \frac {\partial f (\boldsymbol {I} _ {i})}{\partial \boldsymbol {I} _ {i}} \left(\boldsymbol {I} _ {k} - \boldsymbol {I} _ {i}\right) + \pi_ {k} \sum_ {i} \pi_ {i} \frac {1}{2} \frac {\partial f (\boldsymbol {I} _ {k})}{\partial \boldsymbol {I} _ {k}} \left(\boldsymbol {I} _ {k} - \boldsymbol {I} _ {i}\right) \\ = E \left[ \frac {\partial f (\boldsymbol {D})}{\partial \boldsymbol {D}} \frac {\boldsymbol {\pi} _ {k} \left(\boldsymbol {I} _ {k} - \boldsymbol {I} _ {\boldsymbol {D}}\right) + \delta_ {\boldsymbol {D I} _ {k}} \left(\boldsymbol {I} _ {k} - \sum_ {i} \boldsymbol {\pi} _ {i} \boldsymbol {I} _ {i}\right)}{2} \right] = E \left[ \widehat {\nabla} _ {\text {R e i n M a x}, k} \right] \\ \end{array}
$$
E Forward Euler Method and Heun's Method
For simplicity, we consider a simple function  $g(x): \mathcal{R} \to \mathcal{R}$  that is three times differentiable on  $[t_0, t_1]$ . Now, we proceed to a simple introduction to approximate  $\int_{t_0}^{t_1} g'(x) dx$  with the Forward Euler Method and the Heun's Method. For a detailed introduction to numerical ODE methods, please refer to Ascher & Petzold (1998).
Forward Euler Method. Here, we approximate  $g(t_{1})$  with the first-order Taylor expansion, i.e.,  $g(t_{1}) = g(t_{0}) + g^{\prime}(t_{0})\cdot (t_{1} - t_{0}) + O((t_{1} - t_{0})^{2})$ , then we have  $\int_{t_0}^{t_1}g'(x)dx\approx g'(t_0)(t_1 - t_0)$ . Since we used the first-order Taylor expansion, this approximation has first-order accuracy.
Heun's Method. First, we approximate  $g(t_{1})$  with the second-order Taylor expansion:
$$
g \left(t _ {1}\right) = g \left(t _ {0}\right) + g ^ {\prime} \left(t _ {0}\right) \cdot \left(t _ {1} - t _ {0}\right) + \frac {g ^ {\prime \prime} \left(t _ {0}\right)}{2} \cdot \left(t _ {1} - t _ {0}\right) ^ {2} + O \left(\left(t _ {1} - t _ {0}\right) ^ {3}\right). \tag {12}
$$
Then, we show that we can match this approximation by combining the first-order derivatives of two samples. Taylor expanding  $g'(t_1)$  to the first-order, we have:
$$
g ^ {\prime} \left(t _ {1}\right) = g ^ {\prime} \left(t _ {0}\right) + g ^ {\prime \prime} \left(t _ {0}\right) \cdot \left(t _ {1} - t _ {0}\right) + O \left(\left(t _ {1} - t _ {0}\right) ^ {2}\right)
$$
$$
g (t _ {0}) + \frac {g ^ {\prime} (t _ {0}) + g ^ {\prime} (t _ {1})}{2} (t _ {1} - t _ {0}) = g (t _ {0}) + g ^ {\prime} (t _ {0}) \cdot (t _ {1} - t _ {0}) + \frac {g ^ {\prime \prime} (t _ {0})}{2} \cdot (t _ {1} - t _ {0}) ^ {2} + O ((t _ {1} - t _ {0}) ^ {3}).
$$
It is easy to notice that the right-hand side of the above equation matches the second-order Taylor expansion of  $g(t_{1})$  as in Equation 12. Therefore, the above approximation (i.e., approximating  $g(t_{1}) - g(t_{0})$  as  $\frac{g'(t_0) + g'(t_1)}{2} (t_1 - t_0)$ ) has second-order accuracy.
Connection to  $f(\mathbf{I}_i) - f(\mathbf{I}_j)$  in Equation 6. By setting  $g(x) = f(x \cdot \mathbf{I}_i + (1 - x) \cdot \mathbf{I}_j)$ , we have  $g(1) - g(0) = f(\mathbf{I}_i) - f(\mathbf{I}_j)$ . Then, it is easy to notice that the forward Euler Method approximates  $f(\mathbf{I}_i) - f(\mathbf{I}_j)$  as  $\frac{\partial f(\mathbf{I}_j)}{\partial \mathbf{I}_j} (\mathbf{I}_i - \mathbf{I}_j)$  and has first-order accuracy. Also, the Heun's Method approximates  $f(\mathbf{I}_i) - f(\mathbf{I}_j)$  as  $\frac{1}{2} (\frac{\partial f(\mathbf{I}_i)}{\partial \mathbf{I}_i} + \frac{\partial f(\mathbf{I}_j)}{\partial \mathbf{I}_j})(\mathbf{I}_i - \mathbf{I}_j)$  and has second-order accuracy.
F.1Baselines
Here, we consider four methods as our major baselines:
- Straight-Through (ST; Bengio et al., 2013) backpropagate through the sampling function as if it had been the identity function.
- Straight-Through Gumbel-Softmax (STGS; Jang et al., 2017) integrates the Gumbel reparameterization trick to approximate the gradient.
- Gumbel-Rao Monte Carlo (GR-MCK; Paulus et al., 2021) leverages the Monte Carlo method to reduce the variance introduced by the Gumbel noise in STGS. To obtain the optimal performance for this baseline, we set the number of Monte Carlo samples to 1000 in most experiments. Except in our discussions of efficiency, we set the number of Monte Carlo samples to 100, 300, and 1000 for a more comprehensive comparisons.
- Gapped Straight-Through (GST-1.0; Fan et al., 2022) aims to reduce the variance of STGS and constructs a deterministic term to replace the Monte Carlo samples used in GR-MCK. Here, as suggested in (Fan et al., 2022), we set the gap (a hyper-parameter) as 1.0.
GST-1.0 Performance. Despite GST-1.0 achieving good performance on most settings of MNIST-VAE, it fails to maintain this performance on polynomial programming and unsupervised parsing, as discussed before. At the same time, a different variant of GST (i.e., GST-p) achieves a significant performance boost over GST-1.0 on polynomial programming. However, on MNIST-VAE and ListOps, GST-p achieves an inferior performance. Upon discussing with the author of the GST-1.0, we suggest that this phenomenon is caused by different characteristics of GST-1.0 and GST-p.
F.2 Hyper-Parameters
Without specifically, we conduct full grid search for all methods in all experiments, and report the best performance (averaged with 10 random seeds on MNIST-VAE and 5 random seeds on ListOps). The hyper-parameter search space is summarized in Table 7. The search results for Table 2 and Table 1 are summarized in Table 8.
Table 7: Hyper-parameter search space.
Hyperparameters	Search Space
Optimizer	{Adam(Kingma & Ba, 2015), RAdam(Liu et al., 2020)}
Learning Rate	{0.001, 0.0007, 0.0005, 0.0003}
Temperature	{0.1, 0.3, 0.5, 0.7, 1.0, 1.1, 1.2, 1.3, 1.4, 1.5}
Table 8: Hyper-parameters Search Result for Results in Table 1 and Table 2.
STGS	GR-MCK	GST-1.0	ST	ReinMax
MNIST-VAE 8 × 4	Optimizer	Adam	Adam	Adam	Adam	Adam
Learning Rate	0.0003	0.0005	0.0005	0.001	0.0005
Temperature	0.5	0.5	0.7	1.3	1.3
MNIST-VAE 4 × 24	Optimizer	RAdam	RAdam	RAdam	RAdam	RAdam
Learning Rate	0.0005	0.0005	0.0005	0.001	0.0005
Temperature	0.3	0.3	0.5	1.5	1.5
MNIST-VAE 8 × 16	Optimizer	RAdam	RAdam	RAdam	RAdam	RAdam
Learning Rate	0.0005	0.0007	0.0007	0.001	0.0007
Temperature	0.5	0.7	0.5	1.5	1.5
MNIST-VAE 16 × 12	Optimizer	RAdam	Adam	RAdam	Adam	RAdam
Learning Rate	0.0007	0.0005	0.0007	0.0005	0.0007
Temperature	0.7	1.0	0.5	1.5	1.5
MNIST-VAE 64 × 8	Optimizer	RAdam	Adam	RAdam	Adam	RAdam
Learning Rate	0.0007	0.0007	0.0007	0.0005	0.0005
Temperature	0.7	2.0	0.7	1.5	1.5
MNIST-VAE 10 × 30	Optimizer	RAdam	RAdam	RAdam	RAdam	RAdam
Learning Rate	0.0005	0.0005	0.0005	0.0007	0.0005
Temperature	0.5	1.0	0.5	1.4	1.3
ListOps	Optimizer	RAdam	RAdam	RAdam	RAdam	RAdam
Learning Rate	0.0005	0.0005	0.001	0.001	0.0007
Temperature	0.1	0.3	0.1	1.4	1.1
Polynomial Programming. As this problem is relatively simple, we set the learning rate to 0.001 and the optimizer to Adam, and only tune the temperature hyper-parameter.
MNIST-VAE. Following the previous study (Dong et al., 2020b, 2021; Fan et al., 2022), we used 2-layer MLP as the encoder and the decoder. We set the hidden state dimension of the first-layer and the second-layer as 512 and 256 for the encoder, and 256 and 512 for the decoder. For our experiments on MNIST-VAE with 32 latent dimensions and 64 categorical dimensions, we set the batch size to 200, training steps to  $5 \times 10^{5}$ , and activation function to LeakyReLU, in order to be consistent with the literature. For other experiments, we set the batch size to 100, the activation function to ReLU, and training steps to  $9.6 \times 10^{4}$  (i.e., 160 epochs).
Table 9: Test -ELBO on MNIST. Hyper-parameters are chosen based on Train -ELBO.
AVG	8 × 4	4 × 24	8 × 16	16 × 12	64 × 8	10 × 30
STGS	106.89	128.09±0.79	103.60±0.45	99.32±0.33	102.49±0.32	106.20±0.46	101.61±0.54
GR-MCK	109.03	127.90±0.71	102.76±0.33	102.12±0.29	104.23±0.65	113.54±0.50	103.62±0.13
GST-1.0	106.85	128.20±1.12	103.95±0.49	101.44±0.32	101.28±0.59	105.44±0.62	100.78±0.44
ST	118.85	137.06±0.51	113.41±0.49	114.25±0.29	114.48±0.56	115.43±0.29	118.46±0.18
ReinMax	105.74	126.89±0.79	102.40±0.43	100.63±0.41	100.85±0.50	102.91±0.67	100.75±0.50
Table 10: Test -ELBO on MNIST. Hyper-parameters are chosen based on Test -ELBO.
AVG	8 × 4	4 × 24	8 × 16	16 × 12	64 × 8	10 × 30
STGS	107.15	128.09±0.79	103.25±0.22	101.44±0.32	102.29±0.39	106.20±0.46	101.61±0.54
GR-MCK	108.87	127.86±0.54	102.40±0.37	101.59±0.22	104.22±0.63	113.54±0.50	103.62±0.13
GST-1.0	106.55	128.03±1.02	103.63±0.24	100.67±0.34	101.04±0.39	105.44±0.62	100.51±0.37
ST	118.79	137.05±0.36	113.23±0.43	114.11±0.31	114.48±0.56	115.43±0.29	118.46±0.18
ReinMax	105.60	126.29±0.32	102.40±0.43	100.45±0.26	100.84±0.56	102.91±0.68	100.69±0.48
Table 11: Train -ELBO on MNIST. Hyper-parameters are chosen based on Test -ELBO.
AVG	8 × 4	4 × 24	8 × 16	16 × 12	64 × 8	10 × 30
STGS	105.31	126.85±0.85	101.81±0.14	99.32±0.33	100.22±0.47	104.02±0.41	99.63±0.63
GR-MCK	107.37	126.53±0.55	100.47±0.31	99.75±0.29	103.11±0.58	112.34±0.48	102.02±0.18
GST-1.0	104.60	126.63±1.16	102.11±0.24	98.40±0.34	98.76±0.41	102.53±0.57	99.14±0.30
ST	117.76	136.75±0.22	112.09±0.50	113.06±0.26	113.31±0.43	113.90±0.28	117.46±0.09
ReinMax	103.40	124.92±0.38	99.77±0.45	98.06±0.31	98.51±0.54	100.71±0.70	98.40±0.48
Differentiable Neural Architecture Search. We adopt most of the hyper-parameter setting from Dong et al. (2020a). Since GDAS employs a temperature schedule (decaying linearly from 10 to 0.1), and temperature scaling works differently in ReinMax and STGS (as discussed in Section 5 and Section 6.6), we removed the temperature scaling (i.e., set the temperature to a constant 1.0) and increased the weight decay (i.e., from 0.001 to 0.09).
ListOps. We followed the same setting of Fan et al. (2022), i.e., used the same model configuration as in Choi et al. (2017) and set the maximum sequence length to 100.
F.3 Hardware and Environment Setting
Most experiments (except efficiency comparisons) are conducted on Nvidia P40 GPUs. For efficiency comparisons, we measured the average time cost per batch and peak memory consumption on quadratic programming and MNIST-VAE on the same system with an idle A6000 GPU. Also, to better reflect the efficiency of gradient estimators, we skipped all parameter updates in efficiency comparisons.
F.4 Additional Results on Polynomial Programming
We visualized the training curve for polynomial programming with various batch sizes and latent dimensions in Figure 8 (for  $p = 1.5$ ), Figure 9 (for  $p = 2$ ), and Figure 10 (for  $p = 3$ ).
F.5 Additional Results on MNIST-VAE
In our discussions in Section 6, we focused on the training ELBO only. Here, we provide a brief discussion on the test ELBO.
Choosing Hyper-parameter Based on Training Performance. Similar to Table 2, for each method, we select the hyper-parameter based on its training performance. The Test -ELBO in this setting is summarized in Table 9. Despite the model being trained without dropout or other overfitting reduction techniques, ReinMax maintained the best performance in this setting.
(a)  $p = 3$  and  $\boldsymbol {c} = \left[\frac{0.5}{L},\frac{1.5}{L},\dots ,\frac{L - 0.5}{L}\right]$
(c)  $p = 2$  and  $\boldsymbol {c} = [\frac{0.5}{L},\frac{1.5}{L},\dots ,\frac{L - 0.5}{L} ]$
(e)  $p = 1.5$  and  $\boldsymbol {c} = [\frac{0.5}{L},\frac{1.5}{L},\dots ,\frac{L - 0.5}{L} ]$
Figure 7: Training curves of polynomial programming, i.e.,  $\min_{\pmb{\theta}} E_{\mathbf{X}}[\frac{\|\mathbf{X} - \mathbf{c}\|_p^p}{L}]$ , where  $\mathbf{X} \in \{0,1\}^L$ ,  $\mathbf{X}_i \stackrel{\mathrm{iid}}{\sim}$  Multinomial(softmax( $\theta_i$ )),  $\pmb{\theta} = [\theta_1, \dots, \theta_L]^T$ ,  $\pmb{\theta}_i \in \mathbb{R}^2$ , and  $L$  is the number of latent dimensions.
(b)  $p = 3$  and  $\pmb {c} = [0.45,\dots ,0.45]$
(d)  $p = 2$  and  $\pmb {c} = [0.45,\dots ,0.45]$
(f)  $p = 1.5$  and  $\pmb {c} = [0.45,\dots ,0.45]$
Choosing Hyper-parameter Based on Test Performance. We also conduct experiments by selecting hyper-parameters directly based on their test performance. In this setting, the test -ELBO is summarized in Table 10, and the training -ELBO is summarized in Table 11. ReinMax achieves the best performance in all settings except the test performance of the setting with 10 categorical dimensions and 30 latent dimensions.
F.6 More Comparisons with RODEO
To better understand the difference between RODEO and ReinMax, we conduct more experiments on polynomial programming, i.e.,  $\min_{\theta} E_X\left[\frac{|X - c|^p}{L}\right]$ . Specifically, we consider polynomial programming under two different settings that define  $c$  differently:
- In setting A, we have  $c = [0.45, \dots, 0.45]$ . This is the setting we used in the submission.
- In setting B, we have  $c = \left[\frac{0.5}{L}, \frac{1.5}{L}, \dots, \frac{L - 0.5}{L}\right]$ .
As to the difference between the Setting A and the Setting B, we would like to note:
- In setting A, since  $\forall i, c_i = 0.45$  and  $\theta_i \sim \mathrm{Uniform}(-0.01, 0.01)$  at initialization,  $E_{\mathbf{X}_i \sim \mathrm{softmax}(\theta_i)}[\frac{|X_i - c_i|^p}{L}]$  would have similar values. Therefore, the optimal control variates for  $\theta_i$  are similar across different  $i$ .
- In setting B, we set  $c_{i}$  to different values for different  $i$ , and thus the optimal control variate for  $\theta_{i}$  are different across different  $i$ . Therefore, Setting A is a simpler setting for applying control variate to REINFORCE.
As in Figure 7, ReinMax achieves better performance in more challenging scenarios, i.e., smaller batch size, more latent variables, or more complicated problems (Setting B or VAEs). Meanwhile, REINFORCE and RODEO achieve better performance on simpler problem settings, i.e., larger batch size, fewer latent variables, or simpler problems (Setting A).
Figure 8: Polynomial programming training curve, with different batch sizes and random variable counts  $(L)$ , i.e.,  $\min_{\pmb{\theta}} E[\frac{\|\pmb{X} - \pmb{c}\|_{1.5}^{1.5}}{L}]$ , where  $\pmb{\theta} \in \mathcal{R}^{L \times 2}, \pmb{X} \in \{0,1\}^L$ , and  $\pmb{X}_i \stackrel{\mathrm{iid}}{\sim}$  Multinomial(softmax( $\theta_i$ )). More details are elaborated in Section 6.
Figure 9: Quadratic programming training curve, with different batch sizes and random variable counts  $(L)$ , i.e.,  $\min_{\pmb{\theta}} E[\frac{\|\pmb{X} - c\|_2^2}{L}]$ , where  $\pmb{\theta} \in \mathcal{R}^{L \times 2}, \pmb{X} \in \{0,1\}^L$ , and  $\pmb{X}_i \stackrel{\mathrm{iid}}{\sim}$  Multinomial(softmax( $\pmb{\theta}_i$ )). More details are elaborated in Section 6.
Figure 10: Polynomial programming training curve, with different batch sizes and random variable counts  $(L)$ , i.e.,  $\min_{\pmb{\theta}} E[\frac{\|\pmb{X} - \pmb{c}\|_3^3}{L}]$ , where  $\pmb{\theta} \in \mathcal{R}^{L \times 2}, \pmb{X} \in \{0,1\}^L$ , and  $\pmb{X}_i \stackrel{\mathrm{id}}{\sim} \text{Multinomial}(\text{softmax}(\pmb{\theta}_i))$ . More details are elaborated in Section 6.