On-Policy-Destillation von Sprachmodellen: Lernen aus selbstgenerierten Fehlern

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

June 23, 2023

Autoren: Rishabh Agarwal, Nino Vieillard, Yongchao Zhou, Piotr Stanczyk, Sabela Ramos, Matthieu Geist, Olivier Bachem

cs.AI

Zusammenfassung

Wissensdistillation (KD) wird häufig verwendet, um ein Lehrer-Modell zu komprimieren und dadurch dessen Inferenzkosten und Speicherbedarf zu reduzieren, indem ein kleineres Schüler-Modell trainiert wird. Allerdings leiden aktuelle KD-Methoden für autoregressive Sequenzmodelle unter einer Verteilungsdiskrepanz zwischen den während des Trainings gesehenen Ausgabesequenzen und denen, die vom Schüler während der Inferenz generiert werden. Um dieses Problem zu lösen, führen wir die Generalisierte Wissensdistillation (GKD) ein. Anstatt sich ausschließlich auf einen festen Satz von Ausgabesequenzen zu verlassen, trainiert GKD den Schüler auf seinen selbst generierten Ausgabesequenzen, indem es Feedback des Lehrers zu solchen Sequenzen nutzt. Im Gegensatz zu überwachten KD-Ansätzen bietet GKD auch die Flexibilität, alternative Verlustfunktionen zwischen Schüler und Lehrer einzusetzen, was nützlich sein kann, wenn der Schüler nicht die Ausdrucksfähigkeit besitzt, die Verteilung des Lehrers nachzuahmen. Darüber hinaus ermöglicht GKD die nahtlose Integration von Distillation mit RL-Feinabstimmung (RLHF). Wir demonstrieren die Wirksamkeit von GKD für die Distillation autoregressiver Sprachmodelle bei Aufgaben wie Zusammenfassung, Übersetzung und arithmetischem Denken sowie für aufgabenunabhängige Distillation zur Instruktionsfeinabstimmung.

English

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.