On-Policy Distillation of Language Models: Learning from Self-Generated Mistakes

要約

知識蒸留 (KD) は、より小さな生徒モデルをトレーニングすることにより、教師モデルを圧縮して推論コストとメモリ使用量を削減するために広く使用されています。
ただし、自己回帰シーケンス モデルの現在の KD 手法には、トレーニング中に見られる出力シーケンスと推論中に学習者によって生成された出力シーケンスとの間の分布不一致という問題があります。
この問題に対処するために、一般化知識蒸留 (GKD) を導入します。
GKD は、出力シーケンスの固定セットにのみ依存するのではなく、自己生成された出力シーケンスに関する教師からのフィードバックを活用して、そのようなシーケンスについて生徒を訓練します。
教師あり KD アプローチとは異なり、GKD は生徒と教師の間で代替の損失関数を採用する柔軟性も提供します。これは、生徒に教師の分布を模倣する表現力がない場合に役立ちます。
さらに、GKD は蒸留と RL 微調整 (RLHF) のシームレスな統合を促進します。
要約、翻訳、算術推論タスクに関する自己回帰言語モデルの抽出と、命令チューニングのためのタスクに依存しない抽出に対する GKD の有効性を実証します。

要約(オリジナル)

Knowledge distillation (KD) is widely used for compressing a teacher model to reduce its inference cost and memory footprint, by training a smaller student model. However, current KD methods for auto-regressive sequence models suffer from distribution mismatch between output sequences seen during training and those generated by the student during inference. To address this issue, we introduce Generalized Knowledge Distillation (GKD). Instead of solely relying on a fixed set of output sequences, GKD trains the student on its self-generated output sequences by leveraging feedback from the teacher on such sequences. Unlike supervised KD approaches, GKD also offers the flexibility to employ alternative loss functions between the student and teacher, which can be useful when the student lacks the expressivity to mimic the teacher’s distribution. Furthermore, GKD facilitates the seamless integration of distillation with RL fine-tuning (RLHF). We demonstrate the efficacy of GKD for distilling auto-regressive language models on summarization, translation, and arithmetic reasoning tasks, and task-agnostic distillation for instruction-tuning.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google