SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Du Phan* Matthew D. Hoffman* David Dohan† Sholto Douglas Tuan Anh Le Aaron Parisi Pavel Sountsov Charles Sutton Sharad Vikram Rif A. Saurous Google
Large language models (LLMs) solve problems more accurately and interpretably when instructed to work out the answer step by step using a "chain-of-thought" (CoT) prompt. One can also improve LLMs' performance on a specific task by supervised fine-tuning, i.e., by using gradient ascent on some tunable parameters to maximize the average log-likelihood of correct answers from a labeled training set. Naively combining CoT with supervised tuning requires supervision not just of the correct answers, but also of detailed rationales that lead to those answers; these rationales are expensive to produce by hand. Instead, we propose a fine-tuning strategy that tries to maximize the marginal log-likelihood of generating a correct answer using CoT prompting, approximately averaging over all possible rationales. The core challenge is sampling from the posterior over rationales conditioned on the correct answer; we address it using a simple Markov-chain Monte Carlo (MCMC) expectation-maximization (EM) algorithm inspired by the self-taught reasoner (STaR), memoized wake-sleep, Markovian score climbing, and persistent contrastive divergence. This algorithm also admits a novel control-variate technique that drives the variance of our gradient estimates to zero as the model improves. Applying our technique to GSM8K and the tasks in BIG-Bench Hard, we find that this MCMC-EM fine-tuning technique typically improves the model's accuracy on held-out examples more than STaR or prompt-tuning with or without CoT.
For many mathematical, logical, and common-sense reasoning problems, large language models solve problems more accurately when instructed to work out the answer step by step in a chain of thought or a scratchpad (Wei et al., 2022; Nye et al., 2021; Kojima et al., 2022; Rajani et al., 2019; Shwartz et al., 2020). These methods encourage the model to produce a rationale, that is, text describing a sequence of reasoning steps that leads to an answer; the motivation is that it seems to be easier for the model to generate a sequence of correct reasoning steps than to generate the final answer directly. Because of the striking performance of chain-of-thought methods, many variants have been proposed (Wang et al., 2022b; Zhou et al., 2022; Creswell et al., 2022; Ye & Durrett, 2023), but there are still many cases in which the rationales are incorrect.
One way to improve these methods is to fine-tune models to generate better rationales. If gold-standard rationales can be obtained, such as via crowdsourcing (Rajani et al., 2019) or automatically (Nye et al., 2021), then supervised methods can be applied, but obtaining this data can be difficult. An appealing alternative is to start from datasets that contain questions and correct answers only, which are more readily available, and bootstrap rationales during learning. A version of this strategy
*Corresponding authors: {mhoffman,phandu}@google.com. The first two authors contributed equally, and their order was chosen randomly.
†Current affiliation: OpenAI.
was proposed as the self-taught reasoner (STaR) (Zelikman et al., 2022), which generates proposed rationales from an LLM, and then fine-tunes on rationales that lead to the correct answer.
In this paper, we approach the problem of bootstrapping rationales from a different conceptual direction: chain-of-thought methods are probabilistic latent-variable models. The LLM defines a joint probability distribution over questions, rationales, and answers; this joint distribution implies a marginal distribution of answers given questions, averaging over all possible rationales weighted by their probability given the question. The problem of self-training for reasoning then becomes one of learning with incomplete data, a core task in probabilistic machine learning (Murphy, 2022) to which we can apply methods from a large and sophisticated literature.
This perspective raises a technical challenge, because computing the marginal distribution requires averaging over a vast set of potential rationales. To address this, we introduce a learning algorithm for rationale generation, which we call TRICE.<sup>3</sup> TRICE is a simple Markov-chain Monte Carlo (MCMC) expectation-maximization (EM) algorithm combined with a novel control-variate scheme, inspired by ideas from STaR (Zelikman et al., 2022), memoized wake-sleep (Hewitt et al., 2020), Markovian score climbing (Naesseth et al., 2020), and persistent contrastive divergence (Tieleman, 2008).
This view unifies several threads of work in reasoning using LLMs: It provides an alternative interpretation of STaR as a kind of biased stochastic expectation-maximization algorithm (Nielsen, 2000) that underweights difficult examples when its rationalization process fails. Self-consistency (Wang et al., 2022a) can be seen as a Monte Carlo algorithm for computing the most likely answer under the marginal distribution. Compared to self-consistency, the probabilistic learning approach of TRICE allows us to average over rationales not only at inference time, but also at training time. Compared to STaR, TRICE is less likely to ignore difficult examples (which stabilizes convergence and improves performance), and is also able to learn from incorrect rationales as well as correct ones.
We apply our technique to the GSM8K dataset (Cobbe et al., 2021) and to the BIG-Bench Hard benchmark (Suzgun et al., 2022a). We find that TRICE improves the model's performance significantly, outperforming models tuned with STaR, direct tuning with or without CoT, and even supervised fine-tuning on human-generated rationales.
Given a training set of  $N$  questions  $x_{1:N}$  and answers  $y_{1:N}$ , we formalize CoT tuning as optimizing a parameter vector  $\theta$  to maximize the average marginal log-likelihood of answers given questions:
$$
\mathcal {L} (\theta) \triangleq \frac {1}{N} \sum_ {n} \log p _ {\theta} \left(y _ {n} \mid x _ {n}\right) = \frac {1}{N} \sum_ {n} \log \sum_ {z} p _ {\theta} (z \mid x _ {n}) p \left(y _ {n} \mid z, x _ {n}\right), \tag {1}
$$
where  $z$  is an unobserved latent rationale,  $p_{\theta}(z \mid x)$  is the probability<sup>4</sup> of obtaining the rationale  $z$  by prompting an LLM with the question  $x$  and tunable parameters  $\theta$ , and  $p_{\theta}(y \mid z, x)$  is the probability of obtaining the answer  $y$  given rationale  $z$ , question  $x$ , and parameters  $\theta$ . We will be particularly interested in models where the likelihood  $p_{\theta}(y \mid x, z) \in \{0,1\}$ , that is, where the answer  $y$  is a deterministic function of  $z$ . For example, we might say that the model's answer is  $y = \text{"(a)" if } z$  ends with the string "The answer is (a)." For this deterministic model, we define  $p(y \mid z, x) = c(z, y) \in \{0,1\}$ . Details of  $c(z, y)$  for each task can be found in Appendix F. We believe that such a binary likelihood model is appropriate for question-answering tasks where  $z$  is a rationale—a good rationale should leave no ambiguity about the correct answer. The derivations below will therefore assume a binary likelihood function. It is straightforward to generalize our methods to cases where the relationship between  $z$  and  $y$  is weaker and therefore  $p(y \mid x, z)$  is more complicated; Appendix A shows how.
Algorithm 1 outlines the method. A notebook with a reference implementation can be found at https://github.com/google-research/cascades/tree/main/cascades/ examples/notebooks/trice.ipynb.
We start by initializing a memory containing a latent rationale  $z_{n}$  for each example pair  $x_{n}, y_{n}$  by sampling  $z_{n}$  from a hinted guide distribution  $q(z \mid x_{n}, y_{n})$  that may condition on the correct answer  $y_{n}$  as well as the question  $x_{n}$ . For example, the guide might prompt an LLM specifically to give an
$^{3}$ TRICE stands for "Tuning Rationales with Independence-Chain Expectation-maximization."
4 Unless otherwise specified, we sample at temperature 1 throughout.
rationale for the answer; more details about the precise prompts used by the guide are in Appendix F. In some cases sampling from the guide instead of the model  $p_{\theta}(z \mid x_n)$  increases the chances of generating a correct rationale (Zelikman et al., 2022).
We then proceed to the main optimization loop. Each iteration, we sample a minibatch of  $M$  questions and answers from the dataset, and retrieve the rationales associated with those examples from the memory. We then propose new rationales  $\tilde{z}$  from the current model  $p_{\theta}(z \mid x)$ , and whenever the new rationale  $\tilde{z}$  is correct (i.e.,  $c(\tilde{z}, y) = 1$ ) replace the old rationale in memory with the new one.
At this point we have all we need to compute a gradient estimate; we can just average the gradients  $\nabla_{\theta}\log p_{\theta}(z_{i_m}\mid x_{i_m})$  that we obtain from those rationales in the updated memory that are correct (i.e., we ignore examples where both the proposed rationale and the previous rationale were wrong). BASIC_GRADIENT Estimate in Algorithm 1 shows how.
But we can also reduce the variance of our gradient estimator by incorporating a control variate, as in CONTROL_VARIATE_GRADIENT Estimate in Algorithm 1. We first compute leave-one-out estimates  $\beta_{1:M}$  of the average probability of accepting a new rationale. For each example  $m$ , we subtract off a scaled control variate  $\beta_{m}\nabla_{\theta}\log p_{\theta}(\tilde{z}_{m}\mid x_{i_{m}})$  whose expected value is zero (since it is a score function). If the proposed rationale  $\tilde{z}_m$  for example  $m$  is correct, then  $z_{i_m} = \tilde{z}_m$ , and the  $m$ th gradient contribution becomes  $(1 - \beta_{m})\nabla_{\theta}\log p_{\theta}(z_{i_m}\mid x_{i_m})$ , i.e., it is scaled down by  $1 - \beta_{m}$ . If  $\tilde{z}_m$  is incorrect, then we adjust the gradient estimate to try to make  $\tilde{z}_m$  less likely under  $p_{\theta}$ . As the model becomes more accurate (i.e.,  $\beta$  gets closer to 1), we give more weight to incorrect rationales (when they occur) and less weight to correct rationales (most of the time).
CONTROL_VARIATE_GRADIENT Estimate is more expensive than BASIC_GRADIENT Estimate, since we must compute gradients not only for the rationales in memory but also for any incorrect rationales we generate. This may be wasteful, especially if many of the weights on those gradients  $(1 - \beta$  for correct proposals,  $\beta$  for incorrect proposals) are close to zero because  $\beta$  is close to zero or one. To reduce this cost, in SUBSAMPLED_CONTROL_VARIATE_GRADIENT Estimate, we use systematic resampling (Hol et al., 2006) to generate a subsample of  $L$  question-rationale pairs, from which we obtain an unbiased estimate of the output of CONTROL_VARIATE_GRADIENT Estimate. We preferentially sample gradients with higher scalar weights; if  $\beta$  is small, we are less likely to sample incorrect rationales (which have weight  $\beta$ ), and if  $\beta$  is large, we are less likely to sample correct proposed rationales (which have weight  $1 - \beta$ ). This can be seen as a generalization of the strategy of Burda et al. (2015, Section 3) for reducing the cost of computing IwAE gradients.
Below, we derive this variance-reduced stochastic MCMC-EM procedure in more detail.
2.1 Derivation
The true gradient. The gradient of the marginal log-likelihood  $\log p_{\theta}(y \mid x)$  with respect to  $\theta$  is
$$
\nabla_ {\theta} \log \sum_ {z} p _ {\theta} (z, y \mid x) = \sum_ {z} \frac {p _ {\theta} (z , y \mid x) \nabla_ {\theta} \log p _ {\theta} (z , y \mid x)}{\sum_ {z ^ {\prime}} p _ {\theta} \left(z ^ {\prime} , y \mid x\right)} = \sum_ {z} p _ {\theta} (z \mid x, y) \nabla_ {\theta} \log p _ {\theta} (z \mid x), \tag {2}
$$
that is, it is the expectation with respect to the posterior  $p_{\theta}(z \mid x, y)$  of the gradient of the conditional log-prior  $\log p_{\theta}(z \mid x)$ , since the likelihood  $p(y \mid z, x) = c(z, y)$  does not depend on  $\theta$ . So if we can sample from the posterior over rationales  $z$  conditioned on the question-answer pair  $x, y$ , then we can compute an unbiased estimate of the gradient of the marginal log-likelihood  $\log p_{\theta}(y \mid x)$ . We can interpret this as "bootstrapping" rationales  $z$  that are consistent with both the prior on rationales  $p_{\theta}(z \mid x)$  and the observed answer  $y$  (cf. Zelikman et al., 2022).
Independence sampler for  $p_{\theta}(z \mid x, y)$ . We cannot directly sample from  $p_{\theta}(z \mid x, y)$ , so we resort to Markov chain Monte Carlo (MCMC). We maintain a memory (cf. Hewitt et al., 2020) of a single rationale  $z_{n}$  for each question-answer pair  $x_{n}, y_{n}$ , and each iteration we apply a random update to  $z_{n}$  that leaves the posterior  $p_{\theta}(z_{n} \mid x_{n}, y_{n})$  invariant (cf. Tieleman, 2008). Each MCMC update brings the  $z_{n}$ 's closer in distribution to  $p_{\theta}(z_{n} \mid x_{n}, y_{n})$  (Cover, 1999; Murray & Salakhutdinov, 2008). However, updates to  $\theta$  may change the posterior  $p_{\theta}(z_{n} \mid x_{n}, y_{n})$ , so we must keep updating the chains to control the bias of our gradient estimates.
To update the chains, we use a simple, hyperparameter-free independence sampler (Tierney, 1994); a Metropolis-Hastings (Hastings, 1970) update that proposes updating the current state  $z$  with a draw  $\tilde{z}$  from a distribution  $r_{x,y}$  that does not depend on  $z$ , and accepts the update with probability
Algorithm 1 TRICE
Input: Generative model  $p_{\theta}(z,y \mid x)$ , is-correct function  $c(z,y)$ , dataset  $x_{1:N}, y_{1:N}$ , hinted guide distribution  $q(z \mid x,y)$ , initial parameters  $\theta$ , optimizer update function  $h(\theta, g, t)$ , minibatch size  $M$ , gradient minibatch size  $L$ , number of iterations  $T$ .
Output: Tuned parameters  $\theta$ , rationales  $z_{1:N}$ .
1: for  $n\in 1,\dots ,N$  do (in parallel) // Initialize Markov chain states.
2: Sample  $z_{n}$  from  $q(z \mid x_{n}, y_{n})$ . // Sample "fallback" rationale from guide  $q$ .
4: for  $t\in 1,\dots ,T$  do // Main optimization loop.
5: Get next minibatch of  $M$  indices  $i_{1:M}$  into the dataset.
6: for  $m\in 1,\ldots ,M$  do (in parallel) // Take one MCMC step to update Markov chain states.
7: Sample  $\tilde{z}_m$  from  $p_{\theta}(z \mid x_{i_m})$ .
8: if  $c(\tilde{z}_m, y_{i_m})$  then // Accept or reject proposal.
9: Update  $z_{i_m}\gets \tilde{z}_m$
11: Let  $\tilde{c}_m = c(\tilde{z}_m, y_{i_m})$ . // Whether the proposal is correct.
12: Let  $c_{m}^{\prime} = c(z_{i_{m}},y_{i_{m}})$ . // Whether the updated rationale is correct.
14: Compute  $\hat{g}$  using either BASIC_GRADIENT Estimate  $(z,x,c^{\prime})$  CONTROL_VARIATE_GRADIENT Estimate  $(z,x,\tilde{z},\tilde{c},c^{\prime})$  or SUBSAMPLED_CONTROL_VARIATE_GRADIENT Estimate  $(z,x,\tilde{z},\tilde{c},c^{\prime})$
15: Update  $\theta \gets h(\theta ,\hat{g},t)$  //Apply gradient update.
17: return  $\theta$ ,  $z_{1:N}$ .
18: procedure BASIC_GRADIENT Estimate  $(z,x,c^{\prime})$
19: return  $\frac{1}{\sum_{m}c_{m}^{\prime}}\sum_{m}c_{m}^{\prime}\nabla_{\theta}\log p_{\theta}(z_{i_m}\mid x_{i_m})$
20: end procedure
21: procedure CONTROL_VARIATE_GRADIENT Estimate  $(z,x,\tilde{z},\tilde{c},c^{\prime})$
22: for  $m\in 1,\ldots ,M$  do (in parallel)
23: Set  $\beta_{m} = \frac{\sum_{m^{\prime}\neq m}c_{m^{\prime}}^{\prime}\tilde{c}_{m^{\prime}}}{\sum_{m^{\prime}\neq m}c_{m^{\prime}}^{\prime}}$  // Compute leave-one-out control-variate scales.
24: end for
25: return  $\frac{1}{\sum_{m}c_{m}^{\prime}}\sum_{m}c_{m}^{\prime}\left(\nabla_{\theta}\log p_{\theta}(z_{i_m}\mid x_{i_m}) - \beta_{m}\nabla_{\theta}\log p_{\theta}(\tilde{z}_{m}\mid x_{i_m})\right)$
26: end procedure
27: procedure SUBSAMPLED_CONTROL_VARIATE_GRADIENT Estimate(z, x,  $\tilde{z}$ $\tilde{c}, c^{\prime}$
28: for  $m\in 1,\ldots ,M$  do (in parallel)
29: Set  $\beta_{m} = \frac{\sum_{m^{\prime}\neq m}c_{m^{\prime}}^{\prime}\tilde{c}_{m^{\prime}}}{\sum_{m^{\prime}\neq m}c_{m^{\prime}}^{\prime}}$  // Compute leave-one-out control-variate scales.
30: Set  $\tilde{w}_m = c_m'(1 - \tilde{c}_m\beta_m)$ , // Compute unnormalized weights for subsampling.  $\tilde{w}_{M + m} = c_m'(1 - \tilde{c}_m)\beta_m$
31: end for
32: Choose a subset of  $L$  indices  $j_{1:L}$  using systematic resampling with probabilities  $\frac{\tilde{w}}{\sum_{m=1}^{2M} \tilde{w}_m}$ .
33: for  $\ell \in 1,\ldots ,L$  do (in parallel)
34: if  $j_{\ell}\leq M$  then // Selected correct rationale.
35: Let  $\hat{m} = j_{\ell}$ ,  $\hat{z} = z_{i_{\hat{m}}}$ ,  $s = 1$ .
36: else // Selected incorrect rationale.
37: Let  $\hat{m} = j_{\ell} - M, \hat{z} = \tilde{z}_{m}, s = -1$
38: end if
39: Compute  $\hat{g}_{\ell} = s\nabla_{\theta}\log p_{\theta}(\hat{z}\mid x_{i_m})$  // Negate gradient if  $\ell$ th rationale is incorrect.
40: end for
41: return  $\frac{\sum_{m=1}^{2M} \tilde{w}_m}{\sum_m c_m'} \frac{1}{L} \sum_{\ell=1}^L \hat{g}_{\ell}$ .
42: end procedure
$\alpha (\tilde{z}\mid z) = \min \left\{1,\frac{p_{\theta}(\tilde{z},y|x) / r_{x,y}(\tilde{z})}{p_{\theta}(z,y|x) / r_{x,y}(z)}\right\}$ . We choose  $r_{x,y}(z) = p_{\theta}(z\mid x)$ , which simplifies the acceptance probability to  $\alpha (\tilde{z}\mid z) = \min \left\{1,\frac{p_{\theta}(y|x,\tilde{z})}{p_{\theta}(y|x,z)}\right\}$ . This is 1 if  $c(\tilde{z},y) = 1$ , 0 if  $c(\tilde{z},y) = 0$  and  $c(z,y) = 1$ , and ill-defined (implying that we have to reject) if both  $c(z,y) = 0$  and  $c(\tilde{z},y) = 0$ . So we accept whenever the proposal  $\tilde{z}$  is correct, and reject otherwise.
Remarks: Independence samplers can be understood as "Metropolized" importance samplers that spread the work of generating and evaluating proposals over time. In our setting, the update can also be interpreted as attempting to sample from the posterior by rejection sampling, then falling back on an old sample if that fails. The expected number of iterations between successful updates is  $p(y \mid x)^{-1}$ , so mixing will be faster for easier questions  $x$ , and will accelerate as the model improves.
Basic gradient estimator. This MCMC/rejection-sampling procedure lets us approximate the gradient of the marginal log-likelihood in Equation (2). Denoting as  $z$  the state<sup>5</sup> of the Markov chain for an example  $x, y$  before the update, we sample a proposal  $\tilde{z}$  from  $p_{\theta}(z \mid x)$ , accept the new state if it is correct (i.e., if  $c(\tilde{z}, y) = 1$ ), and compute the gradient of the log-probability of the result:
$$
z ^ {\prime} = c (\tilde {z}, y) \tilde {z} + (1 - c (\tilde {z}, y)) z; \hat {g} = \nabla_ {\theta} \log p _ {\theta} \left(z ^ {\prime} \mid x\right); \mathbb {E} _ {z, \tilde {z}} [ \hat {g} \mid \theta ] \approx \mathbb {E} _ {p _ {\theta} (z | x, y)} [ \nabla_ {\theta} \log p _ {\theta} (z \mid x) ], \tag {3}
$$
where  $\mathbb{E}_{z,z}[\cdot \mid \theta ]$  denotes an expectation with respect to both the proposal  $\tilde{z}$  and the previous state  $z$ .
Remarks: The estimate will have low bias if the distribution of  $z'$  is close to the posterior  $p(z \mid x, y)$ , which we expect to be true if the chain is mixing quickly enough relative to how fast  $\theta$  is changing. This will happen if either the probability of getting a correct answer is high, or if  $\theta$  is changing slowly due to a small learning rate and/or gradient. If the model's training-set accuracy improves with training and we use a decaying learning-rate schedule, then as training proceeds both of these factors should work to reduce the bias of the gradient estimate.
Adding a control variate. The mean of an estimator  $\hat{g}$  is not affected by subtracting a zero-mean random variable  $b$  from it. If  $b$  is positively correlated with  $\hat{g}$ , then  $\hat{g} - b$  can have lower variance than  $\hat{g}$ , and we say that  $b$  can be used as a "control variate" (Owen & Zhou, 2000). Since, by the score-function identity,  $\mathbb{E}_{p_{z|x}}[\beta \nabla_{\theta} \log p_{\theta}(z \mid x)] = 0$  (for any scalar  $\beta$  independent of  $z$ ), we can use the proposed samples  $\tilde{z}$  to generate control variates for our gradient estimator:
$$
\begin{array}{l} \mathbb {E} _ {z, \tilde {z}} [ \hat {g} \mid \theta ] = \mathbb {E} _ {z} [ \mathbb {E} _ {\tilde {z}} [ \nabla_ {\theta} \log p _ {\theta} (z ^ {\prime} \mid x) ] \mid \theta ] \\ \begin{array}{l} \mathbb {E} _ {z} \left[ z \left[ \cdot_ {\theta} \log P _ {\theta} (x ^ {\prime} \mid x) \right] \right] \\ = \mathbb {E} _ {z} \left[ \mathbb {E} _ {\tilde {z}} \left[ \nabla_ {\theta} \log p _ {\theta} \left(z ^ {\prime} \mid x\right) - \beta \nabla_ {\theta} \log p _ {\theta} (\tilde {z} \mid x) \right] \mid \theta \right]. \end{array} \tag {4} \\ \end{array}
$$
Remarks: The value of this estimator will depend on whether or not we accept the proposal  $\tilde{z}$ :
$$
\begin{array}{r} \nabla_ {\theta} \log p _ {\theta} \left(z ^ {\prime} \mid x\right) \\ - \beta \nabla_ {\theta} \log p _ {\theta} (\tilde {z} \mid x) = \left\{ \begin{array}{l l} (1 - \beta) \nabla_ {\theta} \log p _ {\theta} \left(z ^ {\prime} \mid x\right) & \text {i f} \tilde {c} = 1, \\ \nabla_ {\theta} \log p _ {\theta} \left(z ^ {\prime} \mid x\right) - \beta \nabla_ {\theta} \log p _ {\theta} (\tilde {z} \mid x) & \text {i f} \tilde {c} = 0, \end{array} \right. \end{array} \tag {5}
$$
where we use the shorthand  $\tilde{c} \triangleq c(\tilde{z}, y)$ .
This control variate can drive the variance of the gradient estimator to zero as the model converges to perfect accuracy on the training set (cf. Roeder et al., 2017). If we set  $\beta = \pi$ , where  $\pi$  is the probability of a correct answer (i.e., that  $\tilde{c} = 1$ ), then as  $\pi$  gets large, most of the time  $\tilde{c} = 1$  and we multiply our gradient estimator by  $1 - \pi$  (multiplying its variance by a factor of  $(1 - \pi)^2$ ). If  $\tilde{c} = 0$ , then we make use of both a correct and incorrect rationale; the weights attached to these updates will not be small, but if incorrect rationales are relatively rare then their contribution to the variance of the gradient estimator will be correspondingly small. On the other hand, if the model has not yet learned to frequently generate good rationales for the training examples, then we should set  $\beta$  closer to 0, since in this case the signal from the incorrect rationale is less informative—in Appendix C.1 we show that the variance of gradient estimators based on incorrect rationales depends strongly on the model's accuracy  $\pi$ . In Appendix B, we show that choosing  $\beta = \pi$  is in fact optimal up to  $O((1 - \pi)^2)$  terms, and that the variance of the resulting estimator is proportional to  $1 - \pi$ .
Estimating  $\beta$ . For each example  $x_{m}, y_{m}$ , we need to compute a  $\beta_{m} \approx \mathbb{E}[\tilde{c}_{m}]$  in a way that ensures that  $\beta_{m}$  is independent of  $\nabla_{\theta} \log p_{\theta}(\tilde{z}_{m} \mid x_{m})$ . We assume that  $\mathbb{E}[\tilde{c}_{m}] \approx \frac{1}{M} \sum_{m} \mathbb{E}[\tilde{c}_{m}]$  (i.e., that the per-example acceptance probability is close to the average acceptance probability across the minibatch<sup>6</sup>), and compute the leave-one-out estimate  $\beta_{m} = \frac{\sum_{m' \neq m} c_{m'}' \tilde{c}_{m'}}{\sum_{m' \neq m} c_{m'}'}$ , where  $c_{m}' := c(z_{m}', y)$ .
There may be some examples (especially early in training) for which we have not yet generated any correct rationales. We omit these examples from our gradient estimate, since they have likelihood 0 and therefore cannot be representative samples from the posterior.
<sup>6</sup>We also tried keeping a running estimate of the average acceptance probability per example, but we did not find that this more complex scheme provided any empirical advantage.