Source-linked AI summary

IIB-LPO: Latent Policy Optimization via Iterative Information Bottleneck

Huilin Deng, Hongchen Luo, Yue Zhu, Long Li, Zhuoyue Chen, Xinghao Zhao, Ming Li, Jihai Zhang, Mengchang Wang, Yang Cao, Yu Kang

arXiv:2601.05870v1cs.LGcs.AI

TL;DR

Semantically homogeneous rollouts can collapse exploration into narrow reasoning templates, while entropy-based approaches have limits in diversifying reasoning. I2B-LPO branches latent reasoning trajectories at high-entropy states and applies the Information Bottleneck as a filter and self-reward; it achieves state-of-the-art accuracy and diversity without excessive generation length.

  • Problem

    Semantically homogeneous rollouts converge on a few reasoning templates, while global entropy regularization can incentivize vacuous verbosity and token-selective updates may not overcome pretrained-model inductive biases.

    Source: §1 · pp. 1–2
  • Method

    I2B-LPO triggers latent branching at high-entropy states and uses the Information Bottleneck to filter trajectories and provide self-reward.

    Source: §1 · p. 2
  • Results

    I2B-LPO achieves state-of-the-art accuracy and diversity without excessive generation length, with margins of 5.3% in accuracy and 7.4% in diversity.

    Source: §1 · p. 2
  • Takeaways & Limitations

    The framework provides structured exploration through entropy-driven branching and a dual-purpose Information Bottleneck for filtering and reward shaping.

    Source: §7 · p. 8

    Entropy may not consistently reflect semantic diversity in open-ended or non-mathematical tasks, and branching is more computationally intensive than single-rollout approaches.

    Source: Limitations and Potential Risk · p. 8

Abstract

from arXiv · show

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.

1 Introduction

RLVR exploration can collapse as semantically similar rollouts converge on a few reasoning templates, leaving weak learning signals. I2B-LPO shifts exploration from token-level perturbations to latent branching and uses the Information Bottleneck to favor concise, informative paths.

1 Introduction

  • Semantic homogeneity makes random rollouts converge on a few reasoning templates, producing vanishing advantage differences and uninformative policy-optimization signals.Rollouts may vary in surface phrasing while sharing nearly identical reasoning patterns.Source: §1 · p. 1
  • Global entropy regularization can reward semantically vacuous verbosity, while sharpening high-entropy tokens may only change synonym choice and fail to overcome pretrained-model inductive biases.Both approaches perturb token probabilities without structurally diversifying reasoning.Source: §1 · pp. 1–2, Figure 1 · p. 2
  • I2B-LPO samples latent variables at high-entropy states to steer distinct reasoning trajectories and uses the Information Bottleneck as a trajectory filter and self-reward.The IB objective favors compact, informative paths while penalizing semantically vacuous verbosity.Source: §1 · p. 2
  • I2B-LPO achieves state-of-the-art reasoning accuracy and semantic diversity without excessive generation length, with margins of 5.3% in accuracy and 7.4% in diversity.The method is presented as a shift from statistical token perturbation to topological branching of reasoning trajectories.Source: §1 · p. 2

2 Related Work

Prior exploration methods use entropy regularization or selective updates to guide token-level behavior, while related work also addresses self-rewarding reasoning. I2B-LPO instead proposes conditional branching for exploration and iterative trajectory refinement.

2 Related Work

  • Prior entropy-based exploration methods use token-selective updates or add entropy regularizers to the loss or advantage function; I2B-LPO proposes conditional branching instead.The cited related work describes regularizers as commonly applied indiscriminately across the token distribution.Source: §2 · p. 2
  • Figure 2 evaluates decoding strategies trained on DeepMath by sampling per-problem truncation points across entropy intervals to simulate different exploration behaviors.The caption describes the evaluation setup but does not report which strategy performs best.Source: Figure 2 · p. 3
  • Self-rewarding approaches use rules, mutual information, semantic-entropy clustering, or majority voting; I2B-LPO uses high-entropy states as gateways for branching and trajectory refinement.Unlike SPINE’s static voting, I2B-LPO actively uses high-entropy states.Source: §2 · p. 3

3 Preliminary Analysis

The preliminary analysis identifies high-entropy states as useful exploration decision points, while standard GRPO shows verbosity and repetition rising after accuracy plateaus. The authors introduce IB self-reward to suppress unproductive expansion.

