cs.CLApr 26, 2025
SaveKETCHUP: K-Step Return Estimation for Sequential Knowledge Distillation
Organizations: Dept. Computing Science, Alberta Machine Intelligence Institute (Amii), University of Alberta, Canada · Canada CIFAR AI Chair, Amii
Abstract
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 values.
Figure 3: Learning curves with different values and model sizes, where the -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 |
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] |
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 |
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 |
| etchup | 32.28 | 16.46 | 28.13 |
Table 6: Distillation results on XSum, Europarl EN–NL, and GSM8K using Qwen1.5 models. Higher is better. The best 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] |
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] |
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.