Knowledge Distillation in Large Language Models: Mathematical Foundations of Forward vs. Reverse KL Divergence, Sequence-Level Distillation (SeqKD), and On-Policy Mode-Seeking Objectives

Deploying frontier large language models in latency-critical and compute-constrained environments remains an active challenge in artificial intelligence systems. While parameter scale correlates directly with downstream generalization and emergent reasoning capabilities, the memory footprint and FLOP intensity of dense multi-hundred-billion parameter architectures impose prohibitive operational costs. Post-training compression strategies such as quantization and structural pruning reduce precision and dimension, but knowledge distillation (KD) provides a foundational mechanism for transferring the learned functional mapping and representational capacity of a large teacher model into an efficient student network.

In classical vision and discriminative architectures, knowledge distillation relies primarily on matching output probability distributions across a fixed label space. In autoregressive generative language models, however, the target space consists of an exponentially large combinatorial sequence domain. The mathematical formulation of the distillation objective, specifically the choice of divergence measure, sampling policy, and sequence-level versus token-level aggregation, dictates whether the student successfully compresses the teacher capabilities or succumbs to distribution shift, exposure bias, and degenerate hallucinations.

1. Mathematical Foundations of Token-Level Forward KL Distillation

The canonical knowledge distillation framework introduced by Hinton et al. (2015) transfers dark knowledge from a teacher model with parameters $\phi$ to a student model with parameters $\theta$ through soft target probability distributions scaled by a temperature hyperparameter $T$.

Let $\mathcal{V}$ denote the discrete vocabulary of size $|\mathcal{V}|$. Given an input prompt sequence $x = (x_1, x_2, \dots, x_N)$ and a target completion sequence $y = (y_1, y_2, \dots, y_T)$, the autoregressive generative probability of the sequence is factorized under the chain rule of probability:

p(y | x) = \prod_{t=1}^{T} p(y_t | y_{<t}, x)

At each decoding step $t$, the teacher and student networks produce unnormalized logit vectors $z_t^{(T)} \in \mathbb{R}^{|\mathcal{V}|}$ and $z_t^{(S)} \in \mathbb{R}^{|\mathcal{V}|}$, respectively. The temperature-scaled probability distribution over the vocabulary is given by:

p_\phi(v | y_{<t}, x; T) = exp(z_{t, v}^{(T)} / T) / \sum_{k \in V} exp(z_{t, k}^{(T)} / T)
p_\theta(v | y_{<t}, x; T) = exp(z_{t, v}^{(S)} / T) / \sum_{k \in V} exp(z_{t, k}^{(S)} / T)

In token-level forward Kullback-Leibler (KL) distillation, the training objective minimizes the cross-entropy between the soft teacher distribution and the soft student distribution at every auto-regressive step along ground-truth or teacher-generated trajectories:

L_{token-KD}(\theta) = E_{(x, y) ~ D} [ \sum_{t=1}^{T} T^2 * D_{KL}( p_\phi(\cdot | y_{<t}, x; T) || p_\theta(\cdot | y_{<t}, x; T) ) ]

Expanding the definition of the Kullback-Leibler divergence:

D_{KL}(p_\phi || p_\theta) = \sum_{v \in V} p_\phi(v | y_{<t}, x; T) * log( p_\phi(v | y_{<t}, x; T) / p_\theta(v | y_{<t}, x; T) )

Because the term $\sum_{v \in \mathcal{V}} p_{\phi}(v) \log p_{\phi}(v)$ is independent of the student parameters $\theta$, minimizing the KL divergence is mathematically equivalent to minimizing the cross-entropy with soft target distributions:

L_{CE-soft}(\theta) = - \sum_{t=1}^{T} \sum_{v \in V} p_\phi(v | y_{<t}, x; T) * log( p_\theta(v | y_{<t}, x; T) )

Gradient Properties and High-Temperature Asymptotics

To understand the mechanics of dark knowledge transfer, consider the gradient of the loss with respect to the pre-softmax student logits $z_{t, i}^{(S)}$ at decoding position $t$:

\partial L_{token-KD} / \partial z_{t, i}^{(S)} = T * ( p_\theta(i | y_{<t}, x; T) - p_\phi(i | y_{<t}, x; T) )

