On-Policy Destillatie van Taalmodellen: Leren van Zelfgegenereerde Fouten

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

June 23, 2023

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

cs.AI

Samenvatting

Kennisdistillatie (KD) wordt veel gebruikt om een leraarmodel te comprimeren, waardoor de inferentiekosten en het geheugengebruik worden verminderd, door een kleiner studentmodel te trainen. Huidige KD-methoden voor auto-regressieve sequentiemodellen hebben echter last van een distributiemismatch tussen uitvoersequenties die tijdens de training worden gezien en die welke door de student tijdens inferentie worden gegenereerd. Om dit probleem aan te pakken, introduceren we Generalized Knowledge Distillation (GKD). In plaats van uitsluitend te vertrouwen op een vaste set uitvoersequenties, traint GKD de student op zijn zelf-gegenereerde uitvoersequenties door gebruik te maken van feedback van de leraar op dergelijke sequenties. In tegenstelling tot begeleide KD-benaderingen biedt GKD ook de flexibiliteit om alternatieve verliesfuncties tussen de student en de leraar te gebruiken, wat nuttig kan zijn wanneer de student niet de expressiviteit heeft om de distributie van de leraar na te bootsen. Bovendien vergemakkelijkt GKD de naadloze integratie van distillatie met RL-finetuning (RLHF). We demonstreren de effectiviteit van GKD voor het distilleren van auto-regressieve taalmodelen voor samenvattings-, vertaal- en rekenkundige redeneertaken, en taakonafhankelijke distillatie voor instructie-afstemming.

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.