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

TL;DR

This paper introduces a new knowledge distillation method for large language models with different tokenizers, improving text generation quality especially on out-of-distribution data.

Contribution

It systematically analyzes existing dual-space KD methods and proposes DSKD-CMA-GA, a novel adversarial approach to address token distribution mismatches.

Findings

Modest but consistent ROUGE-L gains in text generation quality.

Improved performance on out-of-distribution data (+0.37 on average).

Enhanced understanding of attention mechanisms in dual-space KD.

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.