Source-linked AI summary

Reasoning as Compression: Unifying Budget Forcing via the Conditional Information Bottleneck

Fabio Valerio Massoli, Andrey Kuzmin, Arash Behboodi

arXiv:2603.08462v2cs.LG

TL;DR

Verbose CoT improves complex-task reasoning but makes inference expensive, while standard IB is theoretically mismatched to transformers because attention preserves direct prompt access. The paper resolves this with CIB, which yields an RL objective using semantic surprisal costs for reasoning tokens. Experiments report improved or comparable accuracy with substantial compression across model families and task domains.

  • Problem

    Length-based Budget Forcing penalizes essential reasoning and redundant filler alike, while standard IB assumes a Markov structure violated by transformer attention.

  • Method

    CIB treats the prompt as side information and trains an RL policy to maximize answer reward while minimizing reasoning-trace surprisal under a language-model prior.

  • Results

    CIB achieves a more favorable accuracy–length trade-off than naive length penalties and generalizes compression benefits across model families and math, science, knowledge, and code benchmarks.

  • Takeaways & Limitations

    Semantic token costs can compress reasoning while preserving task performance, enabling controlled navigation of the accuracy–compression Pareto frontier.

  • Takeaways & Limitations

    Compression quality depends on prior capacity and can be confounded by tokenization mismatch with cross-family priors; β and prior capacity jointly determine the operating point.

Abstract

from arXiv · show

\ac{CoT} prompting improves LLM accuracy on complex tasks but often increases token usage and inference cost. Existing ``Budget Forcing'' methods reduce cost via fine-tuning with heuristic length penalties, suppressing both essential reasoning and redundant filler. We recast efficient reasoning as a lossy compression problem under the \ac{IB} principle, and identify a key theoretical gap when applying naive \ac{IB} to transformers: attention violates the Markov property between prompt, reasoning trace, and response. To resolve this issue, we model \ac{CoT} generation under the \ac{CIB} principle, where the reasoning trace $Z$ acts as a computational bridge that contains only the information about the response $Y$ that is not directly accessible from the prompt $X$. This yields a general Reinforcement Learning objective: maximize task reward while compressing completions under a prior over reasoning traces, subsuming common heuristics (e.g., length penalties) as special cases (e.g., uniform priors). In contrast to naive token-counting approaches, we introduce a semantic prior that measures token cost by surprisal under a language model. Crucially, the prior is queried only for token-level log-probabilities, adding negligible overhead to the training loop. Empirically, our \ac{CIB} objective prunes reasoning redundancy while preserving fluency and logic, improving accuracy at moderate compression and enabling aggressive compression with minimal accuracy drop. These gains generalize across model families and task domains, confirming \ac{CIB} as a domain-agnostic CoT compression framework.

1 Introduction

CoT improves reasoning on complex tasks but can become excessively verbose, making efficient Budget Forcing difficult. The paper proposes CIB-based semantic compression to preserve useful reasoning while improving the accuracy–compression trade-off.

  • CoT enables test-time computation for complex tasks but increases latency and compute usage through verbose reasoning chains.
  • Naive length penalties and hard token limits charge every token equally, risking deletion of essential reasoning alongside redundant filler.
  • CIB treats the reasoning chain as a minimal sufficient statistic and assigns token costs according to semantic information rather than raw length.
  • The resulting RL objective subsumes length-based penalties as special cases and uses a prior over reasoning traces to control compression.

2 Related Work

Prior efficient-reasoning methods use training or inference budgets to limit excessive computation, but most treat token count as the cost proxy. This paper instead connects Budget Forcing to information theory through a transformer-aware CIB formulation.

  • Budget Forcing methods include length-cost reward shaping, token-level advantages, hard truncation, difficulty-aware allocation, and reference-guided budgeting.
  • These approaches largely rely on naive token counts, which do not measure the semantic information carried by individual tokens.
  • Standard Information Bottleneck applications assume a latent representation mediates all information, unlike transformer attention where the decoder directly accesses the prompt.
  • The paper unifies Budget Forcing and Information Theory under Conditional Information Bottleneck rather than applying the standard Markov-based formulation.

3 Methodology

