Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

GRPO-$λ$: Credit Assignment improves LLM Reasoning

Prasanna Parthasarathi, Mathieu Reymond, Boxing Chen, Yufei Cui, Sarath Chandar

TL;DR

GRPO-$λ$ addresses the credit assignment challenge in RL-finetuned LLM reasoning by reparameterizing generalized advantage estimation with token-level eligibility traces in a critic-free TD framework. The method introduces per-token tracing and several decay styles to propagate rewards more rapidly to earlier tokens, improving learning efficiency and reasoning accuracy. Across 1.5B–7B models and four math-reasoning datasets, GRPO-$λ$ achieves 30-40% faster training and an average gain of over 3 points on benchmarks, including up to 4.5 points on the 7B model; results validate the effectiveness of token-level credit assignment in LLM RL fine-tuning. The work also analyzes bias bounds and stability mechanisms (e.g., negative-advantage clamping) and outlines future directions for trace-type exploration and generalization to broader reasoning tasks.

Abstract

Large language models (LLMs) are increasingly deployed for tasks requiring complex reasoning, prompting significant interest in improving their reasoning abilities through post-training. Especially RL based methods using verifiable reward, like the state-of-the-art GRPO, have shown to tremendously improve reasoning behaviors when applied as post-training methods. However, the lack of an explicit reward or critic model limits GRPO's ability to assign fine-grained credit across token sequences. In this work, we present GRPO-$λ$, a novel extension to GRPO that enhances credit assignment in RL finetuning of LLMs for complex reasoning tasks. We approximate learning from $λ$-return with a reformulation of eligibility traces using token-level log-probabilities applied after each sequence generation, and a novel critic-free approximation of the temporal-difference error. We introduce a few variations for the weighting of the $λ$-return, and their applications to the eligibility-trace, where all the variations provide significant gains over GRPO. We compare GRPO-$λ$ against GRPO by training models from 1.5B to 7B parameters on $4$ different math reasoning datasets. The training plots demonstrate 30-40% improved performance during RL training on both LLaMA-3.1 and Qwen-2.5 architectures. Finally, we show that with GRPO-$λ$, the resulting average performance on AIME24, Math500, OlympiadMath, MinervaMath, and AMC improves over GRPO by over $3$ points and a $4.5$ points improvement on the 7B model.

GRPO-$λ$: Credit Assignment improves LLM Reasoning

TL;DR

GRPO- addresses the credit assignment challenge in RL-finetuned LLM reasoning by reparameterizing generalized advantage estimation with token-level eligibility traces in a critic-free TD framework. The method introduces per-token tracing and several decay styles to propagate rewards more rapidly to earlier tokens, improving learning efficiency and reasoning accuracy. Across 1.5B–7B models and four math-reasoning datasets, GRPO- achieves 30-40% faster training and an average gain of over 3 points on benchmarks, including up to 4.5 points on the 7B model; results validate the effectiveness of token-level credit assignment in LLM RL fine-tuning. The work also analyzes bias bounds and stability mechanisms (e.g., negative-advantage clamping) and outlines future directions for trace-type exploration and generalization to broader reasoning tasks.

Abstract

Large language models (LLMs) are increasingly deployed for tasks requiring complex reasoning, prompting significant interest in improving their reasoning abilities through post-training. Especially RL based methods using verifiable reward, like the state-of-the-art GRPO, have shown to tremendously improve reasoning behaviors when applied as post-training methods. However, the lack of an explicit reward or critic model limits GRPO's ability to assign fine-grained credit across token sequences. In this work, we present GRPO-, a novel extension to GRPO that enhances credit assignment in RL finetuning of LLMs for complex reasoning tasks. We approximate learning from -return with a reformulation of eligibility traces using token-level log-probabilities applied after each sequence generation, and a novel critic-free approximation of the temporal-difference error. We introduce a few variations for the weighting of the -return, and their applications to the eligibility-trace, where all the variations provide significant gains over GRPO. We compare GRPO- against GRPO by training models from 1.5B to 7B parameters on different math reasoning datasets. The training plots demonstrate 30-40% improved performance during RL training on both LLaMA-3.1 and Qwen-2.5 architectures. Finally, we show that with GRPO-, the resulting average performance on AIME24, Math500, OlympiadMath, MinervaMath, and AMC improves over GRPO by over points and a points improvement on the 7B model.

Paper Structure

This paper contains 49 sections, 4 theorems, 20 equations, 15 figures, 6 tables, 2 algorithms.

Table of Contents

  1. Introduction
  2. Background
  3. Related work
  4. Reinforcement Learning
  5. PPO
  6. GRPO
  7. GRPO-L for rapid reward propagation
  8. Bounding the normalized advantage estimation bias
  9. Alternative weighting mechanisms
  10. Traces as per-token weighting
  11. Varying the type of decay
  12. Experiments
  13. Training details
  14. Analyzing the different token weighting schemes
  15. Training Efficiency
  16. ...and 34 more sections

Key Result

Theorem 1

The policy gradient estimate $\hat{g}$ using traces from generalized advantage estimation $A_\text{GAE}$ can be re-parameterized with a critic-free TD-error $\delta_t$ such that $\hat{g} = \sum_{t=0}^\infty A_\text{GAE}(s_t)\nabla_\theta\log\pi_\theta(a_t|s_t) = \sum_{t=0}^\infty\delta_t\sum_{l=0}^t

Figures (15)

  • Figure 1: GRPO-$\lambda$ improves R1-Distill-Qwen-7B's evaluation performance by over 4 points on average across math benchmarks including a 10 point improvement on AIME24.
  • Figure 2: $\epsilon$-trace styles.
  • Figure 3: Comparison of the final accuracy reward (smoothened over the last 20 training iterations) for the token weighting schemes, for both 1.5B models trained on the ORZMath57K dataset. Overall, all token weighting schemes improve training accuracy compared the the GRPO baseline. Interestingly, the both-style trace weight results in higher performance compared to the classic recent-style, showing that alternative token weighting schemes could greatly improve model performance.
  • Figure 4: Comparison of training on different RL Math datasets between the best hyperparameter configuration of GRPO-$\lambda$ and GRPO across the 4 different models used throughout this paper (ordered by row). GRPO-$\lambda$ systematically outperforms GRPO in terms of accuracy reward during training.
  • Figure 5: Comparison of the average generation length during evaluation between GRPO and GRPO-$\lambda$ when post-trained on Qwen2.5-Math-1.5B-Instruct. GRPO-$\lambda$ shows the average length across hyperparameter settings.
  • ...and 10 more figures

Theorems & Definitions (4)

  • Theorem 1
  • Lemma 1
  • Lemma 1
  • Theorem 1