Summary
The work focuses on enhancing the reasoning abilities of large language models (LLMs) during the training phase without relying on external feedback. The authors motivate the work by raising an important question on improving the reasoning ability of LLMs during training phase since most prior works have focussed at achieving this at inference time. To do so, the authors propose to sample diverse reasoning rationales from latent distribution and optimize the LLM through a variational framework using the sampled rationales. The intuition behind application of the variational framework is well motivated through the objective of self-consistency based chain-of-thought prompting. Further, the work proposes to usethe likelihood of generating correct answer conditioned on a rationale as a proxy to explicit reward models to optimise the LLM towards generating better reasoning chains. Results demonstrate that the proposed LaTRO method helps in improving the accuracy achieved for multiple LLMs such as Mistral-7B, Llama-3.1-8B etc. on GSM8K and ARC-Challenge datasets compared to the corresponding base and the SFT versions of the LLMs.
Strengths
1. Enhancing the reasoning ability of LLMs without relying on feedback from any external LLM during the training phase is an important research question.
2. Suitability of applying the variational objective is well motivated, explained and justified in Sections 3 and 4.1.
3. Accuracy gains over the base and the SFT version of the LLMs are shown with some ablation studies on greedy decoding vs. self consistency based sampling. Further, the authors show that inference time scaling of number of reasoning samples obtained using LaTRO can additionally enhance the accuracy.
Weaknesses
1. Presentation of the introduction section of the paper needs improvement - It should better motivate the need behind reducing the reliance on external LLMs/feedback for improving a given LLM. Further, it should provide information about why the variational framework is suitable and elaborate some details about the proposed method in the introduction itself.
2. Discussion on related work is very limited. Even though the proposed method is an effort to reduce reliance on external LLMs, it should still discuss and contrast against existing methods that leverage external LLMs (for example, [1,2]) as well as employ self-play based techniques to improve an LLM by itself with requiring help from any other LLM (for example, [3, 4]). The example references are only indicative (but not exhaustive) of the type of citations that should be included.
3. Lack of appropriate baselines - Even though the work claims and focusses at improving the base/SFT version of the LLM through the proposed LaTRO framework, it should still compare against some of the existing self-play trainable methods (by adapting them for 7B LLMs) as baselines (eg. [3] - on datasets where the ground-truth rationale is available and [4]). Such methods improve the rationales generated by an LLM by exploring the generation space and discriminating better rationales from the inferior ones.
4. More reasoning datasets such as CSQA, Hellaswag etc. should be considered for evaluation studies.
[1] Distilling Step-by-Step! Outperforming Larger Language Models with Less Training Data and Smaller Model Sizes. Cheng-Yu Hsieh, Chun-Liang Li, Chih-kuan Yeh, Hootan Nakhost, Yasuhisa Fujii, Alex Ratner, Ranjay Krishna, Chen-Yu Lee, and Tomas Pfister. In Findings of the Association for Computational Linguistics: ACL 2023, pages 8003–8017, Toronto, Canada. Association for Computational Linguistics.
[2] Zephyr: Direct Distillation of LM Alignment. Lewis Tunstall and Edward Emanuel Beeching and Nathan Lambert and Nazneen Rajani and Kashif Rasul and Younes Belkada and Shengyi Huang and Leandro Von Werra and Clementine Fourrier and Nathan Habib and Nathan Sarrazin and Omar Sanseviero and Alexander M Rush and Thomas Wolf. First Conference on Language Modeling, 2024.
[3] Self-Play Fine-Tuning Converts Weak Language Models to Strong Language Models. Zixiang Chen and Yihe Deng and Huizhuo Yuan and Kaixuan Ji and Quanquan Gu. ICML 2024.
[4] Self-Rewarding Language Models. Weizhe Yuan, Richard Yuanzhe Pang, Kyunghyun Cho, Xian Li, Sainbayar Sukhbaatar, Jing Xu, Jason Weston
Questions
1. Line 290: It is not clear about how the first gradient term in Eq. 4 would lead to optimising the LLM policy to generate higher quality rationales. Further elaboration is needed as to why generating/optimising the likelihood of input question/instruction conditioned on the rationale would lead to better reasoning.
2. Lines 414-415: Please support with examples about why improvements on ARC-Challenge are relatively lesser. The magnitude of improvements is not at all an issue, however, it should be better demonstrated that why better reasoning chains are not leading to higher improvements. Does this make ARC-Challenge ill-suited for this study since it involves limited reasoning scope?
3. Line 430: Why is it needed to generate rationales that are as long as 500 tokens? It would be good to show the usefulness of the information contained in the longer rationales and why does it lead to better performance before saturating.