The methodology models CoT as a conditional information bottleneck in which the prompt remains available while the reasoning trace supplies only additional information needed for the answer. It converts this principle into an RL reward combining answer accuracy with semantic surprisal costs.

  • Attention–Markov Incompatibility: Standard IB fails for transformers because attention lets the answer depend jointly on prompt X and reasoning trace Z, violating the Markov relation.
  • Conditional Information Bottleneck: CIB treats prompt X as side information and requires Z to encode only information additionally needed to predict answer Y.
  • Problem Formulation: The policy πθ(Z | X) generates the CoT bridge, while πθ(Y | X, Z) produces the answer, jointly optimizing sufficiency and minimality.
  • Sufficiency Term: The deterministic verifier is ε-smoothed so its logarithmic score remains defined for incorrect answers.
  • Information Cost: A frozen unconditional base model supplies prior Qϕ, whose token log-probabilities define cumulative surprisal as the reasoning trace information cost.
  • Reward Modeling: The reward combines binary accuracy with a β-controlled minimality term, suppressing redundant tokens that increase cost without improving correctness.

4 Theoretical Analysis: A Unified Framework

The CIB framework shows that common length-based budget penalties are implicit-prior special cases, while a language-model prior assigns semantic token costs rather than relying only on sequence length.

  • A uniform maximum-entropy prior makes every token equally costly, recovering a linear length penalty as a special case of CIB.Each token has surprisal log|V|, so a sequence of length T costs T · log|V|.
  • Target-length penalties correspond to a Laplace-like prior centered on the desired reasoning length.The implied prior makes deviations from the target length exponentially less probable.
  • Length-based implicit priors encode only sequence length, whereas CIB uses a language-model prior with per-token semantic costs.This distinction allows the objective to differentiate token information content rather than treating all tokens equally.

5 Experimental Results

Across small reasoning-model families and diverse tasks, CIB compresses chains of thought while maintaining or improving accuracy. Its semantic prior removes predictable filler, offers smooth accuracy–compression control, and adds limited training overhead, though prior quality depends on model capacity and tokenization compatibility.

  • Main results: CIB improves accuracy while reducing reasoning length across model groups, with gains reported on both easy and hard mathematical benchmarks.On DeepScaleR-1.5B, CIBβ− raises average accuracy by 1.5% while reducing average CoT length by approximately 24%.
  • Main results: On DLER-7B, CIBβ− matches baseline accuracy with approximately 16% shorter CoT, while CIBβ+ cuts length by roughly one-third with a 1.5-point accuracy trade-off.
  • Generalization: CIB generalizes beyond mathematics, improving GPQA Diamond and MMLU-Pro accuracy while matching LiveCodeBench, with token reductions of up to 38.6%.On DLER-7B, it achieves comparable baseline performance while compressing CoT by up to 14.6%.
  • Accuracy–efficiency trade-off: Sweeping β traces a smooth Pareto frontier: lower values prioritize accuracy, whereas higher values produce more aggressive compression.The trade-off is illustrated on AIME24, where intermediate settings retain most baseline accuracy while substantially reducing token count.
  • Compression mechanism: CIB compresses by removing low-surprisal filler while preserving high-information transitions and essential reasoning structure.The semantic prior eliminates scaffolding, redundant verification loops, and tautological checks rather than randomly truncating traces.
  • Training overhead: Scaling the frozen prior to 7B reduces rollout token throughput by less than 6% and effective data collection rate by less than 6.2%.The prior scores generated tokens in a single parallel forward pass and adds no gradient accumulation, backward passes, or optimizer states.
  • Limitations: The semantic signal degrades when the prior is too weak or comes from a cross-family model with mismatched vocabulary and segmentation.These mismatches can penalize native formatting rather than true verbosity, motivating future study of coarser-grained costs.

6 Conclusions

The paper reframes Budget Forcing through Conditional Information Bottleneck, achieving semantic compression with favorable accuracy–length trade-offs while exposing prior-dependent limitations.

  • Up to ∼48% token reduction incurred ≲1.5% accuracy degradation when semantic token costs were used.The regularization coefficient β traverses the accuracy–compression Pareto frontier.
  • CIB subsumes length-based penalties as special cases while replacing raw token counts with surprisal-based semantic costs.The semantic cost is defined using a language-model prior over reasoning traces.
  • Stronger priors enabled more aggressive compression with minimal performance loss and caused approximately 6% throughput reduction during training.Prior quality affects redundancy estimates and compression effectiveness.
  • Compression quality depends on the reference prior, with weak or mismatched priors creating length-like or tokenization-dependent costs.Cross-family priors introduce vocabulary-mismatch confounds, while β and prior capacity jointly determine the operating point.

A Conditional Information Bottleneck

