💬 NLP

SimpleOPD: Simple Tokenizer-Agnostic On-Policy Distillation for Long-Context Reasoning

The paper introduces SimpleOPD, a tokenizer-agnostic on-policy distillation framework that aligns shared text spans and employs a student reference KL loss to effectively transfer long-context proof-reasoning capabilities from strong teachers to short-context students, achieving significant performance gains on mathematical and scientific benchmarks.

Original authors: Haonan He, Haodi Lei, Yun Luo, Haoran Zhang, Shunkai Zhang, Yizhuo Li, Shengji Tang, Zhilin Wang, Runzhe Zhan, Lei Bai, Ganqu Cui, Fangchen Yu, Yafu Li, Peng Ye, Ning Ding, Yu Cheng

Published 2026-08-17

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Haonan He, Haodi Lei, Yun Luo, Haoran Zhang, Shunkai Zhang, Yizhuo Li, Shengji Tang, Zhilin Wang, Runzhe Zhan, Lei Bai, Ganqu Cui, Fangchen Yu, Yafu Li, Peng Ye, Ning Ding, Yu Cheng

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

Imagine you are trying to teach a brilliant, long-winded professor how to explain a complex idea to a quick-witted student who only has a tiny notebook. This is the world of Artificial Intelligence, specifically the field of "Large Language Models" (LLMs). These are computer programs trained to understand and generate human language. Sometimes, we have a "teacher" model that is incredibly smart at solving hard problems, like math proofs, but it thinks in massive, sprawling sentences. Then we have a "student" model that is smaller, faster, and has a shorter memory. The goal is to transfer the teacher's smarts to the student.

For a long time, scientists tried to do this by having the teacher write down answers and the student simply copy them. But this is like asking a student to memorize a dictionary without understanding the words; it doesn't always teach them how to think. A newer, smarter method called "On-Policy Distillation" (OPD) changed the game. Instead of just copying, the student tries to solve a problem, and the teacher watches the student's thought process in real-time, offering corrections and guidance at every single step. It's like a coach standing right next to an athlete, correcting their form as they run, rather than just showing them a video of a perfect run later. However, there's a catch: if the teacher and student speak different "languages" (use different ways of breaking words into pieces, called tokenizers) or if the teacher is used to writing novels while the student only has space for a tweet, the coaching can go wrong. The student might get confused, start rambling endlessly, or crash because the instructions don't fit.

This paper, titled SimpleOPD, tackles exactly that messy situation. The researchers wanted to take a super-smart, long-context teacher model (named SU-01) that excels at mathematical proofs and teach it to various smaller student models, even those that use completely different internal "languages" (tokenizers) and have much shorter memory limits. They discovered that simply letting the teacher coach the student often backfired. The student would get overwhelmed, start writing responses that grew uncontrollably long (a "length explosion"), and eventually run out of space, cutting off their own answers before finishing.

To fix this, the team invented a clever, "tokenizer-agnostic" approach. Instead of forcing the teacher and student to agree on exactly which piece of a word they are looking at, they agreed to look at the actual text on the page. If the teacher says "math" and the student says "mat" and "h," they realized these are just different ways of writing the same sound. They aligned the teaching based on the shared text, ignoring the technical differences in how the computers chopped up the words. But they didn't stop there. To stop the students from rambling, they introduced two safety nets. First, they told the teacher to stop nagging the student about when to stop talking (ignoring specific "stop" tokens), so the student wouldn't get confused about when to finish. Second, they added a gentle "reality check" (a KL loss) that reminded the student to stay close to its original, sensible self, preventing it from drifting too far into wild, endless loops.

The results were impressive. By using this method, the student models didn't just learn to copy; they learned to reason. For example, one student model, Intern-S2-Preview, saw its score on a difficult math proof test (ProofBench) jump by 22.8 points, going from 21.70 to 44.50 in the primary experiments. In a specific evaluation setting using Gemini-2.5-Pro as the judge, the model showed an even larger gain, rising from 34.0 to 55.2. This improvement was so significant that it surpassed the performance of other very powerful, closed-source models like Gemini-2.5-Pro. The paper suggests that this technique works well not just for math, but also helps the students generalize their new skills to science problems they haven't seen before. Essentially, SimpleOPD proved that you can teach a small, short-memory brain to think like a giant, long-winded genius, even if they speak different technical languages, as long as you give them a shared text space and a few rules to keep them from talking forever.

Technical Summary: SimpleOPD

Problem Statement

On-policy distillation (OPD) has emerged as a powerful paradigm for transferring reasoning capabilities from strong teacher models to student models by providing dense, token-level supervision based on student-generated trajectories. However, applying OPD to transfer long-context reasoning capabilities (from models capable of 100K+ token reasoning) to short-context student models across different model families presents significant practical challenges:

Tokenizer Mismatch: Teachers and students often utilize incompatible tokenizers and vocabularies, making direct token-level alignment of probability distributions infeasible.
Distribution Mismatch & Length Explosion: Long-context teachers often generate responses that exceed the context budget of short-context students. Direct OPD training leads to rapid response-length growth, frequent truncation, and incomplete reasoning trajectories, ultimately causing training instability.
Termination Token Suppression: Teacher supervision often discourages students from using termination tokens (e.g., < /think>, < |im end| >) because the teacher's long-context reasoning patterns favor extended generation, causing the student to fail to terminate properly.

