Обучение языковых моделей на собственных данных: извлечение знаний из самостоятельно сгенерированных ошибок

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

June 23, 2023

Авторы: Rishabh Agarwal, Nino Vieillard, Yongchao Zhou, Piotr Stanczyk, Sabela Ramos, Matthieu Geist, Olivier Bachem

cs.AI

Аннотация

Дистилляция знаний (Knowledge Distillation, KD) широко используется для сжатия модели-учителя с целью уменьшения затрат на вывод и объема занимаемой памяти путем обучения более компактной модели-ученика. Однако современные методы KD для авторегрессивных моделей последовательностей страдают от несоответствия распределений между выходными последовательностями, наблюдаемыми во время обучения, и теми, которые генерируются учеником во время вывода. Для решения этой проблемы мы представляем Обобщенную Дистилляцию Знаний (Generalized Knowledge Distillation, GKD). Вместо того чтобы полагаться исключительно на фиксированный набор выходных последовательностей, GKD обучает ученика на его собственных сгенерированных последовательностях, используя обратную связь от учителя на такие последовательности. В отличие от подходов с контролируемой KD, GKD также предоставляет гибкость в использовании альтернативных функций потерь между учеником и учителем, что может быть полезно, когда ученик не обладает достаточной выразительностью для точного воспроизведения распределения учителя. Кроме того, GKD обеспечивает плавную интеграцию дистилляции с тонкой настройкой с использованием обучения с подкреплением (RLHF). Мы демонстрируем эффективность GKD для дистилляции авторегрессивных языковых моделей в задачах суммаризации, перевода и арифметических рассуждений, а также для задач, не зависящих от конкретной задачи, в рамках настройки на инструкции.

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.