IIB-LPO: Latent Policy Optimization via Iterative Information Bottleneck

TL;DR

This work tackles exploration collapse in RLVR for LLM reasoning by shifting from token distribution perturbations to topological diversification of reasoning trajectories. It introduces I2B-LPO, which triggers latent branching at high-entropy states via a CVAE-driven latent code $z$ injected into a Pseudo Self-Attention mechanism and uses the Information Bottleneck as a dual-purpose trajectory filter and self-reward. The method includes entropy-based bifurcation detection, CVAE latent sampling, PSA-guided reasoning, and IB-guided pruning and optimization, with a training objective that combines GRPO loss and an IB term. Empirical results on four mathematical benchmarks show state-of-the-art accuracy and diversity, with improved exploration quality and more concise reasoning, demonstrating the practical potential of structured, information-theoretic exploration in LLM reasoning.

Abstract

Recent advances in Reinforcement Learning with Verifiable Rewards (RLVR) for Large Language Model (LLM) reasoning have been hindered by a persistent challenge: exploration collapse. The semantic homogeneity of random rollouts often traps models in narrow, over-optimized behaviors. While existing methods leverage policy entropy to encourage exploration, they face inherent limitations. Global entropy regularization is susceptible to reward hacking, which can induce meaningless verbosity, whereas local token-selective updates struggle with the strong inductive bias of pre-trained models. To address this, we propose Latent Policy Optimization via Iterative Information Bottleneck (IIB-LPO), a novel approach that shifts exploration from statistical perturbation of token distributions to topological branching of reasoning trajectories. IIB-LPO triggers latent branching at high-entropy states to diversify reasoning paths and employs the Information Bottleneck principle both as a trajectory filter and a self-reward mechanism, ensuring concise and informative exploration. Empirical results across four mathematical reasoning benchmarks demonstrate that IIB-LPO achieves state-of-the-art performance, surpassing prior methods by margins of up to 5.3% in accuracy and 7.4% in diversity metrics.

Figures (13)

• Figure 1: Comparison of exploration paradigms in RLVR. (a) Entropy Regularization globally smooths the probability distribution, leading to high-entropy yet meaningless verbosity. (b) Token-selective Methods locally sharpen the distribution; synonym replacement at these isolated points cannot overcome inductive biases. (c) I2B-LPO introduces topological branching via latent variables $z$, resulting in distinct reasoning trajectories (e.g., Differentiation-based $R_1$ vs. Geometry-based $R_2$).
• Figure 2: Performance of various decoding strategies trained on DeepMath. For each problem, we sample truncation points across entropy intervals to simulate varied exploration behaviors.
• Figure 3: Accuracy and Response Length under GRPO. Notably, gray bars denote 4-gram repetition rate.
• Figure 4: Pipeline of the I2B-LPO. We use a representative path $r_i$ from the initial set $R_o$ to illustrate the workflow, which operates in two phases. (1) Entropy-driven Latent Branching expands $r_i$ into the branching set $R_i$ via Latent Sampling and PSA Injection, which are depicted in the bottom section. (2) Information Bottleneck Regularization applies IB as a dual-purpose filter and self-reward to ensure concise and informative exploration.
• Figure 5: The performance of the trained model on six diversity metrics. We evaluate I2B-LPO using Qwen2.5-7B and Qwen3-14B models across the GSM8K and MATH-500 datasets. For each metric, a higher value indicates greater diversity. And the diversity metrics are calculated across $10$ generated responses per prompt.