Offline Learning and Forgetting for Reasoning with Large Language Models
This work addresses the problem of inefficient inference-time search for researchers and practitioners using large language models in reasoning tasks, offering a more computationally efficient method with significant performance gains, though it is incremental as it builds on existing fine-tuning techniques.
The paper tackles the high computational cost and inference time of using inference-time search to enhance large language models for complex reasoning tasks by proposing an offline fine-tuning approach that integrates search capabilities through learning from successful paths and forgetting failed ones, resulting in a 23% improvement in success rates and a 180x reduction in inference time on arithmetic puzzles.
Leveraging inference-time search in large language models has proven effective in further enhancing a trained model's capability to solve complex mathematical and reasoning problems. However, this approach significantly increases computational costs and inference time, as the model must generate and evaluate multiple candidate solutions to identify a viable reasoning path. To address this, we propose an effective approach that integrates search capabilities directly into the model by fine-tuning it on unpaired successful (learning) and failed reasoning paths (forgetting) derived from diverse search methods. A key challenge we identify is that naive fine-tuning can degrade the model's search capability; we show this can be mitigated with a smaller learning rate. Extensive experiments on the challenging Game-of-24 and Countdown arithmetic puzzles show that, replacing CoT-generated data with search-generated data for offline fine-tuning improves success rates by around 23% over inference-time search baselines, while reducing inference time by 180$\times$. On top of this, our learning and forgetting objective consistently outperforms both supervised fine-tuning and preference-based methods.