DistillLens: Symmetric Knowledge Distillation Through Logit Lens

1. Problem Statement

Standard Knowledge Distillation (KD) compresses Large Language Models (LLMs) by optimizing the student model to match the teacher's final output distribution (logits). This approach treats the teacher's intermediate layers as a "black box," ignoring the evolving "thought processes" that occur across the network's depth.

The authors identify two critical limitations in existing methods that attempt to bridge this gap by distilling intermediate features:

1. Feature Misalignment (MSE): Methods using Mean Squared Error (MSE) assume an isotropic embedding space. However, language embedding spaces are highly anisotropic; identical MSE values can result in vastly different probability divergences depending on whether errors occur on high-probability or low-probability tokens.
2. Asymmetric Divergence Bias: Methods using Kullback-Leibler (KL) divergence are inherently asymmetric.
   • Forward KL ($KL(p||q)$): Exhibits "mean-seeking" behavior, forcing the student to cover all teacher possibilities, often leading to over-smoothed, over-confident outputs.
   • Reverse KL ($KL(q||p)$): Exhibits "mode-seeking" behavior, aggressively penalizing the assignment of probability to low-probability signals, leading to under-confidence and the suppression of nuance.

Consequently, standard KD forces students to reverse-engineer complex reasoning from a sparse signal, leading to poor generalization and a divergence in the internal "thought trajectory" between student and teacher, even if final outputs converge.

2. Methodology: DISTILLLENS

The authors propose DISTILLLENS, a framework that symmetrically aligns the intermediate thought processes of student and teacher models.

Core Mechanism: Logit Lens Projection

Instead of aligning raw hidden states, DISTILLLENS projects the intermediate hidden states ($h^{(l)}$) of both models into the vocabulary space using the model's unembedding matrix ($W_U$). This utilizes the "Logit Lens" technique to interpret intermediate layers as evolving probability distributions:
$$p^{(l)}(y|x) = \text{softmax}(W_U h^{(l)})$$
This allows the distillation process to operate on semantic probability distributions rather than abstract vector spaces.

Symmetric Divergence Objective

To address the biases of asymmetric KL divergence, DISTILLLENS employs a symmetric divergence objective, specifically the Jensen-Shannon Divergence (JSD).

• Dual-Sided Penalty: The JSD loss landscape enforces a convex optimization path that penalizes both overconfidence (student probability $\gg$ teacher) and underconfidence (student probability $\ll$ teacher).
• Theoretical Proof: The authors prove that JSD imposes a linear penalty for overconfidence and a bounded penalty for underconfidence, driving the student's confidence score ($c_\theta = q_\theta / p$) toward 1 (perfect alignment).

Training Framework

The total loss function combines the standard task loss (e.g., Reverse KL on final logits) with the intermediate alignment loss:
$$L_{total} = L_{task} + \lambda \cdot L_{inter}$$
Where $L_{inter}$ is the average JSD computed across a mapped set of intermediate layers. The framework uses a uniform layer mapping strategy to ensure the student captures the evolution of reasoning steps across the full depth of the teacher model.

3. Key Contributions

1. Novel Framework: Introduction of DISTILLLENS, which shifts KD supervision from final outputs to intermediate "thought processes" via Logit Lens projection.
2. Theoretical Insight: A derivation proving that symmetric divergence (JSD) creates a dual-sided penalty landscape, preventing both over- and under-confidence, which is crucial for preserving high-entropy information conduits essential for deduction.
3. Modularity: The method is designed as a modular objective ($L_{inter}$) that can be integrated with any existing distillation strategy (off-policy or on-policy) and any task-specific loss.
4. Empirical Validation: Extensive experiments demonstrating that aligning intermediate distributions yields superior performance compared to standard KD and feature-transfer baselines.

4. Experimental Results

The authors evaluated DISTILLLENS on GPT-2 (120M, 340M) and Llama (1.1B) architectures, using GPT-2-1.5B and Llama-7B as teachers.

• Performance Gains:
  • GPT-2-120M: DISTILLLENS achieved an average Rouge-L score of 21.12, significantly outperforming Standard KD (17.74) and SFT (16.56).
  • GPT-2-340M: Achieved 23.72, marginally surpassing the teacher model's performance (23.52).
  • TinyLlama-1.1B: Achieved 25.48, outperforming the best baseline (24.56) by over 1.6 points.
• Comparison with Feature Transfer: DISTILLLENS consistently outperformed direct feature regression (MSE) and asymmetric feature distillation (FDD) across all benchmarks.
• On-Policy vs. Off-Policy:
  • While on-policy methods (like MiniLLM) reduce exposure bias, they suffer from high computational costs and slow training.
  • DISTILLLENS (off-policy) significantly reduces exposure bias compared to standard KD without the sampling overhead of on-policy methods.
  • Hybrid Approach: Combining DISTILLLENS with MiniLLM (Hybrid) yielded the best results (Avg R-L 21.86), proving the method's modularity and ability to boost on-policy techniques.
• Ablation Studies:
  • Layer Count: Performance peaked at 5 intermediate layers.
  • Selection Pattern: An interleaved selection of layers (spaced throughout the model) performed better than consecutive blocks (first 5 or last 5), highlighting the importance of capturing the full reasoning trajectory.

5. Significance and Impact

• Bridging the Reasoning Gap: DISTILLLENS demonstrates that aligning the internal reasoning steps (intermediate thought processes) is more effective than merely matching final answers. This leads to better generalization and reasoning capabilities.
• Efficiency vs. Robustness: It offers a computationally efficient alternative to on-policy distillation. It shifts the alignment cost entirely to the training phase without impacting inference latency, as the student model architecture remains unchanged.
• Theoretical Advancement: By formalizing the necessity of symmetric divergence for intermediate layers, the paper challenges the dominance of asymmetric KL divergence in feature-based distillation, offering a more robust theoretical foundation for compressing LLMs.

In conclusion, DISTILLLENS represents a paradigm shift in knowledge distillation, moving from "black box" output matching to "white box" structural alignment of evolving thought processes, resulting in more capable and robust compressed language models.