DistillLens: Symmetric Knowledge Distillation Through Logit Lens

TL;DR

DistillLens tackles the limitation of standard KD by explicitly supervising the evolving thought process of the teacher across intermediate layers. It projects hidden states into the vocabulary space using the Logit Lens and enforces symmetric divergence, notably Jensen-Shannon Divergence, to align student and teacher distributions at multiple depths. Theoretical analysis reveals a dual-sided penalty that prevents both overconfidence and underconfidence, while empirical results on GPT-2 and Llama demonstrate consistent gains over standard KD and feature-transfer baselines, with improved Rouge-L and SBERT-based semantic similarity. The approach maintains inference efficiency while incurring additional training cost, and shows promise for combining with on-policy methods to further boost reasoning capabilities.

Abstract

Standard Knowledge Distillation (KD) compresses Large Language Models (LLMs) by optimizing final outputs, yet it typically treats the teacher's intermediate layer's thought process as a black box. While feature-based distillation attempts to bridge this gap, existing methods (e.g., MSE and asymmetric KL divergence) ignore the rich uncertainty profiles required for the final output. In this paper, we introduce DistillLens, a framework that symmetrically aligns the evolving thought processes of student and teacher models. By projecting intermediate hidden states into the vocabulary space via the Logit Lens, we enforce structural alignment using a symmetric divergence objective. Our analysis proves that this constraint imposes a dual-sided penalty, preventing both overconfidence and underconfidence while preserving the high-entropy information conduits essential for final deduction. Extensive experiments on GPT-2 and Llama architectures demonstrate that DistillLens consistently outperforms standard KD and feature-transfer baselines on diverse instruction-following benchmarks. The code is available at https://github.com/manishdhakal/DistillLens.

Figures (8)

• Figure 1: The distilled student (GPT2-120M) notably diverges from the teacher model for the hidden layers, except for the final layer.
• Figure 2: The DistillLens Framework. Comparison between standard Knowledge Distillation (left) and our proposed DistillLens approach (right). Unlike standard KD, which restricts supervision solely to the final output logits, DistillLens aligns the intermediate thought processes of the student and teacher. By projecting intermediate hidden states into the vocabulary space using the Logit Lens, we compute the symmetric divergence loss ($\mathcal{L}_{JSD}$), which is optimized jointly with the standard task loss ($\mathcal{L}_{task}$)
• Figure 3: Exposure bias vs. generated sequence length. Lower bias indicates that the model is more robust to its own generation errors. Please refer to Appendix \ref{['app:sec:exposure_bias']} for calculating the bias.
• Figure 4: DistillLens incorporated with existing KD baselines. DistillLens improves the existing baselines.
• Figure 5: Training Speed Comparison (per GPU): We distill from GPT-2-1.5B to GPT-2-120M using A100 GPUs. Training with DistillLens is slower among the off-policy KDs, but faster when compared against on-policy methods.

Theorems & Definitions (7)

• Definition 3.1: Confidence Score
• Definition 3.2: JSD
• Proposition 3.3: Dual-sided Alignment
• proof
• Definition 1.1: Jeffreys Divergence
• Proposition 1.2: Dual-sided Confidence Penalization
• proof