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

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. Our results 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.

When AI systems solve math problems or answer complex questions, they often "think out loud", i.e., writing out long chains of reasoning before giving an answer. But just like a student who fills exam pages with restated facts, circular logic, and unnecessary commentary, these AI systems frequently ramble far more than needed. A PhD student solving a math problem might use 264 words on average; an AI tackling the same problems uses over 1,600. All that extra "thinking" costs real money and energy at scale, without making answers more accurate. We asked: can we train AI systems to reason more concisely without getting worse at the task? Our approach, IAPO, teaches AI models to identify which parts of their reasoning actually matter. The key idea comes from measuring how much a piece of information reduces uncertainty. We score each word the AI generates based on how much it helps pin down the correct answer, and then use those scores during training to reward informative reasoning and discourage filler. Making this work in practice required solving a computational puzzle: scoring every single word this way naively would be prohibitively slow. We developed efficient techniques — essentially smart caching and batching strategies — that make the computation fast enough to be practical. The results are striking. On standard math benchmarks, AI models trained with IAPO generate up to 47% fewer words while maintaining or even improving accuracy — answering more questions correctly per word written.