Towards Cross-Tokenizer Distillation: the Universal Logit Distillation Loss for LLMs

TL;DR

The paper addresses cross-tokenizer distillation for LLMs by introducing Universal Logit Distillation (ULD) loss, which couples cross-entropy with a Wasserstein distance between teacher and student distributions to align logits across different vocabularies. A fast, closed-form Wasserstein computation under uniform support and cost enables scalable cross-tokenizer distillation, including decoder-to-encoder-decoder scenarios. Empirical results across extractive QA, generative QA, and summarization show that ULD improves over baselines that rely on teacher-generated text, often matching or surpassing teacher performance with less data and smaller models, while stabilizing training. The approach broadens the applicability of KD in real-world settings by removing the requirement of shared tokenizers and offering a practical, architecture-agnostic distillation method with public code and data release.

Abstract

Deploying large language models (LLMs) of several billion parameters can be impractical in most industrial use cases due to constraints such as cost, latency limitations, and hardware accessibility. Knowledge distillation (KD) offers a solution by compressing knowledge from resource-intensive large models to smaller ones. Various strategies exist, some relying on the text generated by the teacher model and optionally utilizing his logits to enhance learning. However, these methods based on logits often require both teacher and student models to share the same tokenizer, limiting their applicability across different LLM families. In this paper, we introduce Universal Logit Distillation (ULD) loss, grounded in optimal transport, to address this limitation. Our experimental results demonstrate the effectiveness of ULD loss in enabling distillation across models with different architectures and tokenizers, paving the way to a more widespread use of distillation techniques.

Figures (5)

• Figure 1: Vocabulary overlap between teacher and student models (\ref{['choices']}), illustrating the challenge of cross-tokenizer distillation. (e.g., Bloomz tokenizer has 30.03% of Mistral's vocabulary).
• Figure 2: Distillation using ULD loss. In block 4, the KL divergence cannot be defined as the two distributions do not have the same support, breaking the absolute continuity of the quotient in the KL logarithmic term. To alleviate this we rely on the ULD loss which leverages a closed form of the Wasserstein distance.
• Figure 3: Student model size ablation with the Pythia family trained by a LLama teacher. Trainings are conducted with ULD loss and teacher-generated text (raw text). Evaluation scores on test sets are depicted on the Y-axis, while Pythia model sizes are on the X-axis.
• Figure 4: Dataset size ablation with a LLama/Pythia-410m pair trained with ULD loss or teacher-generated text. The X-axis indicates the % of data used during training while the y-axis represents the test set score. Minimum and maximum values are represented by the error bars in the plot while the mean is represented by points.
• Figure 5: Evolution of validation ULD and Cross-entropy loss curves during training on SQuAD dataset for a LLama/Pythia-410m and LLama/Bloomz-560m Teacher/Student pair of model. For teacher-generated text models (raw text), the ULD loss was only computed during validation and did not impact the training.

Theorems & Definitions (5)

• Remark
• Remark
• Remark
• Remark
• Remark