Update metadata and improve model card

README.md

@@ -1,72 +1,64 @@
 ---
-license: mit
-datasets:
-- HuggingFaceH4/MATH-500
-- BytedTsinghua-SIA/DAPO-Math-17k
-- openai/gsm8k
 base_model:
 - Qwen/Qwen2.5-0.5B-Instruct
 - Qwen/Qwen2.5-1.5B-Instruct
 - Qwen/Qwen2.5-7B-Instruct
-pipeline_tag: reinforcement-learning
+datasets:
+- HuggingFaceH4/MATH-500
+- BytedTsinghua-SIA/DAPO-Math-17k
+- openai/gsm8k
+license: mit
+pipeline_tag: text-generation
+library_name: transformers
 ---
-# IAPO: Information-Aware Policy Optimization for Token-Efficient Reasoning
 
-  🚀 **Overview**
+# IAPO: Information-Aware Policy Optimization for Token-Efficient Reasoning
 
-  IAPO is an information-theoretic post-training framework that improves the **token efficiency** of Chain-of-Thought
-   (CoT) reasoning. Instead of shaping rewards at the sequence level — as in standard RL methods such as GRPO — IAPO
-  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 tokens and suppressing
-  low-utility exploration, producing shorter reasoning traces without sacrificing accuracy.
+🚀 **Overview**
 
-  Specifically, this checkpoint is a fine-tuned version of `Qwen2.5-0.5B-Instruct` using the IAPO framework on
-  mathematical reasoning datasets (GSM8K, MATH-500, and DAPO-Math-17k).
+IAPO is an information-theoretic post-training framework designed to improve the **token efficiency** of Chain-of-Thought (CoT) reasoning. Instead of shaping rewards at the sequence level—as seen in standard RL methods like GRPO—IAPO assigns **token-wise advantages** based on each token's **conditional mutual information (MI)** with the final answer. This identifies informative reasoning steps and suppresses low-utility exploration, resulting in significantly shorter reasoning traces without sacrificing accuracy.
 
-  **Paper:** [IAPO: Information-Aware Policy Optimization for Token-Efficient 
-  Reasoning](https://arxiv.org/pdf/2602.19049)
-  **Code:** [Official GitHub Repository](https://github.com/YinhanHe123/IAPO)
+- **Paper:** [IAPO: Information-Aware Policy Optimization for Token-Efficient Reasoning](https://huggingface.co/papers/2602.19049)
+- **Code:** [Official GitHub Repository](https://github.com/YinhanHe123/IAPO)
 
-  🎯 **Key Features**
+🎯 **Key Features**
 
-  - 🧠 **Information-Aware Advantage Shaping:** Assigns token-level advantages based on conditional MI w.r.t. the
-  final answer, amplifying informative tokens and suppressing redundant ones.
-  - 🔍 **Exploration Adjustment:** Adds a per-token exploration term that rewards confident tokens in correct
-  trajectories and penalizes them in incorrect ones, preventing premature reasoning collapse while encouraging
-  recovery from failed paths.
-  - ⚡ **Efficient Conditional MI Estimation:** Introduces an *early-exit–based* MI estimator combined with *KV-cache
-   preloading* and *chunk-wise forwarding*, reducing the naive `O(L³d)` cost of per-token MI estimation to a
-  tractable cost suitable for long reasoning.
-  - 📉 **Provable Length Reduction:** Theoretical analysis shows that, under mild assumptions, IAPO monotonically
-  reduces expected completion length without harming correctness.
-  - 🏆 **Strong Empirical Results:** Reduces reasoning length by up to **47%** while improving accuracy across math
-  reasoning datasets, outperforming token-efficient RL baselines such as DAPO, GFPO, GTPO, and S-GRPO.
+- 🧠 **Information-Aware Advantage Shaping:** Assigns token-level advantages based on conditional MI, amplifying informative tokens and suppressing redundant ones.
+- 🔍 **Exploration Adjustment:** Rewards confident tokens in correct trajectories and penalizes them in incorrect ones to prevent reasoning collapse.
+- ⚡ **Efficient Estimation:** Introduces an early-exit–based MI estimator with KV-cache preloading to keep computational costs tractable for long-context reasoning.
+- 📉 **Provable Length Reduction:** Demonstrates monotonic reductions in reasoning verbosity while preserving correctness.
+- 🏆 **Performance:** Reduces reasoning length by up to 36-47% while improving accuracy across various mathematical reasoning benchmarks.
 
-  🔧 **Loading the Checkpoints**
+🔧 **Loading the Checkpoints**
 
-  You can load any of our released IAPO checkpoints directly from the Hugging Face Hub via the `subfolder` argument.
-  Each subfolder corresponds to a `{base_model}_{dataset}` combination. For example:
+This repository contains multiple checkpoints fine-tuned from different base models (Qwen2.5-0.5B, 1.5B, 7B). You can load a specific checkpoint using the `subfolder` argument corresponding to the `{base_model}_{dataset}` combination.
 
-  ```python
-  from transformers import AutoModel
+```python
+from transformers import AutoModelForCausalLM, AutoTokenizer
+import torch
 
-  # Generic usage
-  model = AutoModel.from_pretrained("jonathanhe123/iapo", subfolder="model_name")
+model_id = "jonathanhe123/iapo"
+# Example: Load the Qwen2.5-0.5B-Instruct checkpoint fine-tuned on MATH-500
+subfolder = "Qwen2.5-0.5B-Instruct_MATH-500"
 
-  # Example: load the Qwen2.5-0.5B-Instruct checkpoint fine-tuned on MATH-500
-  model = AutoModel.from_pretrained("jonathanhe123/iapo", subfolder="Qwen2.5-0.5B-Instruct_MATH-500")
-  ```
+tokenizer = AutoTokenizer.from_pretrained(model_id, subfolder=subfolder)
+model = AutoModelForCausalLM.from_pretrained(
+    model_id, 
+    subfolder=subfolder,
+    torch_dtype=torch.bfloat16,
+    device_map="auto"
+)
+```
 
-  **Citation**
+**Citation**
 
-  If you find this work useful, please cite our paper:
+If you find this work useful, please cite the paper:
 
-  ```bibtex
-  @inproceedings{he2026iapo,
-    title={IAPO: Information-Aware Policy Optimization for Token-Efficient Reasoning},
-    author={He, Yinhan and Zhu, Yaochen and Shi, Mingjia and Zheng, Wendy and Su, Lin and Wang, Xiaoqing and Guo, Qi
-  and Li, Jundong},
-    booktitle={International Conference on Machine Learning (ICML 2026)},
-    year={2026}
-  }
-  ```
+```bibtex
+@inproceedings{he2026iapo,
+  title={IAPO: Information-Aware Policy Optimization for Token-Efficient Reasoning},
+  author={He, Yinhan and Zhu, Yaochen and Shi, Mingjia and Zheng, Wendy and Su, Lin and Wang, Xiaoqing and Guo, Qi and Li, Jundong},
+  booktitle={International Conference on Machine Learning (ICML 2026)},
+  year={2026}
+}
+```