弱から強へのオンポリシー蒸留:弱いモデルで強いモデルを強化
本論文では、複数の弱いモデルから蒸留することで強い学生モデルを向上させるフレームワーク、弱から強へのオンポリシー蒸留(W2S-OPD)を提案する。ロジット空間でコントラストペアを用いて代理教師を構築し、数学・コードベンチマークで標準OPDを凌駕する。
大規模言語モデル(LLM)間での能力移転において、オンポリシー蒸留(OPD)は、学生自身の生成系列上で教師のトークンレベルの分布に合わせる効果的なパラダイムである。しかし、既存の手法は通常、教師が学生と少なくとも同等の能力を持つことを前提としており、より大きな教師が存在しないフロンティアでは適用が難しい。また、複数のドメイン専門家を統合するには学生規模での高コストな訓練が必要となる。本研究では、弱から強へのオンポリシー蒸留(W2S-OPD)を提案する。これは、学生よりも小さい複数の弱いモデルから蒸留することで、強い学生を向上させるシンプルかつ効果的なフレームワークである。
W2S-OPDの核心は、ロジット空間でコントラストペア(正モデルと負モデル)を用いて代理教師を構築することにある。これら2つのモデルは学生より小さく、容易に入手できる。それらのロジット差は特定の能力方向を分離し、学生自身のベースモデルに加えることで、その方向を結合しつつ学生と分布的に近い代理教師が得られる。学生は自身の生成系列上のトークンごとの逆KLダイバージェンスを最小化することで、この代理教師を蒸留する。
研究チームは3種類のコントラストペアを実装した。(i) 強化学習後のエキスパートとその強化学習前の初期化の比較は、強化学習が付与するスキルを分離する。(ii) より大きなベースモデルとより小さなベースモデルの比較は、規模による能力を分離する。(iii) 正しいヒントと誤ったヒントを与えた小さなベースモデルの比較は、解決策に向けたインスタンスレベルの方向を分離する。4つの数学ベンチマーク(GSM8K、MATHなど)と3つのコードベンチマーク(HumanEval、MBPPなど)での実験では、W2S-OPDが標準OPDを上回り、学生がドメイン教師を超えることを可能にし、全ての教師が学生より弱い場合でも学生を改善し続けることを示した。分析により、異なるコントラストが異なる信号を生成することが明らかになった。強化学習後およびヒントのコントラストは推論枠組みを強調し、規模のコントラストは解決手順を強調する。コードはGitHubで公開されており、LLM蒸留の新たな方向性を示している。