IAPO: Information-Aware Policy Optimization for Token-Efficient Reasoning

TL;DR

IAPO introduces an information-theoretic framework for token-efficient reasoning in large language models, improving accuracy while reducing reasoning length by focusing on informative reasoning steps through mutual information analysis.

Contribution

It proposes a novel token-wise advantage shaping method based on mutual information, enabling explicit control over reasoning effort distribution during post-training.

Findings

Reduces reasoning length by up to 36%

Improves reasoning accuracy across datasets

Outperforms existing token-efficient RL methods

Abstract

Large language models increasingly rely on long chains of thought to improve accuracy, yet such gains come with substantial inference-time costs. We revisit token-efficient post-training and argue that existing sequence-level reward-shaping methods offer limited control over how reasoning effort is allocated across tokens. To bridge the gap, we propose IAPO, an information-theoretic post-training framework that assigns token-wise advantages based on each token's conditional mutual information (MI) with the final answer. This yields an explicit, principled mechanism for identifying informative reasoning steps and suppressing low-utility exploration. We provide a theoretical analysis showing that our IAPO can induce monotonic reductions in reasoning verbosity without harming correctness. Empirically, IAPO consistently improves reasoning accuracy while reducing reasoning length by up to 36%, outperforming existing token-efficient RL methods across various reasoning datasets. Extensive empirical evaluations demonstrate that information-aware advantage shaping is a powerful and general direction for token-efficient post-training. The code is available at https://github.com/YinhanHe123/IAPO.