Redistributing reward to the reasoning steps that mattered — "RREDCoT" for chain-of-thought

This work (RREDCoT) figures out which reasoning steps mattered and redistributes reward to them when training reasoning LLMs with reinforcement learning (RL).

Recent advances in reasoning language models have been driven by RL fine-tuning, most often the Group Relative Policy Optimization (GRPO) algorithm or variants, which steer models to produce Chain-of-Thought (CoT) traces. The final answer can only be verified — and reward assigned — after the CoT trace is complete, making it a delayed-reward problem. GRPO and its variants correspond to Monte Carlo methods in standard RL, which are known to suffer from high variance.

A possible solution is redistributing reward through credit assignment, emphasizing the segments of the CoT trace that are important for reaching the desirable solution. While Monte Carlo sampling can give an unbiased estimate of intermediate state values, its computational overhead makes it unsuitable for train-time credit assignment in long contexts at high granularity.

The authors introduce RREDCoT (Reward REDistribution for Chain of Thoughts), which uses the model itself to approximate the optimal reward redistribution without additional generation. They study its advantages versus MC sampling and several attribution methods, and analyze aspects of constructing the redistribution such as segmenting CoT traces and estimating state values. The authors include Sepp Hochreiter (JKU Linz), known as a creator of the LSTM.