3 Preliminary Analysis

  • For entropy intervals above 2.5, the hybrid strategy substantially outperforms baselines, with a much larger gap than in intervals below 1.0.The analysis treats high-entropy states as natural decision points where branching can unlock exploration potential.Source: §3 · p. 3
  • Under standard GRPO, accuracy plateaus while response length and 4-gram repetition continue to rise, indicating verbose reasoning without genuine performance gains.Figure 3 marks 4-gram repetition rate with gray bars.Source: §3 · p. 3, Figure 3 · p. 3
  • The authors introduce IB self-reward to suppress unproductive expansion when verbosity and repetition increase without performance gains.This follows the observed GRPO pattern of rising response length alongside plateauing accuracy.Source: §3 · p. 3

4 Methodology

I2B-LPO starts from GRPO rollouts, branches at high-entropy states, and uses IB scores to prune candidates into a training subset. It also incorporates the IB score as an auxiliary training objective.

4.1 Problem Formulation

  • Starting from M initial GRPO rollouts, the method generates K branches per rollout while retaining the original, then keeps N paths after IB pruning.Source: §4.1 · p. 3

4.2.1 Entropy-Driven Bifurcation Detection

  • For each trajectory, the method selects a bifurcation point from the top 5% highest-entropy steps and forms a prefix from the query and preceding path.Source: §4.2.1 · p. 3

4.2.2 Latent Sampling via CVAE

  • Given the prefix, a CVAE models solution trajectories conditioned on context and supplies K independent latent codes from its prior.Source: §4.2.2 · p. 4

4.2.3 Latent-Guided Reasoning via PSA

  • PSA injects latent codes per layer through RMSNorm modulation and augmented attention, steering reasoning while annealing latent influence over time.Source: §4.2.3 · p. 4

4.3.1 IB Score: Theory and Computation

  • The IB objective trades rationale compression against answer informativeness; SIB(r) is defined as the negative IB loss, with a stepwise entropy-and-advantage approximation.Source: §4.3.1 · pp. 4–5

4.3.2 IB-Guided Pruning and Optimization

  • Candidates are ranked by SIB, the top N paths train the policy, and the IB score is an auxiliary objective with a coefficient balancing correctness and reasoning efficiency.Source: §4.3.2 · p. 5

5 Experiment Settings

Experiments use GRPO and veRL, train on DAPO and MATH data, and evaluate two Qwen models against several baseline categories on four mathematical benchmarks. The reported metrics cover response length, diversity, and perplexity.

  • Experiments use GRPO in veRL with Clip-Higher and Group-Sampling, train on DAPO and MATH data, and evaluate four mathematical benchmarks.Source: §5 · p. 5
  • The comparison uses Qwen2.5-37B and Qwen3-14B against entropy-regularization, token-selective, and self-reward baselines.Source: §5 · p. 5
  • Metrics include average token length, Distinct-n, 1 − Score Self-BLEU, and perplexity; higher diversity scores indicate greater diversity, while low perplexity indicates more fluent, coherent responses.Source: §5 · p. 5

6 Experiment

Across four mathematical benchmarks, the experiments assess I2B-LPO’s accuracy and diversity, exploration behavior, IB self-reward, and component ablations.

6 Experiment

  • The experiment frames four questions: overall performance, exploration behavior, IB self-reward efficacy, and the contribution of each component.The stated evaluation covers reasoning accuracy and semantic diversity as well as exploration, self-reward, and ablations.Source: §6 · p. 5

6.1 Quality-diversity balance (RQ1)

  • I2B-LPO consistently outperforms baselines across both backbones in quality and diversity, with Pass@n advantages widening as n increases.Table 1 covers four mathematical benchmarks and two backbones; Figure 5 and Table 2 report diversity metrics.Source: §6.1 · p. 6, Table 1 · p. 6, Figure 5 · p. 6

6.2 Exploration Behavior Analysis (RQ2)

  • I²B-LPO maintains exploratory tokens across a wider entropy spectrum and stabilizes entropy after an initial drop, unlike GRPO’s continuing decay.The analysis also reports that entropy loss can produce anomalously high token entropy (>8) for some tokens.Source: §6.2 · p. 7, Figure 6 · p. 7
  • PSA produces difficulty-sensitive activation in final-layer attention heads, whereas input injection fades in deeper layers and softmax injection yields scattered activations.The compared injection strategies are early input-level fusion, PSA deep fusion, and late softmax-layer fusion.Source: §6.2 · pp. 7–8, Figure 7 · p. 7

6.3 Self-Reward Effectiveness (RQ3)

  • IB self-reward outperforms likelihood, self-consistency, entropy regularization, and LLM-Judge on accuracy, diversity, and confidence, while higher IB scores correspond to concise reasoning.The reported analysis also associates IB scoring with reduced perplexity and improved quality.Source: §6.3 · p. 8, Figure 8 · p. 8

