1 Introduction↩︎

As established modes of language model pretraining hit the boundaries of data availability, the focus shifts further towards maximizing the utility extracted from the pretrained models in subsequent “post-training” stages [1]. Techniques such as Reinforcement Learning from Human Feedback (RLHF) [2] have emerged, relying on a reward signal from possibly vague human preference data. These early attempts to improve the quality of language models were based on the Proximal Policy Optimization (PPO) algorithm [3].

Today, Group Relative Policy Optimization (GRPO) [4] has become the most popular choice for RL-based fine-tuning, because it circumvents the necessity of simultaneously learning an advantage function. This is combined with RL with Verifiable Rewards (RLVR) paradigm, in which model generates a Chain-of-Thought (CoT) before producing the final result, with training often consisting of online tuning followed by distillation on the generated reasoning traces [5]. However, CoT generation poses several challenges: models tend to produce excessively long traces, and answer extraction and reward estimation can yield disparate performance assessments and noisy training signals ([6], [7]).

One prominent issue that RL fine-tuning of reasoning language models still faces is the lack of fine-grained reward signals. In RLVR, the reward for the entire episode is assigned only at the very end of a generated Chain-of-Thought (CoT) and distributed uniformly over the full trajectory, which provides no direct supervision for individual reasoning steps. In on-policy reasoning fine-tuning, this problem has been approached from several directions. Some works propose step-by-step analysis of CoT traces using judge models as a means of credit assignment [8]–[10], while others attempt to extract intermediate utility directly from model statistics [11]. Monte Carlo sampling (MCS) based state value estimation methods have been shown to provide effective intermediate value estimates [12], [13] at the cost of considerable additional token generation.

With RREDCoT, we adapt the core principles of RUDDER [14] to the specific structure of the CoT generation MDP and derive a tractable approximation of the optimal reward redistribution that requires neither additional models nor extra generation steps. Unlike the original RUDDER formulation, which employs an LSTM [15] for return decomposition, we utilize the generating language model itself. This exploits the properties of autoregressive sequence generation together with a novel fast value function estimator inspired by Bayesian methods in natural language generation [16], [17].

Our contributions are as follows:

• We introduce RREDCoT - a tractable credit assignment and reward redistribution algorithm for CoT traces that does not require additional models.

• We devise an entropy-based segmentation strategy for the CoT traces that is specifically tailored for RREDCoT.

• We analyze the relation between our estimator, MC sampling and several attribution methods.

2 CoT Generation as a Reinforcement Learning Problem↩︎

We consider a Markov Decision Process \(\cP = (\cS, \cA, \cR, p, \gamma)\) where \(\cS\) is a set of states, \(\cA\) the set of actions, and \(\cR\) the reward space. A state at step \(t\) is given by \(s_{t} = ( \Bx, \Bu_1, \dots, \Bu_{t-1} )\) where \(\Bx\) is the original query and each \(\Bu_t\) is a generated CoT segment. The language model is then the policy with a discrete categorical action space, i.e., the space of possible next tokens \(a\), also called vocabulary.

Segments are contiguous and non-overlapping sequences of generated tokens that together span the entire sequence. In the corner cases, the whole generated text would be split into segments consisting of individual tokens or kept together as a large single segment. This depends on segmentation strategy discussed further in Sec. 3.1. Importantly, as the state is defined by \(\Bs_{t} = ( \Bx, \Bu_1, \dots, \Bu_{t-1} )\), the dynamics function is deterministic when conditioned on the next segment, i.e., \(p(\Bs_{t+1} \mid \Bs_t,\Bu_t) = 1\). Thus, there is only stochasticity through the policy (the LLM). In the subsequent equations, whenever the transitions (e.g. \(p(\Bs_{t+1} \mid \Bs_{t})\)) are listed without specifying the parameters, \(\Bw\) of the generating model are implied.

2.0.0.1 Reward Modeling for CoT Outputs.

The reward model of text generation used for evaluation is as per Eq. 1 . We call \(\zeta : \By \mapsto \left[0, 1\right]\) a utility function. It maps from the answer space to a real value, which is the inverse of the cost of generating the final output sequence \(\By\). In the simplest case, \(\zeta\) is defined as the correctness of the answer. Often it is a combination of several different cost functions, such as adherence to output format or the overall length of \((\Bu, \By)\), with improperly formatted or overly long sequences being associated with higher cost. The return is then the expected inverse risk, conditioned on the CoT trace \(\Bu\). \[\begin{align} R_{u} = \EXP_{\By \mathop{\mathrm{\sim}}p(\By \mid \Bu, \Bx, \Bw)}[\zeta(\By)] \label{eq:reward95equation95cot} \end{align}\tag{1}\] The CoT trace \(\Bu\) does not affect the reward assigned to the final action \(\By\) directly. Instead, it affects the reward by changing the conditional probability distribution of \(\By\). The VR estimator uses Monte Carlo integration to estimate \(R_{\By}\) with most prior work only sampling a single output together with the original trace (Appx.Eq. 13 ). Several recent works propose using Probability Reward (PR) instead [18], [19]. PR reward estimators use importance sampling (Appx.Eq. 14 ) to reduce variance. This way, the ability of a model to serve as a dynamics function (i.e., world model) is leveraged. It is thus important that the model is reasonably well calibrated to assess the degree of causality between the CoT trace and the answer. PR estimator provides a denser signal but may suffer from increased noise as was pointed out in [20].

