English

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

Computation and Language 2026-02-24 v1 Machine Learning

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.

Keywords

Cite

@article{arxiv.2602.19049,
  title  = {IAPO: Information-Aware Policy Optimization for Token-Efficient Reasoning},
  author = {Yinhan He and Yaochen Zhu and Mingjia Shi and Wendy Zheng and Lin Su and Xiaoqing Wang and Qi Guo and Jundong Li},
  journal= {arXiv preprint arXiv:2602.19049},
  year   = {2026}
}
R2 v1 2026-07-01T10:46:03.827Z