Long Live The Balance: Information Bottleneck Driven Tree-based Policy Optimization

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 (IB) 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 regularization methods 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 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.

When teaching AI models to reason — for example, to solve hard math problems — researchers often use a technique called reinforcement learning, where the model practices solving problems and learns from feedback. A key challenge is striking the right balance between "exploration" (trying out diverse new ideas) and "exploitation" (sticking with what already works). Current popular methods often tip too far one way: either the model quickly locks into a single rigid solution style and stops improving, or it becomes so uncertain that its reasoning falls apart. We introduce a new measurement tool called the IB-Score, inspired by a classical idea from information theory called the Information Bottleneck (IB). At every step of the model's reasoning, IB-Score checks whether the model is being both creative enough and confident in the directions most likely to lead to a correct answer. Building on this, we designed a training method (IB-TPO) that grows reasoning attempts like a branching tree, spending more effort exploring the most promising steps rather than wasting it on dead ends. In experiments on reasoning benchmarks, our method consistently outperformed leading approaches by 2.9% to 3.6%, while also being more efficient — producing 50% more reasoning attempts within the same computing budget.