When the temperature parameter $T$ is high relative to the magnitude of the logits, the exponential terms can be approximated by a first-order Taylor series expansion:

exp(z_{t, i} / T) \approx 1 + (z_{t, i} / T)

Assuming centered logits such that $\sum_{k \in \mathcal{V}} z_{t, k} = 0$, the softmax probability simplifies to:

p(i | y_{<t}, x; T) \approx (1 + z_{t, i} / T) / (|V| + \sum_k z_{t, k} / T) = (1 / |V|) + (z_{t, i} / (|V| * T))

Substituting this approximation into the gradient expression yields:

\partial L_{token-KD} / \partial z_{t, i}^{(S)} \approx T * ( (z_{t, i}^{(S)} - z_{t, i}^{(T)}) / (|V| * T) ) = (z_{t, i}^{(S)} - z_{t, i}^{(T)}) / |V|

At high temperatures, the distillation loss minimizes the mean squared error between teacher and student logit representations directly, penalizing relative differences between non-target tokens. At lower temperatures ($T \to 1$), the gradient focuses heavily on the top-ranking tokens, matching the categorical distribution.

The Logit Communication and Memory Wall

While token-level soft target distillation captures relative entropy across lexical classes, it presents an engineering bottleneck in modern distributed training setups. Vocabulary sizes in state-of-the-art tokenizers range from 128,000 (such as Llama 3) to 256,000 tokens (such as Gemma and Qwen). Transmitting full logit matrices of shape [batch_size, seq_len, |V|] across tensor-parallel and pipeline-parallel interconnects induces massive memory overhead and network bandwidth saturation.

To mitigate this, industrial implementations utilize top-$K$ or nucleus top-$p$ truncated logit distillation:

\tilde{p}_\phi(v) = p_\phi(v) / \sum_{k \in K_t} p_\phi(k)   if v \in K_t, else 0

where $\mathcal{K}_t \subset \mathcal{V}$ contains only the top $K$ most probable tokens ($K \in [8, 64]$), slashing inter-node communication volume by over 99%.

2. Sequence-Level Knowledge Distillation (SeqKD)

To bypass the memory and compute bottlenecks of dense token-level logit transfer, Kim & Rush (2016) introduced Sequence-Level Knowledge Distillation (SeqKD). Instead of transferring per-step distribution slices, SeqKD optimizes the student against the teacher complete sequence-level probability distribution over the combinatorial space of all sequences $\mathcal{Y} = \mathcal{V}^*$.

The sequence-level forward KL objective is formulated as:

L_{seq}(\theta) = E_{x ~ D} [ D_{KL}( p_\phi(\cdot | x) || p_\theta(\cdot | x) ) ]
               = - E_{x ~ D} [ \sum_{y \in Y} p_\phi(y | x) * log p_\theta(y | x) ] + const

Summing over all possible sequences $\mathcal{Y}$ is computationally intractable due to exponential complexity $\mathcal{O}(|\mathcal{V}|^T)$. SeqKD resolves this by approximating the full teacher distribution $p_{\phi}(y \mid x)$ with its mode (mode approximation):

p_\phi(y | x) \approx \delta(y - \hat{y})

where $\hat{y}$ is the sequence generated by the teacher via greedy decoding or beam search:

\hat{y} = \arg\max_{y \in Y} p_\phi(y | x) \approx BeamSearch(p_\phi, x)

Under this Dirac delta approximation, the sequence-level distillation loss reduces to standard negative log-likelihood (cross-entropy) over pseudo-labeled teacher rollouts:

L_{SeqKD}(\theta) = - E_{x ~ D} [ \sum_{t=1}^{|\hat{y}|} log p_\theta(\hat{y}_t | \hat{y}_{<t}, x) ]

SeqKD offers three decisive operational advantages:

• Decoupled Architecture: The student model does not need to share a tokenizer, hidden dimension, or vocabulary alignment with the teacher network.
• Offline Data Generation: Teacher sequences can be pre-generated, filtered, and cached, decoupling teacher inference from student backpropagation.
• Complexity Reduction: Beam-searched teacher outputs exhibit lower conditional entropy and simpler syntactic structures than raw human web scrape text, accelerating student convergence.

3. Information-Theoretic Divergence Dynamics: Forward KL vs. Reverse KL

