自分の言葉で説明する: トークン選択的デュアル知識蒸馏による推論の改善

要約: Knowledge Distillation（KD）は、大規模モデルの推論能力を小型のモデルへ転送することができ、推論タスクの思考過程（Chain-of-Thoughts）を生成するコストを削減する可能性があります。KD手法は通常、学生に教師の出力全体にわたる分布を模倣させることを求めます。しかし、容量が限られた学生はこのような広範な監督によって圧倒され、特に複雑な推論タスクにおいて分布の不一致を生じさせることがあります。私たちは、学生中心の蒸留を目指す Token-Selective Dual Knowledge Distillation（TSD-KD）というフレームワークを提案します。TSD-KDは、推論のために重要なトークンの蒸留に焦点を当て、学生が自分の言葉で推論を説明することを奨励します。TSD-KDは間接蒸留と直接蒸留を組み合わせます。間接蒸留は、好みに基づくランキングという弱い形のフィードバックを使用します。学生は自分で生成した候補応答を提案します；教師はそれらの候補を間接フィードバックとして再ランク付けしますが、全体の分布を強制することはありません。直接蒸留は分布の整合を用いますが、教師と学生の相対的な確信度に基づいてトークンを選択的に蒸留します。最後に、蒸留中の学生の自信を維持するためにエントロピー正則化を追加します。全体として、私たちの手法は学生に対して推論過程を支援し自己改善を促進するための、的を絞った間接フィードバックを提供します。実験は、10個の難易度の高い推論ベンチマークにおけるTSD-KDの最先端性能を示し、ベースラインおよびランナーアップをそれぞれ最大で54.4%および40.3%上回りました。特筆すべきことに、TSD-KDで訓練された学生は4例で自分の教師モデルを上回り、最大で20.3%上回りました。ソースコードは https://github.com/kmswin1/TSD-KD で公開されています。