The Conditional Information Bottleneck treats reasoning traces as compressed information about answers beyond what prompts already provide. It resolves the incompatibility between standard Information Bottleneck assumptions and transformer attention through a tractable variational reinforcement-learning formulation.

  • Reasoning generates Z so the model can maximize answer probability while compressing the trace.Budget Forcing is formulated as compressing the reasoning trace Z for each query–answer pair.
  • Standard Information Bottleneck assumes Y ↔ X ↔ Z, but transformer attention lets answer prediction depend directly on both X and Z.This violates the Markov relation required by the classical formulation.
  • CIB maximizes conditional information I(Y; Z|X), making Z informative about Y given X while minimizing redundant prompt information I(X; Z).The formulation interprets reasoning as a computational bridge conditioned on the prompt.
  • The model is parameterized by πθ(z|x) and πθ(y|x,z), which reparameterize the original bottleneck problem without changing its optimal value.The factorization yields a valid conditional distribution over reasoning traces.
  • The formulation faces unknown P(y|x), approximated information quantities, and intractable marginalization over P(z).These challenges motivate sample-based and variational approximations.
  • Variational distributions Qϕ and Qρ provide bounds for the intractable prior and answer-prediction terms, yielding a policy-optimization objective.The final policy πθ(z|x) maximizes returns derived from these approximate bounds.

B.1 Domain Generalization

CIB generalizes beyond mathematical reasoning to science, knowledge, and code benchmarks while usually preserving or improving accuracy and reducing inference compute. Its compressed models also retain competitive inference-time scaling under tight length limits.

  • Reliability: CIBβ− improves accuracy on 25 of 32 model×benchmark evaluations, while L1-based penalties show clearer degradations.On DeepScaleR AIME24, CIBβ− rises from 38.2±1.5 to 40.2±1.9, whereas L3L1-EXACT falls from 38.2±1.5 to 24.4±1.3.
  • Domain generalization: On DLER-7B, CIBβ− matches or slightly improves out-of-domain accuracy while reducing tokens by 10.9%–17.4%.GPQA rises from 50.1 to 50.5, MMLU-Pro remains 58.7, and LiveCodeBench changes from 35.9 to 35.6.
  • Domain generalization: CIBβ− reduces DLER-1.5B GPQA Diamond by only 1.5 points, compared with a 4.1-point drop for L1-Exact.The cited comparison is 40.9±0.7 to 39.4±0.6 for CIBβ− and 40.9±0.7 to 36.8±0.6 for L1-Exact.
  • Inference-time compute: CIB models show on-par or better inference-time scaling than baselines, especially with 3K or 2K maximum generation lengths and a 7B prior.This comparison concerns Pass@k accuracy on AIME24 and AIME25.

C CoT Qualitative Comparison: Pruning Reasoning Redundancy

Qualitative comparisons show that CIB removes redundant verification, parsing, and trial-and-error computation rather than merely truncating traces. The resulting traces preserve necessary logic while sometimes selecting shorter solution strategies.

  • General pattern: CIB changes which computations appear in reasoning traces instead of merely truncating outputs.The semantic prior imposes information cost while the sufficiency objective preserves task success.
  • Geometry reasoning: In geometry reasoning, CIB replaces redundant input parsing and self-correction with a concise trigonometric identity.The baseline favors lengthy coordinate calculation, whereas CIB converges on sin T = cos R.
  • Arithmetic search: In arithmetic search, CIB avoids calculating the incorrect candidate 983 and converges directly on 97.The figure reports an approximately 78% token reduction without loss of accuracy.
  • Constraint satisfaction: In constraint satisfaction, CIB retains the derived constraint c < 16 while eliminating tautology checks and repeated self-verification.The baseline exhibits verification bloat after deriving the necessary constraint.

D Training Details

The experiments use GRPO-based training and standardized evaluation infrastructure across specified model and hardware settings. Reproducibility details include library versions, sampling parameters, and shared hyperparameters.

  • Software: Training uses the trl library version 0.26.2 for Group Relative Policy Optimization, while lighteval version 0.8.1 supports evaluation.The implementation uses lighteval for downstream benchmarks.
  • Training configuration: All models are fine-tuned with the hyperparameters listed in Table 6.Table 6 is identified as the shared GRPO training configuration.
  • Inference and hardware: Inference uses vLLM 0.10.2, temperature 0.6, top-p 0.95, 32K maximum completion length, and 16 generations per prompt.Training runs on 8 NVIDIA H100 GPUs with 80GB each.

E Results from literature

The paper reports additional compression results for DeepScaler-1.5B from prior work. These results are presented separately from the paper’s main experiments.

  • Additional results: Table 7 reports compression results for DeepScaler-1.5B from reference [46].The supplied passage identifies the table but does not provide its numerical results.

F The list of datasets and models

This section provides tables listing the datasets and models used in the paper.

  • Table 8 lists the datasets used in the paper.
  • Table 9 lists the models used in the paper.
Loading 2603.08462v2…