💬 NLP

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

This paper proposes IAPO, an information-theoretic post-training framework that optimizes token-efficient reasoning by assigning token-wise advantages based on conditional mutual information, thereby reducing inference costs by up to 36% while improving or maintaining accuracy.

Original authors: Yinhan He, Yaochen Zhu, Mingjia Shi, Wendy Zheng, Lin Su, Xiaoqing Wang, Qi Guo, Jundong Li

Published 2026-06-01

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Yinhan He, Yaochen Zhu, Mingjia Shi, Wendy Zheng, Lin Su, Xiaoqing Wang, Qi Guo, Jundong Li

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

The Big Problem: The "Over-Explainer" Robot

Imagine you ask a very smart robot a math question. You want the answer, but the robot starts talking. It doesn't just give the answer; it talks through every single thought, checks its work three times, says things like "Okay, let me think about this," and then accidentally circles back to a thought it already had.

This is what happens with current AI models (Large Language Models). They are great at solving problems, but they are verbose. They use way too many "tokens" (chunks of words) to get there.

The Cost: Every extra word costs money and time. It's like paying for a taxi ride where the driver takes a scenic route, stops to smell the flowers, and checks the map five times before dropping you off.
The Goal: We want the robot to be smart and concise. We want it to cut the fluff without losing the correct answer.

The Old Way: The "Length Limit" Coach

Previously, researchers tried to fix this by acting like a strict coach who says, "If you write more than 50 words, you get no points."

The Flaw: This is like telling a student, "Don't write more than 50 words," without caring what they wrote. The student might cut out the most important math steps just to stay under the limit, or they might keep the boring parts and cut the good parts. The old methods didn't understand which words were actually helpful and which were just noise.

The New Solution: IAPO (The "Smart Editor")

The authors propose a new system called IAPO. Instead of just counting words, IAPO acts like a Smart Editor that understands the value of every single word the robot writes.

Here is how it works, broken down into three simple parts:

1. The "Value Meter" (Information Awareness)

Imagine the robot is writing a story to solve a puzzle. IAPO looks at every sentence and asks: "Does this sentence actually help solve the puzzle?"

High Value: "The answer is 42." (This is crucial!)
Low Value: "Wait, let me double-check my math... hmm, I think I was right before." (This is just noise/redundancy).

IAPO uses a mathematical concept called Conditional Mutual Information. In plain English, this is a way to measure: "If I remove this specific word, how much more confused would I be about the final answer?"

If removing the word makes you confused, that word is high value. IAPO rewards it.
If removing the word changes nothing, that word is low value (fluff). IAPO penalizes it.

2. The "Exploration Safety Net"

There is a risk: If you only reward the robot for being short, it might get lazy and guess the answer too quickly, skipping the hard thinking needed to get it right.
To fix this, IAPO has a second rule: The Exploration Adjustment.

If the robot gets the answer right, IAPO says, "Great job! You were confident and correct. Keep doing exactly what you did." (It stops the robot from wandering off).
If the robot gets the answer wrong, IAPO says, "Oops. You were too confident in a bad path. Let's try something different next time." (It encourages the robot to explore new ideas).

3. The "Speed Trick" (Efficient Estimation)

Calculating the "value" of every single word in a long story is usually very slow and expensive for computers. It's like trying to weigh every grain of sand on a beach individually.
The authors invented a Speed Trick (using something called "Early-Exit" and "KV-Cache").

The Analogy: Instead of re-reading the whole story from the beginning every time to check a specific word, the system saves the "memory" of the story as it goes. It can then jump straight to the part it needs to check. This makes the process fast enough to actually use on real computers.

The Results: Shorter, Smarter, Faster

The paper tested this new "Smart Editor" on math problems (like those found in school or competitions).

The Outcome: The AI trained with IAPO solved the problems just as accurately as before, but it used up to 47% fewer words.
The Comparison: In one example, a standard AI took 105 words to solve a simple math problem. The IAPO-trained AI solved the exact same problem in just 15 words, cutting out all the "umms," "ahhs," and repetitive checks.

Summary

Think of IAPO as teaching an AI to be a great editor rather than just a fast typist.

