SlowGuess/NeurIPS_2023_OCR huggingface dataset:
We restrict the estimate to consider only examples for which we have a correct rationale (i.e., where  $c_{m}^{\prime} = 1$ ), since these are the only examples that influence our gradient estimate. Leaving out  $\tilde{c}_{m}$  and  $c_{m}^{\prime}$  from the estimate  $\beta_{m}$  ensures that  $\beta_{m}$  is independent of  $\tilde{z}_{m}$ .
Gradient subsampling. Finally, as described above, we can reduce the cost of our gradient estimator by using systematic resampling to select a subset of rationales. This does not affect the expected value of the estimator as long as the marginal probability of selecting a rationale is proportional to the corresponding weight  $\tilde{w}_m$ , and the averaged gradient is reweighted by  $\frac{\sum_{m=1}^{2M} \tilde{w}_m}{\sum_m c_m'}$ .
2.2 Why not variational inference, reweighted wake-sleep, or rejection sampling?
We considered three alternatives to the MCMC-EM approach that we pursue in this paper: variational EM (e.g., Bishop, 2006), reweighted wake-sleep (RWS; Bornschein & Bengio, 2015; Le et al., 2019), and rejection sampling.
Variational expectation-maximization is a common strategy for training latent-variable models, but variational inference with discrete latent variables is challenging (e.g., Tucker et al., 2017).
RWS is an attractive alternative that avoids high-variance score-function gradients; it proceeds by sampling  $M$  samples  $z_{1:M}$  from a guide model  $q_{\phi}(z \mid x, y)$ , assigning the samples weights  $w_{m} \propto \frac{p_{\theta}(y, z \mid x)}{q_{\phi}(z \mid x, y)}$ , and updating both the model parameters  $\theta$  and the guide parameters  $\phi$  to maximize the reweighted log-probabilities  $\sum_{m} w_{m} \log p_{\theta}(z_{m} \mid x)$  and  $\sum_{m} w_{m} \log q_{\phi}(z_{m} \mid x, y)$ . Unfortunately, we found that RWS training sometimes led to degenerate zero-length rationales  $z$ . Figure 1 suggests a partial explanation: shorter sequences get higher weights, so the model and guide learn to produce shorter and shorter sequences until they consistently produce empty rationales.
Why do longer sequences tend to get lower weights? We can write the unnormalized weights as  $\tilde{w}_m = c(y,z_m)\frac{p_\theta(z_m|x)}{q_\phi(z_m|x,y)} = c(y,z_m)\prod_{t=1}^{T_m}\frac{p_\theta(z_{m,t}|x,z_{m,1:(t-1)})}{q_\phi(z_{m,t}|x,y,z_{m,1:(t-1)})}$ , where  $T_m$  is the length of  $z_m$  and  $\epsilon$  is added to address the case where none of the samples are correct. If there is a mismatch between  $q(z_{m,t}|x,z_{m,1:(t-1)})$  and  $p(z_{m,t}|x,z_{m,1:(t-1)})$ , then  $\frac{p_\theta(z_{m,t}|x,z_{m,1:(t-1)})}{q_\phi(z_{m,t}|x,y,z_{m,1:(t-1)})}$  will usually be less than one, with rare high-weight exceptions that ensure that  $\mathbb{E}_q[p(z|x)/q(z|x)] = 1$ .
If these exceptions are rare enough to not typically appear in a sample of  $M$  sequences  $z_{1:M}$ , then the normalized weights  $w_{1:M} = \frac{\tilde{w}_{1:M}}{\sum_m\tilde{w}_m}$  will tend to assign higher mass to shorter sequences unless those shorter sequences are much less likely to be correct.
With careful initialization and learning-rate tuning, we could sometimes get RWS to avoid this problem
of empty rationales. But this led to a new problem: the guide  $q_{\phi}(z \mid x, y)$  learned to closely mimic the prior  $p(z \mid x)$  until the very end of the rationale, and then simply paste in the correct answer whether or not it had anything to do with the rationale up to that point (cf. Turpin et al., 2023). Figure 5 in Appendix E shows a representative example in which the guide model ignores the answer it arrived at through incorrect reasoning and pastes in the correct answer.
Quantitatively, denoting by  $t$  the index of the token at which the "final answer" section of the rationale begins, in one run we found that the average KL between  $q(z_{1:t} \mid x, y)$  and  $p(z_{1:t} \mid x)$  was about 0.61 nats, while the conditional KL between  $q(z_{(t+1):T} \mid x, y, z_{1:t})$  and  $p(z_{(t+1):T} \mid x, z_{1:t})$  was about 42.5 nats, confirming that the guide was not "reasoning backwards", just copying the correct answer.
Finally, we considered a rejection-sampling scheme in which we sample  $K$  proposal rationales  $z_{1:K}$  from  $p(z \mid x)$ , and average the gradients from those rationales that lead to correct answers. We will present the quantitative results in Section 4; our main finding is that, while this scheme can work, it requires reducing the minibatch size by a factor of  $K$  to keep the per-iteration cost constant compared to TRICE, which in turn leads to slower convergence and/or worse final results.
A number of methods have proposed rationale generation for problem-solving tasks in neural sequence models, including both fully supervised and few-shot approaches (Wei et al., 2022; Nye et al., 2021; Kojima et al., 2022; Rajani et al., 2019; Shwartz et al., 2020; Wang et al., 2022b; Zhou et al., 2022; Creswell et al., 2022; Ye & Durrett, 2023). Particularly relevant to our approach is self-consistent chain-of-thought (Wang et al., 2022b), because this can be approximately viewed as marginalizing over rationales at test time. This technique has been successfully applied for a range of quantitative reasoning tasks (Lewkowycz et al., 2022). There is relatively much less work that does imputation or averaging over rationales at training time; perhaps the main instance is STaR (Zelikman et al., 2022), which we discuss in Section 3.1.
Dohan et al. (2022) present a position paper which advocates representing a composition of language model interactions via probabilistic programming. Our treatment of rationales as latent variables is inspired by that work. Lievin (2022) offers another example of interpreting LLMs with CoT as latent-variable models.
Variational inference (e.g., Kingma & Welling, 2013) and wake-sleep methods (e.g., Bornschein & Bengio, 2015) are workhorses of the latent-variable-modeling community, but as we discuss in Section 2.2 we found the bias of these methods to cause serious problems. MCMC-EM is a less-common strategy these days, although a version of it based on Gibbs sampling (Geman & Geman, 1984) it has been widely applied to training undirected graphical models (Tieleman, 2008). TRICE can also be cast as an instance of Markovian score climbing (Naesseth et al., 2020).
ReAct (Yao et al., 2023) demonstrated that injecting reasoning into an RL-style observe-and-act loop significantly increases performance. This approach was extended in Reflexion (Shinn et al., 2023), where an agent can conditionally reflect on an RL trajectory, augmenting the resulting examples which can be used as few-shot examples in subsequent rollouts. These approaches reported significant improvements on their respective evaluation tasks but still rely on the model being able to produce useful and actionable feedback through pure few-shot prompting, whereas our method actively tunes the model to produce thoughts amenable to the task.
Recent work on tool-use within language models also works via imputation, inferring where to insert calls to tools (Parisi et al., 2022; Schick et al., 2023). Their loss functions are similar in spirit to ours, filtering out trajectories which do not lead to valid answers. In this paper, we have treated rationales as latent variables; one could also treat tool-use as a latent variable.
3.1 Self-Taught Reasoner
The most closely related work is the self-taught reasoner (STaR; Zelikman et al., 2022). Besides the arguments in their derivations, there are three significant differences between TRICE and STaR. First, STaR uses greedy decoding, which reduces the diversity of the rationales it trains on. The authors made this choice to reduce the danger of the model getting the right answer despite having a bad rationale. While we do find that our procedure sometimes generates correct answers for the wrong reasons, this did not seem to stand in the way of the model improving on most tasks. One reason may be that our base models are more powerful than the 6B-parameter GPT-J model used in the STaR paper, so they are more likely to generate good rationales from the beginning.
A second difference is that TRICE resamples rationales every iteration, so it are less likely to overfit to any particular rationale. STaR has an inner loop that runs many training iterations on a single set of rationales, meaning it uses stale rationales to estimate the gradient of the marginal likelihood.
<sup>7</sup>We also considered optimizing an importance-weighted bound (Burda et al., 2015) using the prior  $p(z \mid x)$  as a proposal distribution, but instead opted for a simple rejection sampling scheme since this is less biased and equally feasible in our setting.
In our experiments, we observed that this leads to the model effectively memorizing a fixed set of rationales for the training set. Once this happens, the greedy decoding procedure will almost certainly reproduce exactly the same rationales at the beginning of the next outer loop. If these rationales all lead to the correct answer, and STaR has a rationale for each question, then this is a global optimum of the marginal likelihood on the training set! But empirically, STaR often does not find a good rationale for each question, and so it ignores some fraction of the training set (see Section 4).
This tendency to ignore the most difficult questions in the training set follows from STaR's derivation as an approximate policy-gradient algorithm trying to directly minimize the 0-1 loss  $\mathbb{E}_p[1 - c(z,y)] = 1 - p_{\theta}(y\mid x)$ . The derivative of this marginal likelihood is  $p_{\theta}(y\mid x)\nabla_{\theta}\log p_{\theta}(y\mid x)$ , that is, it is the derivative of the marginal log-likelihood (which TRICE tries to maximize) weighted by  $p_{\theta}(y\mid x)$ . This weighting causes difficult examples to contribute little to the gradient used to update the model, so the model may "give up" on questions that it cannot yet solve. This is one argument for trying to maximize log-likelihoods instead of likelihoods.
A final, minor difference is that when STaR updates its rationales, it may replace a rationale from the model  $p(z \mid x)$  with a rationale from a surrogate  $q_{\theta}(z \mid x, y)$ . As the model memorizes a set of correct rationales for the training set, STaR becomes less likely to fall back on the surrogate, but this choice could affect early training dynamics.
We evaluate TRICE on the GSM8K (Cobbe et al., 2021) dataset and the 27 BigBench-Hard (BBH) tasks (Suzgun et al., 2022b) using the medium-size PaLM 2-M (Anil et al., 2023) Transformer-based LLM (Vaswani et al., 2017). For the BBH experiments, we used the Flan instruction-tuned (Chung et al., 2022) version of PaLM 2; for GSM8K, we used the base PaLM 2 model, since GSM8K is included in the Flan training datasets. All experiments were run on TPU v4 and v5e chips (Jouppi et al., 2023). Examples of generated rationales can be found in Appendix E.
Rather than fine-tune the model weights, we use prompt tuning (Lester et al., 2021); we preprocess a sequence of embedding vectors  $\theta$  (a "soft prompt") to the embeddings corresponding to the tokenized CoT prompt used to condition the model. Prompt tuning can achieve similar accuracy gains to full fine-tuning, but using a small fraction of the parameters. We initialize the soft prompt with the embedding sequence obtained from a series of three (for BBH) or five (for GSM8K) exemplar CoT prompts, each of the form "Question: <QUESTION>\nAnswer: Let's think step by step.\n<RATIONALE>". We consider two initialization schemes: one where we use the standard few-shot CoT prompts that are provided with BBH, and one where we try to bootstrap a few-shot CoT prompt by sampling random questions from the training set, generating random rationales from the base model, and picking three or five examples where the random rationales lead to correct answers. The first scheme can be seen as a way of fine-tuning a good initial few-shot prompt, but it does require a small amount of detailed CoT supervision, while the second scheme only requires label supervision.
On each BBH task, we split the examples into  $60\%$  train and  $40\%$  test sets. For all but three tasks, this is 150 training and 100 test examples. For GSM8K, we use the standard 7473-example training set and 1319-example test set. We evaluate CoT models' accuracy in two ways: first, using greedy (temperature-0) decoding, and second, using "self-consistency" (Wang et al., 2022b). In self-consistency evaluation, we draw 40 samples and check whether the most common answer is correct; this is a plug-in estimator for the prediction  $\arg \max_y p(y \mid x)$  that minimizes 0-1 loss under the model (although this is not how Wang et al. (2022b) originally motivated the procedure).
We compare against four baseline prompt-tuning methods: direct prompt tuning, CoT prompt tuning, rejection sampling, and STaR (Zelikman et al., 2022). All methods are evaluated against the same validation sets, and use the same training labels, few-shot prompts (except for direct tuning, where we only use question-answer pairs), and initialization strategies as appropriate. Details for each method and its corresponding experimental hyperparameters can be found in Appendix F.
Table 1 and Table 2 summarize the results; more detailed task-by-task BBH summaries are in Appendix D. Even with no human-generated exemplar rationales, TRICE is able to learn to generate rationales that lead to the correct answer. TRICE also outperforms a model trained directly on human-generated rationales on GSM8K (cf. Uesato et al., 2022), perhaps because the cross-entropy
Table 2: Average accuracies (columns 2 and 3) and fraction of training examples for which we can generate correct rationales (column 4) on GSM8K. Direct prompt tuning is initialized with an embedding sequence obtained from a few-shot prompt containing example question-answer pairs ("Q-A"). All remaining prompt-tuning methods are initialized with an embedding sequence obtained from a few-shot prompt containing example question-rationale-answer triples ("Q-R-A") obtained randomly from the GSM8K training set or bootstrapped as described in the text.
loss used in supervised fine-tuning may place more weight on style than substance; it takes far more bits to encode how one expresses a chain of reasoning than it does to encode the reasons themselves.
Initializing the soft prompt with a human-generated 3-shot exemplar question-rationale-answer prompt slightly improves performance on BBH, as does evaluating with self-consistency. By the end of training, TRICE has managed to generate at least one valid rationale for almost all training examples, while STaR fails to generate valid rationales for almost  $10\%$  of training examples. Unlike in the experiments done on Commonsense QA (Talmor et al., 2019) by Zelikman et al. (2022), STaR does not outperform the direct-prompted prompt-tuned model on BBH. This may be because each BBH task includes relatively little training data (150 examples as opposed to CommonsenseQA's 9,741), and so in its inner loop STaR overfits to its relatively small set of bootstrapped rationales. TRICE, on the other hand, can overfit to the small set of questions but at least has a chance to generate a somewhat diverse set of rationales from those questions.
One piece of evidence for this overfitting-rationales hypothesis is that on the final step of its final inner loop, STaR (with bootstrapped initialization) achieves a training sequence-level (not per-token) cross-entropy loss of less than 0.06 on all tasks, and of less than 0.01 on 19 out of 27 tasks. This implies that it has learned to exactly reproduce a single set of rationales with very high probability, which makes it very likely that it will generate those same rationales in the next iteration.
Figure 2 compares estimates for GSM8K of the average training marginal likelihood (i.e., how often a proposal is accepted) and the validation accuracy with greedy decoding as a function of number of
training steps $^{8}$  for rejection sampling and for TRICE with and without the control-variate scheme. The control-variate scheme improves average convergence speed, particularly towards the end of training as the probability of generating correct answers on the training set increases. Both versions of TRICE converge to high training accuracy much faster than rejection sampling.
We proposed TRICE, a method for tuning LLMs to be better at solving question-answering tasks using chain-of-thought (CoT) prompting. By framing the CoT-prompted LLM as a latent-variable model, we were able to derive a principled and effective fine-tuning method. When applied to GSM8K and BIG-Bench Hard (BBH) tasks, TRICE outperforms three strong baselines: direct prompt-tuning, STaR, and rejection sampling. While we derived TRICE in the context of CoT question-answering, its basic MCMC-EM strategy could be employed more broadly, for example to tool-use problems.
Limitations: We only evaluated TRICE with prompt-tuning on a medium-size LLM; it may be that it behaves differently on smaller models, larger models, or when using other fine-tuning strategies. TRICE is a gradient-based tuning algorithm, but many of the most capable LLMs are proprietary, and their owners often do not provide any public mechanism for gradient-based fine-tuning. This makes it hard to evaluate how well TRICE would work when used with, say, GPT-4 (OpenAI, 2023). Finally, our quantitative evaluations focused on whether the LLM could produce the right answer; we did not formally evaluate the quality of the reasoning in the rationales themselves (cf. Uesato et al., 2022).
Broader Impacts: This work aims to improve the capabilities of LLMs by making them better able to answer questions accurately and transparently. However, more-capable LLMs may be used in malicious or unsafe ways, fine-tuning on uncurated question-answering datasets may introduce biases into the models, and more widely used LLMs will contribute a larger carbon footprint.
Rationales may make it easier for motivated users to judge the trustworthiness of LLM outputs. But many users may not read and critique an LLM's rationales, taking the mere existence of a rationale as evidence of truth. If chain-of-thought rationales promote uncritical trust, they could lead to harm.
<sup>8</sup>We set the cost per iteration of rejection sampling and TRICE with and without the control-variate scheme to be directly comparable: for rejection sampling, we reduce the minibatch size by a factor of four and generate four times as many proposals per example; for TRICE with the control-variate scheme, we set the gradient minibatch size  $L$  equal to the number of examples per minibatch  $M$  (note that this does still involve subsampling, since each example could potentially contribute both a correct and an incorrect rationale to the gradient estimate).
Acknowledgements: We appreciate Daniel Freeman and Enrique Piqueras' contributions to the infrastructure that we used in our experiments. We thank Kevin Murphy, Ben Lee, Brian Patton, and Jascha Sohl-Dickstein for helpful discussions.
Anil, R., Dai, A. M., First, O., Johnson, M., Lepikhin, D., Passos, A., Shakeri, S., Taropa, E., Bailey, P., Chen, Z., et al. Palm 2 technical report. arXiv preprint arXiv:2305.10403, 2023.
Bishop, C. M. Pattern recognition and machine learning. Springer, 2006.
Bornshein, J. and Bengio, Y. Reweighted wake-sleep, 2015.
Burda, Y., Grosse, R., and Salakhutdinov, R. Importance weighted autoencoders. arXiv preprint arXiv:1509.00519, 2015.
Chung, H. W., Hou, L., Longpre, S., Zoph, B., Tay, Y., Fedus, W., Li, E., Wang, X., Dehghani, M., Brahma, S., et al. Scaling instruction-finetuned language models. arXiv preprint arXiv:2210.11416, 2022.
Cobbe, K., Kosaraju, V., Bavarian, M., Chen, M., Jun, H., Kaiser, L., Plappert, M., Tworek, J., Hilton, J., Nakano, R., Hesse, C., and Schulman, J. Training verifiers to solve math word problems. arXiv preprint arXiv:2110.14168, 2021.
Cover, T. M. Elements of information theory. John Wiley & Sons, 1999.
Creswell, A., Shanahan, M., and Higgins, I. Selection-Inference: Exploiting large language models for interpretable logical reasoning. arXiv preprint arXiv:2205.09712, May 2022.
Dohan, D., Xu, W., Lewkowycz, A., Austin, J., Bieber, D., Lopes, R. G., Wu, Y., Michalewski, H., Saurous, R. A., Sohl-dickstein, J., Murphy, K., and Sutton, C. Language model cascades, 2022.
Geman, S. and Geman, D. Stochastic relaxation, Gibbs distributions, and the Bayesian restoration of images. IEEE Transactions on pattern analysis and machine intelligence, 6:721-741, 1984.
Hastings, W. K. Monte Carlo sampling methods using Markov chains and their applications. Biometrika, 57(1):97-109, 04 1970. ISSN 0006-3444. doi: 10.1093/biomet/57.1.97. URL https://doi.org/10.1093/biomet/57.1.97.
Hewitt, L. B., Le, T. A., and Tenenbaum, J. B. Learning to learn generative programs with memoised wake-sleep. In Uncertainty in Artificial Intelligence, 2020.
Hol, J. D., Schon, T. B., and Gustafsson, F. On resampling algorithms for particle filters. In 2006 IEEE Nonlinear Statistical Signal Processing Workshop, pp. 79-82, 2006. doi: 10.1109/NSPW.2006.4378824.
Jouppi, N., Kurian, G., Li, S., Ma, P., Nagarajan, R., Nai, L., Patil, N., Subramanian, S., Swing, A., Towles, B., et al. Tpu v4: An optically reconfigurable supercomputer for machine learning with hardware support for embeddings. In Proceedings of the 50th Annual International Symposium on Computer Architecture, pp. 1-14, 2023.
Kingma, D. P. and Ba, J. Adam: A method for stochastic optimization. In Bengio, Y. and LeCun, Y. (eds.), 3rd International Conference on Learning Representations, ICLR 2015, San Diego, CA, USA, May 7-9, 2015, Conference Track Proceedings, 2015. URL http://arxiv.org/abs/1412.6980.
Kingma, D. P. and Welling, M. Auto-encoding variational bayes. arXiv preprint arXiv:1312.6114, 2013.
Kojima, T., Gu, S. S., Reid, M., Matsuo, Y., and Iwasawa, Y. Large language models are Zero-Shot reasoners. arXiv preprint arXiv:2205.11916, May 2022.
Le, T. A., Kosiorek, A. R., Siddharth, N., Teh, Y. W., and Wood, F. Revisiting reweighted wake-sleep for models with stochastic control flow. In Uncertainty in Artificial Intelligence, 2019.