Multi-Level Optimal Transport for Universal Cross-Tokenizer Knowledge Distillation on Language Models

TL;DR

This work tackles cross-tokenizer knowledge distillation for large language models by introducing MultiLevelOT, which aligns teacher and student logits through both token-level and sequence-level optimal transport, avoiding strict token-by-token vocabulary matching. The framework uses three losses—Holistic Absolute Difference Loss, Sequential Logarithmic Loss, and Sinkhorn Distance Loss—combined with cross-entropy to distill knowledge across diverse vocabularies and model families. Empirical results on QED, FairytaleQA, and DIALOGSum show that MultiLevelOT consistently outperforms state-of-the-art cross-tokenizer KD methods and generalizes across student sizes, architectures, and teacher models. The work highlights the practical potential of geometry-aware logit distribution matching for scalable, versatile KD in real-world multilingual or multi-model settings.

Abstract

Knowledge distillation (KD) has become a prevalent technique for compressing large language models (LLMs). Existing KD methods are constrained by the need for identical tokenizers (i.e., vocabularies) between teacher and student models, limiting their versatility in handling LLMs of different architecture families. In this paper, we introduce the Multi-Level Optimal Transport (MultiLevelOT), a novel approach that advances the optimal transport for universal cross-tokenizer knowledge distillation. Our method aligns the logit distributions of the teacher and the student at both token and sequence levels using diverse cost matrices, eliminating the need for dimensional or token-by-token correspondence. At the token level, MultiLevelOT integrates both global and local information by jointly optimizing all tokens within a sequence to enhance robustness. At the sequence level, we efficiently capture complex distribution structures of logits via the Sinkhorn distance, which approximates the Wasserstein distance for divergence measures. Extensive experiments on tasks such as extractive QA, generative QA, and summarization demonstrate that the MultiLevelOT outperforms state-of-the-art cross-tokenizer KD methods under various settings. Our approach is robust to different student and teacher models across model families, architectures, and parameter sizes. Codes and models are available at https://github.com/2018cx/Multi-Level-OT.

Figures (3)

• Figure 1: An illustration of vocabulary mismatch resulting from cross-tokenizer discrepancies. Unlike strict token-wise distillation methods that may lead to token misalignment, we employ sequence-level and sequence-aware token-level optimal transport to facilitate effective knowledge transfer.
• Figure 2: Illustration of our pipeline. MultiLevelOT computes sequence-aware token-level and sequence-level optimal transport distances between the output logits of the teacher and student models. This approach effectively transfers local and global information within the logits distribution, accommodating vocabulary differences and enabling cross-tokenizer distillation.
• Figure 3: Performance at different student scales (Pythia 160M, 410M, and 1B) on QED and FairytaleQA.