adu2021 Reasoning as Compression: Information-Theoretic Budget Optimization
Large language models generate lengthy chain-of-thought (CoT) reasoning that improves accuracy but increases inference latency and cost. Standard approaches use heuristic length penalties: if reasoning is too long, penalize it. But this ignores the semantic content—some reasoning steps are valuable, others are redundant.
This framework reframes efficient reasoning as a compression problem. Rather than counting tokens, we use information theory: measure the "cost" of reasoning by its semantic complexity (surprisal under a language model prior). The goal is to maximize task reward while compressing the reasoning trace below a budget. This yields both heuristic penalties and more sophisticated information-theoretic solutions automatically.
Core Concept
The framework applies the Conditional Information Bottleneck (CIB) principle to reasoning:
Information Bottleneck: Given input (prompt) and reasoning trace, the trace should compress information while still predicting the output accurately.
Semantic Cost: Cost of reasoning is measured by surprisal (negative log probability) under a language model, not token count. This allows intelligent pruning: rare/surprising tokens are costly, redundant tokens are cheap.
Optimization Target: Maximize task reward while keeping information cost below a budget, trading off reasoning quality and latency.
The result unifies several existing methods (length penalties, token budgets, etc.) as special cases of the general CIB optimization.
Architecture Overview
- Language Model Prior: Reference model that estimates surprisal (semantic cost)
- Reward Model: Task-specific model assessing reasoning quality
- Information Bottleneck Objective: Balances reward and compression
- Semantic Cost Estimation: Computes cost via language model probability
- RL Formulation: Optimize via policy gradient with information constraints
- Approximate CIB Solver: Tractable algorithm for constrained optimization
Implementation Steps
Step 1: Define Information-Theoretic Cost
Measure reasoning complexity using surprisal.
import torch
import torch.nn.functional as F
from typing import List, Tuple
import numpy as np
class InformationCostComputer:
"""
Compute information-theoretic cost of reasoning traces.
Uses language model prior to estimate semantic complexity.
"""
def __init__(self, reference_lm, temperature: float = 1.0):
self.reference_lm = reference_lm
self.temperature = temperature
def compute_surprisal(self, text: str) -> float:
"""
Compute surprisal (negative log probability) of text under LM prior.
Surprisal measures how unexpected/informative text is.
Higher surprisal = more information = higher cost.
"""
tokens = self.reference_lm.tokenize(text)
log_probs = []
with torch.no_grad():
for i in range(1, len(tokens)):
# Get probability of token i given tokens 0..i-1
context = tokens[:i]
next_token = tokens[i]
logits = self.reference_lm(context)
log_prob = F.log_softmax(logits, dim=-1)[next_token]
log_probs.append(-log_prob.item()) # Negative = surprisal
# Average surprisal per token
if log_probs:
return np.mean(log_probs)
return 0.0
def compute_segment_cost(self, reasoning_trace: str, segmentation: List[Tuple[int, int]]) -> List[float]:
"""
Compute cost of each segment in reasoning.
Segments: list of (start, end) character positions.
"""
segment_costs = []
for start, end in segmentation:
segment_text = reasoning_trace[start:end]
cost = self.compute_surprisal(segment_text)
segment_costs.append(cost)
return segment_costs
def segment_by_sentences(self, text: str) -> List[Tuple[int, int]]:
"""Split reasoning into sentence-level segments."""
import re
sentences = re.split(r'(?<=[.!?])\s+', text)
segmentation = []
pos = 0
for sentence in sentences:
start = pos
end = pos + len(sentence)
segmentation.append((start, end))
pos = end + 1 # +1 for space
return segmentation
Step 2: Formulate the Conditional Information Bottleneck
Define the optimization objective.
class ConditionalInformationBottleneck:
"""
CIB principle applied to reasoning:
The reasoning trace should be maximally informative about the answer
while being minimally informative beyond what the answer requires.
"""
def __init__(self, reference_lm, reward_model, information_budget: float = 10.0):
self.reference_lm = reference_lm
self.reward_model = reward_model
self.information_budget = information_budget
self.cost_computer = InformationCostComputer(reference_lm)
def compute_task_reward(self, reasoning: str, answer: str, ground_truth: str) -> float:
"""
Compute reward for this reasoning-answer pair.
Typically: is the answer correct?
"""
if answer.strip().lower() == ground_truth.strip().lower():
return 1.0
else:
# Partial credit based on semantic similarity
return self.reward_model.compute_similarity(answer, ground_truth)
def compute_information_cost(self, prompt: str, reasoning: str, answer: str) -> float:
"""
Compute information-theoretic cost of reasoning.
Cost = surprisal of reasoning under reference LM
This measures: how much new information does this reasoning contribute?
"""
# Surprisal of the full reasoning
full_trace_cost = self.cost_computer.compute_surprisal(reasoning)
# Surprisal of just the prompt (baseline)
# If LM expects reasoning given prompt, cost is low
# If LM is surprised by reasoning, cost is high
# Conditional cost: information in reasoning beyond what prompt implies
prompt_surprise = self.cost_computer.compute_surprisal(prompt)
conditional_cost = full_trace_cost - prompt_surprise
return max(0.0, conditional_cost) # Cost is non-negative
def compute_cib_loss(self, prompt: str, reasoning: str, answer: str,
ground_truth: str, beta: float = 0.1) -> float:
"""
Conditional Information Bottleneck loss:
L = -reward + beta * information_cost
This trades off:
- Reward: is the answer correct?
- Information cost: how much information in reasoning?
beta controls the tradeoff (higher = prefer compressed reasoning)
"""
reward = self.compute_task_reward(reasoning, answer, ground_truth)
info_cost = self.compute_information_cost(prompt, reasoning, answer)
loss = -reward + beta * info_cost
return loss
def check_information_budget(self, reasoning: str) -> bool:
"""Verify reasoning stays within information budget."""
cost = self.cost_computer.compute_surprisal(reasoning)
return cost <= self.information_budget
Step 3: Optimize Reasoning via Policy Gradient
Use reinforcement learning to find efficient reasoning.
import torch
import torch.optim as optim
from torch.distributions import Categorical
class ReasoningOptimizer:
"""
Optimize reasoning under information budget via RL.
"""
def __init__(self, reasoning_generator, cib_objective):
self.reasoning_generator = reasoning_generator
self.cib_objective = cib_objective
def optimize_reasoning_trajectory(self, prompt: str, ground_truth: str,
num_steps: int = 100,
beta: float = 0.1) -> Tuple[str, float]:
"""
Generate reasoning that balances quality and efficiency.
Strategy: Progressive shortening
- Start with full verbose reasoning
- Iteratively remove/compress low-value steps
- Stop when budget exceeded
"""
# Generate initial reasoning (verbose)
initial_reasoning = self.reasoning_generator.generate_verbose(prompt)
initial_answer = self.reasoning_generator.extract_answer(prompt + initial_reasoning)
# Segment reasoning into steps
segmentation = self.cib_objective.cost_computer.segment_by_sentences(initial_reasoning)
segment_costs = self.cib_objective.cost_computer.compute_segment_cost(
initial_reasoning, segmentation
)
# Iteratively remove low-value segments
current_reasoning = initial_reasoning
current_answer = initial_answer
for iteration in range(num_steps):
# Evaluate current reasoning
current_cost = self.cib_objective.compute_information_cost(
prompt, current_reasoning, current_answer
)
current_reward = self.cib_objective.compute_task_reward(
current_reasoning, current_answer, ground_truth
)
# Check budget
if current_cost <= self.cib_objective.information_budget:
break # Within budget, done
# Remove lowest-value segment
if segmentation:
# Find segment with minimum contribution to understanding
value_scores = self._score_segment_importance(
prompt, current_reasoning, current_answer, segmentation
)
min_value_idx = np.argmin(value_scores)
start, end = segmentation[min_value_idx]
# Remove segment
new_reasoning = current_reasoning[:start] + current_reasoning[end:]
new_answer = self.reasoning_generator.extract_answer(prompt + new_reasoning)
# Only keep removal if it doesn't hurt reward too much
new_reward = self.cib_objective.compute_task_reward(
new_reasoning, new_answer, ground_truth
)
if new_reward >= 0.9 * current_reward: # Allow 10% degradation
current_reasoning = new_reasoning
current_answer = new_answer
segmentation.pop(min_value_idx)
else:
# Can't remove this segment, try next
segmentation.pop(min_value_idx)
final_loss = self.cib_objective.compute_cib_loss(
prompt, current_reasoning, current_answer, ground_truth, beta=beta
)
return current_reasoning, final_loss
def _score_segment_importance(self, prompt: str, reasoning: str, answer: str,
segmentation: List[Tuple[int, int]]) -> np.ndarray:
"""Score importance of each reasoning segment."""
scores = []
for start, end in segmentation:
# Remove this segment and measure impact
modified_reasoning = reasoning[:start] + reasoning[end:]
modified_answer = self.reasoning_generator.extract_answer(prompt + modified_reasoning)
# Importance = impact on reward if removed
original_reward = self.cib_objective.compute_task_reward(reasoning, answer, answer)
modified_reward = self.cib_objective.compute_task_reward(
modified_reasoning, modified_answer, answer
)
importance = original_reward - modified_reward
scores.append(importance)
return np.array(scores)
Step 4: Semantic Prior for Cost Estimation
Implement efficient cost computation.
class SemanticPrior:
"""
Faster approximation of information costs using language model probabilities.
"""
def __init__(self, lm_model):
self.lm = lm_model
self.cache = {}
def estimate_cost(self, text: str) -> float:
"""Fast cost estimation via cached computations."""
if text in self.cache:
return self.cache[text]
# Tokenize
tokens = self.lm.tokenize(text)
if len(tokens) == 0:
return 0.0
# Get per-token probabilities
with torch.no_grad():
logits = self.lm.get_logits(text)
log_probs = F.log_softmax(logits, dim=-1)
# Compute perplexity as proxy for cost
# (perplexity ~= average surprisal)
avg_log_prob = log_probs.mean().item()
cost = -avg_log_prob / np.log(2) # Convert to bits
self.cache[text] = cost
return cost
def semantic_prior_objective(self, prompt: str, reasoning_segments: List[str],
answer: str) -> float:
"""
Objective that prefers reasoning that is likely under semantic prior.
High-likelihood reasoning (given prompt) has lower cost.
"""
total_cost = 0.0
for segment in reasoning_segments:
# Cost reduced if segment is likely given prompt
segment_cost = self.estimate_cost(prompt + segment)
prior_cost = self.estimate_cost(prompt) # Baseline
# Relative cost
relative_cost = max(0, segment_cost - prior_cost)
total_cost += relative_cost
return total_cost
Step 5: End-to-End Reasoning System
Integrate all components into a practical system.
class EfficientReasoningSystem:
"""
Complete system for budget-constrained reasoning optimization.
"""
def __init__(self, base_lm, reward_model, information_budget: float = 10.0):
self.base_lm = base_lm
self.cib_objective = ConditionalInformationBottleneck(
base_lm, reward_model, information_budget
)
self.optimizer = ReasoningOptimizer(base_lm, self.cib_objective)
self.semantic_prior = SemanticPrior(base_lm)
def generate_efficient_reasoning(self, prompt: str, ground_truth: str,
budget: float = None) -> Dict[str, any]:
"""
Generate reasoning optimized for both quality and efficiency.
"""
if budget:
self.cib_objective.information_budget = budget
# Optimize reasoning under budget
optimized_reasoning, loss = self.optimizer.optimize_reasoning_trajectory(
prompt, ground_truth, num_steps=50, beta=0.1
)
# Extract answer
answer = self.base_lm.generate_answer(prompt + optimized_reasoning)
# Compute metrics
reward = self.cib_objective.compute_task_reward(
optimized_reasoning, answer, ground_truth
)
info_cost = self.cib_objective.compute_information_cost(
prompt, optimized_reasoning, answer
)
budget_ratio = info_cost / self.cib_objective.information_budget
return {
'reasoning': optimized_reasoning,
'answer': answer,
'reward': reward,
'information_cost': info_cost,
'budget_utilization': budget_ratio,
'loss': loss
}
def compare_reasoning_efficiency(self, prompt: str, ground_truth: str) -> Dict:
"""
Compare different reasoning strategies.
"""
# Verbose reasoning (no compression)
verbose_result = self.generate_efficient_reasoning(
prompt, ground_truth, budget=float('inf')
)
# Compressed reasoning (tight budget)
compressed_result = self.generate_efficient_reasoning(
prompt, ground_truth, budget=5.0
)
return {
'verbose': verbose_result,
'compressed': compressed_result,
'accuracy_trade': verbose_result['reward'] - compressed_result['reward'],
'cost_reduction': verbose_result['information_cost'] - compressed_result['information_cost']
}
Practical Guidance
Hyperparameters:
- Information budget: 5-20 bits (task-dependent; higher = more reasoning)
- Beta (tradeoff weight): 0.01-1.0 (higher = prefer compression)
- Segmentation granularity: sentence-level or step-level
- Reference LM temperature: 1.0 (standard; controls cost smoothness)
When to Use:
- Tasks where reasoning is helpful but costs matter (API pricing, latency)
- Scenarios requiring transparency into cost-quality tradeoffs
- Domains where reasoning steps vary dramatically in usefulness
- Settings where you can measure ground truth (supervised learning)
When NOT to Use:
- Tasks requiring full verbose reasoning (e.g., math proofs need all steps)
- Online learning where reference LM isn't available
- Streaming inference (buffering reasoning adds latency)
- Domains where reasoning structure is unpredictable
Pitfalls:
- Reference LM quality critical: poor prior leads to meaningless costs
- Information budget setting requires tuning: no universal default
- Semantic cost can be noisy: use ensemble or smoothing
- Answer extraction fragile: must correctly parse reasoning output
Reference
Paper: arxiv.org/abs/2603.08462