Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

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

Stella Eva Tsiapali, Cong-Thanh Do, Kate Knill

Abstract

Large language models (LLMs) achieve state-of-the-art (SOTA) performance across language tasks, but are costly to deploy due to their size and resource demands. Knowledge Distillation (KD) addresses this by training smaller Student models to mimic larger Teacher models, improving efficiency without significant performance loss. Dual-Space Knowledge Distillation with Cross-Model Attention (DSKD-CMA) has emerged as a SOTA method for KD between LLMs with distinct tokenizers, yet its internal workings remain largely opaque. In this work, we systematically analyse the attention mechanism of DSKD-CMA through manual token alignment probing and heatmap visualisations, revealing both strengths and limitations. Building on this, we introduce a novel method, DSKD-CMA-GA, based on Generative Adversarial (GA) learning, to address the mismatched distributions between the keys and queries computed from distinct models. Experiments show modest but consistent ROUGE-L gains in text generation quality, particularly on out-of-distribution data (+0.37 on average), narrowing the gap between cross- and same-tokenizer KD.

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

Abstract

Large language models (LLMs) achieve state-of-the-art (SOTA) performance across language tasks, but are costly to deploy due to their size and resource demands. Knowledge Distillation (KD) addresses this by training smaller Student models to mimic larger Teacher models, improving efficiency without significant performance loss. Dual-Space Knowledge Distillation with Cross-Model Attention (DSKD-CMA) has emerged as a SOTA method for KD between LLMs with distinct tokenizers, yet its internal workings remain largely opaque. In this work, we systematically analyse the attention mechanism of DSKD-CMA through manual token alignment probing and heatmap visualisations, revealing both strengths and limitations. Building on this, we introduce a novel method, DSKD-CMA-GA, based on Generative Adversarial (GA) learning, to address the mismatched distributions between the keys and queries computed from distinct models. Experiments show modest but consistent ROUGE-L gains in text generation quality, particularly on out-of-distribution data (+0.37 on average), narrowing the gap between cross- and same-tokenizer KD.
Paper Structure (15 sections, 16 equations, 5 figures, 1 table)

This paper contains 15 sections, 16 equations, 5 figures, 1 table.

Table of Contents

  1. Introduction
  2. Background and Related Work
  3. Knowledge Distillation (KD)
  4. Cross-Tokenizer KD
  5. Methods
  6. Original DSKD-CMA Method
  7. Proposed Cross-Model Attention Analysis
  8. Proposed Key-Query Matching Distillation
  9. Experimental Setup
  10. Results and Discussion
  11. Chunk-Based Probing Insights
  12. Key-Query Matching Improvements
  13. Conclusion
  14. Acknowledgments
  15. Compliance with Ethical Standards

Figures (5)

  • Figure 1: The sequence of matrix multiplications employed by the attention mechanism of DSKD-CMA.
  • Figure 2: Chunk alignment of token sequences by distinct tokenizers.
  • Figure 3: Flowchart of the proposed manual chunk alignment process.
  • Figure 4: The sequence of matrix multiplications involved in Chunk-Level Projection (CLP) and Chunk-Level Attention (CLA).
  • Figure 5: Alignment weights assigned to pairs of tokens in the Student and Teacher sequences, using CMA and its chunk-based ablations: Chunk-Level Projection (CLP) and Chunk-Level Attention (CLA).