2.0.0.2 CoT Generation MDP and its Bellman Equation.

A peculiar feature of the particular MDP at hand is that the action space \(\cA\) is partitioned into \(\{\cY, \cU\}\) subspaces. The \(\cY\) subspace contains all the actions that result in the beginning of the final output, while the \(\cU\) subspace contains all of the actions that do not trigger the output and instead continue the generation of the CoT trace.

In the commonly used models, the first subset of actions is initiated when end of thinking token (EOT) is produced, i.e. </think>. Upon selecting the EOT token, the model can then produce the answer for which the reward can be evaluated as per Eq. 1 . Therefore, we can rewrite the Bellman equation for such MDP in the following fashion: \[\begin{align} v_{\Bw} (\Bs_{t}) = & \sum_{a \in \cY} p(a \mid \Bs_t) \; \zeta(\Bs_{t+1}) + \sum_{a \in \cU} p(a \mid \Bs_{t}) \;\gamma \;v_{\Bw}(\Bs_{t+1}) \nonumber \\ = & \: \EXP_{\By \mathop{\mathrm{\sim}}p(\By \mid \Bs_{t})} [ \zeta(\By) ]+ \EXP_{\Bu_t \mathop{\mathrm{\sim}}p(\Bu_t \mid \Bs_{t})} [ v_{\Bw}(\Bs_{t+1}) ] \label{eq:cot-mdp-bellman} \end{align}\tag{2}\]

In Eq. 2 , we use the knowledge of the CoT MDP structure to split the value function on the action space \(\cA\) into two parts: one that corresponds to the model selecting an answer and the other that corresponds to the model selecting to continue CoT generation with fragment \(\Bu_t\). With CoT generation, the model only gets rewarded for generating useful output in the end. The immediate reward for generating another thought is zero. CoT traces from commonly used models can reach thousands of tokens in length [5], [21], resulting in highly delayed rewards. We can further derive the corresponding action-value function \(q^{\Bw} (\Bs_{t}, a)\): \[\begin{align} q^{\Bw} (\Bs_{t}, a) = \begin{cases} \EXP_{a \mathop{\mathrm{\sim}}p(a \mid \Bs_t, \Bx, \Bw)} [ \zeta(\Bs_{t}) ] & \text{if } a \in \cY \\ v_{\Bw}(\Bs_{t+1}) & \text{if } a \notin \cY \end{cases} \label{eq:cot-mdp-action-value} \end{align}\tag{3}\]

Alternatively, the value term of the CoT MDP can be interpreted as an expectation under the predictive distribution of the reasoning language model: \[\begin{align} v_{\Bw} (\Bs_{t}) = & \; \EXP_{\By, \Bu \mathop{\mathrm{\sim}}p(\By, \Bu \mid \Bx, \Bw)} \bigg[ \zeta(\By) \bigg] \nonumber \\ = & \; \sum_{\By \in \cY, \Bu \in \cU} p(\By, \Bu \mid \Bx, \Bw) \cdot \zeta(\By) \nonumber \\ = & \; \sum_{\By \in \cY, \Bu \in \cU} p(\By \mid \Bu, \Bx, \Bw) p(\Bu \mid \Bx, \Bw) \cdot \zeta(\By) \label{eq:cot-mdp-bellman-bayesian} \end{align}\tag{4}\] In this formulation, the reward function \(\zeta\) is integrated over the space of \(\BY \times \BU\) that consists of the final outputs \(\By\) and chains of thought \(\Bu\). What enables such a transformation is that (a) the dynamics function is trivial; (b) the intermediate rewards for transitions within \(\Bu\) are zero; (c) \(\gamma\) is normally ignored and set to one in the CoT setting. This view enables us to use some of the insights from Bayesian methods applied to LLMs to estimate the value term.

3 RREDCoT↩︎

