SimpleOPD: 長文脈推論のためのシンプルなトークナイザー非依存オン方策蒸留

要旨

オン・ポリシー蒸留（OPD）は、より強力な教師モデルから推論能力を転移する有望な手法であるが、これを長文脈推論の教師モデルと短文脈の学生モデルに適用すると、トークナイザの不一致、教師と学生の分布不一致、応答長の爆発、学習の不安定性といった実際上の課題が生じる。本研究では、長文脈推論モデルSU-01から短文脈の学生モデルへ証明推論能力を転移することにより、この設定を検討する。トークナイザの違いに対処するため、共有テキスト空間でOPDを実行し、学生トークナイザと教師トークナイザの下で同一のテキスト区間を占めるトークンのみを整列させる。過剰な生成長と頻繁なトランケーションの問題を緩和するため、学生参照KL損失を導入し、`</think>`や`<|im_end|>`といった特殊な終了トークンのアドバンテージをマスクする。この戦略は、学生が初期方針から過度に逸脱するのを抑えることで、教師と学生の分布不一致の問題を緩和し、安定した長さの成長を促す。同一系列および異なる系列の学生モデル（Qwen3、Qwen3.5、Intern-S2、GLM-4.7、Gemma-4を含む）での実験では、数学的推論、特に自然言語による数学証明において一貫した改善が見られた。特に、Intern-S2-PreviewはProofBenchで21.2ポイント向上し、55.2に達してGemini-2.5-Proを上回った。また、HLEやHiPhOといった科学ベンチマークでも改善が見られ、OPDが数学の学習領域を超えて一般化する推論能力を転移することを示唆している。

English

On-policy distillation (OPD) offers a promising way to transfer reasoning capabilities from stronger teacher models, but applying it to long-context reasoning teachers and short-context students introduces practical challenges, including tokenizer mismatch, teacher-student distribution mismatch, response length explosion, and training instability. In this work, we study this setting by transferring proof-reasoning capabilities from the long-context reasoning model SU-01 to short-context student models. To handle tokenizer differences, we perform OPD in a shared text space and align only tokens that occupy identical text spans under the student and teacher tokenizers. To mitigate the problem of excessive generation length and frequent truncation, we introduce a student reference KL loss and mask the advantages of special termination tokens such as </think> and <|im_end|>. This strategy constrains the student from drifting excessively from its initial policy, thereby mitigating the teacher-student distribution mismatch problem and fostering steady length growth. Experiments on both same-family and different-family student models, including Qwen3, Qwen3.5, Intern-S2, GLM-4.7, Gemma-4, show consistent gains in mathematical reasoning, especially natural-language math proving. Notably, Intern-S2-Preview improves by 21.2 points on ProofBench, reaching 55.2 and surpassing Gemini-2.5-Pro. It also improves on science benchmarks such as HLE and HiPhO, suggesting that OPD transfers reasoning capabilities that generalize beyond the mathematical training domain.