Exploiting Tree Structure for Credit Assignment in RL Training of LLMs

• URL: http://arxiv.org/abs/2509.18314v2
• Date: Thu, 02 Oct 2025 20:38:22 GMT
• Title: Exploiting Tree Structure for Credit Assignment in RL Training of LLMs
• Abstract summary: Reinforcement learning improves reasoning, yet sparse delayed reward over long sequences makes token-level credit assignment the key bottleneck.<n>We study the verifiable-reward setting, where the final answer is checkable and multiple responses can be drawn per prompt.<n>We propose textbfTEMPO (emphtextbfTree-textbfEstimated textbfMean Prefix Value for textbfPolicy textbfOptimization), a critic-free algorithm that augments the group
• Score: 11.64053639889468
• License: http://creativecommons.org/licenses/by-sa/4.0/
• Abstract: Reinforcement learning improves LLM reasoning, yet sparse delayed reward over long sequences makes token-level credit assignment the key bottleneck. We study the verifiable-reward setting, where the final answer is checkable and multiple responses can be drawn per prompt. Reasoning tasks in math and medical QA align with this setup, where only a few decision tokens significantly impact the outcome. PPO offers token-level advantages with a learned value model, but it is complex to train both the actor and critic models simultaneously, and it is not easily generalizable, as the token-level values from the critic model can make training prone to overfitting. GRPO is critic-free and supports verifiable rewards, but spreads a single sequence-level return across tokens and ignores branching. We introduce \textbf{Prefix-to-Tree (P2T)}, a simple procedure that converts a group of responses into a prefix tree and computes \emph{nonparametric} prefix values \(V(s)\) by aggregating descendant outcomes. Built on P2T, we propose \textbf{TEMPO} (\emph{\textbf{T}ree-\textbf{E}stimated \textbf{M}ean Prefix Value for \textbf{P}olicy \textbf{O}ptimization}), a critic-free algorithm that augments the group-relative outcome signal of GRPO with \emph{branch-gated} temporal-difference corrections derived from the tree. At non-branch tokens, the temporal-difference (TD) term is zero, so TEMPO reduces to GRPO; at branching tokens, it supplies precise token-level credit without a learned value network or extra judges/teachers. On Qwen3-1.7B/4B, TEMPO outperforms PPO and GRPO on in-distribution (MATH, MedQA) and out-of-distribution (GSM-HARD, AMC23, MedMCQA, MMLU-Medical) benchmarks, and reaches higher validation accuracy with roughly the same wall-clock time.

Related papers

• TRACE: A Unified Rollout Budget Allocation Framework for Efficient Agentic Reinforcement Learning [52.67756371231985]
  Reinforcement learning with verifiable rewards (RLVR) is a promising approach for enhancing reasoning and agentic behavior in large language models.<n>We introduce Tree Rollout Allocation for Contrastive Exploration (TRACE), a unified rollout allocation framework.<n>Technically, TRACE allocates rollout budget to both prompt roots and intermediate prefixes that are most likely to yield mixed terminal rewards.
  arXiv  Detail & Related papers  (2026-06-09T17:16:03Z)
• Unleashing Implicit Rewards: Prefix-Value Learning for Distribution-Level Optimization [74.91418266859297]
  Implicit process reward models (PRMs) provide fine-grained reward signals along the reasoning process.<n>PRMs suffer from a train-inference mismatch: training only constrains a sequence-level aggregate, whereas inference requires token-level scores to reflect local step quality.<n>We propose a novel Implicit Prefix-Value Reward Model (IPVRM), which directly learns a prefix-conditioned value function estimating the probability of eventual correctness.<n>We also propose Distribution-Level RL (DistRL), which computes TD advantages for both sampled tokens and high-probability candidate tokens.
  arXiv  Detail & Related papers  (2026-04-14T18:19:54Z)
• Balanced Aggregation: Understanding and Fixing Aggregation Bias in GRPO [70.38763678943648]
  Reinforcement learning with verifiable rewards (RLVR) has become a central paradigm for improving reasoning and code generation in large language models.<n>Standard GRPO uses sequence aggregation, while recent work has advocated token aggregation as a better alternative.<n>We show that these two rules induce different optimization biases; token aggregation introduces sign-length coupling, while sequence aggregation implicitly downweights longer responses.
  arXiv  Detail & Related papers  (2026-04-14T09:48:46Z)
• AdaPonderLM: Gated Pondering Language Models with Token-Wise Adaptive Depth [23.442686851761298]
  AdaPonderLM is a self-supervised recurrent language model that learns token-wise early exiting during pretraining.<n>AdaPonderLM reduces inference compute at about 10% while maintaining comparable language modeling perplexity and competitive downstream accuracy.
  arXiv  Detail & Related papers  (2026-03-02T14:28:16Z)
• Speculative Decoding Speed-of-Light: Optimal Lower Bounds via Branching Random Walks [39.54576236079211]
  Speculative generation has emerged as a promising technique to accelerate inference in large language models.<n>In this work, we establish the first tight'' lower bounds on the runtime of any deterministic speculative generation algorithm.
  arXiv  Detail & Related papers  (2025-12-12T16:54:33Z)
