Learning a Dense Reasoning Reward Model from Expert Demonstration via Inverse Reinforcement Learning

2025-10-02

AI Summary

This work addresses the lack of fine-grained supervision signals for multi-step reasoning in language models. We propose a token-level reasoning reward modeling method based on inverse reinforcement learning (IRL), which implicitly learns a dense, reusable process-level reward function from expert reasoning trajectories—contrasting with style-mimicking supervised fine-tuning. Our approach enables both policy optimization during training and path resampling during inference. Crucially, we introduce adversarial IRL to large-model reasoning supervision for the first time, achieving unified token-level alignment across training signals, inference-time selection, and error localization. Experiments on GSM8K using Llama-3 and Qwen2.5 demonstrate that our reward model significantly improves reasoning accuracy, accurately predicts answer validity, and precisely localizes intermediate reasoning errors.

Technology Category

• Machine Learning: Large Multimodal Models (LMMs)
• Natural Language Processing: (Large) Language Models
• Multiagent Systems: Adversarial Agents

Application Category

• User Modeling, Personalization and Recommendation: Large Language Models (LLM) for user modeling and recommendation
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Semantics and Knowledge: Data modeling to support human-machine intelligence, including LLMs agents, intelligent system behavior, explanations, and user-friendly interactions

Abstract

We reframe and operationalise adversarial inverse reinforcement learning (IRL) to large language model reasoning, learning a dense, token-level reward model for process supervision directly from expert demonstrations rather than imitating style via supervised fine-tuning. The learned reasoning reward serves two complementary roles: (i) it provides step-level feedback to optimise a reasoning policy during training; and (ii) it functions at inference as a critic to rerank sampled traces under fixed compute budgets. We demonstrate that our approach prioritises correctness over surface form, yielding scores that correlate with eventual answer validity and enabling interpretable localisation of errors within a trace. Empirically, on GSM8K with Llama3 and Qwen2.5 backbones, we demonstrate: (i) dense reasoning rewards can be used as a learning signal to elicit reasoning, and (ii) predictive performance is improved from reward-guided reranking (notably for Llama-based policies). By unifying training signals, inference-time selection, and token-level diagnostics into a single reasoning reward, this work suggests reusable process-level rewards with broad potential to enhance multi-step reasoning in language models.

Problem

• Learning dense token-level rewards from expert demonstrations
• Providing step-level feedback for reasoning policy optimization
• Enabling inference-time reranking and error localization in reasoning traces

Innovation

• Learning dense token-level reward via inverse reinforcement learning
• Using reward for step-level feedback and inference-time reranking
• Unifying training and diagnostics into single reasoning reward