Semantic EOS Aggregation：オンポリシー蒸留の長文化を終了トークン不整合から捉える

大規模言語モデルの能力を小型モデルに引き継がせる技術として蒸留が広く用いられています。中でも「オンポリシー蒸留（OPD）」は、生徒モデルが自ら言葉を生成し、その生成されたトークンを教師モデルがその場で採点して学習を進める枠組みです。あらかじめ用意された固定テキストを学ぶオフポリシー方式と異なり、生徒自身が出力した内容に基づいて誤りを修正できるため、高い性能を引き出しやすい利点があります。しかし、このオンポリシー蒸留では、学習が進むにつれて生徒の回答が必要以上に長くなり、設定された上限トークン数（本研究では7168トークン）まで生成が止まらなくなる「長文化」が頻発することが課題でした。生成が長文化すると計算負荷や推論コストが急激に跳ね上がります。従来の研究では、この原因として逆KL目的の性質や、長い系列での教師信号の劣化、確率分布が偏るエントロピー崩壊などが指摘されてきました。本論文はこれらと並行して存在する、より低水準で具体的な要因として、終了トークンの扱いを検証しています。

言語モデルはテキスト生成を終える際、「EOS（End of Sequence）トークン」と呼ばれる特殊な記号を出力します。システムはこの記号を検知して出力を停止します。これを日常の業務に例えてみましょう。ある職場で、部下（生徒モデル）は報告書の末尾に「以上」と書き、上司（教師モデル）は「完了」と書く習慣を持っているとします。人間から見ればどちらも「報告書をここで終える」という同じ意図を表しています。しかしLLMの世界では、事前学習段階のベースモデル（生徒）と、指示学習などを施された追加学習済みモデル（教師）の間で、これと似た食い違いが頻発します。双方が同じ単語一覧を共有し、設定上はどちらの記号も停止対象に含まれていても、モデル内部でどちらの記号に高い確率を割り当てているかが異なるのです。論文ではこれを「終了トークン不整合」と呼びます。ただし、この報告書の例えは単一回答の終了という同じ事象を表す表記揺れにのみ当てはまります。対話での発言交代や外部ツールへの処理引き渡しなど、異なる制御フローを担う終了記号の区別には当てはまらない点に留意が必要です。

オンポリシー蒸留では、生徒が実際にサンプルしたトークンのみを教師が直接評価し、両者の対数確率の差を学習係数として生徒を更新します。生徒が自身の好むEOS（例での「以上」）を出力して回答を終えたとき、教師が別のEOS（「完了」）に停止確率を集中させていると、生徒が出した「以上」に対する教師の確率は極めて低くなります。すると学習の計算上、生徒は「ここで終了したこと」に対して負の係数（ペナルティ）を受け取ってしまいます。一方で、生徒が教師の好む「完了」へ乗り換えることも容易ではありません。著者らがQwen3で観測したところ、生成終端付近における教師優先EOSの生徒確率は約10のマイナス11乗と極小でした。生徒が自発的にその記号を出力する機会がほぼないため、正の学習信号を受けることもできません。この結果、生徒は自身が持つ終了トークンを出すたびに罰せられて確率を削られ、Qwen3の通常蒸留では生徒固有のEOS確率が約0.8からほぼゼロまで低下しました。停止する手段を失った生徒は、生成上限の7168トークンに達するまで出力し続けることになります。

終了トークンの食い違いに対し、最も手軽な対症療法は「どちらのEOSが出ても生成を止めるよう、推論時の停止集合に両方を登録しておく」ことです。著者らはこれを「Fix 1」として検証しました。しかしこれは生成時の停止条件を変えるだけで、学習時の負の信号を解消できません。著者らはこれに加え、教師の全EOS確率を生徒側の正準EOSへ写像する「Fix 2」、教師確率を集約した上で生徒のサンプリング確率からも余分なEOSを除いて単一トークン空間に制約する「Fix 4」を比較しました。そして本研究の主提案が「Fix 3」にあたる「Semantic EOS Aggregation」です。この手法は、機能的に等価なEOS群を単一の「意味的な停止行動」として扱います。教師と生徒のそれぞれにおいて等価EOS群の総確率を計算し、生徒がEOSを出力した際には、個別の単語ではなく総停止確率の対数比を用いて学習更新を行います。通常の非EOSトークンの更新は変更せず、すべての等価EOSをロールアウトの停止記号として登録します。これにより、表層的な単語表記の違いを無視し、「止まるか続けるか」という行動単位で正しく蒸留を行えるようにします。

