🤖 AI

RREDCoT: Segment-Level Reward Redistribution for Reasoning Models

The paper introduces RREDCoT, a novel method that leverages the reasoning model itself to approximate optimal segment-level reward redistribution for Chain-of-Thought traces, thereby addressing the high variance and computational inefficiency of traditional Monte Carlo credit assignment in reinforcement learning fine-tuning.

Original authors: Mykyta Ielanskyi, Kajetan Schweighofer, Lukas Aichberger, Sepp Hochreiter

Published 2026-06-05

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Mykyta Ielanskyi, Kajetan Schweighofer, Lukas Aichberger, Sepp Hochreiter

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

The Big Problem: The "Black Box" of Reasoning

Imagine you are teaching a student to solve a very difficult math problem. The student writes down a long, step-by-step thought process (a "Chain of Thought") before finally writing the answer.

In current AI training methods, the teacher only gives feedback at the very end.

The Student: Writes 50 pages of thinking, makes a mistake on page 10, but keeps going, and eventually guesses the right answer on page 50.
The Teacher: Says, "Great job! You got the right answer."
The Result: The student thinks, "Wow, my mistake on page 10 was actually helpful!" or "I don't know which part of my 50 pages was actually useful."

This is called a delayed reward problem. Because the AI doesn't know which specific thoughts led to the success, it learns slowly and inefficiently. It's like trying to learn to drive a car by only getting a "Good Job" or "Bad Job" signal after you've parked, without knowing if you turned the wheel too early or hit the brakes too late.

The Solution: RREDCoT (The "Rewind and Re-distribute" Tool)

The authors created a new method called RREDCoT (Reward REDistribution for Chain of Thoughts).

Think of RREDCoT as a smart editor that watches the student's 50-page draft. Instead of just grading the final answer, the editor goes back and assigns a "score" to every single paragraph.

Paragraph 1-5: "Good setup, but not critical." (Low score)
Paragraph 10: "This was a wrong turn, but you recovered." (Negative score)
Paragraph 25: "This was the key insight that solved the puzzle!" (High score)
Paragraph 50: "Correct answer." (Bonus score)

By giving credit (or blame) to specific parts of the thinking process, the AI learns much faster which thoughts are actually useful.

How It Works (Without the Math)

The paper explains that previous methods tried to solve this in two ways, both of which had flaws:

The "Guessing Game" (Monte Carlo Sampling): The AI would generate the same problem 100 times to see which steps usually lead to success. This is accurate but takes forever (like running a marathon 100 times just to figure out the best route).
The "Blame Game" (Attribution): Looking at the AI's internal "attention" to guess what was important. The paper argues this is often misleading because it looks at what the AI looked at, not what actually worked.

RREDCoT's Trick:
Instead of running the AI 100 times or guessing blindly, RREDCoT uses the AI's own knowledge to estimate the value of each step.

It looks at the "Reference Solution" (the correct path).
It asks: "If I had taken this specific step, how much closer did I get to the correct answer?"
It uses a clever math shortcut (inspired by how we predict the next word in a sentence) to calculate this instantly, without needing to generate 100 extra versions of the story.

The "Hybrid Segmentation" (Cutting the Cake)

To make this work, the AI needs to know where one "thought" ends and another begins. You can't just look at every single letter (token) because that's too much data.

The Paper's Idea: They use a "Hybrid Segmentation" strategy. Imagine cutting a long cake.
- First, they cut at obvious places (like new paragraphs or keywords like "Wait" or "Therefore").
- Then, they look at the "confusion" (entropy) of the text. If the AI was very unsure about what to write next, that's a good place to cut the cake.
- This creates logical "chunks" of thinking that are easy to grade.

What They Found (The Results)

The researchers tested this on math problems (like the AIME and MATH datasets).

Faster Learning: Models using RREDCoT learned to solve problems better and faster than models using the standard method (GRPO).
Better Efficiency: They got better results even when the AI was generating very long chains of thought (up to 25,000 tokens).
No Extra Models Needed: Unlike some other methods that require a second "judge" AI to grade the work, RREDCoT uses the main AI itself to do the grading, saving time and resources.

The Catch (Limitations)

The paper is honest about where this method might struggle:

You Need the Answer Key: RREDCoT works best when you already know the correct solution path (or at least a good hint of it). If you are trying to solve a problem where the solution is unknown or the "right way" is vague, this method can't help much.
It Costs a Bit More: It requires about 1.5 to 2 times more computer power than the standard method, but the authors say this is a fair trade-off for the speed of learning it provides.

Summary

RREDCoT is a new way to train AI to think better. Instead of waiting until the end to say "Good job," it breaks the thinking process into small chunks and tells the AI exactly which thoughts were helpful and which were distractions. It does this quickly and efficiently, allowing AI to learn complex reasoning skills much faster than before.

Technical Summary: RREDCoT – Segment-Level Reward Redistribution for Reasoning Models

Problem Statement

Recent advancements in reasoning language models rely heavily on Reinforcement Learning (RL) fine-tuning, particularly using the Group Relative Policy Optimization (GRPO) algorithm combined with the Reinforcement Learning with Verifiable Rewards (RLVR) paradigm. In this setting, models generate a Chain-of-Thought (CoT) trace before producing a final answer. However, a critical bottleneck exists: the reward signal is typically assigned only after the entire CoT trace is complete and the final answer is verified. This creates a delayed reward problem.

