SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Suzgun, M., Scales, N., Schärli, N., Gehrmann, S., Tay, Y., Chung, H. W., Chowdhery, A., Le, Q. V., Chi, E. H., Zhou, D., et al. Challenging big-bench tasks and whether chain-of-thought can solve them. arXiv preprint arXiv:2210.09261, 2022b.
Talmor, A., Herzig, J., Lourie, N., and Berant, J. CommonsenseQA: A question answering challenge targeting commonsense knowledge. In Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume 1 (Long and Short Papers), pp. 4149-4158, Minneapolis, Minnesota, June 2019. Association for Computational Linguistics. doi: 10.18653/v1/N19-1421. URL https://aclanthology.org/N19-1421.
Tieleman, T. Training restricted boltzmann machines using approximations to the likelihood gradient. In Proceedings of the 25th international conference on Machine learning, pp. 1064-1071, 2008.
Tierney, L. Markov chains for exploring posterior distributions. Annals of Statistics, 22(4):1701-1728, December 1994.
Tucker, G., Mnih, A., Maddison, C. J., Lawson, J., and Sohl-Dickstein, J. Rebar: Low-variance, unbiased gradient estimates for discrete latent variable models. Advances in Neural Information Processing Systems, 30, 2017.
Turpin, M., Michael, J., Perez, E., and Bowman, S. R. Language models don't always say what they think: Unfaithful explanations in chain-of-thought prompting. arXiv preprint arXiv:2305.04388, 2023.
Uesato, J., Kushner, N., Kumar, R., Song, F., Siegel, N., Wang, L., Creswell, A., Irving, G., and Higgins, I. Solving math word problems with process-and outcome-based feedback. arXiv preprint arXiv:2211.14275, 2022.
Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., Kaiser, L., and Polosukhin, I. Attention is all you need. Advances in neural information processing systems, 30, 2017.
Wang, X., Wei, J., Schuurmans, D., Le, Q., Chi, E., Narang, S., Chowdhery, A., and Zhou, D. Self-Consistency improves chain of thought reasoning in language models. arXiv preprint arXiv:2203.11171, March 2022a.
Wang, X., Wei, J., Schuurmans, D., Le, Q., Chi, E., and Zhou, D. Self-consistency improves chain of thought reasoning in language models. arXiv preprint arXiv:2203.11171, 2022b.
Wei, J., Wang, X., Schuurmans, D., Bosma, M., Chi, E., Le, Q., and Zhou, D. Chain of thought prompting elicits reasoning in large language models. "arXiv preprint arXiv:2201.11903", January 2022.
Yao, S., Zhao, J., Yu, D., Du, N., Shafran, I., Narasimhan, K., and Cao, Y. React: Synergizing reasoning and acting in language models, 2023.
Ye, X. and Durrett, G. Explanation selection using unlabeled data for In-Context learning. "arXiv preprint arXiv:2302.04813", February 2023.
Zelikman, E., Wu, Y., and Goodman, N. D. Star: Bootstrapping reasoning with reasoning. arXiv preprint arXiv:2203.14465, 2022.
Zhou, D., Scharli, N., Hou, L., Wei, J., Scales, N., Wang, X., Schuurmans, D., Bousquet, O., Le, Q., and Chi, E. Least-to-Most prompting enables complex reasoning in large language models. "arXiv preprint arXiv:2205.10625", May 2022.
Supplemental Material for "Training Chain-of-Thought via Latent-Variable Inference"
A Generalizing TRICE to Nondeterministic Likelihood Models
To apply TRICE beyond question-answering problems, we might want to use a nondeterministic likelihood model. For example, our desired output  $y$  might be a summary of a text  $x$ , and  $z$  might be a scratchpad or outline. In situations like this, there might be many  $y$ 's that are appropriate for a given  $x$  and  $z$ . So from a modeling perspective, it could make sense to make  $p(y \mid x, z)$  have nonzero entropy. But there is also a computational reason to prefer such a model: as the number of reasonable values that  $y$  could take given  $x$  increases, the probability of sampling the precise  $y$  that was observed goes down at a rate that might be exponential in the size of the output space.
Fortunately, we can easily extend TRICE to accommodate nondeterministic likelihoods. The differences are:
1. As discussed in Section 2.1, the acceptance probability for an independence sampler where the proposals are generated from  $p(z \mid x)$  is  $\alpha(\tilde{z} \mid z) = \min\{1, \frac{p(y|\tilde{z},x)}{p(y|z,x)}\}$ . So instead of deterministically accepting correct proposed latents  $\tilde{z}$ , we update the memory probabilistically, always accepting proposals that make the observed  $y$  more likely, but sometimes also accepting proposals that make  $y$  a little less likely.
2. The control-variate weights  $\beta_{m}$  are now a function of the acceptance probabilities  $\alpha$ :  $\beta_{m} = \frac{1}{M - 1}\sum_{m^{\prime}\neq m}\alpha (\tilde{z}_{m^{\prime}}\mid z_{i_{m^{\prime}}})$
The value of the control variate in this setting may be less than it is in the deterministic-likelihood setting. Even if we learn a model that consistently produces good latents  $z$  (in the sense that they lead to valid outputs  $y$ ), this does not guarantee that it will consistently generate latents that are consistent with the particular  $y$  that was observed. For example, there might be multiple reasonable ways to outline a long text, some of which lead to different summary paragraphs. In this scenario, the acceptance probability may not converge to something close to 1, and the variance-reduction effect from the control variate will be modest.
B Derivation of the Control Variate Scaling Heuristic
Our analysis will focus on the variance of the  $k$ th element  $\hat{g}_k$  of our gradient estimator. To minimize clutter, we suppress subscript indices and define  $g' = \nabla_{\theta_k} \log p(z' \mid x)$ ,  $\tilde{g} = \nabla_{\theta_k} \log p(\tilde{z} \mid x)$  and  $g = \nabla_{\theta_k} \log p(z \mid x)$ , where  $z'$  is the updated rationale,  $\tilde{z}$  is the proposed rationale, and  $z$  is the previous rationale in memory. Note that  $g' = \tilde{g}$  if the proposal  $\tilde{z}$  is accepted, and  $g' = g$  if  $\tilde{z}$  is rejected. We assume that the previous rationale  $z$  is a sample from  $p(z \mid x, y)$ . We also define the shorthands  $g_+ = \mathbb{E}[\tilde{g} \mid c(\tilde{z}, y) = 1]$ ,  $g_- = \mathbb{E}[\tilde{g} \mid c(\tilde{z}, y) = 0]$ ,  $v = \mathrm{Var}(\tilde{g}) = \mathbb{E}[\tilde{g}^2]$ ,  $v_+ = \mathrm{Var}(\tilde{g} \mid c(\tilde{z}, y) = 1)$ ,  $v_- = \mathrm{Var}(\tilde{g} \mid c(\tilde{z}, y) = 0)$ . Finally, we define the probability of accepting a proposal  $\pi = \mathbb{E}[c(\tilde{z}, y)]$ .
Our gradient estimator is
$$
\hat {g} = g ^ {\prime} - \beta \tilde {g}; \quad \mathbb {E} [ \hat {g} ] = g _ {+}. \tag {6}
$$
The variance of  $\hat{g}$  is
$$
\begin{array}{l} \mathbb {E} \left[ \left(g ^ {\prime} - g _ {+} - \beta \tilde {g}\right) ^ {2} \right] = \mathbb {E} \left[ \left(g ^ {\prime} - g _ {+}\right) ^ {2} \right] + \beta^ {2} \mathbb {E} \left[ \tilde {g} ^ {2} \right] - 2 \beta \mathbb {E} \left[ \left(g ^ {\prime} - g _ {+}\right) \tilde {g} \right] (7) \\ = v _ {+} + \beta^ {2} v - 2 \beta \mathbb {E} \left[ g ^ {\prime} \tilde {g} \right]. (8) \\ \end{array}
$$
In Equation (8) we use the fact that  $\mathbb{E}[\tilde{g}] = 0$  and  $g_{+}$  is nonrandom and therefore independent of  $\tilde{g}$ . Breaking out the cases where  $c(\tilde{z},y) = 0$  (so  $g^{\prime} = g$ ) and  $c(\tilde{z},y) = 1$  (so  $g^{\prime} = \tilde{g}$ ), the rightmost expectation is
$$
\begin{array}{l} \mathbb {E} \left[ g ^ {\prime} \tilde {g} \right] = \pi \mathbb {E} \left[ \tilde {g} ^ {2} \mid c (\tilde {z}, y) = 1 \right] + (1 - \pi) \mathbb {E} \left[ g \tilde {g} \mid c (\tilde {z}, y) = 0 \right] (9) \\ = \pi \left(v _ {+} + g _ {+} ^ {2}\right) + (1 - \pi) g _ {+} g _ {-}. (10) \\ \end{array}
$$
We can simplify the term on the right. Since  $0 = \mathbb{E}[\tilde{g}] = \pi g_{+} + (1 - \pi)g_{-}$ ,  $(1 - \pi)g_{+}g_{-} = -\pi g_{+}^{2}$ . Plugging this into Equation (10), we have
$$
\mathbb {E} \left[ g ^ {\prime} \tilde {g} \right] = \pi \left(v _ {+} + g _ {+} ^ {2}\right) - \pi g _ {+} ^ {2} = \pi v _ {+}. \tag {11}
$$
So the variance of  $\hat{g}$  simplifies to
$$
\operatorname {V a r} (\hat {g}) = (1 - 2 \beta \pi) v _ {+} + \beta^ {2} v. \tag {12}
$$
Taking the derivative with respect to  $\beta$  shows that this is minimized when
$$
2 \left(\beta v - \pi v _ {+}\right) = 0; \quad \beta^ {\star} = \pi \frac {v _ {+}}{v}. \tag {13}
$$
Plugging this back into Equation (12) gives the optimal variance  $v^{\star}$ :
$$
v ^ {\star} = v _ {+} + \pi^ {2} \frac {v _ {+} ^ {2}}{v} - 2 \pi^ {2} \frac {v _ {+} ^ {2}}{v} = v _ {+} \left(1 - \pi^ {2} \frac {v _ {+}}{v}\right). \tag {14}
$$
We can expand  $v$  by considering the case where  $c(\tilde{z},y) = 1$  (which happens with probability  $\pi$  and the case where  $c(\tilde{z},y) = 0$  (which happens with probability  $1 - \pi$ ):
$$
\begin{array}{l} v = \mathbb {E} [ \tilde {g} ^ {2} ] = \pi v _ {+} + \pi g _ {+} ^ {2} + (1 - \pi) v _ {-} + (1 - \pi) g _ {-} ^ {2} \\ = \pi v _ {+} - (1 - \pi) g _ {+} g _ {-} + (1 - \pi) v _ {-} - \pi g _ {+} g _ {-} \\ = \pi v _ {+} + (1 - \pi) v _ {-} - g _ {+} g _ {-} \\ = \pi v _ {+} + (1 - \pi) v _ {-} + \frac {1 - \pi}{\pi} g _ {-} ^ {2} \tag {15} \\ = \pi v _ {+} + (1 - \pi) v _ {-} + (1 - \pi) (2 - \pi) g _ {-} ^ {2} + O ((1 - \pi) ^ {3}) \\ = \pi v _ {+} + (1 - \pi) v _ {-} + (1 - \pi) (1 + 1 - \pi) g _ {-} ^ {2} + O ((1 - \pi) ^ {3}) \\ = \pi v _ {+} + (1 - \pi) \left(v _ {-} + g _ {-} ^ {2}\right) + O \left((1 - \pi) ^ {2}\right), \\ \end{array}
$$
where in the second line we again use the fact that  $\pi g_{+} = -(1 - \pi)g_{-}$ , and in the third-to-last line we approximate  $\frac{1}{\pi}$  with the first-order Taylor approximation  $\frac{1}{\pi} = 2 - \pi + O((1 - \pi)^2)$ . Thus, we can write
$$
\begin{array}{l} \pi \frac {v _ {+}}{v} = \frac {\pi v _ {+}}{\pi v _ {+} + (1 - \pi) \left(v _ {-} + g _ {-} ^ {2}\right) + O ((1 - \pi) ^ {2})} \\ = \left(1 + \frac {1 - \pi}{\pi v _ {+}} \left(v _ {-} + g _ {-} ^ {2}\right) + O \left(\left(1 - \pi\right) ^ {2}\right)\right) ^ {- 1} \tag {16} \\ = 1 - \frac {1 - \pi}{\pi v _ {+}} (v _ {-} + g _ {-} ^ {2}) + O ((1 - \pi) ^ {2}). \\ \end{array}
$$
So the optimal variance  $v^{\star}$  is
$$
\begin{array}{l} v ^ {\star} = v _ {+} \left(1 - \pi^ {2} \frac {v _ {+}}{v}\right) \\ = v _ {+} \left(1 - \pi + \frac {1 - \pi}{v _ {+}} \left(v _ {-} + g _ {-} ^ {2}\right) + O \left((1 - \pi) ^ {2}\right)\right) \tag {17} \\ = (1 - \pi) \left(v _ {+} + v _ {-} + g _ {-} ^ {2}\right) + O \left(\left(1 - \pi\right) ^ {2}\right)). \\ \end{array}
$$
By contrast, plugging our heuristic value of  $\beta = \pi$  into Equation (12) gives the suboptimal variance  $v^{\pi}$ :
$$
\begin{array}{l} v ^ {\pi} = (1 - 2 \pi^ {2}) v _ {+} + \pi^ {2} v \\ = (1 - 2 \pi^ {2}) v _ {+} + \pi^ {2} (\pi v _ {+} + (1 - \pi) (v _ {-} + g _ {-} ^ {2}) + O ((1 - \pi) ^ {2})) \tag {18} \\ = (1 - 2 \pi^ {2} + \pi^ {3}) v _ {+} + \pi^ {2} (1 - \pi) \left(v _ {-} + g _ {-} ^ {2}\right) + O ((1 - \pi) ^ {2}) \\ = (1 - \pi) \left(v _ {+} + v _ {-} + g _ {-} ^ {2}\right) + O \left((1 - \pi) ^ {2}\right), \\ \end{array}
$$
where we use the approximation  $\pi^k = (1 - (1 - \pi))^k = 1 - k(1 - \pi) + O((1 - \pi)^2)$  to simplify the  $\pi^2$  and  $\pi^3$  terms. Thus, we conclude that  $v^{\star}$  and  $v^{\pi}$  are the same up to  $O((1 - \pi)^2)$ , and so as the probability  $\pi$  of getting the correct answer increases, the suboptimality of setting  $\beta = \pi$  goes down faster than the variance does.
C On Gradient Estimators Based Solely on Incorrect Rationales
We adopt the same shorthands as in Appendix B.
Our analysis centers on the identity
$$
0 = \mathbb {E} [ \tilde {g} ] = \pi g _ {+} + (1 - \pi) g _ {-}, \tag {19}
$$
which relates the gradient we want to estimate  $g_{+}$  (the expected gradient given that the rationale is correct) to  $g_{-}$  (the expected gradient given that the rationale is incorrect).
C.1 Variance of incorrect-rationale gradient estimators
Equation (19) implies that
$$
g _ {+} = - \frac {1 - \pi}{\pi} g _ {-}. \tag {20}
$$
Deferring for the moment the difficulty in estimating  $\pi^{-1}$  (see Appendix C.2 below), we can consider the variance of an estimator based on the right hand side of Equation (20):
$$
\begin{array}{l} \mathbb {V} \left[ - \frac {1 - c (\tilde {z} , y)}{\pi} \tilde {g} \right] = \mathbb {E} \left[ \frac {(1 - c (\tilde {z} , y)) ^ {2}}{\pi^ {2}} \tilde {g} ^ {2} \right] - g _ {+} ^ {2} \\ = \frac {1 - \pi}{\pi^ {2}} \left(v _ {-} + g _ {-} ^ {2}\right) - g _ {+} ^ {2} \tag {21} \\ = \frac {1 - \pi}{\pi^ {2}} \left(v _ {-} + g _ {-} ^ {2}\right) - \frac {(1 - \pi) ^ {2}}{\pi^ {2}} g _ {-} ^ {2} \\ = \frac {1 - \pi}{\pi^ {2}} (v _ {-} + \pi g _ {-} ^ {2}). \\ \end{array}
$$
If  $\pi$  is small, then this becomes
$$
\frac {1 - \pi}{\pi^ {2}} \left(v _ {-} + \pi g _ {-} ^ {2}\right) = \frac {v _ {-}}{\pi^ {2}} + O \left(\pi^ {- 1}\right), \tag {22}
$$
so that unless the variance  $v_{-}$  of incorrect rationales is very low, the variance of this estimator will be  $O(\pi^{-2})$ , which is very high. By contrast, the variance of a gradient estimator based purely on correct rationales is simply  $v_{+}$ , so unless the gradient variance for incorrect rationales is dramatically lower than that for correct rationales, then if  $\pi$  is small then incorrect rationales will lead to much noisier gradient estimates.
On the other hand, if  $1 - \pi$  is small, then we have
$$
\frac {1 - \pi}{\pi^ {2}} \left(v _ {-} + \pi g _ {-} ^ {2}\right) = (1 - \pi) \left(v _ {-} + g _ {-} ^ {2}\right) + O \left(\left(1 - \pi\right) ^ {2}\right), \tag {23}
$$
which is likely a significant improvement on the variance  $v_{+}$  of the correct-rationale estimator; in particular, it goes to zero as  $\pi$  approaches 1.
C.2 TRICE control variate as a debiased estimator based on incorrect rationales
Equation (20) implies that, in principle, given unbiased estimates for  $1 - \pi$  and  $\pi^{-1}$ , we could compute gradient updates purely based on rationales that fail to obtain the correct answer. Unbiased estimates of  $1 - \pi$  are easy to obtain; we need only compute  $1 - c(\tilde{z},y)$ . But unbiased estimates of  $\pi^{-1}$  are harder to come by.
Instead, we can compute a biased estimator that ignores the  $\pi^{-1}$  term and then correct for the bias:
$$
\begin{array}{l} \mathbb {E} \left[ - (1 - c (\tilde {z}, y)) \tilde {g} \right] = - (1 - \pi) g _ {-} \\ = \pi g _ {+} \tag {24} \\ = g _ {+} - (1 - \pi) g _ {+}. \\ \end{array}
$$
So the expected value of the estimator  $-(1 - c(\tilde{z},y))\tilde{g}$  is too small by a factor of  $(1 - \pi)g_{+}$ . We can correct this bias by adding in an unbiased estimator of it that uses  $g$  (the gradient for a correct rationale from memory):
$$
\mathbb {E} [ - (1 - c (\tilde {z}, y)) (\tilde {g} - g) ] = g _ {+} - (1 - \pi) g _ {+} + (1 - \pi) g _ {+} = g _ {+}. \tag {25}
$$
This is precisely the gradient estimator that TRICE uses when the sampled rationale  $\tilde{g}$  is incorrect and  $\beta = 1$ . Smaller values of  $\beta$  interpolate between this estimator and an estimator based purely on correct rationales.
D BBH Per-Task Experimental Results
Table 3 summarizes our experimental results for each task in BBH.
E Example Rationales
Figure 3 illustrates some examples of rationales generated by the TRICE-tuned LLM on the BBH task Logical Deduction Three Objects.
Figure 4 illustrates some examples of rationales generated by the TRICE-tuned LLM on GSM8K. Although we did find examples where the LLM got the answer right for the wrong reasons, this was much less common on GSM8K than on BBH, since the numeric output space for GSM8K is much larger than that for the typical multiple-choice BBH task.
Figure 5 shows an example where the guide model  $q_{\phi}(z \mid x, y)$  in reweighted wake-sleep learns to closely mimic the prior model  $p_{\theta}(z \mid x)$  until the very end of the rationale, at which point it pastes in the correct answer.
F Method and Template Details
In this section, we present more details on the methods and templates that we used in the experiments.
To sample from  $p_{\theta}(z \mid x)$ , we prompt the LLM with the template "Question: <QUESTION>\nAnswer: Let's think step by step.\n". We cap the length of the generated rationales at 1.25 times the length of the longest of the exemplar rationales used to initialize the soft prompt. To initialize the memory (i.e., to sample from  $q(z \mid x, y)$  in line 2 of Algorithm 1), on BBH we sample from the base model with a "guide" prompt of the form "Question: <QUESTION>\nAnswer: The answer is <ANSWER>. Let's think step by step.\n". We use the same guide prompt to generate rationalizations in STaR, but with temperature 0 (see below). On GSM8K, we instead initialize the memory with samples from  $p_{\theta}(z \mid x)$ , since we found that initializing the memory using a prompt that includes the answer led to
Example Rationales for Logical Deduction Three Objects BBH Task
An example where TRICE gets the answer right with a good rationale:
Question: On a shelf, there are three books: a brown book, a yellow book, and an orange book. The yellow book is the leftmost. The orange book is to the right of the brown book.
(A) The brown book is the rightmost
(B) The yellow book is the rightmost
(C) The orange book is the rightmost
TRICE answer: Brown book is to the right of yellow book. And orange book to the right of brown book. So, the answer is (C).
An example where TRICE gets the answer right with a bad rationale:
Question: A fruit stand sells three fruits: plums, apples, and loquats. The loquats are the cheapest. The plums are less expensive than the apples.
(A) The plums are the second-most expensive
(B) The apples are the second-most expensive
(C) The loquats are the second-most expensive
TRICE answer: Apples cost more than plums, but less than loquat. So, the answer is (A).
An example where TRICE gets the answer wrong:
Question: A fruit stand sells three fruits: mangoes, watermelons, and kiwis. The watermelons are less expensive than the kiwis. The kiwis are the second-most expensive.
(A) The mangoes are the most expensive
(B) The watermelons are the most expensive
(C) The kiwis are the most expensive
TRICE answer: The Mangoes are more expensive than the Watermelon. And are more expensive than kiwis. So, the answer is (C).
Figure 3: Examples of rationales where TRICE gets the answer right, right but for the wrong reasons, and wrong.
Example Rationales for GSM8K
Question: Gabrielle sells eggs. On Monday she sells 5 crates of eggs. On Tuesday she sells 2 times as many crates of eggs as Monday. On Wednesday she sells 2 fewer crates than Tuesday. On Thursday she sells half as many crates of eggs as she sells on Tuesday. How many total crates of eggs does she sell for the 4 days?
Question: Samantha's last name has three fewer letters than Bobbie's last name. If Bobbie took two letters off her last name, she would have a last name twice the length of Jamie's. Jamie's full name is Jamie Grey. How many letters are in Samantha's last name?
TRICE answer: Jamie's full name consists of 5 letters.  
If Bobbie took 2 letters off her last name, she would have a last name twice the length of Jamie's, which is 10 letters.  
So, Bobbie's last name has 10 letters.  
If Bobbie's last name has 10 letters, and Samantha's last name has three fewer letters than Bobbie's, Samantha's last name has  $10 - 3 = 7$  letters.  
<answer>  
7  
</answer>
Question: John drives for 3 hours at a speed of 60 mph and then turns around because he realizes he forgot something very important at home. He tries to get home in 4 hours but spends the first 2 hours in standstill traffic. He spends the next half-hour driving at a speed of 30 mph, before being able to drive the remaining time of the 4 hours going at 80 mph. How far is he from home at the end of those 4 hours?
TRICE answer: First, he drives 3 hours * 60 mph = 180 miles away from home.  
Then, he turns around and spends 2 hours in traffic (0 mph)  
Then, he spends 0.5 hours at 30 mph = 15 miles  
Then, he spends 1.5 hours at 80 mph = 120 miles  
So, at the end of the 4 hours, he is 180 - 15 + 120 = 285 miles away from home  
<answer>  
285  
</answer>
Figure 4: Examples of rationales where TRICE gets the answer right, right but for the wrong reasons, and wrong.
Figure 5: Example where the prompt-tuned RWS guide model pastes in the correct answer at the end, contradicting the rationale up to that point. The rationales generated by the guide and model are almost identical up to the final answer block.
slower convergence and worse results; it may be that including the answer in the prompt sometimes encourages the model to produce untrustworthy explanations (Turpin et al., 2023).
To evaluate the correctness  $c(z, y)$  of a rationale  $z$  given the answer  $y$ , in BBH we search the end of the rationale for final answers in the form "the answer is <ANSWER>". In GSM8K, we initialize the soft prompt to encourage the model to wrap its answers in "<answer>" and "<answer>" tags, and then extract the answer from those tags. To encourage the bootstrapped few-shot examples in GSM8K to follow this template, we the following example to the CoT prompt: "Question: What is 1 plus 1? \nAnswer: Let's think step by step. \n1 plus 1 is 2. \n<answer>\n2 \n</answer>\n\n
Figure 6 and Figure 7 show the bootstrapped few-shot CoT examples that we used to initialize the soft prompt in the experiments.
TRICE. For all BBH tasks, we run TRICE for 500 steps with batch size  $M = 32$  and do not use subsampling (i.e., compute  $L = 64$  gradients per batch). We use the Adam optimizer (Kingma & Ba, 2015) with an initial learning rate 1.0 and a cosine decay schedule (Loshchilov & Hutter, 2017) that reduces the learning rate by 10x over the first 450 steps. For GSM8K, we run TRICE for 5000 steps with a constant learning rate of 1.0, batch size  $M = 128$ , and compute  $L = 128$  gradients per batch.
STaR. We use an adaptation of the STaR strategy proposed by Zelikman et al. (2022), where we do prompt-tuning rather than fine-tuning on all weights. The method alternates between updating its memory and retuning the model from scratch on the updated memory in an inner loop. We apply this procedure for 10 outer-loop steps. Following Zelikman et al. (2022), we start with 40 inner-loop optimization steps, increasing the number of inner-loop steps by  $20\%$  each outer-loop iteration up to a maximum of 200 steps. If the training loss goes below 0.01 we break out of the inner loop. For STaR's inner-loop optimization, we use the same prompt-tuning initialization, Adam hyperparameters as above, but with cosine decay from 1.0 to 0.1 over the course of each inner loop. To update the
STaR memory, we first try generating a rationale from the prompt-tuned model by greedy decoding, then if that rationale is incorrect fall back on a rationalization generated by greedy decoding from the same guide model we use in TRICE to initialize the MCMC memory, and finally if neither procedure generates a valid rationale we omit the example from the memory.
Rejection Sampling. We reduce mini-batch size by 4 and draw 4 rationales for each example in the mini-batch. We use the same mini-batch size, train steps, and optimizer as in TRICE for all BBH and GSM8K experiments. In BBH, we use the initial learning rate 1.0 as in TRICE. In GSM8K, we use the learning rate 0.1 because it achieved better results than learning rate 0.3, and the training procedure became unstable with learning rate 1.0.
CoT Prompt Tuning. To do supervised CoT tuning, we prompt-tune the model to maximize the log-likelihoods of the training rationales given questions. The BBH datasets include very few exemplar rationales, so we cannot apply this strategy to BBH. On GSM8K, we use the same hyperparameters as in TRICE except that we early-stop the algorithm after only 1000 train steps, since the model overfits badly when we run longer.