クロストークナイザー蒸留では、網羅性より監督信号の信頼性が重要

オンポリシー蒸留（OPD）は、生徒モデル自身が生成した応答を学習対象にし、教師モデルの予測分布からフィードバックを与える手法である。固定データを模倣するだけでなく、生徒が実際に訪れる生成軌跡を直接改善できる点が特徴だ。しかし教師と生徒が異なるトークナイザーを使う場合、両者の予測を単純に比較することはできない。系列の位置だけでなく、語彙レベルでも対応関係を構築する必要がある。

本論文が問い直すのは、対応範囲を広げることが本当に学習改善につながるのか、という点である。著者らは、数学推論とコード生成を対象に、異なるトークナイザーを持つ三つの教師・生徒組み合わせを調べた。その結果は、「監督する範囲が広いほどよい」という直感を単純には支持しなかった。

主な結果は次の通りである。

• 厳密な1対1対応グループは、語彙に大きな違いがあっても、生徒が生成するトークンの大部分をカバーした。
• 蒸留前に生徒からサンプリングした応答では、厳密に対応する位置において、共有語彙が教師・生徒双方の確率質量のほぼ全体を保持していた。
• 各位置で生徒が選んだ共有語彙Top-16だけに反復KLを適用すると、完全な共有語彙OPDと同程度の精度となり、評価対象のクロストークナイザー基線を上回った。
• 不一致グループに対して、スパンの対数確率を使った平均二乗誤差を追加すると監督範囲は完全になるが、精度は低下した。

追加監督が逆効果になり得る理由も分析されている。厳密損失だけで学習したチェックポイントでは、不一致スパン損失の勾配と厳密損失の勾配の方向一致が弱く、一部では反対方向になった。また、スパン側の勾配は厳密損失に対して相対的に大きくなった。つまり、形式上の監督の空白を埋める信号が、最も有効な最適化方向と協調するとは限らない。

この研究が示す設計上の教訓は、カバレッジだけを指標にしないことである。監督信号が十分な確率質量を含むか、安定して計算できるか、既存の損失と同じ方向に学習を進めるかも確認すべきだ。クロストークナイザーOPDでは、すべての不一致スパンに損失を課すより、対応が明確な位置で小さく信頼できる候補集合を使う方が有効な場合がある。

実験対象は限られたモデル組み合わせとタスクであり、結論をすべての状況に一般化することはできない。それでも、追加損失を導入する前に、対応率、確率質量、勾配の方向と大きさを診断するという実践的な指針を与えている。網羅性の追求から、監督信号の信頼性を重視する設計への転換が、今後の重要な論点になりそうだ。

出典：Hugging Face Daily Papers