We propose a novel k-step return estimation method (called KETCHUP) for Reinforcement Learning(RL)-based knowledge distillation (KD) in text generation tasks. Our idea is to induce a K-step return by using the Bellman Optimality Equation for multiple steps. Theoretical analysis shows that this K-step formulation reduces the variance of the gradient estimates, thus leading to improved RL optimization especially when the student model size is large. Empirical evaluation on three text generation tasks demonstrates that our approach yields superior performance in both standard task metrics and large language model (LLM)-based evaluation. These results suggest that our K-step return induction offers a promising direction for enhancing RL-based KD in LLM research.
Figures & tables
Figure 1: Average predicted return of different approaches.
Model
XSum
Europarl EN-NL
GSM8K
ROUGE-1 ↑
ROUGE-2 ↑
ROUGE-L ↑
BLEU4 ↑
chrF ↑
TER ↓
Accuracy(%) ↑
Teacher
41.32
18.86
33.79
25.36
51.11
63.17
40.71
Student
19.60
3.19
13.72
0.95
24.80
100.21
0.00
Distilled Student
SeqKD Kim and Rush (2016)
33.54
11.90
26.67
22.09
48.33
66.18
20.02
KL Hinton et al. (2015)
34.36
12.86
27.38
22.35
48.58
65.93
23.96
JS Wen et al. (2023)
34.87
13.18
27.84
22.55
48.71
65.74
24.72
Table 1: Main results on XSum, Europarl EN–NL, and GSM8K datasets. The best student result is in bold . ↑/↓ The higher/lower, the better. We prompt the teacher and off-the-shelf student in a zero-shot manner to gain the first two rows. We select the best checkpoint based on the performance of the held-out validation set and report the performance of these checkpoints on the test set for all distilled students.
Figure 2: Variance and bias with different K values.
Figure 3: Learning curves with different K values and model sizes, where the x -axis is the number of training steps.
Dataset
Method
Overall
Informativeness
Coherence
XSum
TVD
67.50%
68.15%
65.90%
LLMR
69.95%
70.55%
66.30%
Ketchup
73.50%
73.90%
70.40%
Europarl
TVD
53.80%
54.15%
54.85%
LLMR
56.45%
55.85%
56.30%
Ketchup
58.85%
57.95%
58.45%
Table 2: LLM-based evaluation on the summarization and translation tasks. Europarl refers to the Europarl EN–NL dataset. We show the winning rates of each method over the KL distillation baseline in terms of overall quality, informativeness, and coherence.
Appendix figures & tables6 assets
Supplementary material from the paper’s appendix.
Appendix
Dataset
Task
# of Samples
Train
Dev
Test
XSum ( Narayan et al., 2018 )
Summarization
202,926
11,332
11,333
Europarl EN-NL ( Koehn, 2005 )
Machine Translation
1,167,808
10,014
10,016
GSM8K ( Cobbe et al., 2021 )
Arithmetic reasoning
6,705
768
1,319
Appendix
Table 3: Statistics of our datasets.
Hyperparameter
Value
Training Epochs
3
Train Batch size
8
Eval Batch size
32
Optimizer
AdamW
Grad Accumulation Steps
32
Reward Clip Range
[-100, 100]
Appendix
Table 4: Hyperparameter Details for experiments on RL-based approaches ( Ketchup , LLMR, LLMR with mean baseline, and LLMR with min-variance baseline).
Hyperparameter
Value
Training Epochs
2
Train Batch size
32
Eval Batch size
32
Optimizer
AdamW
Grad Accumulation Steps
16
Dropout
0.25
Appendix
Table 5: Hyperparameter Details for experiments on divergence-based KD approaches (seqKD, KL, JSD, TVD).
Model
XSum (ROUGE-1 ↑ )
Europarl (BLEU4 ↑ )
GSM8K (Acc. (%) ↑ )
Teacher (Qwen1.5-4B)
38.15
21.32
42.08
Student (Qwen1.5-0.5B)
8.80
0.02
0.00
KL Hinton et al. (2015)
31.29
15.76
26.31
TVD Wen et al. (2023)
31.18
16.22
26.99
LLMR Li et al. (2024)
31.61
15.90
27.29
K etchup
32.28
16.46
28.13
Appendix
Table 6: Distillation results on XSum, Europarl EN–NL, and GSM8K using Qwen1.5 models. Higher ↑ is better. The best K values are 2, 2, and 16 for the three datasets, respectively.
Please evaluate the overall quality of the following summaries given the document.
Evaluation Criteria:
Overall Quality: A good summary should be both precise and concise, summarizing the most important points in the given document,
without including unimportant or irrelevant details
Document: [Source]
Summary [ID1] : [Summary-A]
Summary [ID2] : [Summary-B]
Appendix
Table 7: Prompt templates for LLM evaluation on the summarization task in terms of overall quality, informativeness, and coherence. Here, “ Source ” is the document to be summarized. The choices of IDs are “A” and “B”; “ Summary-A ” and “ Summary-B ” are replaced with model-generated texts. Since LLMs are not robust to ID and order ( Zheng et al., 2023 ; Shen et al., 2023 ) , we enumerate different combinations for a given pair, resulting in four LLM queries.
Please evaluate the overall quality of the following translations from English to Dutch.
Evaluation Criteria:
Overall Quality: A good translation should: 1) faithfully reflect the meaning of the source text; 2) avoid adding unnecessary or irrelevant
details. 3) use natural and fluent Dutch.
Source: [Source]
Translation [ID1] : [Translation-A]
Translation [ID2] : [Translation-B]
Appendix
Table 8: Prompt templates for LLM evaluation on the machine translation task in terms of overall quality, informativeness, and coherence. Here, “ Source ” is the source sentence to be translated. The choices of IDs are “A” and “B”; “ Translation-A ” and “ Translation-B ” are replaced with model-generated texts. We still enumerate different combinations for a given pair, resulting in four LLM queries.
Knowledge distillation (KD) is a key technique for compressing Large Language Models (LLMs), yet methods relying on a single KL objective often fail to balance primary distribution fitting with long-tail probability modeling, limiting both generation quality and generalization. To address this, we analyze the complementary roles of forward and reverse KL divergence (FKL/RKL) in distribution alignment from theoretical and empirical perspectives. We then propose a reinforcement-learning-based adaptive KL-weighted distillation framework, in which a policy network dynamically assigns weights to FKL and RKL based on teacher-student distributional characteristics, guided by immediate reward signals to achieve dual alignment on principal and long-tail modes. Extensive experiments demonstrate consistent improvements across Rouge-L and BertScore metrics, surpassing greedy heuristics by 0.4-0.6 points and outperforming other baseline methods on diverse benchmarks.
Zilong Liu, Xuewen Zhang, Jinrui Xing +3
Li Auto, Beijing, China · School of Software and Microelectronics, Peking University, China
Large language models have become increasingly popular and demonstrated remarkable performance in various natural language processing (NLP) tasks. However, these models are typically computationally expensive and difficult to be deployed in resource-constrained environments. In this paper, we propose LLMR, a novel knowledge distillation (KD) method based on a reward function induced from large language models. We conducted experiments on multiple datasets in the dialogue generation and summarization tasks. Empirical results demonstrate that our LLMR approach consistently outperforms traditional KD methods in different tasks and datasets.
Dongheng Li, Yongchang Hao, Lili Mou
Dept. Computing Science & Alberta Machine Intelligence Institute (Amii), University of Alberta · Canada CIFAR AI Chair, Amii
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.