Training Chain-of-Thought via Latent-Variable Inference

Summary

This paper introduces a principled and practical approach to boost the training of the "chain-of-thought" through the inference of latent variables. Instead of the typical variational inference, the authors opt for a MCMC-EM method to circumvent the issue of posterior collapse, a common occurrence in auto-regressive models with latent variables, which could be exacerbated when dealing with large language models. By treating rationales as latent variables, this methodology enables training without the need for detailed rationales, which can be expensive to produce. Moreover, the authors propose a control-variate technique that reduces training variance while maintaining unbiasedness. The theoretical underpinnings of their approach are validated on the BIG-Bench Hard dataset.

Strengths

1. Although the soft-prompt approach is widely used, this proposal introduces an innovative perspective on 'prompt-finetuning' utilizing a chain-of-thought. 2. The assumptions and derivations presented are straightforward and intuitive with only minor confusion. 3. The authors thoroughly explore the links between their work and previous studies, while also critically examining the limitations of their own findings. 4. The authors validate their theories by applying them to the significant benchmark of BBH.

Weaknesses

See questions.

Questions

1. For accept rate in Metroplis-Hastings, why the transition kernel $p(\tilde{z},y|x)$ is not conditional on previous state $z$? 2. For sampling from intractable posterior $p(z|x,y)$, what if you directly sampling from this distribution using finite-step MCMC as $p(z|x,y)\propto p(z|x)p(y|z,x)$? You do not need to add more experiments. Could you explain the difference? 3. For control variate, in appendix you mentioned it may not be that useful in the stochastic setting. For figure 1, I want to know that suppose you train long enough, will the training estimate becomes the same for two settings (with or without control variate) based on your claim in Eq. 3? Meanwhile, what's the notation meaning for expectation over MCMC? 4. How many tokens you use for latent variable $z$? 5. While I commend the authors for validating their approach on an important benchmark, an aspect I deeply appreciate, I believe that testing on a more diverse dataset would bolster the paper's credibility and appeal. This is particularly true when making a comparison with [1]. 6. Do you have failure case analysis and do you need human evaluation? I would be more than willing to raise my review scores if the authors could provide further clarification and deeper insights about the paper. [1] STaR: Self-Taught Reasoner Bootstrapping Reasoning With Reasoning

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

The authors has discussed this paper's limitations.

Summary

This paper explores considering Chain-of-Thoughts (CoT) as a latent variable and introduces TRICE: an MCMC-EM algorithm for optimizing CoT-latent-latent variable models using a Metropolis-Hastings algorithm (rejection sampling) coupled with a control variate. The authors tested the model by fine-tuning a proprietary LLM on 27 BIG-bench hard tasks and report improvements when compared to direct prompt tuning and STaR. The authors tested the effectiveness of the control variate but did not test the method against more conventional optimisation methods.

Strengths

Overall, This paper discusses a well-defined issue and explores innovative solutions to the optimization problem, it was a refreshing and educative read, thank you! Main strengths: 1. Optimizing LLMs as latent variable models has great potential and the authors present the problem well 2. The authors introduce an innovative gradient estimator (MCMC-EM + control variate) which is computationally efficient (single sample per data point) and seems to work very well 3. The authors demonstrate progress on a subset of the MMLU dataset 4. Theoretical and experimental limitations are thoroughly discussed

Weaknesses

