Exploiting Tree Structure for Credit Assignment in Reinforcement Learning with Large Language Models.

Area of Science:

• Artificial Intelligence
• Machine Learning
• Natural Language Processing

Background:

• Reinforcement learning (RL) enhances large language model (LLM) reasoning.
• Token-level credit assignment is challenging due to sparse, delayed rewards in LLMs.
• Existing methods like PPO (actor-critic) and GRPO (critic-free) have limitations in complexity or reward distribution.

Purpose of the Study:

• To address the challenge of token-level credit assignment in RL for LLMs.
• To develop a critic-free RL algorithm that improves upon existing methods.
• To enhance the reasoning abilities and performance of LLMs through more effective reward propagation.

Main Methods:

• Proposed Prefix-to-Tree (P2T) to organize sampled responses into a prefix tree.
• Developed TEMPO (Tree-Estimated Mean Prefix Value for Policy Optimization), a critic-free algorithm.
• Integrated branch-aware temporal-difference (TD) corrections into GRPO using P2T-derived prefix values.

Main Results:

• TEMPO consistently improved convergence and final performance over PPO and GRPO on Qwen3 models.
• Performance gains were observed on both in-distribution (MATH, MedQA) and out-of-distribution benchmarks (GSM-HARD, AMC23, MedMCQA, MMLU-Medical).
• TEMPO achieved higher validation accuracy within comparable wall-clock time.

Conclusions:

• TEMPO offers an effective critic-free approach for RL in LLMs, overcoming limitations of previous methods.
• The prefix tree structure and TD corrections enable more nuanced credit assignment.
• TEMPO demonstrates significant improvements in LLM reasoning and performance across diverse tasks.