The choice of mathematical divergence objective exerts a structural influence on the behavior of distilled generative models. In probabilistic machine learning, the directionality of the Kullback-Leibler divergence determines whether the fitted distribution exhibits mode-covering or mode-seeking characteristics.

+--------------------------------------------------------------------------------------------------+
|                   Information Divergence Property Comparison in Language Models                  |
+--------------------------------------------------------------------------------------------------+
| Feature                | Forward KL: D_KL(p_teacher || p_student) | Reverse KL: D_KL(p_student || p_teacher) |
|------------------------+------------------------------------------+------------------------------------------|
| Mathematical Formula   | \sum_y p_teacher(y) * log(p_T / p_S)     | \sum_y p_student(y) * log(p_S / p_T)     |
| Statistical Property   | Zero-avoiding / Mode-covering            | Zero-forcing / Mode-seeking              |
| Asymptotic Penalty     | Infinite when p_S(y) -> 0 and p_T(y) > 0 | Infinite when p_S(y) > 0 and p_T(y) -> 0 |
| Capacity Mismatch Risk | Blurry distributions, hallucinations     | Selective mode dropping, sharp peaks     |
| Sampling Trajectory    | Off-policy (teacher rollouts / dataset)  | On-policy (student self-rollouts)        |
| Optimization Engine    | Supervised Maximum Likelihood / SFT      | Policy Gradients / RL Fine-Tuning        |
+--------------------------------------------------------------------------------------------------+

Forward KL and the Mode-Covering Pathology

In Forward KL minimization:

D_{KL}(p_\phi || p_\theta) = \sum_{y \in Y} p_\phi(y | x) * log( p_\phi(y | x) / p_\theta(y | x) )

The objective is weighted by $p_{\phi}(y \mid x)$. If the teacher assigns non-zero probability to a sequence $y$ ($p_{\phi}(y \mid x) > 0$), the student must assign non-zero probability ($p_{\theta}(y \mid x) > 0$) to avoid the ratio diverging to infinity. Consequently, Forward KL is zero-avoiding.

When a student model possesses significantly fewer parameters or narrower representational bandwidth than the teacher, it lacks the capacity to represent the rich multimodal distribution of the teacher. Forced by the zero-avoiding penalty to cover all modes, the student averages across distinct trajectories. In natural language generation, this distribution averaging forces probability mass into improbable intermediate states, yielding grammatical drift, incoherent hallucinations, and soft, uncommitted probability distributions.

Reverse KL and Mode-Seeking Concentration

In Reverse KL minimization:

D_{KL}(p_\theta || p_\phi) = \sum_{y \in Y} p_\theta(y | x) * log( p_\theta(y | x) / p_\phi(y | x) )

The objective is weighted by the student own output probability $p_{\theta}(y \mid x)$. If the student assigns probability to a region where the teacher density is near zero ($p_{\phi}(y \mid x) \to 0$), the term diverges to infinity. Conversely, if $p_{\phi}(y \mid x) > 0$ but $p_{\theta}(y \mid x) = 0$, the expression evaluates to $0 \log 0 = 0$.

Hence, Reverse KL is zero-forcing. The student is not penalized for ignoring entire modes of the teacher distribution, provided that the modes it does choose to represent are strictly supported by the teacher. In generative modeling, the student concentrates its limited capacity on the highest-confidence pathways of the teacher, producing sharp, accurate, and faithful generations.

4. On-Policy Distillation and Exposure Bias (MiniLLM and GKD)

Supervised knowledge distillation algorithms (both token-level KD and SeqKD) suffer from exposure bias and distribution mismatch. During training, the loss is computed along trajectories $y$ sampled from either the ground-truth dataset or the teacher model. During inference, however, the student generates tokens autoregressively conditioned on its own prior predictions:

\hat{y}_t \sim p_\theta(\cdot \mid \hat{y}_{<t}, x)

If the student makes an early compounding error, it enters a state space out-of-distribution (OOD) relative to the training trajectories, triggering cascading failures.

To solve distribution mismatch, modern frameworks employ On-Policy Distillation (OPD), where the training distribution matches the student inference rollout policy.

       [ Context Prompt x ]
               │
       ┌───────┴───────┐
       ▼               ▼
 [ Student Policy ]   [ Teacher Policy ]
  q_θ(y | x) rollouts    p_φ(y | x) evaluation
       │               │
       └───► [ Reverse KL / GKD Loss ] ◄───┘
               │
        [ Policy Gradient Update ]
               ▼
        [ Updated Student θ ]

