Offline Learning and Forgetting for Reasoning with Large Language Models

TL;DR

The paper tackles the problem that inference-time search in large language models, while effective for reasoning, incurs high computational costs. It proposes an offline three-phase pipeline: generate diverse CoT-style reasoning data from multiple reasoners (including classic BFS/DFS and inference-time solvers), fine-tune the base LLM with a combined learning-and-forgetting objective called unlikelihood fine-tuning (UFT), and evaluate across CoT and search-based reasoning. The core idea is to leverage unpaired positive and negative reasoning paths with a small learning-rate to preserve search capability, using objectives $J_{NLL}$ and $J_{UL}$ and the combined loss $J(\theta;\\mathcal{D},\\alpha)=(1-\\alpha)J_{NLL}(\\theta;\\mathcal{D}^+) +\\alpha J_{UL}(\\theta;\\mathcal{D}^-)$, along with a rule-based verifier for rewards. Experiments on Game-of-24 and Countdown show that high-quality, diverse training data and the unlikelihood term substantially improve CoT reasoning and retain inference-time search efficiency, reducing inference cost by up to $180\times$ and achieving around a $23\%$ absolute rise in success rates over inference-time baselines. The results consistently indicate that data quality and the forgetting objective are pivotal, with CoT data often better preserving search ability than mixed datasets. This approach offers a practical, offline path to more reliable and faster reasoning in LLMs, with potential applicability beyond arithmetic puzzles to broader reasoning tasks.

Abstract

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.

Figures (7)

• Figure 1: Trade-off between inference time and success rate in challenging arithmetic puzzles, Countdown. We find that fine-tuning on CoT-style data from diverse reasoners substantially enhances CoT inference over the base LLM (Qwen2.5-Math 7B) while preserving inference-time search (ToT yao2023tree, RAP hao2023reasoning) using a smaller learning rate. Among fine-tuning methods including SFT and preference optimization, our method (UFT) achieves the best CoT and search performance.
• Figure 2: Our method for reasoning tasks. We first generate synthetic reasoning data using multiple LLM reasoners (e.g., CoT, ToT, RAP) and classic algorithms (e.g., BFS, DFS). This data is converted to a unified CoT-style format and labeled as correct or failed by a ground-truth verifier to form an unpaired dataset. We then fine-tune the base LLM with negative log-likelihood (NLL) loss on correct paths and unlikelihood loss on failed paths, which we refer to as unlikelihood fine-tuning (UFT). Finally, we evaluate the improved LLM with multiple LLM reasoners on a test set.
• Figure 3: Impact of (peak) learning rate and data sources on CoT vs. search-based reasoning for standard SFT ($\alpha$=0) using Qwen2.5-Math 7B as the base model. Learning rate mediates a trade-off between CoT and search performance; CoT data better preserves search capability in most cases.
• Figure 4: Effects of unlikelihood loss and data sources on CoT vs. search-based reasoning using Qwen2.5-Math 7B as the base model, with the peak learning rate as 1e-6. Unlikelihood loss leads to greater improvements in search performance than in CoT inference in most cases.
• Figure 5: Results of standard SFT ($\alpha$=0) using Qwen2.5-Math 7B as the base model for Game-of-24. The same conclusion in \ref{['fig:lr_sft-cd4']} holds.