Source-linked AI summary
Long Live The Balance: Information Bottleneck Driven Tree-based Policy Optimization
Hao Jiang, Shurui Li, Tianpeng Bu, Bowen Xu, Xin Liu, Qihua Chen, Hongtao Duan, Lulu Hu, Bin Yang, Minying Zhang
TL;DR
Online RL for LLMs lacks a consistently balanced exploration–exploitation mechanism, limiting stable optimization. The paper introduces IB-Score and IB-TPO with IB-guided tree sampling, achieving 2.9-3.6% gains over GRPO baselines and other state-of-the-art online RL methods. IBTree also improves sampling effectiveness while supporting IB-Score estimation.
Problem
Online RL approaches for LLMs struggle with inadequate exploration, insufficient exploitation, and unstable exploration–exploitation balance.
Method
IB-TPO incorporates IB-Score, which measures reasoning-step diversity and mutual information with the correct answer, into policy optimization and IB-guided tree sampling.
Results
2.9-3.6%: IB-TPO outperforms GRPO baselines and state-of-the-art online RL methods across multiple in- and out-of-domain benchmarks.
Takeaways & Limitations
IB-TPO maintains a better exploration–exploitation trade-off and achieves higher sampling effectiveness than compared online RL approaches.
Takeaways & Limitations
IBTree remains slightly slower than independent sampling under the same token budget, and future work extends IB-TPO to multimodal reasoning and function calling.
Abstract
from arXiv · showhide
Recent advances in online reinforcement learning (RL) for large language models (LLMs) have demonstrated promising performance in complex reasoning tasks. However, they often exhibit an imbalanced exploration-exploitation trade-off, resulting in unstable optimization and sub-optimal performance. We introduce IB-Score, a novel metric grounded in Information Bottleneck theory that evaluates policy's exploration-exploitation balance by quantifying the trade-off between step-level reasoning diversity and mutual information shared with the correct answer. Analysis based on IB-Score shows that popular online RL approaches (e.g., GRPO) with common regularizers fail to consistently maintain balance during training with suboptimal results. To address this, we propose Information Bottleneck-driven Tree-based Policy Optimization (IB-TPO), a principled framework that formulates IB-Score as a fine-grained optimization objective and utilizes a novel IB-guided tree sampling strategy that not only improves the efficiency of online sampling with 50% more trajectories under the same token budget, but also reuses the tree structure for effective IB-Score Monte Carlo estimation. Extensive experiments across standard benchmarks show that our method significantly outperforms GRPO baseline by 2.9% to 3.6% and also outperforms other state-of-the-art online RL approaches. Our code is available at https://github.com/alibaba/EfficientRL.
1. Introduction
Online RL for LLMs struggles to balance exploration and exploitation, motivating IB-Score and IB-TPO, which integrate information-bottleneck guidance with tree-based optimization.
- GRPO faces inadequate exploration from structurally inefficient independent sampling and insufficient exploitation from sparse outcome-based rewards.
- Existing entropy, process-reward, and tree-based approaches still leave learning signals vulnerable to early convergence toward high-certainty local optima.
- IB-Score quantifies fine-grained exploration–exploitation balance using reasoning-step diversity and mutual information with the correct answer.
- IB-TPO incorporates IB-Score into policy optimization and uses IBTree to guide diverse reasoning directions while estimating IB-Score efficiently.
- 2.9-3.6%: IB-TPO outperforms GRPO baselines and state-of-the-art online RL methods across multiple in- and out-of-domain benchmarks.
- The paper’s contributions comprise IB-Score, IB-TPO with IBTree, and comprehensive comparisons against GRPO and other state-of-the-art approaches.
2. Preliminaries
GRPO independently samples trajectories, collects verifier rewards, normalizes group-relative advantages, and applies clipped policy optimization. Information Bottleneck theory frames balancing information discarded from representations against information retained about targets.
- GRPO independently samples G trajectories from the policy to explore alternative solutions.
- Verifier rewards are normalized within each group to compute relative advantages against a group baseline.
- GRPO applies a clipping-based policy optimization using trajectory-token importance weights and advantages.
- Information Bottleneck theory extracts task-relevant information from input X through representation Z to predict output Y.
- The IB objective minimizes I(X; Z) while maximizing I(Z; Y), and mutual information can be expanded as I(X; Y) = H(X)−H(X|Y).
3. Diagnosing Exploration-Exploitation Imbalance: An IB Perspective
The paper diagnoses online RL imbalance through Information Bottleneck theory, showing that GRPO’s exploration and exploitation signals deteriorate during training. It introduces IB-Score to quantify step-level diversity and answer-related information using entropy and rollout-based estimates.
- Exploration–Exploitation Dilemma: GRPO variants with entropy-oriented regularization suffer performance stagnation because premature entropy decline makes policies overly deterministic and reduces effective sampled groups.Entropy regularization and asymmetric clipping may alleviate entropy degradation but do not prevent declining Eff-Rate; excessive entropy can instead destabilize training.
- Exploration–Exploitation Dilemma: The paper identifies a need for a unified mechanism that quantifies and encourages a balanced exploration–exploitation trade-off in online RL.Existing methods address exploration, process-level feedback, or tree-based estimation from different perspectives but do not provide this unified balance mechanism.
- IB-Score: IB-Score measures step-level exploration diversity together with mutual information between each reasoning step and the correct answer.The metric is derived from an Information Bottleneck objective and is designed to evaluate the trade-off in a fine-grained manner.
- IB-Score: IB-Score depends on the covariance Cov(η1, η2), linking environmental-feedback information gain with model confidence in each reasoning branch.This covariance makes IB-Score more nuanced than entropy alone because effective balance requires confidence to align with informative reasoning paths.
- IB-guided Tree Procedure: The proposed IB-guided tree procedure repeatedly expands nodes with the highest IB-Scores, then computes global and local advantages for step-level policy optimization.The method uses existing rollouts to calculate node scores and reuses the tree structure for IB-Score estimation.
- IB-Score: GRPO baselines show a rapid decline in Cov(η1, η2) toward zero, indicating increasingly uniform confidence across reasoning paths.The resulting conflict between high confidence in valuable paths and sufficient exploration makes regularization yield marginal improvement or training collapse.
4. Information Bottleneck Driven Tree-based Policy Optimization
IB-TPO integrates IB-Score into online RL and uses IBTree to guide branching toward promising reasoning directions while reusing the tree for efficient estimation. Its training pipeline combines tree-structured sampling with local and global step-level advantages for policy optimization.
- Information Bottleneck Guided Tree Search: IBTree selects branching nodes with the highest IB-Score, favoring directions with greater exploration diversity and promising confidence allocation.
- Information Bottleneck Guided Tree Search: IBTree shares prefix responses during branching, enabling more trajectory attempts under the same token budget than independent sampling.
- Information Bottleneck Guided Tree Search: IBTree also serves as a Monte Carlo estimator of IB-Score while adapting branching-node selection as IB-Score updates.
- Online RL Training with IBTree: The overall procedure samples G trajectories with fine-grained IB-Scores, computes supervision signals from the tree, and updates the policy using the combined advantage.
- IB-TPO incorporates IB-Score into the online RL objective to encourage balanced exploration and exploitation during reasoning.
- Online RL Training with IBTree: The optimization combines local and global advantages, with the local term comparing each step with its parent and the global term providing a complementary signal.
5. Experiments
Experiments evaluate IB-TPO across Qwen3 scales and mathematical benchmarks, showing stronger performance, sampling efficiency, and exploration–exploitation dynamics than compared methods. Ablations attribute gains to IBTree and IB-based advantage estimation, while moderate delimiter noise remains tolerable.
- Setup: IB-TPO is evaluated on Qwen3-1.7B-Base and Qwen3-8B-Base using DAPO-Math-17K, with avg@32 reported across in-domain and out-of-domain benchmarks.Benchmarks include MATH-500, AIME 24/25, AMC 23/24, GPQA Diamond, and IFEval.
- Main Results: 2.9% and 3.6% average improvements over GRPO are reported, while IB-TPO also surpasses all compared methods across benchmarks.Compared methods include GRPO variants, IBRO, TreeRL, and TreePO.
- More Analysis: 6.8% higher average Eff-Rate and 50% more trajectories under a similar token budget demonstrate the efficiency of IB-guided tree sampling over independent sampling.IB-Score-guided branching achieves the best Eff-Rate and improves Avg-Rate among compared branching strategies.
- More Analysis: IB-TPO shows more consistent validation improvement and stronger Cov(η1, η2) and IB-Score dynamics than other approaches.These results are reported in the Qwen3-8B-base training-dynamics comparison.
- Ablation Studies: Both IBTree and IB-based advantage estimation improve GRPO, while their full combination achieves the best ablation results.The IB-based objective also outperforms TreeRL, and IBTree cannot be readily replaced by alternative tree structures.
- Ablation Studies: β = 5 produces better Eff-Rate and Avg-Rate with less token consumption, while λ = 0.1 achieves the best advantage-estimation performance.Training with only global advantage is reported as sub-optimal.
- Ablation Studies: Moderate step-boundary perturbations do not significantly change performance, indicating robustness to the tested delimiter noise.The perturbation relocates 10% of existing split points randomly to simulate under-segmentation and over-segmentation.
6. Related Work
Related work combines online RL with information-theoretic regularization and tree search to improve LLM reasoning. These approaches motivate IB-TPO while leaving efficiency and search-space challenges for the paper to address.
- Information Bottleneck in Online RL: Prior work shows online RL can improve LLM post-training, while other studies indicate that balanced exploration–exploitation benefits RL training.IB-based advantage-weighted entropy regularization has been proposed for this setting.
- Tree Search in Online RL: Tree search supports deeper reasoning and process supervision, but integrating it into online RL can involve large search spaces that hinder efficiency.Prior online-RL tree methods motivate more selective tree sampling.
7. Conclusion
The paper introduces IB-Score to evaluate exploration–exploitation balance and IB-TPO to integrate it into policy optimization with IBTree. Experiments report superiority over GRPO and other state-of-the-art approaches.
- Conclusion: IB-Score measures fine-grained exploration–exploitation balance, while IB-TPO integrates it into policy optimization through IBTree.IBTree provides IB-guided exploration and efficient IB-Score estimation.
- Conclusion: Comprehensive experiments demonstrate IB-TPO’s superiority over GRPO baselines and state-of-the-art approaches.
8. Limitation
The remaining limitation is computational: multi-iteration tree sampling increases time consumption, and IBTree remains slightly slower than independent sampling under the same token budget. Future work targets multimodal reasoning and function calling.
- Efficiency Limitation: IBTree remains slightly slower than independent sampling under the same token budget despite implementation optimizations and parallel sampling.The limitation arises from time consumption introduced by multi-iteration tree sampling.
- Future Scope: Future extensions are planned for multimodal reasoning and function calling.
Impact Statement
The paper presents IB-Score and IB-TPO as tools for more stable and effective online RL, linking the objective to confidence allocation and current policy confidence. Training analysis reports that the relevant success-rate product remains close to 1 during IB-TPO training.
- IB-TPO uses IB-Score to improve sampling efficiency and reasoning quality for more reliable LLM post-training on complex reasoning tasks.
- The step-wise IB objective favors allocating confidence to the most feedback-informative path while reducing the policy model’s current confidence.
- The average product of sub-tree success rates remains close to 1 during IB-TPO training.
A.2. Difference between IB-Score and Entropy
IB-Score is an Information Bottleneck-derived surrogate for the exploration–exploitation trade-off that combines policy confidence with reward gains across branches. Unlike token-level entropy, it evaluates uncertainty at the reasoning-step level.
- IB-Score evaluates exploration–exploitation balance using both model confidence and reward gains from different branches.
- IB-Score computes uncertainty at the step level rather than the token level to better reflect model decision-making.
- Step-level entropy is less susceptible to outlier tokens that can arise from inference–training precision discrepancies in RL.
B. More Implementation Details
The implementation uses consistent sampling and optimization settings across methods, with IBTree branching guided by IB-Score. Experiments compare sampling strategies, pass@K, runtime, and matched-time performance under token-budget-aligned evaluation.
- The framework uses T=0.7, top p=0.95, top k=20, 2K truncation, learning rate 1 × 10^-6, and KL weight 0.001 across compared methods.
- IBTree samples 4 initial trajectories and expands the highest-IB-Score non-leaf node with one trajectory per iteration, collecting 12 trajectories for optimization.
- The comparison includes independent, random, fixed-width, entropy-guided, and IB-Score-guided branching strategies on Qwen3-8B-Base.
- IBTree samples 50% more trajectories than independent sampling under the same token budget, increasing 8 trajectories to 12.
- IBTree’s runtime gap versus independent sampling narrows as tree parallelism increases, while future work targets lower wall-clock sampling time.
- Under matched sampling time, IBTPO has lower sampling runtime than GRPO while achieving better performance.
C.4. Extended experiments under longer-context settings
Longer-context experiments increase sampling time and memory burden, but IBTree retains sampling-efficiency advantages and IBTPO’s performance improvements remain consistent at 4K and 8K context lengths. The extended evaluation also examines broader model and task settings.
- Longer-context settings: Only 6% of Qwen3-8B-Base responses are truncated at 2K tokens, supporting 2K as the main experiment’s practical maximum length.
- Longer-context settings: Increasing maximum truncation length from 2K to 4K or 8K raises wall-clock sampling time and memory burden.
- Longer-context settings: IBTree maintains a time-efficiency advantage at G=8 and samples 50% more trajectories at G=12 with a smaller incremental time increase.
- Longer-context settings: IBTPO continues to outperform GRPO and TreeRL when maximum context length increases from 2K to 4K or 8K.
- Generalization: The extended evaluation includes Qwen3-14B-Base and Llama3.1-8B-Instruct on code generation.