KETCHUP: K-Step Return Estimation for Sequential Knowledge Distillation

2025-04-26

AI Summary

This work addresses the high gradient variance and training instability of reinforcement learning (RL)-based knowledge distillation (KD) for text generation—particularly with large-scale student models. We propose a sequence-level reward estimation method grounded in the K-step Bellman optimality equation, the first to incorporate K-step returns into RL-based KD frameworks. Theoretically, we prove that this approach substantially reduces the variance of policy gradient estimates. By modeling multi-step reward signals along the student’s autoregressive generation process and integrating them into policy gradient optimization, our method enhances both training stability and generalization. Empirically, it achieves state-of-the-art performance across three text generation tasks—abstractive summarization, dialogue generation, and data-to-text generation—under both automatic metrics (e.g., BLEU, ROUGE) and large language model–based human evaluation (LLM-as-a-judge), with especially pronounced gains for large-parameter student models.

Technology Category

• Machine Learning: Reinforcement Learning
• Natural Language Processing: Generation
• Search and Optimization: Learning to Search

Application Category

• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs
• Economics, Online Markets and Human Computation: Cost models of using LLMs in production systems

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.

Problem

• Improves RL-based knowledge distillation in text generation
• Reduces gradient variance via K-step return estimation
• Enhances performance in LLM evaluation metrics

Innovation

• K-step return estimation using Bellman Optimality Equation
• Reduces gradient estimate variance for better optimization
• Enhances RL-based knowledge distillation in LLMs