6.4 Ablation of Main Components (RQ4)

  • The full IB loss-plus-pruning mechanism reduces perplexity by 11.7, while MATH Pass@1 rises from 54.4% to 81.5%.The ablations also report that entropy-based branching outperforms other strategies, PSA injection depth peaks at the last 12 layers, and max_length of 8192 tokens provides sufficient capacity for complex reasoning.Source: §6.4 · p. 8

7 Conclusion

The conclusion presents I²B-LPO as a response to exploration collapse, combining entropy-driven latent branching with dual-purpose information bottleneck. Experiments are reported to show state-of-the-art accuracy and diversity.

7 Conclusion

  • The paper targets exploration collapse associated with semantically homogeneous reasoning paths.It identifies the problem as models losing exploration amid narrow reasoning behaviors.Source: §7 · p. 8
  • I²B-LPO uses entropy-driven latent branching and a dual-purpose IB mechanism for filtering and reward shaping.The framework is described as enabling structured exploration.Source: §7 · p. 8
  • Experiments are reported to confirm state-of-the-art accuracy and diversity, and the paper presents the framework as addressing the exploration bottleneck.The authors characterize it as a principled, extensible framework for exploration-aware policy optimization in LLM reasoning.Source: §7 · p. 8

Limitations and Potential Risk

The authors identify uncertainty about entropy as a diversity signal beyond mathematics and higher computational cost than single-rollout approaches. Their evaluation focuses on mathematical reasoning benchmarks and abstract symbolic domains.

Limitations

  • Entropy may not consistently reflect semantic diversity in open-ended or non-mathematical tasks.This limits how confidently the uncertainty signal can be interpreted beyond the stated task setting.Source: Limitations and Potential Risk · p. 8
  • Branching remains more computationally intensive than single-rollout approaches despite IB-based pruning.The authors identify more efficient algorithms as future work.Source: Limitations and Potential Risk · p. 8

Potential Risk

  • The method focuses on established public mathematical benchmarks and abstract symbolic domains, without sensitive data or content generation.The authors assess its ethical and safety risks as minimal within this stated scope.Source: Limitations and Potential Risk · p. 8

Ethical Considerations

The paper reports ethics safeguards and categorizes RLVR tokens by their roles in reasoning. It also describes the scope of component-ablation and hyperparameter-sensitivity analyses.

Ethical Considerations

  • Table 3 covers branching triggers, latent injection, and IB self-rewards; Table 4 analyzes injection depth, branching factor, and maximum response length.Source: Table 3 · p. 9, Table 4 · p. 9
  • The study uses public, anonymized data, reports results truthfully, provides reproducible resources, and recommends oversight and filtering in practical applications.The authors state they followed the ACL Code of Ethics and submitted the Responsible NLP Research Checklist.Source: Ethical Considerations · pp. 8–9

A Categorizing High-Entropy Tokens

B Implementation Details of the CVAE

The appendix specifies the CVAE architecture and training, derives an IB-aware token score, and describes how attention-head contributions to difficulty perception are localized.

B Implementation Details of the CVAE

  • The CVAE uses a context-conditioned prior, a DeBERTa-v2large posterior encoder, and a decoder that fuses sampled z into LLaMA; it is trained on MATH and GSM8K.The encoder mean-pools token embeddings to produce variational parameters for the latent representation.Source: B Implementation Details of the CVAE · pp. 12–13

C IB-Aware Scoring Metric

  • The IB score balances token-level entropy as an exploration proxy against relevance, using likelihood among validity-verified correct trajectories to approximate informativeness.The derivation links token-level entropy to the complexity term and correct-path likelihood to predictive informativeness.Source: C IB-Aware Scoring Metric · p. 13

D Localization of Difficulty-Sensitive Attention Heads

D.1 Difficulty Scoring

  • Difficulty scoring projects each head’s representation onto a pretrained difficulty direction, then aggregates attributions across samples with the same ground-truth difficulty.Source: §D.1 · p. 14

D.2 Differentiation Score and Localization

  • The differentiation score subtracts easy-cohort attribution from hard-cohort attribution: positive values indicate heads associated with difficult problems, and negative values with simple ones.Source: §D.2 · p. 15

E Additional Implementation Details

The appendix gives training and data-filtering details, then reports qualitative comparisons of generated responses and a latent-injection decay setting.

Additional implementation details

Loading 2601.05870v1…