著者らは提案手法の妥当性を確かめるため、体系的な実験を行いました。主実験では、生徒モデルにQwen3-1.7B-Base、固定教師に非思考モードのQwen3-4Bを用い、数学データセット「DAPO-Math-17K」とTTRL形式のもとで200学習ステップを実行しました。応答上限は7168トークンです。さらにモデル系列横断の確認として「Llama 3.2」および「Gemma 3」のbase-to-post-trained蒸留でも検証しました。加えて、学習の進捗段階による影響を調べるため「K2-Horizon」を用い、事前学習段階から教える「Pretrain-to-Final」（400ステップ）、中間段階からの「Midtrain-to-Final」（200ステップ）、指示学習後からの「SFT-to-Final」（200ステップ）を比較しました。追跡指標には、平均応答長、7168トークン上限に達した割合である「clipping ratio」、終端位置のEOS確率を採用しました。また数学性能の評価として、「AMC23」「AIME24」「AIME25」の各問題から16回答を生成して問題内正解率を求め、3ベンチマークを等重みで平均する「Avg@16」を用いました。

この評価枠組みは、推論時の設定変更（Fix 1）と学習目的の補正（Fix 3）を切り離し、長文化の主要因が損失関数側の学習信号にあることを明確に示せます。また、応答長だけでなくclipping ratioや終端EOS確率を併用することで、単なるテキスト短縮ではなく停止機能の回復を捉えられます。さらにK2-Horizonにより、不整合の影響がモデルの学習段階にどう依存するかを分離して示せます。一方で、この評価では確立できない領域もあります。本検証は単一ターンの数学推論に集中しており、発言交代や外部ツール呼び出しが別の終了記号で制御されるマルチターン環境へそのまま拡張できるかは未検証です。また、Avg@16による評価は数学的正確さを測るものであり、停止行動の改善が論理的思考力の向上を直接もたらすとは言えません。実際の推論遅延やシステム運用の金銭的コスト削減についても、本実験では測定されていません。

著者らが報告した実験結果において、まず注目すべきはFix 1の失敗です。推論時の停止集合に複数のEOSを追加したFix 1は、Qwen3主実験で通常の未補正蒸留とほぼ完全に一致する平均応答長およびclipping ratioの軌跡を示しました。推論側の出口を揃えても、学習損失側での負の信号が残る限り長文化は防げないことが実証されました。これに対し、Semantic EOS AggregationはQwen3、Llama 3.2、Gemma 3の3系列すべてにおいて、通常の蒸留と比べて平均応答長とclipping ratioを大幅に低減しました。生成上限である7168トークンへ突き進むロールアウトが抑えられ、正常な停止行動が回復したと著者らは報告しています。ただし、系列ごとに回復度合いには差がありました。Qwen3やGemma 3で顕著な改善が見られた一方、Llama 3.2では教師モデルの参照長との間に、より大きな差が残りました。不整合の解消が有効である一方、モデル系列固有の性質による差異が残存することを示しています。

K2-Horizonを用いた段階別検証は、本手法の有効な境界と限界を浮き彫りにしました。中間段階（Midtrain）やSFT後の初期化では、生徒がすでに教師EOSを十分支持しており、Semantic EOS Aggregationの有無による応答長、clipping ratio、停止確率、下流正解率の一貫した差は見られませんでした。これは本手法が長さを無条件に削る一般的なペナルティではなく、不整合がある場合に特化して働くことを裏付けます。しかし、事前学習直後から教えるPretrain-to-Finalでは深刻な現象が確認されました。初期の長さ増加と高止まりする停滞期（plateau）を経て、学習後期に再び応答長が急増したのです。後期には停止確率がゼロ付近へ低下し、ほぼ全ロールアウトが7168トークンの上限へ達しました。そしてこの後期の再長文化は、Semantic EOS Aggregationを適用した後にも解消されずに残存しました。この結果は、終了トークン不整合が長文化の重要因ではあっても網羅的な唯一の原因ではなく、学習後期には別の停止破壊機構が働いていることを示しています。

本研究の知見を実務に生かす上では、評価の感度と適用条件を冷静に見極める必要があります。まず下流の数学正解率（Avg@16）について、停止挙動の改善が一貫した精度向上につながるとは示されませんでした。さらに評価器の形式感度も極めて重要です。同一のQwen3-4B教師出力に対し、元のDAPO評価器ではAvg@16が6.28%にとどまった一方、等価なLaTeX配置を扱える互換パーサでは24.34%へと大きく変化しました（18.06ポイントの差）。出力形式の許容度だけでスコアが大きく変わるため、精度の解釈には注意を要します。実務的な適用境界として、本手法が強く推奨されるのは、同一の語彙を共有し、単一ターン回答の終了を表す複数EOSが存在するbase-to-post-trained蒸留です。生徒固有のEOSが低下し推論時の停止設定だけでは改善しない場合、損失側での集約が有効な処方箋となります。一方で、会話交代やツール連携など終了記号ごとに意味が異なるマルチターン環境や、学習後期の停止崩壊に対しては、本手法だけで完結させず別の機構を切り分けて診断する必要があります。