Old AI: Writes a 10-page essay to say "The answer is 5."
IAPO AI: Writes a one-sentence note: "The answer is 5."

It achieves this by rewarding the AI only for the words that actually matter, while using a clever speed trick to do the math behind the scenes. This saves money, time, and computing power, making AI more efficient without making it less smart.

`) to intermediate prefixes. This forces the model to immediately predict the final answer distribution from the current state, allowing the calculation of entropy reduction $H(y|q, o_{<t}) - H(y|q, o_{\le t})$ without generating the full reasoning chain.
2. KV-Cache Preloading & Chunk-wise Forwarding:
* KV-Cache Preloading: The model performs a single forward pass over the full completion to cache Key-Value (KV) states. Subsequent entropy estimations reuse these cached states for the shared prefix, requiring computation only for the short postfix prompt. This reduces complexity from $O(L^3d)$ to approximately $O((K^3 + L^2)d)$ , where $K$ is the postfix length ( $K \ll L$ ).
* Chunk-wise Forwarding: To further reduce overhead, completions are partitioned into chunks, and conditional MIs for all tokens within a chunk are estimated simultaneously via batched forward passes.

3. Theoretical Analysis

The paper provides a theoretical analysis demonstrating that IAPO induces monotonic reductions in reasoning verbosity without harming correctness.

Length Reduction: The authors show that the difference in expected completion length between IAPO and vanilla GRPO is proportional to the covariance between the completion length $L(o)$ and the informativeness-weighted gradient signal $S(o)$ . Under the assumption that longer completions tend to accumulate more low-informativeness tokens (negative correlation), IAPO is proven to reduce expected completion length.
Exploration Control: The exploration adjustment term is shown to regulate policy entropy. It suppresses entropy for correct trajectories (stabilizing the solution) and increases entropy for incorrect trajectories (encouraging exploration of new paths), thereby maintaining reasoning accuracy.

4. Empirical Results

Experiments were conducted on three mathematical reasoning datasets (GSM8K, MATH-500, DAPO-Math-17k) using Qwen2.5 models (0.5B, 1.5B, 7B). IAPO was compared against state-of-the-art baselines including DAPO, GFPO, GTPO, and S-GRPO.

Token Efficiency: IAPO consistently achieved the highest token efficiency (Ratio@k = Pass@k / Length@k). On the 7B model, IAPO reduced reasoning length by up to 47% compared to the base model and significantly outperformed existing token-efficient RL baselines.
Accuracy Preservation: Despite the drastic reduction in tokens, IAPO maintained or improved reasoning accuracy (Pass@k). For instance, on GSM8K with the 7B model, IAPO achieved a Pass@32 of 0.9805 with an average length of ~84 tokens, whereas the base model required ~156 tokens for similar accuracy.
Inference Time: Wall-clock inference time was reduced by more than 11% due to shorter sequences.
Ablation Studies: Removing the informativeness term (IAPO-NI) or replacing it with next-token entropy (IAPO-NE) degraded performance, confirming that conditional MI is the critical factor for identifying informative tokens.
Case Studies: Visualizations of reasoning trajectories showed that IAPO successfully suppressed redundant meta-commentary and failed derivations, focusing generation on high-informativeness tokens (e.g., key numerical values and logical conjunctions).

5. Significance and Claims

The paper claims that information-aware advantage shaping is a powerful and general direction for token-efficient post-training. By moving beyond heuristic length or position penalties, IAPO provides a principled, information-theoretic mechanism to identify and retain only the reasoning steps that contribute to answer correctness.

The authors emphasize that their method:

Formalizes token-efficient post-training as maximizing the ratio of task accuracy to completion length.
Offers a scalable solution to the computational bottleneck of token-level MI estimation via early-exit estimators and KV-cache optimization.
Demonstrates that it is possible to significantly reduce inference costs (up to 47% token reduction) while preserving, and in some cases enhancing, the reasoning capabilities of LLMs.

The work suggests that this framework can lead to more sustainable AI deployment by lowering compute and energy requirements, while also offering more interpretable reasoning chains by filtering out redundant content.