RUDDER [14] is an algorithm for decomposing and redistributing the rewards for delayed reward problems. In Deep Reinforcement Learning (DRL), it manages to mitigate the bias of the TD methods and the variance of the MCMC methods by learning an additional LSTM model [15] that is then used for credit assignment to intermediate actions, speeding up the convergence of on-policy methods. [14] describe the optimal reward redistribution \(\tilde{\cP}\) as such, that any future reward \(\kappa(T-t-1, t) = 0\) for \(0 \le t \le T - 1\). Under such redistribution, the following conditions are equivalent and sufficient for optimal reward redistribution: \[\begin{align} \kappa(T-t-1, t) &= 0 \quad \nonumber \\ \Leftrightarrow \nonumber \\ \quad \EXP[R_{t+1} \mid \Bs_{t-1}, a_{t-1}, \Bs_{t}, a_{t}] &= \tilde{q}^{\pi} (\Bs_{t}, a_{t}) - \tilde{q}^{\pi} (\Bs_{t-1}, a_{t-1}) \label{eq:rudder-q-form} \end{align}\tag{5}\] We can use the action values derived for CoT MDP in Eq. 3 together with Eq. 2 and the second definition of optimal reward redistribution in Eq. 5 to derive an optimal reward redistribution for the CoT trace. Let us consider the scenario where the model is in state \(\Bs_t\), \(a_t \notin \cY\), and \(a_{t-1} \notin \cY\): \[\begin{align} \tag{6} \EXP[R_{t+1} & \mid \Bs_{t-1}, a_{t-1}, \Bs_{t}, a_{t}] = \\ \tag{7} = & \; q^{\Bw} (\Bs_{t}, a_{t}) - q^{\Bw} (\Bs_{t-1}, a_{t-1}) \\ \tag{8} = & \; \EXP_{\By \mathop{\mathrm{\sim}}p(\By \mid \Bs_{t+1})} [ \zeta(\By) ] - \EXP_{\By \mathop{\mathrm{\sim}}p(\By \mid \Bs_{t})} [ \zeta(\By) ] \\ & + \; \EXP_{\Bu_{t+1} } [ v^{\Bw}(\Bs_{t+2}) ] - \EXP_{u_{t} } [ v^{\Bw}(\Bs_{t+1}) ] \nonumber \end{align}\] The part of the sum in Eq. 7 can be computed by direct estimation of expectations in Eq. 1 using the model with parameters \(\Bw\). To estimate these expectations under the output distribution of an autoregressive models one could use multinomial sampling as was done by [22]. At the same time, under \(\{0,1\}\) outcome reward \(\zeta\) and a given small set of reference answers is known, we can estimate the expectation in Eq. 1 using a PR estimator.

The value part of the reward redistribution (Eq. 8 ) requires estimating an expectation over the plausible subsequent steps. Estimating this quantity would require sampling for every segment, which would render it either too slow to be practically applied to online training or require a considerable reduction in segment number, yielding coarser credit assignment. However, this expectation would generally be low if the probabilities of the actions \(p(\Bu_t \mid \Bs_t, \Bw)\) are small. This gives rise to our segmentation strategy in the following section, which aims to ensure that the fragments from existing CoT are as decoupled as possible.

3.1 Hybrid Segmentation Strategy↩︎

While our algorithm can be used for token-level attribution, this would require thousands of evaluations for every trace, rendering it impractical. The segmentation strategy may influence the subsequent credit assignment and additional computational costs. Several prior works consider the segmentation of CoT traces. For instance, [22] define subthoughts based on special sequences, e.g. “Wait”, “But”, etc. More generic substring segmentation approaches are also possible, such as splitting on double newlines. [13] and [23] propose segmentation strategies based on token probabilities. While keyword based approaches are inherently anthropomorphic and may lead to suboptimal segmentation from the algorithmic standpoint, the purely entropy or probability based approaches can fall into the trap of synonyms, where the high immediate entropy does not lead to actual plurality of the subsequent strings.

We propose using a hybrid keyword-entropy segmentation approach (Fig. 2 (Left)). The hybrid segmentation starts with splitting by generic keywords, such as a newline character, and later iteratively merging consecutive segments with the lowest combined entropy until the desired consolidated number of exit points is reached. High entropy tokens have high plurality and are therefore opportune ‘exit points’ for estimating the immediate contributions. Additionally, we can then homogenize segment likelihood and allow better estimation of the value term in Eq. 8 .

3.2 Credit Assignment↩︎

Once the segmentation is done, we must assign intermediate advantages or proportions thereof to every segment. The goal of this stage is to construct an efficient estimator for the Eq. 7 and 8 . Here, it is important to delineate credit assignment from the attribution analysis. The question of attribution analysis is “what contributed to obtaining this trajectory” whereas the question of credit assignment is "what contributed to getting this reward". In other words, credit assignment introduces additional information about the reward, while the attribution analysis operates merely on the sampled trajectory and environment dynamics.

With the PR estimator we can estimate not only the immediate term in Eq. 7 , but also get an estimate of the value term. We can achieve that by using a reference solution path, which is usually provided in the datasets. Immediate reward estimates in this case can be viewed as a distance to the goal state, where the generation of the desired high utility output is inevitable.

Our PR-style estimator for the value function is then as follows: \[\begin{align} \hat{v}^\text{our}_{\Bw} (\Bs_{t}) = & \sum_{\By, \Bu \mathop{\mathrm{\sim}}q} \frac{p(\By \mid \Bu, \Bx, \Bw) p(\Bu \mid \Bx, \Bw)}{q(\By, \Bu \mid \Bx, \Bw)} \cdot \zeta(\By) \nonumber \\ = & \frac{1}{N} \sum_{\substack{ \By \in \{\cY^\star\} \\ \Bu \in \{\cU^\star\} }} p(\By \mid \Bu, \Bx, \Bw) p(\Bu \mid \Bx, \Bw) \cdot \zeta(\By) \label{eq:value-estimator-our} \end{align}\tag{9}\] Where \(q\) is a proposal distribution which we define as a uniform distribution over selected reference answer-solution pairs \(\cY^\star\) and \(\cU^\star\). These sequences are a predetermined set of answers and corresponding solution paths accordingly. The crucial quality that the members of the set must possess is non-zero utility \(\zeta\). In the simplest case, this set consists of a single reference answer and solution. Alternatively, it can feature multiple combinations of answers, solutions coming from different sources, including the model’s own high utility solution answer pairs and solutions generated by teacher models.

