RREDCoT: Segment-Level Reward Redistribution for Reasoning Models

2026-06-04

AI Summary

This work addresses the high variance in reinforcement learning fine-tuning caused by delayed rewards, which hinders effective credit assignment to intermediate steps in chain-of-thought (CoT) reasoning. To overcome this challenge, the authors propose a segment-level reward redistribution mechanism that requires neither additional generation nor Monte Carlo sampling, enabling endogenous and fine-grained credit assignment using the model’s own internal signals. By estimating state values at the segment level, the method achieves near-optimal reward guidance without incurring extra inference overhead. Experimental results demonstrate that, compared to existing attribution and sampling-based approaches, the proposed technique yields more efficient and stable training on long-context reasoning tasks and significantly enhances the model’s reasoning performance.

Abstract

Recent advancements in reasoning language models have been driven by Reinforcement Learning (RL) fine-tuning. Most often, these rely on the Group Relative Policy Optimization (GRPO) algorithm or modifications thereof to steer the models to produce Chain-of-Thought (CoT) traces. The final answer can only be verified, and the reward assigned, after the CoT trace is complete, making it a delayed reward problem. GRPO and its modifications correspond to Monte Carlo methods in standard RL, which are known to suffer from high variance. A possible solution to this problem is the redistribution of rewards through credit assignment, where segments of the CoT trace that are important for arriving at the desirable solution are emphasized by assigning a higher reward. While Monte Carlo sampling can be used to provide an unbiased estimate of intermediate state values, its computational overhead makes it unsuitable for train-time credit assignment in long contexts at high granularity. We introduce RREDCoT (Reward REDistribution for Chain of Thoughts), which utilizes the model itself to approximate the optimal reward redistribution without additional generation. We investigate the advantages of our method compared to MC sampling and several attribution methods. We further analyze several aspects relevant to the construction of the redistribution such as segmentation of CoT traces and state value estimation.

Problem

• reward redistribution
• credit assignment
• Chain-of-Thought
• delayed reward
• reinforcement learning

Innovation

• Reward Redistribution
• Chain-of-Thought
• Reinforcement Learning
• Credit Assignment
• Segment-Level Reward