Distillazione On-Policy di Modelli Linguistici: Apprendimento dagli Errori Auto-Generati

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

June 23, 2023

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

cs.AI

Abstract

La distillazione della conoscenza (KD) è ampiamente utilizzata per comprimere un modello insegnante al fine di ridurre i costi di inferenza e l'impronta di memoria, addestrando un modello studente più piccolo. Tuttavia, gli attuali metodi di KD per modelli sequenziali auto-regressivi soffrono di una mancata corrispondenza tra le sequenze di output osservate durante l'addestramento e quelle generate dallo studente durante l'inferenza. Per affrontare questo problema, introduciamo la Distillazione della Conoscenza Generalizzata (GKD). Invece di fare affidamento esclusivamente su un insieme fisso di sequenze di output, la GKD addestra lo studente sulle sequenze di output auto-generate, sfruttando il feedback dell'insegnante su tali sequenze. A differenza degli approcci di KD supervisionati, la GKD offre anche la flessibilità di impiegare funzioni di perdita alternative tra lo studente e l'insegnante, il che può essere utile quando lo studente non ha l'espressività necessaria per imitare la distribuzione dell'insegnante. Inoltre, la GKD facilita l'integrazione senza soluzione di continuità della distillazione con l'ottimizzazione tramite RL (RLHF). Dimostriamo l'efficacia della GKD per la distillazione di modelli linguistici auto-regressivi in compiti di riassunto, traduzione e ragionamento aritmetico, nonché per la distillazione indipendente dal compito per l'ottimizzazione tramite istruzioni.

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.