TokAlign: Efficient Vocabulary Adaptation via Token Alignment
This addresses inefficiencies in tokenization for LLMs, enabling better adaptation to new domains or languages, though it is incremental as it builds on existing vocabulary adaptation techniques.
The paper tackles the problem of vocabulary mismatch in Large Language Models (LLMs) for new domains or languages, which slows training and hinders knowledge transfer, by proposing TokAlign, an efficient method that aligns vocabularies via token co-occurrence and rearranges model parameters, resulting in a perplexity decrease from 3.4e² to 1.2e² after initialization and a 4.4% improvement in token-level distillation over sentence-level distillation.
Tokenization serves as a foundational step for Large Language Models (LLMs) to process text. In new domains or languages, the inefficiency of the tokenizer will slow down the training and generation of LLM. The mismatch in vocabulary also hinders deep knowledge transfer between LLMs like token-level distillation. To mitigate this gap, we propose an efficient method named TokAlign to replace the vocabulary of LLM from the token co-occurrences view, and further transfer the token-level knowledge between models. It first aligns the source vocabulary to the target one by learning a one-to-one mapping matrix for token IDs. Model parameters, including embeddings, are rearranged and progressively fine-tuned for the new vocabulary. Our method significantly improves multilingual text compression rates and vocabulary initialization for LLMs, decreasing the perplexity from 3.4$\text{e}^2$ of strong baseline methods to 1.2$\text{e}^2$ after initialization. Experimental results on models across multiple parameter scales demonstrate the effectiveness and generalization of TokAlign, which costs as few as 5k steps to restore the performance of the vanilla model. After unifying vocabularies between LLMs, token-level distillation can remarkably boost (+4.4% than sentence-level distillation) the base model, costing only 235M tokens.