Overall, the paper lacks structure and more extensive comparison with the existing literature on Monte Carlo methods. I believe not addressing these points would greatly reduce the impact of this paper. Nevertheless, I would happily accept this paper if these limitations are resolved! 1. The paper overall lacks structure. Although I appreciated the lengthy discussions, I think the main storyline could be made more clear by separating the main ideas from the discussions. 2. The optimization method is only presented using full text and Algorithm 1. This paper would convey its idea more clearly using a set of two additional equations: A) expression of the *true* gradients and B) Monte Carlo estimate (which translates into Algorithm 1). 3. Classic gradient estimation methods are not covered sufficiently and some of the statements in section 2.2 *Why not variational inference?* are misleading. 1. Posterior collapse (discussed in Lucas et al., 2019) is an issue inherent to the optimization of the proposal distribution $q(z | x, y)$ and $q$ is in this paper never optimized. 2. The author discusses the limitations of the ELBO, but the ELBO only arises when using a single-sample estimate of the log-likelihood. When working with discrete latent variables, it is more common to use the Importance Weighted Bound (IWB). I.e., use a K-samples importance-weighted estimate of the likelihood (e.g., $\log p(y | x) \approx \log \frac{1}{K} \sum_{j=1}^K \frac{p(y, z^{(j)} | x)}{q(z^{(j)}|x,y)})$ -- as done in ["Revisiting Reweighted Wake-Sleep for Models with Stochastic Control Flow", Let et al., 2019](https://arxiv.org/abs/1805.10469)) which variance scales as $\sqrt{1/K}$. See ["Monte Carlo Gradient Estimation in Machine Learning", Mohamed et al., 2020] for more background on optimization methods. 6. The experimental comparison with RWS is not included in the paper.

Questions

1. How exactly did you compute the gradients in the RWS experiment?

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

1. I stumbled upon this PhD thesis: [Deep Latent Variable Models for Natural Language Processing, Lievin 2022](http://vlievin.github.io/deep-lvms-for-nlp.pdf), which discusses similar ideas. You might not have been the first ones to consider that "CoT methods are latent-variable probabilistic models"

Summary

This work introduces a new method for training models for chain-of-thought prompting, aimed at maximizing the marginal probability of generating a correct answer using Markov-chain Monte Carlo, expectation maximization, and a novel control variate technique, along with prompt tuning. Much of the work is spent deriving and analyzing the method (e.g. compared to related methods like self-taught reasoner STaR), with a final section dedicated to testing the method and demonstrating very significant performance gains over vanilla prompt tuning (no CoT) and STaR.

Strengths

- Principled, novel method along with extensive derivation. Particularly, this well motivates the need for taking a marginalization framing rather than the greedy approach often used before - It is generally very useful to see more theoretical justification provided in this area. Past works (e.g. STaR) provide many novel ideas, but it is useful to take a more comprehensive (and marginalization-based) approach to the problem - extensive theoretical/technical rationale for method over other options (STaR, variational inference) - very strong experimental results, demonstrating significant gains over STaR, another recent method for training for CoT prompting

Weaknesses

- Fairly limited analysis of experimental results (mostly aggregate results on Big Bench tasks in the main paper, without too much specific analysis) - It would also be useful/interesting to see what the generated rationales look like: how do they differ from e.g. STaR, and how often are they correct in a human-interpretable way rather than just leading to a correct answer? Adding some information about this in the main paper would be useful - Fairly limited discussion of limitations

Questions

None

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

4 excellent

Presentation

4 excellent

Contribution

3 good

Limitations

Limitations section is quite limited. It would be useful to see e.g. a discussion of the effects of rationales on user trust (will having rationales make users trust models more even if the rationales are wrong?)

Summary

This paper proposes treating rationales as latent-variables and considers the marginal distribution over answers, averaging over possible rationales. To do so, a MCMC procedure is proposed, in which rationales are proposed via an independence sampler. The approach is compared to STaR (Zelikman et al., 2022) as a baseline on 27 BigBench-HARD tasks. Rather than fine-tune all the model weights, prompt tuning is used a parameter-efficient alternative. The proposed approach is found to outperform STaR on the subset of tasks considered.

Strengths

* The derivations the authors present are detailed and comprehensive. The design choices are explained well and justified in context. * Compared to alternate methods, the approach advocated in this paper seems more principled. * The authors discuss a potential alternate approach (variational Bayes) and justify why it is not appropriate.

Weaknesses

* As the authors acknowledge, TRICE is only tested on small datasets. The behavior of the comparison model is significantly different than the expected behavior; the authors speculate this is due to overfitting of the comparison model, but do not compare TRICE’s performance on Commonsense QA or reduce the number of inner loop steps to prevent overfitting. * While the overall approach is novel, it’s not clear that the improvements and alterations proposed present significant innovations on existing methods. Furthermore, I have some concerns about the sensitivity of the approach to hyper-parameters, such as the # of MCMC steps, initialization scheme, etc. * Some implementation details are unclear. The methods as-is would not be reproducible, and it seems that the authors do not plan to release their code, raising concerns about replicability.

Questions

* What is the specific $c(z,y)$ used in the experiments? It never seems to be defined— an example is given on line 71, but it’s unclear whether that’s the specific correctness function used in the experiments. * Line 80 says that “For example, the guide might prompt an LLM specifically to give an rationale for the answer”— was this the case for all experiments? * Line 296 says that "an adaptation of the STaR strategy" is tested-- what aspects specifically were adapted?

Rating

5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.

Confidence

5: You are absolutely certain about your assessment. You are very familiar with the related work and checked the math/other details carefully.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

There is a nominal discussion of limitations but I think this could be significantly expanded. For example, MCMC can be sensitive to hyperparameters and mixing is hard to diagnose.

Summary

The paper proposes a prompt-tuning strategy that tries to maximize the marginal log-likelihood of generating a correct answer using CoT prompting.

Strengths

The paper presents a prompt-tuning strategy that aims to optimize the marginal log-likelihood of producing accurate answers through the utilization of CoT prompting. And the paper uses simple Markov chain Monte Carlo (MCMC) expectation-maximization (EM) algorithm, memoized wake-sleep, and persistent contrastive divergence to address the challenge of sampling from the posterior. The results look like promising over other baselines.

Weaknesses

The author did not provide specific information about the model details. If the model size is close to 100 billion, it would be beneficial for the paper to include a baseline comparison with text-davinci-003 (Zero-shot or few shot learning). Otherwise, I feel it unnecessary to conduct prompt tuning on a large-scale model if its performance is still inferior to that of zero/few-shot learning using the same scale model but from GPT family. If the author's intention is to compare with open-source models, it is essential to include Alpaca-CoT (https://github.com/PhoebusSi/alpaca-CoT), regardless of the specific model or fine-tuning technique they employ.

Questions

See weakness part.

Rating

4: Borderline reject: Technically solid paper where reasons to reject, e.g., limited evaluation, outweigh reasons to accept, e.g., good evaluation. Please use sparingly.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

2 fair

Contribution

2 fair

Limitations

See weakness part.

Improving clarity and presenting main equations more clearly

### 1. Structural improvement Enclosing remarks would indeed help. ### 2. Providing true gradient and Monte Carlo estimate Yes, the true gradient is provided in Eq 2. However, the expression of the gradient approximation is spread across sections "**2.1 derivations**" and Eq 3&4. Equations 3&4 (which are the core equations defining TRICE) are enclosed in a section titled "**6 Adding a control variate**". A better structure, better choice of section titles and an improved framing of the equations would help. 1. Gradient $$\eta_\theta = E_{p_\theta(z | x, y)} \left[ \nabla_\theta \log p_\theta(z | x) \right]$$ 2. Gradient approximation via MCMC $$\eta_\theta \approx E_{\mathrm{MCMC}} \left[ \nabla_\theta \log p_\theta (z | x) \right] $$ 3. Gradient approximation via MCMC with control variate $$\eta_\theta \approx \ldots $$ **NB**: Similarly to Reviewer zJeM, I find the notation $E_{\mathrm{MCMC}}\left[ \cdot \right]$ confusing.

VI, RWS, and the IWB - great insights

Thank you for the clarification. It makes sense you obtain zero-length rationales, and I understand why you relate it to posterior collapse. > How exactly did you compute the gradients in the RWS experiment? Thank you for the details. What are your thoughts on optimizing $\theta$ using the "M-sample IWB" only? Meaning, using the expression of the RWS gradients you provided and using, for instance, $p(z|x,y) \approx p_\theta(z|x) c(x,y)$ as sampling distribution? How does it relate to TRICE?