$\mathcal{X}$-KD: General Experiential Knowledge Distillation for Large Language Models

1. Problem Statement

Knowledge Distillation (KD) for Large Language Models (LLMs) aims to transfer knowledge from a large "teacher" model to a smaller "student" model. However, existing approaches suffer from a fundamental limitation:

• Behavioral Imitation vs. Experiential Learning: Current methods (e.g., Sequence-level KD, Supervised KD, GKD, MiniLLM) primarily focus on imitating the teacher's behavior (output distributions or token probabilities). They treat the teacher as a static expert to be cloned.
• Missing the "Why": These methods overlook the original learning environment (the reward function) that shaped the teacher's knowledge. By only mimicking the output, students fail to understand the underlying principles or rewards that guided the teacher's decisions.
• Limitations of Existing Solutions:
  • Behavioral Cloning (BC): Simple imitation often leads to distribution mismatch and poor generalization.
  • MiniLLM: While it incorporates Inverse Reinforcement Learning (IRL) concepts to model the teacher's reward, it relies on policy gradient methods. This deviates from the standard supervised learning paradigm, introduces training instability, and limits compatibility with existing LLM training frameworks and various divergence functions.

2. Methodology: Experiential Knowledge Distillation (X-KD)

The authors propose X-KD, a framework that enables student models to learn in the teacher's original learning environment by modeling the teacher's underlying reward function.

Core Theoretical Foundation

• Inspiration: The method is inspired by Experiential Learning Theory (learning by experiencing the same environment) and Bayesian Inverse Reinforcement Learning (IRL).
• Framework: X-KD utilizes the Approximated Variational Reward Imitation Learning (AVRIL) framework. AVRIL allows for the joint learning of a reward function and a policy.
• Goal: Instead of just minimizing the divergence between student and teacher outputs, X-KD encourages the student to maximize the expected reward of the teacher's original environment.

Technical Formulation

The approach models the Natural Language Generation (NLG) task as a Markov Decision Process (MDP).

1. Original Reward Modeling (ORM):
   • The teacher's reward function $R_\pi$ is unknown. X-KD approximates the posterior distribution of this reward, $p(R|\pi)$, using a parameterized distribution $q_\phi(R)$.
   • It optimizes an Evidence Lower Bound (ELBO) to learn this reward distribution jointly with the student policy.
2. Loss Function Decomposition:
   The total loss is decomposed into two parts:
   • Knowledge Distillation Term ($L_{seq}$ or $L_{kd}$): Standard distillation objectives (Sequence-level, Forward-KL, or Generalized KD) that ensure the student mimics the teacher's output distribution.
   • Experiential Regularization Term ($L_{ex}$): A novel term derived from AVRIL that enforces consistency with the learned reward function. It involves:
     • Minimizing the KL divergence between the learned reward distribution and a prior.
     • Enforcing a Temporal Difference (TD) error constraint, ensuring the student's Q-values (estimated via the policy) align with the learned reward structure.

Generalization and Flexibility

X-KD is designed as a general framework that can be integrated into various existing KD pipelines:

• Sequence-Level X-KD: Integrates with standard sequence-level distillation (maximizing likelihood of teacher sequences).
• Supervised X-KD: Integrates with Forward-KL distillation (token-level divergence).
• Generalized X-KD: Integrates with GKD (using $\beta$-JS divergence and on-policy data).
• Black-Box Support: A specific pipeline (Algorithm 2) allows X-KD to work with black-box teachers where only final outputs are accessible, using sequence-level objectives without requiring access to the teacher's internal logits.

Implementation Simplicity

Unlike MiniLLM, X-KD maintains the supervised learning paradigm. The Q-value decoder is inversely formulated as a log-softmax of the student's action probabilities, allowing it to be trained using standard gradient descent without complex policy gradient stabilization techniques.

3. Key Contributions

1. Novel Framework: Proposes X-KD, the first general KD framework that explicitly models the teacher's original learning environment (reward function) rather than just the behavior.
2. Theoretical Unification: Demonstrates that X-KD can be derived from the AVRIL framework and seamlessly integrated into sequence-level, supervised, and generalized KD methods while retaining the simplicity of supervised learning.
3. Empirical Superiority: Shows that X-KD outperforms strong baselines (including GKD and MiniLLM) across multiple tasks.
4. Efficiency and Trade-offs: Achieves better performance-diversity trade-offs and higher data efficiency compared to existing methods.

4. Experimental Results

The authors evaluated X-KD on three NLG tasks: Abstractive Summarization (XSum), Machine Translation (WMT14 en-de), and Arithmetic Reasoning (GSM8k).

• Performance:
  • X-KD variants (SeqXKD, S-XKD, G-XKD) consistently outperformed their non-experiential baselines (SeqKD, SKD, GKD) and the concurrent MiniLLM in most settings.
  • G-XKD (Generalized X-KD) achieved the best overall results, surpassing GKD and MiniLLM on summarization and translation.
• Performance-Diversity Trade-off:
  • X-KD demonstrated superior robustness. As sampling temperature increased (higher diversity), X-KD maintained higher quality scores (ROUGE/BLEU) compared to baselines, whose performance dropped more sharply. This suggests X-KD learns more generalizable strategies rather than overfitting to specific teacher outputs.
• Data Efficiency:
  • X-KD achieved comparable performance to baselines trained on 100% of the data while using only 50% to 75% of the training data. This indicates that modeling the reward function helps the student capture fundamental task principles more efficiently.
• Black-Box Distillation:
  • In black-box settings (using Llama-2-7b as teacher), X-KD (SeqXKD) achieved a higher win rate against the baseline (SeqKD) when evaluated by GPT-4.

5. Significance and Impact

• Paradigm Shift: X-KD shifts the KD paradigm from "copying the teacher" to "learning the environment the teacher learned in." This aligns machine learning more closely with human experiential learning theories.
• Practicality: By avoiding policy gradient methods and maintaining a supervised learning structure, X-KD is easier to integrate into existing LLM training pipelines than MiniLLM.
• Resource Efficiency: The improved data efficiency is crucial for LLMs, where training data and computational resources are expensive.
• Limitations: The authors acknowledge that X-KD is sensitive to the tuning of the "experiential weight" hyperparameter ($\lambda$) and that while black-box applicability is demonstrated, it requires further systematic study.

In conclusion, X-KD provides a robust, flexible, and theoretically grounded approach to distilling LLMs that significantly improves model quality, diversity, and training efficiency by leveraging the concept of experiential learning.