PaperlayerView exact revision on arXiv

Source-linked AI summary

Training Chain-of-Thought via Latent-Variable Inference

Du Phan, Matthew D. Hoffman, David Dohan, Sholto Douglas, Tuan Anh Le, Aaron Parisi, Pavel Sountsov, Charles Sutton, Sharad Vikram, Rif A. Saurous

arXiv:2312.02179v1cs.LGcs.AIcs.CL

TL;DR

CoT tuning can require detailed rationales that are difficult to obtain, despite the availability of question-answer data. TRICE treats rationales as latent variables and uses MCMC-EM to tune models; on GSM8K and BIG-Bench Hard, it outperforms several tuning baselines.

  • Problem

    Detailed rationales needed for supervised CoT tuning can be difficult to obtain, even though question-answer datasets without rationales are more readily available.

    Source: §1 · p. 1

  • Method

    TRICE treats rationales as latent variables and uses an MCMC expectation-maximization algorithm with a control-variate scheme to tune rationale generation.

    Source: §1 · p. 2

  • Results

    TRICE improves performance on GSM8K and BIG-Bench Hard, outperforming STaR, direct tuning with or without CoT, and supervised fine-tuning on human-generated rationales.

    Source: §1 · p. 2

  • Takeaways & Limitations

    Compared with STaR, TRICE is less likely to ignore difficult examples and can learn from incorrect as well as correct rationales.

    Source: §1 · p. 2

    Evaluation used prompt-tuning on a medium-size LLM and assessed answer correctness, not the quality of the rationales themselves.

    Source: §5 · p. 10

Abstract

