Усиление осведомленности о подкреплении при дистилляции знаний для логического вывода в больших языковых моделях

Reinforcement-aware Knowledge Distillation for LLM Reasoning

February 26, 2026

Авторы: Zhaoyang Zhang, Shuli Jiang, Yantao Shen, Yuting Zhang, Dhananjay Ram, Shuo Yang, Zhuowen Tu, Wei Xia, Stefano Soatto

cs.AI

Аннотация

Последующее обучение с подкреплением (RL) недавно позволило достичь значительного прогресса в крупных языковых моделях (LLM), способных к длинным цепочкам рассуждений (chain-of-thought), однако высокая вычислительная стоимость вывода (inference) таких моделей стимулирует их дистилляцию в модели-ученики меньшего размера. Большинство существующих методов дистилляции знаний (KD) разработаны для обучения с учителем (SFT) и основаны на фиксированных траекториях учителя или регуляризации с использованием дивергенции Кульбака-Лейблера (KL) между учителем и учеником. При сочетании с RL эти подходы часто страдают от рассогласования распределений и интерференции целей: контроль со стороны учителя может не соответствовать эволюционирующему распределению траекторий ученика, а KL-регуляризатор может конфликтовать с максимизацией вознаграждения и требовать тщательного балансирования потерь. Для решения этих проблем мы предлагаем RL-осознанную дистилляцию (RLAD), которая осуществляет выборочное имитирование в процессе RL — направляя ученика в сторону учителя только тогда, когда это улучшает текущее обновление политики. Наш ключевой компонент, дистилляция на основе отношения доверительной области (TRRD), заменяет KL-регуляризатор "учитель-ученик" на целевой функционал, основанный на отношении правдоподобия в стиле PPO/GRPO и заякоренный на смеси политик учителя и старой политики, что обеспечивает учитывающую преимущества (advantage-aware), ограниченную доверительной областью дистилляцию на траекториях ученика и естественным образом балансирует исследование, использование и имитацию. На различных наборах данных для проверки логических и математических рассуждений RLAD последовательно превосходит оффлайн-дистилляцию, стандартный GRPO и основанную на KL дистилляцию знаний "учитель-ученик" на политике (on-policy).

English

Reinforcement learning (RL) post-training has recently driven major gains in long chain-of-thought reasoning large language models (LLMs), but the high inference cost of such models motivates distillation into smaller students. Most existing knowledge distillation (KD) methods are designed for supervised fine-tuning (SFT), relying on fixed teacher traces or teacher-student Kullback-Leibler (KL) divergence-based regularization. When combined with RL, these approaches often suffer from distribution mismatch and objective interference: teacher supervision may not align with the student's evolving rollout distribution, and the KL regularizer can compete with reward maximization and require careful loss balancing. To address these issues, we propose RL-aware distillation (RLAD), which performs selective imitation during RL -- guiding the student toward the teacher only when it improves the current policy update. Our core component, Trust Region Ratio Distillation (TRRD), replaces the teacher-student KL regularizer with a PPO/GRPO-style likelihood-ratio objective anchored to a teacher--old-policy mixture, yielding advantage-aware, trust-region-bounded distillation on student rollouts and naturally balancing exploration, exploitation, and imitation. Across diverse logic reasoning and math benchmarks, RLAD consistently outperforms offline distillation, standard GRPO, and KL-based on-policy teacher-student knowledge distillation.