TLDR: A new research paper introduces an Inverse Reinforcement Learning (IRL) approach to train large language models (LLMs) to reason more effectively. Instead of just imitating expert style, the method learns a dense, token-level reward model from demonstrations. This reward model serves two purposes: providing step-by-step feedback for policy optimization during training and acting as a critic to re-rank sampled reasoning traces at inference time. Experiments on GSM8K show that this approach prioritizes correctness, enables error localization, and improves predictive performance, particularly for Llama-based models.
A new research paper introduces a novel approach to enhance the reasoning capabilities of large language models (LLMs) by learning a dense, token-level reward model directly from expert demonstrations. This method, termed ‘Expert Reasoning via Inverse Reinforcement Learning,’ moves beyond simply imitating a teacher’s style, aiming instead to optimize the reasoning process itself.
Traditionally, LLMs are trained to reason through supervised fine-tuning (SFT) on human or stronger model-generated reasoning traces. While effective, this approach primarily teaches models to mimic a style rather than truly understand and optimize the underlying decision-making process of reasoning. The new paper reframes this challenge as an inverse reinforcement learning (IRL) problem, where the goal is to infer a reward function that explains expert behavior, rather than being explicitly given one.
The learned reasoning reward model serves two crucial functions. Firstly, it provides step-level feedback during training, acting as a signal to optimize the LLM’s reasoning policy. This dense, token-level feedback offers more granular guidance than traditional outcome-based rewards, which only indicate final correctness. Secondly, at inference time, this reward model acts as a ‘critic’ to re-rank sampled reasoning traces, selecting higher-quality outputs even under fixed computational budgets.
The researchers focused on four key objectives for their approach:
A Usable Training Signal
The dense, token-level rewards learned by the model provide a powerful training signal, guiding the LLM to develop better reasoning strategies. Unlike training solely on final outcome correctness, which can lead to shortcut learning, this method encourages a more robust and step-wise understanding.
Inference-Time Performance Improvement
By scoring and re-ranking sampled reasoning traces, the learned reward model significantly improves predictive performance. This means that even with a limited number of generated samples, the model can identify and select the most likely correct answers, leading to substantial gains without needing to retrain the policy.
Prioritizing Correctness Over Style
A critical aspect of this work is ensuring that the reward model values actual correctness rather than just mimicking the stylistic elements of expert demonstrations. The empirical results show that the learned reward scores correlate positively with the eventual validity of the answer, with much weaker correlations to superficial formatting or structural cues. This preference for correctness is further sharpened by introducing ‘perturbed’ examples during training, which expose near-miss errors that might otherwise be overlooked.
Also Read:
- Guiding Language Models to Smarter Reasoning with Abstractions
- GRAD: Enhancing LLM Reasoning with Dynamic, Generated Examples
Interpretable Error Localization
The dense nature of the reward allows for interpretable diagnostics. By visualizing token-level rewards along a reasoning trace, it becomes possible to pinpoint exactly where an error occurs. Correct solutions show contiguous positive reward bands on decisive computations, while incorrect solutions exhibit early negative bands at the point of deviation, with penalties propagating forward. This capability enables auditing and could potentially facilitate self-correction mechanisms in future LLMs.
The experiments were conducted on the GSM8K dataset, a benchmark for grade school mathematics problems, using Llama3 and Qwen2.5 backbones. The findings demonstrate that the dense reasoning rewards effectively elicit reasoning capabilities and improve predictive performance, particularly for Llama-based policies. While the raw accuracy (pass@1) might sometimes trail strong supervised fine-tuning baselines, the dual role of the learned reward model for both training and inference offers significant capability advantages.
This work represents a significant step towards creating reusable, process-level reward models that can unify training signals, inference-time selection, and token-level diagnostics into a single, coherent framework. The researchers acknowledge limitations, such as potential biases inherited from expert traces and the computational cost, but highlight promising future directions, including improving optimization, refining reward calibration, and integrating the critic at generation time for self-correction. For more details, you can read the full paper here.