from arXiv · show

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 \emph{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.

1 Introduction

TRICE treats chain-of-thought tuning as learning with latent rationales, avoiding the need to supervise detailed rationales directly. Its MCMC-EM approach improves performance on GSM8K and BIG-Bench Hard compared with several tuning baselines.

1 Introduction

  • Gold rationales can be difficult to obtain, while question-answer datasets without rationales are more readily available for bootstrapping.Existing CoT rationales can also be incorrect.Source: §1 · p. 1
  • TRICE frames CoT methods as latent-variable models, where answer probabilities marginalize over possible rationales conditioned on the question.This recasts rationale bootstrapping as learning with incomplete data.Source: §1 · p. 2
  • TRICE uses MCMC expectation-maximization with a control-variate scheme to address the challenge of averaging over a vast set of potential rationales.The algorithm draws on STaR, memoized wake-sleep, Markovian score climbing, and persistent contrastive divergence.Source: §1 · p. 2
  • Compared with STaR, TRICE is less likely to ignore difficult examples and can learn from incorrect as well as correct rationales.The paper characterizes STaR as underweighting difficult examples when rationalization fails.Source: §1 · p. 2
  • TRICE significantly outperforms STaR, direct tuning with or without CoT, and supervised fine-tuning on human-generated rationales on GSM8K and BIG-Bench Hard.The paper also reports that TRICE outperforms direct prompt-tuning, STaR, and rejection sampling on these tasks.Source: §1 · p. 2

2 Method

The method treats rationales as latent variables and tunes the model to maximize the marginal likelihood of correct answers. TRICE approximates the required posterior over rationales with MCMC-EM, using accepted correct proposals and a control variate to reduce gradient variance.

2 Method

  • CoT tuning maximizes the average marginal log-likelihood of answers given questions, averaging over latent rationales.The model assumes a binary likelihood in which the answer is a deterministic function of the rationale, while noting that the method can generalize to weaker relationships.Source: §2 · p. 2
  • TRICE initializes one rationale per question-answer pair from a guide that can use the correct answer, then proposes model-generated rationales and replaces a stored rationale whenever the proposal is correct.The algorithm updates parameters using gradients from the updated memory of rationales.Source: §2 · pp. 2–3, §2.1 · p. 4

2.1 Derivation

  • The marginal-likelihood gradient is an expectation over rationales drawn from the posterior conditioned on the question and answer, motivating posterior sampling with an MCMC chain.The independence sampler proposes from the model and falls back to the previous state when a proposal is incorrect; mixing is faster for easier questions and as the model improves.Source: §2.1 · pp. 3, 5
  • A control variate subtracts a zero-mean score-function term from the gradient estimate, scaling correct and incorrect proposal contributions according to an estimated acceptance probability.The variance can approach zero as training accuracy approaches perfection; subsampling gradients by their scalar weights reduces computation while preserving an unbiased estimate.Source: §2 · p. 3, §2.1 · p. 5

2.2 Why not variational inference, reweighted wake-sleep, or rejection sampling?

  • RWS can produce empty rationales because shorter sequences receive higher weights, and its guide can instead mimic the prior before appending the correct answer regardless of the preceding reasoning.In one run, the reported KL was about 0.61 nats before the final-answer section and about 42.5 nats for the conditional final-answer section.Source: §2.2 · p. 6
  • Rejection sampling can work, but keeping per-iteration cost constant requires reducing the minibatch size by a factor of K, leading to slower convergence and/or worse final results.The scheme averages gradients from K sampled rationales that produce correct answers.Source: §2.2 · p. 7

3 Related Work

The related work situates TRICE among methods that average or impute rationales during training, and contrasts it with prompting and latent-variable approaches. Its closest comparison is STaR, which differs in rationale decoding and updating and can omit difficult training examples.

3 Related Work

  • Prior work includes supervised and few-shot rationale generation, while self-consistent CoT approximately marginalizes over rationales at test time; comparatively less work averages or imputes rationales during training.STaR is identified as a main example of training-time rationale imputation.Source: §3 · p. 7

3.1 Self-Taught Reasoner

  • TRICE differs from STaR’s greedy decoding by sampling more diverse rationales, and the authors report that occasional correct answers for wrong reasons did not prevent improvement on most tasks.They suggest stronger base models than the 6B-parameter GPT-J used in STaR may help explain this pattern.Source: §3.1 · p. 7
  • TRICE resamples rationales each iteration, whereas STaR’s inner loop reuses stale rationales to estimate the marginal-likelihood gradient.STaR may also use a surrogate rationale distribution during updates, potentially affecting early training dynamics.Source: §3.1 · pp. 7–8
  • STaR trains repeatedly on a fixed rationale set and can memorize it; empirically, it often fails to find a good rationale for every question and consequently ignores part of the training set.The paper connects this tendency to an objective that weights examples by answer likelihood, giving difficult examples less influence on the gradient.Source: §3.1 · p. 8

4 Experiments

The experiments evaluate TRICE against prompt-tuning baselines on GSM8K and BIG-Bench Hard (BBH). TRICE learns valid rationales without human-generated exemplars and outperforms supervised tuning on human-generated rationales on GSM8K.

4 Experiments

  • Experiments use medium-size PaLM 2-M on GSM8K and 27 BBH tasks, with base PaLM 2 for GSM8K and Flan PaLM 2 for BBH.BBH examples are split 60%/40% for training and testing; GSM8K uses its standard 7,473-example training and 1,319-example test sets.Source: §4 · p. 8
  • Comparisons include direct prompt tuning, CoT prompt tuning, rejection sampling, and STaR; TRICE uses soft prompts initialized from few-shot exemplars or bootstrapped rationales.The standard BBH initialization uses a small amount of detailed CoT supervision, while bootstrapped initialization requires label supervision only.Source: §4 · p. 8
  • Without human-generated exemplar rationales, TRICE learns to generate rationales leading to correct answers, and on GSM8K it outperforms a model trained directly on human-generated rationales (cf. Uesato et al., 2022).Table 2 reports average accuracies (columns 2 and 3) and the fraction of GSM8K training examples for which correct rationales can be generated (column 4).Source: §4 · p. 8, Table 2 · p. 9
  • On BBH, TRICE generates a valid rationale for almost all training examples, while STaR fails to do so for almost 10%.The authors suggest STaR may overfit its bootstrapped rationales with only 150 training examples per task.Source: §4 · p. 9
  • STaR’s final inner-loop sequence-level cross-entropy is below 0.06 on every task and below 0.01 on 19 of 27, consistent with reproducing one rationale set.The authors present this as evidence for their overfitting-rationales hypothesis.Source: §4 · p. 9
  • The control-variate scheme improves average convergence speed, and both TRICE variants reach high training accuracy much faster than rejection sampling.Figure 2 compares training-set marginal-likelihood estimates and greedy-decoding validation accuracy over training steps on GSM8K.Source: §4 · p. 10, Figure 2 · p. 10

5 Discussion

TRICE outperforms direct prompt-tuning, STaR, and rejection sampling on GSM8K and BIG-Bench Hard. The paper notes limits in evaluation scope and discusses potential risks from more capable models and users' treatment of rationales.

5 Discussion

  • On GSM8K and BIG-Bench Hard, TRICE outperforms direct prompt-tuning, STaR, and rejection sampling.The authors suggest its basic MCMC-EM strategy could also be used for tasks such as tool use.Source: §5 · p. 10
  • Evaluation covered prompt-tuning on a medium-size LLM and answer correctness, not formal rationale quality.The authors note behavior may differ with other model sizes or fine-tuning strategies; gradient-based tuning also limits evaluation on proprietary models such as GPT-4.Source: §5 · p. 10
  • More capable LLMs may be used maliciously or unsafely, while uncurated question-answering data may introduce bias and wider use may increase carbon footprint.Source: §5 · p. 10
  • Rationales may help motivated users judge trustworthiness, but users who mistake their presence for evidence of truth could be harmed.Source: §5 · p. 10
  • For cost-matched iterations, rejection sampling uses one-quarter the minibatch and four times as many proposals, while TRICE sets L equal to M.The TRICE setting still subsamples because each example may contribute both a correct and an incorrect rationale.Source: §5 · p. 10

A Generalizing TRICE to Nondeterministic Likelihood Models

For nondeterministic likelihoods, many outputs may fit a question and rationale, making the probability of sampling the observed output fall as the output space grows. TRICE can accommodate this setting by probabilistically updating rationale memory, though its control-variate benefit may be modest.

A Generalizing TRICE to Nondeterministic Likelihood Models

  • As the number of reasonable outputs for an input and rationale grows, sampling the particular observed output can become exponentially less likely.Nondeterministic likelihoods can be appropriate when several outputs are reasonable for the same input and rationale.Source: A Generalizing TRICE to… · p. 14
  • For nondeterministic likelihoods, TRICE updates rationale memory probabilistically, always accepting proposals that make the observed output more likely and sometimes accepting those that make it less likely.Source: A Generalizing TRICE to… · p. 14
  • The control variate may reduce variance only modestly when valid rationales do not consistently match the particular observed output.With multiple reasonable outputs, acceptance probability may not approach 1 even when the model generates valid rationales.Source: A Generalizing TRICE to… · p. 14

B Derivation of the Control Variate Scaling Heuristic

The appendix analyzes the variance of TRICE's gradient estimators and the control-variate scaling heuristic. It shows the heuristic's gap from optimal scaling shrinks faster than the variance as the probability of a correct proposal increases, while incorrect-rationale estimates can be especially noisy when that probability is small.

B Derivation of the Control Variate Scaling Heuristic

  • As correct-proposal probability π increases, choosing β=π has suboptimality O((1−π)^2), which decreases faster than the variance.The heuristic and optimal variances agree up to O((1−π)^2).Source: B Derivation of the Control Variate… · p. 15

C On Gradient Estimators Based Solely on Incorrect Rationales

  • The identity 0 = E[˜g] = πg+ + (1−π)g− relates the expected gradients for correct and incorrect rationales.Here g+ and g− are the conditional expected gradients for correct and incorrect proposals, respectively.Source: C On Gradient Estimators Based Solely… · p. 16

C.1 Variance of incorrect-rationale gradient estimators

  • When π is small, an incorrect-rationale estimator has variance O(π^-2) unless the variance of incorrect-rationale gradients is very low, making estimates much noisier than correct-rationale estimates.Source: §C.1 · p. 16

C.2 TRICE control variate as a debiased estimator based on incorrect rationales

The estimator can use incorrect sampled rationales by correcting the bias from dropping a difficult posterior-probability term. The correction uses the stored gradient for a correct rationale, with β controlling interpolation.

C.2 TRICE control variate as a debiased estimator based on incorrect rationales

  • An estimator based only on incorrect rationales would need unbiased estimates of 1 − π and π^-1; the latter is harder to obtain.Source: §C.2 · p. 16
  • Ignoring the difficult π^-1 term produces an estimator whose expectation is too small by a factor of (1 − π)g+.Source: §C.2 · p. 16
  • Adding an unbiased correction based on the stored correct-rationale gradient yields TRICE’s estimator for an incorrect sample when β = 1.Source: §C.2 · p. 16
  • Smaller β values interpolate toward an estimator based purely on correct rationales.Source: §C.2 · p. 16

D BBH Per-Task Experimental Results

The section points readers to Table 3 for per-task BIG-Bench Hard experimental results.

D BBH Per-Task Experimental Results

E Example Rationales

The examples show TRICE-generated rationales on BIG-Bench Hard and GSM8K, including cases where an answer is correct despite faulty reasoning. A separate example depicts a guide model inserting the correct answer at the end of a rationale.

E Example Rationales

  • Figure 3 illustrates TRICE-generated rationales for the BIG-Bench Hard Logical Deduction Three Objects task.Source: E Example Rationales · p. 17
  • On GSM8K, getting the answer right for the wrong reasons was much less common than on typical multiple-choice BIG-Bench Hard tasks, which have smaller numeric output spaces.Source: E Example Rationales · p. 17
  • Figure 5 shows an RWS guide model closely mimicking the prior rationale before inserting the correct answer at the end.Source: E Example Rationales · p. 17

F Method and Template Details

The appendix details the prompts, answer extraction, and training configurations used for TRICE and comparison methods. It also supplies example prompts and rationales, alongside a warning about answer-conditioned prompting.

F Method and Template Details

  • The generation prompt asks the model to think step by step, caps rationale length at 1.25 times the longest exemplar, and initializes BBH memory with an answer-conditioned guide prompt.Source: F Method and Template Details · p. 17
  • An RWS guide example pastes in the correct answer at the end, contradicting the preceding rationale; answer-in-prompting may also encourage untrustworthy explanations.The passage also associates this setting with slower convergence and worse results.Source: Figure 5 · p. 20, F Method and Template Details · p. 20
  • BBH correctness is extracted from a final “the answer is <ANSWER>” form, while GSM8K uses answers enclosed in <answer> tags.Source: F Method and Template Details · p. 20
  • Figures 6 and 7 show the bootstrapped few-shot CoT examples used to initialize the soft prompt: three shots for formal fallacies and five for GSM8K.Source: F Method and Template Details · p. 20, Figure 6 · p. 22, Figure 7 · p. 23
  • TRICE runs for 500 steps on BBH and 5,000 on GSM8K, using batch sizes 32 and 128, respectively.The BBH setup uses 64 gradients per batch; GSM8K uses 128.Source: F Method and Template Details · p. 20
  • The STaR adaptation alternates memory updates with prompt retuning from scratch over 10 outer-loop steps, falling back to guide-generated rationales when needed.Source: F Method and Template Details · pp. 20–21
  • Rejection sampling draws four rationales per example, while supervised CoT tuning is unavailable for BBH because its datasets have very few exemplar rationales.On GSM8K, supervised CoT tuning stops after 1,000 steps because longer runs overfit badly.Source: F Method and Template Details · p. 21
  • Direct prompt tuning trains without generating a rationale and uses 150 optimization steps, with hyperparameter adjustments intended to reduce overfitting.Source: F Method and Template Details · p. 21
Loading 2312.02179v1…