• Multi-GRPO: Multi-Group Advantage Estimation for Text-to-Image Generation with Tree-Based Trajectories and Multiple Rewards [18.89733659101796]
  Group Relative Policy Optimization has shown promising potential for aligning text-to-image (T2I) models.<n>Existing GRPO-based methods suffer from two critical limitations.<n>We propose textbfMulti-GRPO, a multi-group advantage estimation framework.
  arXiv  Detail & Related papers  (2025-11-30T05:44:35Z)
• SPINE: Token-Selective Test-Time Reinforcement Learning with Entropy-Band Regularization [24.699488237384134]
  We propose SPINE, a token-selective test-time reinforcement learning framework.<n>We show that SPINE consistently improves Pass@1 over TTRL while avoiding response-length collapse.<n>These results indicate that aligning updates with chain-of-thought branch points is a simple and label-free mechanism for stable and effective test-time adaptation.
  arXiv  Detail & Related papers  (2025-11-22T06:32:34Z)
• Rethinking Thinking Tokens: LLMs as Improvement Operators [80.12087211785949]
  Reasoning training incentivizes LLMs to produce long chains of thought (long CoT), which allows them to explore solution strategies with self-checking.<n>This results in higher accuracy, but inflates context length, token/compute cost, and answer latency.<n>We ask: Can current models leverage their metacognition to provide other combinations on this Pareto frontier?<n>We identify an interesting inference family Parallel-Distill-Refine (PDR), which performs the following: (i) generate diverse drafts in parallel; (ii) distill them into a bounded, textual workspace; and (iii) refine conditioned on this workspace
  arXiv  Detail & Related papers  (2025-10-01T17:08:59Z)
• Tree Reward-Aligned Search for TReASURe in Masked Diffusion Language Models [13.433506313486701]
  Tree search has emerged as a powerful framework for aligning generative models with task-specific rewards at test time.<n>We propose TReASURe, a tree-search test-time alignment method that addresses these issues.<n>TReASURe achieves state-of-the-art results on perplexity, linguistic acceptability, and control of sentiment and toxicity.
  arXiv  Detail & Related papers  (2025-09-27T06:22:45Z)
• Pipeline Parallelism is All You Need for Optimized Early-Exit Based Self-Speculative Decoding [73.67253077506672]
  Large language models (LLMs) deliver impressive generation quality, but incur very high inference cost.<n>Early-exit based self-speculative decoding (EESD) has emerged to mitigate this cost.<n>We propose Pipeline-Parallel Self-Speculative Decoding (PPSD) that fully pipelines the draft and verification work.
  arXiv  Detail & Related papers  (2025-09-19T04:51:41Z)
• TreePO: Bridging the Gap of Policy Optimization and Efficacy and Inference Efficiency with Heuristic Tree-based Modeling [65.46347858249295]
  TreePO is a self-guided rollout algorithm that views sequence generation as a tree-structured searching process.<n>TreePO essentially reduces the per-update compute burden while preserving or enhancing exploration diversity.
  arXiv  Detail & Related papers  (2025-08-24T16:52:37Z)
• GTPO and GRPO-S: Token and Sequence-Level Reward Shaping with Policy Entropy [0.0]
  Reinforcement learning (RL) with algorithms like Group Relative Policy Optimization (GRPO) improves Large Language Model (LLM) reasoning.<n>This paper solves this with textbfDynamic Entropy Weighting.<n>Our core idea is that high-entropy tokens in correct responses can guide the policy toward a higher performance ceiling.
  arXiv  Detail & Related papers  (2025-08-06T11:42:47Z)
• TreeRPO: Tree Relative Policy Optimization [55.97385410074841]
  name is a novel method that estimates the mathematical expectations of rewards at various reasoning steps using tree sampling.<n>Building on the group-relative reward training mechanism of GRPO, name innovatively computes rewards based on step-level groups generated during tree sampling.
  arXiv  Detail & Related papers  (2025-06-05T15:56:38Z)
• Reinforcing Video Reasoning with Focused Thinking [65.85683941058916]
  We propose TW-GRPO, a novel framework that enhances visual reasoning with focused thinking and dense reward granularity.<n>Specifically, we employ a token weighting mechanism that prioritizes tokens with high informational density.<n>We also reformulate RL training by shifting from single-choice to multi-choice QA tasks.
  arXiv  Detail & Related papers  (2025-05-30T15:42:19Z)
• Prune Spatio-temporal Tokens by Semantic-aware Temporal Accumulation [89.88214896713846]
  STA score considers two critical factors: temporal redundancy and semantic importance. We apply the STA module to off-the-shelf video Transformers and Videowins. Results: Kinetics-400 and Something-Something V2 achieve 30% overshelf reduction with a negligible 0.2% accuracy drop.
  arXiv  Detail & Related papers  (2023-08-08T19:38:15Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.