Standard approaches like GRPO correspond to Monte Carlo methods in RL, which suffer from high variance because the reward is distributed uniformly across the entire trajectory, providing no direct supervision for individual reasoning steps. While Monte Carlo Sampling (MCS) can estimate intermediate state values, it requires generating multiple completions for every segment, leading to prohibitive computational overhead for high-granularity credit assignment in long contexts. Existing alternatives, such as using judge models for step-by-step analysis or extracting intermediate utility from model statistics, often introduce noise or require additional computational resources.

Methodology: RREDCoT

The authors propose RREDCoT (Reward REDistribution for Chain of Thoughts), a method that adapts the principles of the RUDDER algorithm to the specific structure of CoT generation. The core objective is to redistribute the final reward to intermediate segments of the CoT trace to provide denser learning signals without requiring additional models or extra generation steps.

1. MDP Formulation

The CoT generation process is modeled as a Markov Decision Process (MDP) where the state $s_t$ consists of the original query and the sequence of generated segments up to $t-1$ . The action space is partitioned into:

$Y$ (Answer): Actions that trigger the final output.
$U$ (Continuation): Actions that generate further CoT segments.
The immediate reward for generating a thought is zero; the reward is only realized upon generating a correct final answer.

2. Optimal Reward Redistribution

RREDCoT derives an optimal reward redistribution $\sigma$ based on the difference in expected returns between consecutive states. The method utilizes the generating language model itself to approximate the value function, avoiding the need for an auxiliary LSTM (as used in original RUDDER) or extensive sampling.

The redistribution is calculated as:
$\sigma_t \propto \hat{R}_{t+1} - \hat{R}_t + \hat{v}_{t+1} - \hat{v}_t$
Where:

$\hat{R}_t$ is the immediate reward estimate (probability of the reference answer given the current trace).
$\hat{v}_t$ is the estimated value of the state (expected utility of future completions).

To estimate these values efficiently, the authors employ a Probability Reward (PR) estimator using importance sampling over a set of reference solution paths ( $Y^\star, U^\star$ ). This allows the model to act as a dynamics function (world model) to estimate the likelihood of reaching a high-utility state from the current segment.

3. Hybrid Segmentation Strategy

To make the credit assignment tractable, RREDCoT introduces a hybrid keyword-entropy segmentation strategy:

Initial Split: The trace is split by generic keywords (e.g., newlines).
Iterative Merging: Consecutive segments with the lowest combined entropy are merged until a target number of segments is reached.
Rationale: High-entropy tokens represent points of high plurality (decision points) and serve as natural "exit points" for estimating immediate contributions. This strategy ensures segments are as decoupled as possible, minimizing the bias in the value function estimation.

4. Integration with RL Objectives

The redistributed rewards are integrated into the standard policy gradient objective. The method ensures return-equivalence, meaning the sum of redistributed rewards equals the original episode return, preserving the optimal policy while accelerating convergence.

Key Contributions

RREDCoT Algorithm: A tractable credit assignment and reward redistribution algorithm for CoT traces that does not require additional models or extra generation steps, leveraging the generating model itself for value estimation.
Entropy-Based Segmentation: A novel segmentation strategy tailored for RREDCoT that balances keyword-based splitting with entropy-driven merging to optimize the estimation of intermediate contributions.
Theoretical and Empirical Analysis: A detailed analysis comparing the proposed estimator against Monte Carlo sampling and various attribution methods (e.g., Shapley values, Leave-One-Out), demonstrating that attribution methods unconditioned on correct solutions capture different aspects of generation than reward redistribution.

Experimental Results

The authors evaluated RREDCoT on several benchmarks, including AIME, MATH, and Minerva datasets.

Long-Context Performance: On the Numina-CoT dataset with generation lengths up to 25k tokens, RREDCoT demonstrated superior performance compared to standard GRPO. For instance, on the Qwen3-4B Instruct model, RREDCoT achieved an AIME24 score of 0.908 compared to GRPO's 0.850.
Small-Scale Refinement: When applied to small reasoning models (1.5B parameters) for online refinement, RREDCoT showed competitive or superior results compared to models tuned with GRPO and hyperparameter tuning, particularly on the AMC23 and MATH500 benchmarks.
Variance and Bias: The study showed that while truncated Monte Carlo sampling leads to significant loss of valid completions (introducing bias), RREDCoT's estimator aligns well with MC sampling results at a fraction of the computational cost.
Correlation Analysis: The paper found that gradient-based and Leave-One-Out (LOO) attribution methods are highly correlated with each other but less correlated with RREDCoT. This suggests that attribution methods capture different signal properties than reward redistribution, and using attribution alone for training weights may be redundant or suboptimal.

Significance and Claims

The paper claims that RREDCoT addresses the high variance and delayed reward issues inherent in RL fine-tuning of reasoning models. By redistributing rewards to intermediate segments based on the model's own predictive capabilities, the method accelerates convergence and improves final performance without the computational prohibitions of Monte Carlo sampling.

The authors emphasize that their approach is a tractable approximation of optimal reward redistribution. They note that the method relies on the availability of reference solutions and some subgoal structure (e.g., unique solutions in math problems). While the method introduces a computational overhead of 1.5–2x compared to naive GRPO, it is significantly more efficient than full Monte Carlo sampling. The work positions RREDCoT as both a training signal amplifier and an analysis tool, offering a new perspective on credit assignment that is distinct from traditional attribution techniques.