\(\hat{v}^\text{our}_{\Bw}\) is an importance sampling estimator that uses the biased proposal distribution of the available solution paths. Let us derive the extent of bias of this point estimator using the unbiased MC estimator (for detailed transformations refer to Appx. 9.5): \[\begin{align} \text{Bias} & \left[ v_{\Bw}(\Bs_{t}), \hat{v}^\text{our}_{\Bw} (\Bs_{t})\right] = \EXP[\hat{v}^\text{our}_{\Bw} (\Bs_{t})] - v_{\Bw}(\Bs_{t}) \nonumber \\ = & - \sum_{\By \in \cY \setminus \cY^\star, \Bu \in \cU \setminus \cU^\star} \frac{1}{N} \underbrace{p(\By, \Bu \mid \Bx, \Bw)}_{\text{ answer \& solution prob.}} \cdot \underbrace{\zeta(\By)}_{\text{utility}} \label{eq:value95estimator95bias} \end{align}\tag{10}\] where \(Z\) is the cardinality of the set of all possible solutions and answers for the autoregressive model and \(N\) is the cardinality of the reference solution set. We note that \(Z\), being the number of all possible sequences generated by the model, is finite since the maximum length must be specified for the autoregressive decoding algorithm [16]. Essentially, this means that for every pair \(\By^\star\) and \(\Bu^\star\) we sum the integrand over all continuations that are not this specific one. Under the practically sensible assumption that \(\zeta\) is non-negative, this bias is less or equal to zero and will lead to underestimation of the value function. In the best case scenario, our proposal distribution \(q\) captures all sequences that have non-zero probability and non-zero utility leading to zero bias.

While explicitly estimating the bias in Eq. 10 is intractable, we can draw some hypotheses about it by leveraging the empirical insights from Bayesian methods for LLMs. It is known that the predictive distribution of Language Models is structured and will generally contain clusters that are correlated [24], [25]. With this in mind, if we take a difference between \(\hat{v}^\text{our}_{\Bw}(\cdot)\) evaluated at \(\Bs_{t}\) and \(\Bs_{t+1}\) as is required to estimate the Eq. 8 , we can expect the difference of the probabilities of the reference solution to be indicative of the dynamics for the whole cluster. One scenario where this bias could be substantial is when there are multiple diverse solutions that are plausible under the model \(\Bw\) and have high utility for the given problem \(\Bx\). If the solutions are diverse, they might be parts of different clusters that are disjoint in terms of probability. In this case, the mass in Eq. 10 would have a high magnitude and it would be possible that the estimates obtained through the difference of \(\hat{v}^\text{our}_{\Bw}(\cdot)\) could even be nonsensical.

3.2.0.1 Subgoal structure of reference solution.

We assume that the regions of the predictive distribution of the reasoning model that are disjoint from the region containing the reference solution, yet contain sequences with positive utility are rare. Given the body of work in Bayesian Language Modeling (i.e., [17]) we feel such an assumption is safe. [5] and the majority of subsequent RLVR works consider mathematics datasets (e.g. [26], [27]) which have the advantage of having unique solutions (down to symbolic transformation). Non-mathematical problems, such as logic and programming [28], [29], also provide unique solutions down to invariances. In many problems with composite answers (i.e., answers that contain multiple parts that can be independently correct or wrong), such as Sudoku puzzles or Python programs, the subgoal structure is revealed even without an explicit solution path. For example, in a Sudoku puzzle, the placement of each number is a subtask and when we see the filled-out answer to it, we can assess the achievement of individual subgoals.

Often a reference solution path is not available. As we show in the later section, the reward redistribution requires more information than just the produced sequence itself. The manner in which such information is obtained (whether the reference solution or an extensive tree search) is beyond the scope of our work. In this case one can consider starting with applying a search algorithm of choice which attempts to find a solution for a given problem. This solution path can then be summarized and used for redistribution. We assume that some hint of the subgoal structure is available during training.

3.3 Integrating Reward Redistribution into Commonly Used RL Objectives↩︎

The RREDCoT redistribution approximation derived so far depends only on the properties of the CoT generation MDP. This means that it can be integrated into any RL objective so long as it is applied to a problem of that MDP structure. One important condition of reward redistribution is return-equivalence, meaning that the original episode return must be preserved in the reward redistribution SDP. We further reformulate the generalized policy gradient objective from [4] to include the redistribution of rewards: \[\begin{align} \nabla_\theta & \cJ_{\cA}(\theta) = \EXP \underbrace{_{(q,o) \mathop{\mathrm{\sim}}\cD}}_{\text{Data Source}} \bigg[ \sum_{t=0}^{|o|} \underbrace{ \sigma(t, q, o, \pi_\text{rf})}_{\text{Redistribution}} \underbrace{GC_{\cA} (q,o,\pi_\text{rf})}_{\text{Gradient Coefficient}} \underbrace{\nabla_\theta \log \pi_\theta (o_t \mid q, o_{<t})}_{\text{Token-Wise Gradient}} \label{eq:deepseek-math-objective-with-redistribution} \bigg] \end{align}\tag{11}\] Where \(\sigma (t, q, o, \pi_\text{rf})\) is the token-wise reward coefficient that depends on the reference model and question-answer pair. The policy gradient formulations attribute reward uniformly over the whole sequence Eq. 12 : \[\begin{align} \label{eq:deepseek-math-objective-uniform-sigma} \sigma(t, q, o, \pi_\text{rf}) = \frac{1}{|o|} \end{align}\tag{12}\] \(\sigma\) must sum up to \(1\) in order to satisfy the return equivalence from [14]. We note, that this formulation does not bias the original objective according to Theorem 1 in [14] if the condition holds. This leads to the same optimal policy. In practice, the diversity of data and breadth of autoregressive predictive distribution mean that full convergence i.e. perfect knowledge of the general reasoning, is unattainable in practice. As a consequence, speeding up the convergence rate would mean de facto improvement of final performance.

