Toward Adaptive Reasoning in Large Language Models with Thought Rollback

TL;DR

This paper introduces Thought Rollback (TR), a novel framework enabling large language models to adaptively revise their reasoning steps through error correction, significantly improving problem-solving accuracy in complex tasks.

Contribution

The paper proposes Thought Rollback, a new error correction mechanism that allows LLMs to revise their reasoning by rolling back to previous thoughts, enhancing reasoning flexibility and accuracy.

Findings

GPT-4 with TR outperforms previous methods by 9% on MATH dataset.

TR achieves state-of-the-art problem-solving rates in mathematical and multi-task reasoning.

TR reduces interaction costs while improving reasoning reliability.

Abstract

Large language models (LLMs) have been routinely used to solve various tasks using step-by-step reasoning. However, the structure of intermediate reasoning steps, or thoughts, is rigid and unidirectional, such as chains, trees, or acyclic-directed graphs. Consequently, the resulting inflexible and forward-only reasoning may not address challenging tasks and fail when the LLM frequently gives false responses, i.e., ``hallucinations''. This paper proposes a new reasoning framework, called Thought Rollback (TR), allowing LLMs to adaptively build thought structure while maintaining effective reasoning toward problem-solving under ``hallucinations''. The core mechanism of TR is rolling back thoughts, which allows LLMs to perform error analysis on thoughts, and thus roll back to any previously mistaken thought for revision. Subsequently, by including such trial-and-error in the prompt to guide the LLM, each rollback leads to one more reliable reasoning path. Therefore, starting with a simple prompt without human annotations, LLM with TR adaptively and gradually explores thoughts for a correct solution. Comprehensive experiments on mathematical problems and multi-task reasoning demonstrate the state-of-the-art performance of TR in terms of problem-solving rate and interaction cost. For instance, the solving rate of GPT-4 with TR outperforms the current best by $9\%$ on the MATH dataset.