Destilación de Modelos de Lenguaje en Política Activa: Aprendiendo de Errores Autogenerados
On-Policy Distillation of Language Models: Learning from Self-Generated Mistakes
June 23, 2023
Autores: Rishabh Agarwal, Nino Vieillard, Yongchao Zhou, Piotr Stanczyk, Sabela Ramos, Matthieu Geist, Olivier Bachem
cs.AI
Resumen
La destilación de conocimiento (KD, por sus siglas en inglés) se utiliza ampliamente para comprimir un modelo maestro con el fin de reducir su costo de inferencia y huella de memoria, entrenando un modelo estudiante más pequeño. Sin embargo, los métodos actuales de KD para modelos de secuencias auto-regresivos sufren de un desajuste en la distribución entre las secuencias de salida observadas durante el entrenamiento y aquellas generadas por el estudiante durante la inferencia. Para abordar este problema, introducimos la Destilación de Conocimiento Generalizada (GKD, por sus siglas en inglés). En lugar de depender únicamente de un conjunto fijo de secuencias de salida, GKD entrena al estudiante en sus propias secuencias de salida generadas, aprovechando la retroalimentación del maestro sobre dichas secuencias. A diferencia de los enfoques de KD supervisados, GKD también ofrece la flexibilidad de emplear funciones de pérdida alternativas entre el estudiante y el maestro, lo que puede ser útil cuando el estudiante carece de la capacidad expresiva para imitar la distribución del maestro. Además, GKD facilita la integración fluida de la destilación con el ajuste fino mediante aprendizaje por refuerzo (RLHF). Demostramos la eficacia de GKD para destilar modelos de lenguaje auto-regresivos en tareas de resumen, traducción y razonamiento aritmético, así como en la destilación independiente de tareas para el ajuste por instrucciones.
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.