Внутриполитическая дистилляция в доверительной области

Аннотация

Дистилляция по политике (On-Policy Distillation, OPD) является фундаментальным методом эффективного пост-тренинга больших языковых моделей (LLM) с широким спектром применений в обучении агентов, мультизадачном улучшении и сжатии моделей. Однако обучение OPD становится нестабильным при значительном расхождении распределений учителя и студента, поскольку супервизия учителя на токенах, сгенерированных студентом, может давать ненадёжные градиенты политики и даже приводить к сбою оптимизации. Данная работа решает проблему надёжной супервизии на уровне токенов по политике с помощью стратегий назначения кредитов и предлагает дистилляцию по политике в доверительной области — TrOPD. Она обладает следующими характеристиками: 1) Обучение по политике в доверительной области: TrOPD выполняет OPD только в областях, где учитель предоставляет надёжную супервизию, смягчая трудности оптимизации оценки обратного KL-расхождения K1 при несоответствии распределений. 2) Оценка выбросов: Для областей-выбросов мы исследуем клиппирование градиентов, маскирование и оценку прямого KL-расхождения, чтобы уменьшить негативное влияние ненадёжной супервизии. 3) Руководство вне политики: Студент продолжает генерацию с префиксов учителя и использует прямое KL-расхождение для имитации руководства вне политики, поощряя исследование по политике в сторону надёжных областей. Эксперименты показывают, что TrOPD стабильно превосходит современные базовые методы OPD, включая OPD, EOPD и REOPOLD, в задачах математического рассуждения, генерации кода и бенчмарках общего назначения.

English

On-Policy Distillation (OPD) is a fundamental technique for efficient post-training of large language models (LLMs), with broad applications in agent learning, multi-task enhancement, and model compression. However, OPD training becomes unstable when the teacher and student distributions differ substantially, as teacher supervision on student-generated tokens may yield unreliable policy gradients and even cause optimization failure. This work addresses reliable on-policy token-level supervision through credit assignment strategies, and proposes Trust Region On-Policy Distillation, TrOPD. It features the following characteristics: 1) Trust-Region On-Policy Learning: TrOPD performs OPD only in regions where the teacher provides reliable supervision, mitigating the optimization difficulty of the K1 reverse-KL estimator under distribution mismatch. 2) Outlier Estimation: For outlier regions, we explore gradient clipping, masking, and forward-KL estimation to reduce the adverse effects of unreliable supervision. 3) Off-Policy Guidance: The student continues generation from teacher prefixes and uses forward KL to imitate off-policy guidance, encouraging on-policy exploration toward reliable regions. Experiments show that TrOPD consistently outperforms SoTA OPD baselines, including OPD, EOPD, and REOPOLD, across mathematical reasoning, code generation, and general-domain benchmarks.