MiniLLM: Formulations and Policy Gradients

Gu et al. (2024) formulated MiniLLM, an on-policy distillation framework minimizing the sequence-level Reverse KL divergence between student $q_{\theta}$ and teacher $p$:

L_{MiniLLM}(\theta) = E_{x ~ D} [ D_{KL}( q_\theta(\cdot | x) || p(\cdot | x) ) ]
                    = E_{x ~ D, y ~ q_\theta(\cdot | x)} [ log( q_\theta(y | x) / p(y | x) ) ]

Because the expectation is taken over the parameter-dependent student distribution $y \sim q_{\theta}(\cdot \mid x)$, we cannot differentiate the objective with standard backpropagation. Applying the log-derivative trick (REINFORCE algorithm):

\nabla_\theta L_{MiniLLM}(\theta) = E_{x ~ D} [ \sum_{y \in Y} \nabla_\theta q_\theta(y | x) * log( q_\theta(y | x) / p(y | x) ) + q_\theta(y | x) * \nabla_\theta( log q_\theta(y | x) - log p(y | x) ) ]

Since $\sum_{y} q_{\theta}(y \mid x) \nabla_{\theta} \log q_{\theta}(y \mid x) = \sum_{y} \nabla_{\theta} q_{\theta}(y \mid x) = \nabla_{\theta}(1) = 0$, and using $\nabla_{\theta} q_{\theta} = q_{\theta} \nabla_{\theta} \log q_{\theta}$:

\nabla_\theta L_{MiniLLM}(\theta) = E_{x ~ D, y ~ q_\theta(\cdot | x)} [ \nabla_\theta log q_\theta(y | x) * ( log( q_\theta(y | x) / p(y | x) ) + 1 ) ]

In practice, to stabilize policy gradient variance, MiniLLM employs token-level credit assignment with baseline subtraction and length normalization:

\nabla_\theta L(\theta) = - E_{x ~ D, y ~ q_\theta(\cdot | x)} [ \sum_{t=1}^{T} \nabla_\theta log q_\theta(y_t | y_{<t}, x) * R_t ]

where the per-token pseudo-reward $R_t$ reflects the relative advantage of the token under the teacher distribution:

R_t = ( log( p(y_t | y_{<t}, x) / q_\theta(y_t | y_{<t}, x) ) ) - b(y_{<t}, x)

Generalized Knowledge Distillation (GKD)

Agarwal et al. (2024) introduced Generalized Knowledge Distillation (GKD), unifying supervised KD and on-policy optimization into an interpolatable continuum. GKD controls trajectory sampling through a mixture distribution $\pi_{\text{mix}}$:

y \sim \pi_{mix}(\cdot | x) = \lambda * q_\theta(\cdot | x) + (1 - \lambda) * p_{data}(\cdot | x)

where $\lambda \in [0, 1]$ parameterizes the fraction of student on-policy rollouts versus dataset ground truth.

Furthermore, GKD generalizes the loss to arbitrary divergences $D_f$ at the token level along the generated trajectories:

L_{GKD}(\theta) = E_{x ~ D, y ~ \pi_{mix}(\cdot | x)} [ \sum_{t=1}^{|y|} D_f( p(\cdot | y_{<t}, x) || q_\theta(\cdot | y_{<t}, x) ) ]

Setting $\lambda = 0$ with Forward KL recovers classical token-level KD. Setting $\lambda = 1$ with Reverse KL or Jensen-Shannon Divergence (JSD) executes fully on-policy exploration, directly teaching the student how to recover from self-generated errors.

5. Reasoning Trajectory Distillation and Long-Form Chain-of-Thought

A major breakthrough in open-weight models has been the distillation of reasoning traces from frontier reasoning architectures (exemplified by DeepSeek-R1-Distill-Qwen and DeepSeek-R1-Distill-Llama).

Instead of distilling standard question-answer pairs, reasoning distillation transfers extended exploration paths containing intermediate verification steps, backtracking, and algorithmic decomposition.

