Distillation en politique active des modèles de langage : Apprentissage à partir des erreurs auto-générées

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

Résumé

La distillation de connaissances (KD) est largement utilisée pour compresser un modèle enseignant afin de réduire ses coûts d'inférence et son empreinte mémoire, en entraînant un modèle étudiant plus petit. Cependant, les méthodes actuelles de KD pour les modèles de séquence auto-régressifs souffrent d'un décalage de distribution entre les séquences de sortie observées pendant l'entraînement et celles générées par l'étudiant lors de l'inférence. Pour résoudre ce problème, nous introduisons la Distillation de Connaissances Généralisée (GKD). Au lieu de s'appuyer uniquement sur un ensemble fixe de séquences de sortie, GKD entraîne l'étudiant sur ses propres séquences de sortie générées en exploitant les retours du modèle enseignant sur ces séquences. Contrairement aux approches supervisées de KD, GKD offre également la flexibilité d'utiliser des fonctions de perte alternatives entre l'étudiant et l'enseignant, ce qui peut être utile lorsque l'étudiant manque de la capacité expressive pour imiter la distribution de l'enseignant. De plus, GKD facilite l'intégration transparente de la distillation avec le réglage fin par apprentissage par renforcement (RLHF). Nous démontrons l'efficacité de GKD pour distiller des modèles de langage auto-régressifs sur des tâches de résumé, de traduction et de raisonnement arithmétique, ainsi que pour la distillation agnostique aux tâches dans le cadre de l'ajustement par instructions.

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.