Learning to Refine: Self-Refinement of Parallel Reasoning in LLMs
This addresses the problem of improving LLMs' multi-step reasoning for complex tasks, offering a model-agnostic enhancement that is robust across scales and generalizes to out-of-distribution tasks, though it is incremental as it builds on existing test-time scaling methods.
The paper tackles the limitation of existing test-time scaling methods for LLMs, which fail when all candidate responses are incorrect, by introducing Generative Self-Refinement (GSR), a framework that enables a unified model to generate and refine candidate responses in parallel, achieving state-of-the-art performance across five mathematical benchmarks.
To further enhance the ability of Large Language Models (LLMs) to solve complex, multi-step reasoning problems, test-time scaling (TTS) methods have gained widespread attention. Existing approaches such as Best-of-N and majority voting are limited as their performance depends on the quality of candidate responses, making them unable to produce a correct solution when all candidates are incorrect. Introducing an additional model to select the best response also incurs significant deployment costs. To this end, we introduce Generative Self-Refinement (GSR), a novel parallel test-time scaling framework where a unified model first generates a set of candidate responses in parallel and then performs self-refinement to synthesize a new superior solution based on a prompt consisting of the problem and these candidates. However, LLMs struggle to perform refinement effectively when prompted directly. Therefore, we design a hybrid training pipeline by jointly optimizing for two complementary objectives, solving problems directly and refining candidate responses. Experimental results demonstrate that our method achieves state-of-the-art performance across five mathematical benchmarks. We further show that this learned self-refinement skill is a model-agnostic enhancement, robust across different model scales and generalizing to out-of-distribution reasoning tasks.