Several GRPO improvements use the normalizer \(\sigma\) for reward shaping without adhering to these conditions, therefore changing the optimal policy. For example, BNPO [32] uses \(\sigma = \frac{1}{|\cB|}\) where \(\cB\) is the number of tokens in a batch. DR-GRPO [33] uses \(\sigma = \frac{1}{|\cM|}\) where \(\cM\) is the maximum number of tokens allowed in the batch.

4 Experiments↩︎

In this section, we investigate some of the empirical properties relating to the proposed RREDCoT. The implementation details of the experiments in this section are provided in Appx. 10.

4.0.0.1 Variance and Bias of Truncated MC Value Estimator.

We conducted an analysis of the variability of the estimates of intermediate values by MC sampling. In Appx.Fig. 8, we show how the number of MC samples impacts the standard deviation of the estimates. We observe SD of the estimator falling below the \(0.1\) mark after approximately \(5\) samples. In this example, we have computed the MC values by generating sequences with a fixed total horizon, so the effects of increased variance at the early stages of CoT generation are not visible. Such computation would be impractical at these levels of granularity for on-policy optimization. In Fig. 4 we show that truncating the MC sampling completions can lead to loss of a large number of completions, including ones that lead to correct answers. This introduces a bias similar to that in Eq. 10 .

4.0.0.2 Correlation analysis between attribution and credit assignment methods.

A number of attribution techniques exist, such as Shapley values [34], Leave One Out (LOO) [35], [36] and gradient-based techniques [37]. Attribution methods have their advantages and drawbacks, and there is no one-size-fits-all. Generally, attribution in neural networks is a computationally intensive procedure. We want to test the hypothesis that attribution unconditioned on additional information is disjoint from credit assignment/reward redistribution and would be ill-suited for explicitly weighing the tokens during training. The results of the correlation analysis are presented in Fig. 5. The gradient attribution is highly correlated with LOO attribution, meaning that using the LOO attribution values as an explicit weighing signal at training time would be redundant. This shows that using LOO style attribution analysis for the purposes of assigning intermediate rewards would be redundant at training time, since much of its allocation would be implicitly performed by the gradient descent. We speculate that this would also negate the effect of any other purely attribution explicit weight assignment, including, for example, attention weight-based attribution. At the same time, RREDCoT reward redistribution shows relatively high correlation, especially towards the later parts of the trajectories, where the MC estimates might be more precise.

4.0.0.3 On Policy LM Fine-Tuning.

We have utilized the objective in Eq. 11 to assess the utility of our reward redistribution. We used 4B parameter Qwen3 Instruct model as a starting point with maximum generation length set to \(25\)k tokens. The optimization was performed for \(500\) steps with hyperparameters detailed in Appx.10. The \(\sigma\) values for redistribution were normalized using L1 norm. This aligns conveniently with using PR approaches [18], [20] to estimate the reward, as the sum of all intermediate rewards in log space equals the log PR. For GRPO, standard verifiers along were used along with several beneficial adjustments, such as sequence level importance [38] and implementation level importance weights as is provided in the Transformer Reinforcement Learning library [39].

For each question in evaluation datasets \(8\) independent rollouts have been produced. Each output was evaluated using an ensemble of the LLM-as-a-Judge models with two prompts \(8\) samples each [40] for robustness. Tab. 1 shows a confident increase in efficiency when using our method for longer generation lengths.

4.0.0.4 Redistributing Reward Using Model’s Own Answer.

In order to check if our objective works with bootstrapping the model on its own answers, we have utilized the objective in Eq. 11 but used the models own answer for redistribution. The advantage terms were estimated as in GRPO (Eq. 13 ). Since the advantage values were computed from the models own answers, we could normalize the \(\sigma\) values using the softmax, making the \(\sigma\) values positive and sum up to \(1\). Log-differences of transitions were used as attribution values to estimate the optimal redistribution (Eq. 7 ).

The results are presented in Tab. 2. Best of \(2\) checkpoints were taken for RREDCoT and best of \(3\) for open-rs over a short training run (300 steps). The resulting model compares favorably to those obtained by [31] using GRPO and hyperparameter tuning.

5 Limitations↩︎

Our proposed method relies on the knowledge of the solution to the problem and some degree of subgoal structure. Normally this would be an existing solution trace that is not atypical under the distribution of general text. While a vast majority of the post-training datasets known to us provide such information, it cannot be taken for granted. Our method cannot be simply applied to the problems for which no hint of solution path is available or the provided reference solution is uninformative. One category of such problems is constraint satisfaction problems. Mitigation of this could be attempted with the use of bootstrapping on the models own correct answers (as was assessed in the last paragraph (Sec.4/P.4.0.0.4) of the preceding section). This, however, will only work if a positive signal is attainable by the model and precludes the use of PR based approaches to sequence level advantage estimation.

