Flexible Multi-Level Optimal Transport for Universal Cross-Tokenizer Knowledge Distillation on LLMs and Beyond.

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, as they assume a consistent semantic correspondence across logit dimensions, 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 logit distribution structures using a flexible Sinkhorn loss with relaxed marginal constraints, allowing partial transport and absorbing unmatched tokens, thereby serving as an efficient OT-based surrogate while tolerating unaligned mass. Extensive experiments on tasks of extractive QA, generative QA, summarization, instruction- following, math reasoning, and Big-Bench-Hard demonstrate that the MultiLevelOT outperforms state-of-the-art cross-tokenizer KD methods. Our approach is robust to different student and teacher models across model families, architectures, and parameter sizes. We further validate MultiLevelOT in multimodal-input settings and introduce task-specific technical extensions for both multi-teacher text distillation and structured visual-token distillation.