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 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.

📋 한 줄 요약

**[LLM RL / Information Bottleneck]** IB-TPO가 IB-Score(reasoning diversity vs mutual information 균형)를 fine-grained 목표로 정식화, IB-guided tree sampling으로 같은 토큰 예산에 50% 더 많은 trajectory·tree 재사용 MC 추정, GRPO 대비 2.9~3.6% 향상.

🎯 핵심 기여도

LLM 온라인 RL의 최근 진보가 복잡 reasoning에서 promising 성능을 보이지만 imbalanced exploration-exploitation trade-off로 unstable optimization·sub-optimal 결과를 보임을 진단.
IB-Score 도입 — Information Bottleneck 이론 기반의 새 metric으로 정책의 exploration-exploitation 균형을 step-level reasoning diversity와 correct answer와 공유하는 mutual information 사이 trade-off로 정량화.
IB-Score 분석으로 GRPO 등 인기 온라인 RL 방법이 일반 regularizer와 함께도 학습 중 균형을 일관 유지하지 못하며 suboptimal 결과를 낳음을 보임.
IB-TPO(Information Bottleneck-driven Tree-based Policy Optimization) 제안 — IB-Score를 fine-grained 최적화 목표로 정식화, 새 IB-guided tree sampling 전략으로 동일 토큰 예산에서 50% 더 많은 trajectory 효율화·tree 구조 재사용으로 효과적 IB-Score Monte Carlo 추정.

💡 핵심 아이디어

LLM RL의 안정·강력한 학습은 reward만 최적화하는 것이 아니라 reasoning diversity와 답 정보량의 Information Bottleneck 균형을 fine-grained 목표로 두는 데서 비롯되며, tree-based sampling은 같은 예산으로 더 많은 trajectory를 만들고 tree 구조를 재사용해 IB-Score를 효율적으로 추정한다.

🔬 기술적 접근법

**방법론**: IB-Score metric + IB-TPO(Tree-based Policy Optimization).
**핵심 기법**: (1) IB-Score: step-level reasoning diversity와 correct answer와의 mutual information의 trade-off로 exploration-exploitation 균형 정량화, (2) IB-Score를 fine-grained 최적화 목표로 정식화, (3) IB-guided tree sampling 전략으로 동일 토큰 예산에서 50% 더 많은 trajectory 효율화, (4) Tree 구조 재사용으로 IB-Score MC 추정 효율화, (5) GRPO 등 기존 RL과 standard reasoning benchmark에서 비교.

📊 주요 결과

광범위 표준 벤치마크에서 GRPO baseline 대비 2.9%~3.6% 향상.
다른 SOTA 온라인 RL 방법들도 능가.
Tree sampling으로 동일 토큰 예산에서 50% 더 많은 trajectory 효율화.
Tree 구조 재사용으로 효과적 MC 추정.
Code: github.com/alibaba/EfficientRL.

💭 의의 및 한계

**의의**: LLM RL의 exploration-exploitation 균형을 IB 이론으로 원리적 정량화, GRPO의 단점(balance 유지 실패) 진단·해결, tree 구조의 sampling·추정 두 측면 동시 활용. **한계**: IB-Score 추정의 tree 구조 의존성, 50% 효율 향상의 task·모델 의존, IB-Score 외 다른 균형 metric과 비교 부재.

🚀 실용적 활용

수학·코딩·과학 reasoning LLM 학습.
GRPO 대체 RL 최적화기.
LLM의 reasoning trace 다양성 진단·개선 도구.