Furthermore our method involves additional computational overhead compared to the simpler GRPO variants. In practice our optimization runs with RREDCoT required \(1.5-2\)x the amount of compute as the runs without redistribution. This is still much faster compared to the MC sampling based intermediate value estimates and we believe that it is a reasonable trade-off. Our implementation of value function estimation for the intermediate values was aggressively optimized for reuse of KV values and batched inference. The peak GPU memory utilization was unchanged from naive GRPO.

6 Conclusion and Outlook↩︎

In this work, we introduce a new reward redistribution and credit assignment method RREDCoT which is intended to improve the sample efficiency of reasoning models fine-tuning. We have described the method and its versatile compatibility with the existing RLVR / RLPR pipelines and optimization objectives, conducted evaluations regarding its empirical benefits and discussed its relation to alternatives. The estimator agrees well with the values provided by our thorough MC sampling estimate while being considerably less computationally expensive. We additionally show that traditional attribution methods capture different aspects of CoT generation than reward redistribution and MC-based credit assignment. These insights show that attribution methods that are not conditioned on the correct solutions are bound to capture aspects not related to arriving at the desired state. We anticipate future work to further explore the application of Bayesian methods to credit assignment problems in this increasingly important setting. The RREDCoT can function as both a training signal amplifier and an analysis tool.

Acknowledgments↩︎

The ELLIS Unit Linz, the LIT AI Lab, the Institute for Machine Learning, are supported by the Federal State Upper Austria. We thank the projects FWF AIRI FG 9-N (10.55776/FG9), AI4GreenHeatingGrids (FFG- 899943), Stars4Waters (HORIZON-CL6-2021-CLIMATE-01-01), FWF Bilateral Artificial Intelligence (10.55776/COE12). We thank NXAI GmbH, Audi AG, Silicon Austria Labs (SAL), Merck Healthcare KGaA, GLS (Univ. Waterloo), TÜV Holding GmbH, Software Competence Center Hagenberg GmbH, dSPACE GmbH, TRUMPF SE + Co. KG.

7 Datasets and Models↩︎

7.1 Datasets Used↩︎

7.1.1 Dataset derived from open-rs↩︎

The models were trained on open-rr dataset. open-rr dataset is derived from open-rs dataset [31]. Despite the dataset having been compiled from other math datasets, most of which are supposed to be curated to any extent, our manual inspection has determined that it contained an error rate of \(15\). Almost all of the relatively trivial-looking questions designated as ‘Hard’ and having no step-by-step solution were found to have incorrect labels. This could additionally explain why these questions were difficult for the model to learn: the reference solutions for them were incorrect! We have filtered the dataset further from \(7,000\) examples to include only the examples containing the solution path, as well as on the length of the reference solution.

7.1.2 Evaluation Datasets↩︎

We used the following evaluation datasets: AIME24/25 [26], AMC, MATH-500 [27], Minerva [41] and OlympiadBench [42]. The verification was done with [43]. The maximum generation length was chosen to be \(4096\) for the experiments with 1.7B models and \(25k\) tokens for the Qwen3-4B based experiments.

7.2 Models Used↩︎

We have made extensive use of Qwen2.5 distilled to Deepseek R1 traces [5] and Qwen3 [44] families of models. We avoided MoE models due to the increased complexity of training dynamics with them.

8 Additional Figures↩︎

9 Additional Equations↩︎

9.1 Sequence Return Estimation↩︎

The RLPR and RLVR estimators are as follows: \[\begin{align} R^{VR}_{\Bu} & \approx \frac{1}{n} \sum_{i=0}^{n} \zeta(\By_{i}); \quad \By_{i} \mathop{\mathrm{\sim}}p(. \mid \Bu, \Bx, \Bw) \tag{13} \\ R^{PR}_{\Bu} & \approx \sum_{i=0}^{n} \zeta(\By_{i}) \; p(\By_{i} \mid \Bu, \Bx, \Bw) ; \quad \By_{i} \mathop{\mathrm{\sim}}\{\By^\star\}^{n} \tag{14} \end{align}\]

Theorem : if \(\By^\star\) contains all \(\By\) such that \(\zeta(\By_{i}) > 0\) then Eq. 13 and Eq. 14 converge to the same \(R_{\By}\).

This can be an obstacle for the practical application of RLPR when, e.g., the number of solutions is large or unbounded.

9.2 Transformations of Bellman Equation for CoT MDP↩︎