Frontier Reasoning Model (Teacher)
               │
               ▼  [ High-Compute Search & RL Verification ]
  [ Long-Form Chain-of-Thought Trajectories: <think> ... </think> ]
               │
               ▼  [ Rejection Sampling & Verification Filtering ]
  [ Verified Accurate Reasoning Traces ]
               │
               ▼  [ Sequence-Level SFT / On-Policy GKD ]
    Compact Dense Student Model (1.5B - 70B)

Why Reasoning Distillation Overcomes the Capacity Wall

Under classical task settings, distilling a 671B mixture-of-experts model into a 1.5B dense student results in substantial capability loss. In reasoning distillation, however, smaller students achieve competitive performance on benchmarks like MATH-500 and AIME. The mathematical explanation lies in test-time compute substitution:

1. Explicit Search Path Linearization: The teacher amortizes internal latent search into explicit tokenized reasoning tokens. The student learns the procedural transition probabilities:

p(r_{k+1} | r_{\le k}, x)

converting implicit hidden-state search into sequential autoregressive inference.

1. Verification Filtering: By executing rejection sampling over teacher traces:

D_{distill} = { (x, y_{CoT}) | Verify(y_{CoT}, x) = 1 }

the training distribution eliminates low-density error paths, presenting the student with a unimodal target manifold well within its representational capacity.

6. Comprehensive Algorithmic Comparison

The following matrix compares primary knowledge distillation strategies across their mathematical objectives, implementation complexities, and operational properties:

+----------------------------------------------------------------------------------------------------------------------+
|                                     LLM Knowledge Distillation Framework Comparison                                  |
+----------------------------------------------------------------------------------------------------------------------+
| Framework         | Mathematical Objective     | Target Space        | Sampling Policy    | Operational Cost | Use Case        |
|-------------------+----------------------------+---------------------+--------------------+------------------+-----------------|
| Token-Level KD    | \sum_t D_KL(p_T || p_S)    | Full soft logits    | Teacher/Dataset    | Very High (I/O)  | Same-tokenizer  |
| SeqKD (Kim & Rush)| -\sum_t log p_S(\hat{y}_t) | Teacher mode \delta | Teacher beam/greedy| Low (SFT format) | Synthetic data  |
| MiniLLM           | D_KL(p_S || p_T)           | Teacher likelihood  | Student rollouts   | High (Policy RL) | Mode-seeking RL |
| GKD               | D_f(p_T || p_S) on \pi_mix | Arbitrary divergence| Hybrid student/data| Medium-High      | Instruction KD  |
| CoT Reasoning KD  | -\sum_t log p_S(r_t)       | Verified traces     | Rejection-sampled  | Low-Medium       | Math and code   |
+----------------------------------------------------------------------------------------------------------------------+

Summary

Knowledge distillation in large language models has evolved from simple logit-matching heuristics into a rigorous discipline grounded in divergence dynamics and reinforcement learning. While Forward KL remains the standard for offline imitation, its mode-covering nature inherently penalizes capacity-constrained students, leading to hallucination and distribution blurring.

By leveraging Sequence-Level KD for scalable cross-architecture transfer, Reverse KL for sharp mode-seeking optimization, and on-policy sampling frameworks like MiniLLM and GKD to eliminate exposure bias, engineers can systematically compress frontier LLM capabilities into efficient, high-throughput architectures suitable for real-time edge and datacenter serving.

Sources

• Hinton, G., Vinyals, O., & Dean, J. (2015). Distilling the Knowledge in a Neural Network. arXiv:1503.02531
• Kim, Y., & Rush, A. M. (2016). Sequence-Level Knowledge Distillation. Proceedings of EMNLP 2016, arXiv:1606.07947
• Gu, Y., Dong, L., Wei, F., & Huang, M. (2024). MiniLLM: Knowledge Distillation of Large Language Models. Proceedings of ICLR 2024, arXiv:2306.08543
• Agarwal, R., Vieillard, N., Zhou, Y., Stanczyk, P., Ramos, S., Geist, M., & Bachem, O. (2024). On-Policy Distillation of Language Models: Learning from Self-Generated Mistakes. Proceedings of ICLR 2024, arXiv:2306.13649
• Xu, X., Li, M., Tao, C., Shen, T., Cheng, R., Li, J., Xu, C., Tao, D., & Zhou, T. (2024). A Survey on Knowledge Distillation of Large Language Models. arXiv:2402.13116