Methodology: SimpleOPD

The authors propose SimpleOPD, a tokenizer-agnostic on-policy distillation framework designed to transfer reasoning capabilities from the long-context model SU-01 to diverse short-context student models. The method consists of three core components:

1. Cross-Tokenier Alignment in Shared Text Space

Instead of forcing a mapping between incompatible token vocabularies, SimpleOPD performs distillation in a shared text space.

Process: The student generates a response sequence $y_{1:n}$ using its own tokenizer, which is decoded into a surface string $s$ . The teacher then evaluates this exact same string $s$ using its own native tokenizer and chat template.
Alignment: Tokens are aligned only if they occupy identical text spans in the shared response string. The authors define a set of aligned pairs $\mathcal{M}$ where the cumulative text prefix and the incremental text span of a teacher token $z_i$ match those of a student token $y_t$ .
Loss Construction: For aligned positions, the student is supervised by the teacher's log-probability. For unaligned positions (where token boundaries do not match), the student's own log-probability is used as a fallback. This creates a token-aligned surrogate for the reverse KL divergence.

2. Termination Token Advantage Masking

To prevent the student from learning to suppress termination tokens (a common failure mode in direct OPD), the authors mask the advantages of specific structural tokens.

Mechanism: The OPD loss is masked (set to zero) for tokens such as < /think> and < |im end| >.
Rationale: These tokens control output format and termination behavior rather than reasoning content. By not forcing the student to match the teacher's distribution on these specific tokens, the student retains its ability to terminate generation properly, preventing infinite loops or excessive length growth.

3. Student-Reference KL Loss

To mitigate the distribution mismatch and prevent the student policy from drifting too far from its initial capabilities (which leads to instability and truncation), a student-reference KL divergence loss is introduced.

Mechanism: A KL penalty is added between the current student policy and the initial student policy (the reference).
Effect: This constrains the student from deviating excessively, stabilizing the training dynamics and ensuring steady growth in response length within the student's context budget. The coefficient for this loss is tuned based on the magnitude of the teacher-student gap (e.g., 0.5 for same-family, 1.0 for cross-family).

Key Contributions

Tokenizer-Agnostic Alignment: Demonstrated that effective OPD can be achieved across different model families by aligning tokens based on shared text spans rather than requiring identical vocabularies.
Stabilization Techniques: Identified that direct OPD causes length explosion and truncation. Proposed and validated two simple but effective stabilization techniques: termination token masking and student-reference KL regularization.
Cross-Family Transfer: Successfully transferred proof-reasoning capabilities from a 30B long-context teacher (SU-01) to diverse short-context students (Qwen, Intern-S2, GLM, Gemma) without supervised fine-tuning on teacher trajectories.

Experimental Results

The authors evaluated SimpleOPD on mathematical reasoning benchmarks, including ProofBench (natural-language proofs), AnswerBench, AIME25, and AMOBench.

Same-Family & Cross-Family Gains: SimpleOPD consistently improved performance across all tested student models.
- Intern-S2-Preview: Achieved a massive 21.2-point improvement on ProofBench (from 34.0 to 55.2 when evaluated with Gemini-2.5-Pro as the judge, and from 21.7 to 44.5 in the primary evaluation), surpassing strong frontier models like Gemini-2.5-Pro and GPT-5.
- Qwen3-30B-A3B: Improved ProofBench score by 22.67 points (13.80 $\to$ 36.47).
- GLM-4.7 & Gemma-4: Showed consistent gains on ProofBench, demonstrating the method's applicability to models with significantly different tokenizers (e.g., SentencePiece vs. BPE).
Stability: The combination of token masking and KL loss effectively reduced the truncation rate to nearly zero and prevented the "length explosion" observed in baseline OPD training.
Out-of-Domain Generalization: Despite being trained exclusively on mathematical proof data, the distilled models showed improvements on scientific reasoning benchmarks (HLE, HiPhO), suggesting the transferred reasoning capabilities generalize beyond the training domain.
Distillation Length: Longer distillation sequences (32k vs 6k) yielded better performance on complex proof tasks, indicating that preserving long reasoning traces is crucial for transferring multi-step reasoning behaviors.

Significance and Claims

The paper claims that SimpleOPD provides a robust and generalizable solution for transferring long-context reasoning capabilities to short-context models, even across different architectures and tokenizers.

Efficiency: It offers a cost-effective alternative to training long-context models from scratch or using complex RL pipelines, leveraging existing strong teachers to boost smaller, faster models.
Stability: The proposed stabilization techniques (masking and KL regularization) are essential for making OPD viable in heterogeneous settings, solving the critical issues of training instability and truncation that previously hindered such transfers.
Generalization: The results suggest that reasoning capabilities distilled via this method are not limited to the specific training domain (math proofs) but can enhance performance in broader scientific reasoning tasks.

The authors conclude that their approach effectively bridges the gap between long-context teachers and short-context students, enabling the democratization of advanced reasoning capabilities across the model ecosystem.