Value function decomposition for CoT MDP: \[\begin{align} v_{\Bw} (\Bs_{t}) = & \; \EXP[G_t \mid S_t = \Bs_{t}] \nonumber \\ = & \; \EXP[R_{t+1} + \gamma G_{t+1} \mid S_t = \Bs_{t}] \nonumber \\ = & \; \sum_{a} \pi(a \mid \Bs_t) \sum_{\Bs_{t+1}, r} p(\Bs_{t+1}, r \mid \Bs_{t}, a) \left[r + \gamma v_{\pi}(\Bs_{t+1})\right] \nonumber \\ = & \sum_{a} \pi(a \mid \Bs_t) \sum_{\Bs_{t+1}, r} p(\Bs_{t+1}, r \mid \Bs_{t}, a) \;r + \sum_{a} \pi(a \mid \Bs_t) \sum_{\Bs_{t+1}, r} p(\Bs_{t+1}, r \mid \Bs_{t}, a) \;\gamma \;v_{\pi}(\Bs_{t+1}) \nonumber \\ = & \sum_{\By} p(\By \mid \Bs, \Bw) \; r(\By) + \sum_{\Bu_t} p(\Bu_t \mid \Bs_{t}, \Bw) \sum_{\Bu_n} p(\Bs_{t+1} \mid \Bs_{t}, \Bu_t) \; \gamma \; v_{\pi}(\Bs_{t+1}) \nonumber \\ = & \sum_{a \in \cY} p(a \mid \Bs_t, \Bw) \; r(\By) + \sum_{a \in \cU} p(a \mid \Bs_{t}, \Bw) \;\gamma \;v_{\Bw}(\Bs_{t+1}) \end{align}\]

This uses the return at time step \(t\), defined as \(G_t = \sum^{\infty}_{k=0} \gamma^{k} R_{t+k+1}\).

Alternatively, we can reformulate the Bellman equation as estimation of a quantity under the Bayesian posterior of the reasoning language model \(\Bw\):

\[\begin{align} \label{appx:eq:bellman95to95bayesian95mdp} v_{\Bw} (\Bs_{t}) \; = & \; \EXP[G_t \mid S_t = \Bs_{t}] \\ = & \; \EXP[R_{t+1} + \gamma G_{t+1} \mid S_t = \Bs_{t}] \nonumber \\ = & \; \sum_{a} \pi(a \mid \Bs_t) \sum_{\Bs_{t+1}, r} p(\Bs_{t+1}, r \mid \Bs_{t}, a) \left[r + \gamma v_{\pi}(\Bs_{t+1})\right] \nonumber \end{align}\tag{15}\] From this point onward, given that (a) the dynamics function is trivial and (b) that \(\gamma\) is set to \(1.\) we can continue as follows: \[\begin{align} \label{appx:eq:cot-mdp-bellman-p2} = & \; \EXP_{\By, \Bu \mathop{\mathrm{\sim}}p(\By, \Bu \mid \Bx, \Bw)} \bigg[ \zeta(\By) \bigg] \nonumber \\ = & \; \sum^{\BY, \BU} p(\By, \Bu \mid \Bx, \Bw) \cdot \zeta(\By) \nonumber \\ = & \; \sum^{\BY, \BU} p(\By \mid \Bu, \Bx, \Bw) p(\Bu \mid \Bx, \Bw) \cdot \zeta(\By) \end{align}\tag{16}\]

9.3 CoT Optimal Reward Redistribution Identity↩︎

Detailed derivation of the optimal reward redistribution for CoT generation: \[\begin{align} \EXP[R_{t+1} \mid \Bs_{t-1}, a_{t-1}, \Bs_{t}, a_{t}] \; = & \; q^{\Bw} (\Bs_{t}, a_{t}) - q^{\Bw} (\Bs_{t-1}, a_{t-1}) \nonumber \\ = & \; v^{\Bw}(\Bs_{t+1}) - v^{\Bw}(\Bs_{t}) \tag{17} \\ = & \; \EXP_{\By \mathop{\mathrm{\sim}}p(\By \mid \Bs_{t+1}, \Bw)} [ r(\By) ] + \EXP_{u_{t+1} \mathop{\mathrm{\sim}}p(u_t \mid \Bs_{t+1}, \Bw)} [ v_{\pi}(\Bs_{t+2}) ] - \nonumber \\ & \; - \EXP_{\By \mathop{\mathrm{\sim}}p(\By \mid \Bs_{t}, \Bw)} [ r(\By) ] - \EXP_{\Bu_{t} \mathop{\mathrm{\sim}}p(\Bu_t \mid \Bs_{t}, \Bw)} [ v_{\pi}(\Bs_{t+1}) ] \nonumber \\ = & \; \left[ \EXP_{\By \mathop{\mathrm{\sim}}p(\By \mid \Bs_{t+1}, \Bw)} [ r(\By) ] - \EXP_{\By \mathop{\mathrm{\sim}}p(\By \mid \Bs_{t}, \Bw)} [ r(\By) ] \right] + \tag{18} \\ & \; + \left[ \EXP_{\Bu_{t+1} \mathop{\mathrm{\sim}}p(\Bu_{t+1} \mid \Bs_{t+1}, \Bw)} [ v^{\Bw}(\Bs_{t+2}) ] - \EXP_{\Bu_{t} \mathop{\mathrm{\sim}}p(\Bu_t \mid \Bs_{t}, \Bw)} [ v^{\Bw}(\Bs_{t+1}) ] \right] \tag{19} \end{align}\]

9.4 Other Equations↩︎

General policy gradient [4]: \[\begin{align} \nabla_\theta \cJ_{\cA} = \EXP \underbrace{[(q,o) \mathop{\mathrm{\sim}}\cD]}_{\text{Data Source}} \left( \frac{1}{|o|} \sum_{t=0}^{|o|} \underbrace{GC_{\cA} (q,o,t,\pi_{rf})}_{\text{Gradient Coefficient}} \underbrace{\nabla_\theta \log \pi_\theta (o_t \mid q, o_{<t})}_{\text{Token-Wise Gradient}} \right) \label{appx:eq:deepseek-math-objective-generalization} \end{align}\tag{20}\]

