On Reward Functions For Self-Improving Chain-of-Thought Reasoning Without Supervised Datasets (Abridged Version)

Thomas Foster, Eltayeb Ahmed, Jonathan Cook, Shalev Lifshitz, Tim Rocktäschel, Jakob Nicolaus Foerster

Keywords: large language models, self-improvement, CoT reasoning, reinforcement learning

TL;DR: An analysis of reward functions for self-improving CoT reasoning on general, unstructured text.

Abstract: Prompting a Large Language Model (LLM) to output Chain-of-Thought (CoT) reasoning improves performance on complex problem-solving tasks. Moreover, several popular approaches exist to "self-improve'' the CoT reasoning abilities of LLMs on tasks where supervised (question, answer) datasets are already available. An emerging line of work explores whether self-improvement is possible without these supervised datasets, instead utilizing the same large, unstructured text corpora as used during pretraining. This would overcome the data availability bottleneck present in current self-improvement methods, and open the door towards *compute-only scaling* of language model reasoning ability. We investigate a fundamental question in this line of work: What constitutes a suitable reward function for learning to reason during general language model pretraining? We empirically demonstrate how different functions affect what reasoning is learnt and where reasoning is rewarded. Using these insights, we introduce a novel reward function called Reasoning Advantage (RA) that facilitates self-improving CoT reasoning on free-form question-answering (QA) data, where answers are unstructured and difficult to verify. We also explore the optimization of RA on general unstructured text using offline RL, and our analysis indicates that future work should investigate more powerful optimization algorithms, potentially moving towards more online algorithms that better explore the space of CoT generations.

Submission Number: 33