Dual-Space Knowledge Distillation with Key-Query Matching for Large Language Models with Vocabulary Mismatch

1. Problem Statement

Large Language Models (LLMs) offer state-of-the-art (SOTA) performance but are computationally expensive to deploy. Knowledge Distillation (KD) is a standard solution where a smaller "Student" model learns from a larger "Teacher" model. However, a significant barrier exists in modern LLMs: Vocabulary Mismatch.

• Different LLMs use distinct tokenizers and vocabularies, resulting in sequences of varying lengths and different predictive spaces.
• Existing KD methods struggle to align the internal states (logits or hidden representations) of models with different tokenizers.
• The current SOTA method for this challenge, Dual-Space Knowledge Distillation with Cross-Model Attention (DSKD-CMA), achieves good results but operates as a "black box." Its internal alignment mechanisms are opaque, making it difficult to diagnose failures or improve upon them.

2. Methodology

The authors propose a two-pronged approach: first, a systematic analysis of the existing DSKD-CMA mechanism, and second, the introduction of a novel enhancement using adversarial learning.

A. Analysis of DSKD-CMA

The authors deconstructed DSKD-CMA to understand how it aligns Teacher and Student hidden states.

• Mechanism: DSKD-CMA uses a Cross-Model Attention (CMA) mechanism where Student embeddings are projected to form Queries ($Q$) and Teacher embeddings serve as Keys ($K$). Attention matrices are computed to project states into each other's spaces.
• Probing Strategy: To make this opaque process interpretable, the authors introduced manual token alignment probes. They segmented token sequences into "chunks" (minimal text spans shared by both tokenizers) and created explicit alignment matrices.
• Ablation Studies: They replaced the learned CMA with two explicit alternatives:
  1. Chunk-Level Projection (CLP): A uniform projection based strictly on aligned chunks (no attention).
  2. Chunk-Level Attention (CLA): Attention restricted only to tokens within aligned chunks.
• Findings: The analysis revealed that CMA implicitly learns chunk-level structures effectively but suffers from noise (especially with repeated tokens) and lacks precise localization. Interestingly, CLP performed similarly to CMA, suggesting that the complex attention mechanism might be over-engineered for simple chunk alignment, though CLP is computationally heavier.

B. Proposed Solution: DSKD-CMA-GA

The core limitation identified was that the Keys ($K$) and Queries ($Q$) in CMA originate from distinct models with mismatched distributions ($p_K$ and $p_Q$). This distributional shift weakens the attention mechanism. To fix this, the authors introduced Key-Query Matching objectives:

1. Generative Adversarial Alignment (GA):
   • A discriminator ($D$) is trained to distinguish between Teacher Keys and Student Queries.
   • The Student's query projector ($P_Q$) acts as a generator, trained to fool the discriminator (minimizing the ability to distinguish the distributions).
   • This aligns the distributions of $K$ and $Q$ before the attention step.
2. Conditional Transport (CT):
   • An alternative using a constrained version of Optimal Transport to encourage bidirectional matching between distributions via a trainable critic.

The final training objective combines the original DSKD-CMA loss with the new alignment loss ($L_{KQ}$):
$$L_{train} = L_{DSKD-CMA} + L_{KQ}$$

3. Key Contributions

1. Interpretability Framework: The first systematic analysis of DSKD-CMA using manual chunk alignment and heatmap visualizations, revealing that the method implicitly captures chunk structures but struggles with localization.
2. Novel Distillation Objectives: The proposal of DSKD-CMA-GA and DSKD-CMA-CT, which explicitly align the distributions of Keys and Queries using adversarial learning and transport methods.
3. Performance Benchmarking: Comprehensive evaluation showing that the proposed methods narrow the performance gap between cross-tokenizer and same-tokenizer KD.

4. Experimental Results

• Setup:

  • Student: GPT-2 (124M, 50K vocab).
  • Teachers: GPT-2 XL (Same tokenizer) and Qwen1.5 (1.8B, 150K vocab - Cross-tokenizer).
  • Datasets: Dolly (In-Distribution) and Self-Instruct, Vicuna, Super-Natural Instructions, Unnatural Instructions (Out-of-Distribution).
  • Metric: ROUGE-L.

• Key Findings:

  • GA Superiority: DSKD-CMA-GA achieved the strongest performance, yielding consistent ROUGE-L gains over the baseline DSKD-CMA (0.15–1.04 points).
  • Out-of-Distribution (OOD) Gains: The GA method showed particular strength on OOD data, with an average improvement of +0.37 ROUGE-L points.
  • Closing the Gap: DSKD-CMA-GA significantly narrowed the gap between cross-tokenizer and same-tokenizer KD. In some cases, it surpassed the performance of same-tokenizer distillation.
  • Divergence Interaction: The choice of divergence function (e.g., Skewed KL) interacts with the alignment method. GA performed best with Skewed KL (SKL), while the baseline worked best with Skewed Reverse KL (SRKL).
  • Ablation Insight: The "Chunk-Level Attention" (CLA) variant degraded performance, suggesting that DSKD benefits from stable, coarse-grained alignments rather than over-sensitive, fine-grained token weighting.

5. Significance

This work is significant for several reasons:

• Demystifying KD: It moves cross-tokenizer KD from a "black box" approach to an interpretable framework, identifying distributional mismatch in attention mechanisms as a key bottleneck.
• Practical Efficiency: By using Generative Adversarial learning to align distributions, the method improves generation quality without the massive computational overhead of explicit token mapping or Chain-of-Thought alignment (unlike previous attempts).
• Scalability: The results suggest that high-quality distillation is possible even between models with vastly different vocabularies (e.g., 50K vs. 150K tokens), making it easier to deploy efficient student models trained on diverse, large-scale teacher models.
• Future Directions: The paper opens avenues for dynamic loss weighting and broader applications of adversarial learning in model alignment.

In summary, the paper demonstrates that by explicitly aligning the statistical distributions of attention keys and queries via adversarial learning, one can significantly enhance the efficiency and quality of knowledge distillation across models with incompatible tokenizers.