9.5 Value Estimator↩︎

Deriving the bias of the \(\hat{v}^{our}(\Bs_t)\): \[\begin{align} \text{Bias} \left[ v_{\Bw}(\Bs_{t}), \hat{v}^{our}_{\Bw} (\Bs_{t})\right] \; = & \; \EXP[\hat{v}^{our}_{\Bw} (\Bs_{t})] - v_{\Bw}(\Bs_{t}) \nonumber \\ = & \; \frac{1}{N} \sum_{\By, \Bu \mathop{\mathrm{\sim}}q} p(\By \mid \Bu, \Bx, \Bw) p(\Bu \mid \Bx, \Bw) \cdot \zeta(\By) \nonumber \\ & - \frac{1}{Z} \sum_{\By \in \cY, \Bu \in \cU} p(\By \mid \Bu, \Bx, \Bw) p(\Bu \mid \Bx, \Bw) \cdot \zeta(\By) \label{appx:eq:bias:true95val} \\ = & - \sum_{\By \in \cY \setminus \cY^\star, \Bu \in \cU \setminus \cU^\star} \frac{Z-N}{NZ} p(\By \mid \Bu, \Bx, \Bw) p(\Bu \mid \Bx, \Bw) \cdot \zeta(\By) \nonumber \\ \approx & - \sum_{\By \in \cY \setminus \cY^\star, \Bu \in \cU \setminus \cU^\star} \frac{1}{N} p(\By \mid \Bu, \Bx, \Bw) p(\Bu \mid \Bx, \Bw) \cdot \zeta(\By) \end{align}\tag{21}\] For these transformations, we use the true value of the value function in Eq. 21 , which is summed over the space of all possible final outputs and solution paths. The penultimate transition is made by moving the factor inside the sum and essentially means that for every pair \(\By^\star\) and \(\Bu^\star\) we sum the integrand over all sequences that are not the specific sequence from the reference pool. The final transition is made by using the fact that \(N \ll Z\). This bias is non-positive. It is zero if our \(q\) covers the entirety of \(\cY, \cU\).

10 Details on Experimental Settings↩︎

10.1 MC Sampling Variance Experiment↩︎

We have computed a single rollout for each question in the MATH-500 dataset. The maximum completion length was \(8196\). The generation and evaluation were performed with the Qwen3-4B-Thinking-2507 model. Segmentation was performed using the hybrid method with a maximum number of segments set to \(40\). \(20\) completions were generated at every exit point, each keeping the maximum completion lengths constraint. This procedure took approximately \(80\) GPU-hours on RTX 5000 PRO. \(100\) samples were subsampled to bootstrap the estimator.

10.2 Attribution and Credit Assignment Method Correlation↩︎

The same rollouts and model as in the previous subsection were used. In addition, we have computed gradient, LOO and RREDCoT values for the segments.

The LOO values were computed as pointwise mutual information between the segment and all of the subsequent segments under the model’s predictive distribution. The gradient attribution was computed by computing the gradient of logprobs of the selected completion tokens with respect to input embeddings. To get per segment value out of that, the first-order norm of these gradients was taken and averaged per segment. This is reminiscent of the gradient propagation occurring during the model update.

The loo value computation took \(0.5\) GPU-hours, while RREDCoT and gradient took about \(0.3\) GPU-hours. The considerably more computationally expensive MC sampling was used as a proxy for ground truth value segment-wise advantage estimates.

10.3 On Policy Model Training↩︎

10.3.0.1 Small Scale Experiments with open-rs dataset.

The setting and hyperparameters from [31] were used. Hybrid segmentation was performed with the maximum segment number of \(120\). Completion length of \(1024\) were used in training.

For validation during training, we used avg@4 strategy with 4 datasets: Minerva (first 50), AIME24 (full), AIME25 (full) and MATH-500 (first 50). The generation length was limited to \(4096\) tokens for performance reasons. Answer extraction and validation were performed using the math-verify package. This experiment is estimated to have consumed on the order of \(500\) GPU-hours.

10.3.0.2 Experiments with Qwen3-4B-Instruct-2507

Numina-CoT [30] dataset was used to train reasoning starting with Qwen3-4B-Instruct-2507 [44]. The shared hyperparameters of GRPO and RREDCoT run were identical, with the key difference being the learning rate that was \(1.e-6\) for GRPO and \(5.e-7\) for RREDCoT. Note, that the learning rate was kept at the given value for the GRPO as it was found to perform worse with \(5.e-7\). In other words, the LR is individually tailored for each of the two algorithm. The group size was set to \(4\), gradient accumulation to \(1\), total step size to \(500\) with cosine LR scaling.

For validation we used avg@\(8\) strategy with 5 datasets: AIME24, AIME25, AIME26, Minerva-Math and MATH-500. The generation length was limited to \(25K\), similar to training, to assess the method performance at larger scale. Answer extraction and validation were performed using the math-verify package. This experiment is estimated to have consume on the order of \(1000\) GPU-hours.

References↩︎

---

1. Work done while at Johannes Kepler University Linz.↩︎