Adapting Language Models via Token Alignment

Contemporary large language models (LLM) are trained with fixed tokenizers, which are created in a separate process before the LLM training stage. This two-stage procedure leads to several challenges when the models are evaluated on out-of-distribution (OOD) datasets or other modalities like protein sequences, for example, it can incur worse compression ratio and semantic misalignment in the vocabularies. In this work, we propose a simple token alignment technique via sparse Sinkhorn projections, the aligned tokens can serve as good initializations for continual finetuning on downstream tasks. We demonstrate the effectiveness of the proposed algorithm on protein sequences. In addition, the proposed alignment present \textit{weak-to-strong model transferability} -- tokens aligned on a weak model can be transferred to a stronger model -- provided that both models share the same tokenizer.