Unifying Vocabulary of Large Language Model with Statistical Token-level Alignment

Authors: Chong Li, Jiajun Zhang, Chengqing Zong
Primary area: generative models
Keywords: Vocabulary Adaptation, Large Language Model, Efficient NLP

Abstract:
Large Language Models (LLMs) achieve great success across many general tasks, but the mismatch among different vocabularies hinders further applications like token-level distillation and inference with various models. To align the vocabularies of LLMs, we propose a simple yet effective method named **UnifyVocab** to replace the vocabulary of an LLM at a limited cost. A new vocabulary alignment method is devised first to align the source vocabulary to the target one. We then rearrange the corresponding parameters like embeddings, and progressively fine-tune the model. Experimental results on models across multiple parameter scales demonstrate the effectiveness and generalization of UnifyVocab, which costs as few as 10B tokens to recover 98.02\% performance of the vanilla models on average. We further find that unifying the vocabularies significantly facilitates the token-level distillation which remarkably boosts (+4.4\%) the model with only 235M tokens. Moreover, our method provides a better initialization of multilingual vocabulary for LLMs to adapt to new languages.

Decision: reject
Reviewer scores: 5, 6, 6

Reviews and discussion:
<thread>
creation date:
1729608358962

thread_type:
official_review

<thread_object>
creator: reviewer_8lxg

<content>
summary:
This paper wants to address the mismatch among different vocabularies used by various LLMs. UnifyVocab is proposed. The high-level idea is to use the embeddings of the tokens from the source tokenizer to initialize the embeddings of tokens from the target tokenizer. To achieve this, the authors train GloVe embeddings for the tokens in the source and target vocabularies respectively, and then align the tokens using the similarity between the source and target tokens.

strengths:
- The paper is generally easy to follow.

- The experiments are extensive.

- UnifyVocab seems to be simple and effective in aligning tokens among the vocabularies of two tokenizers.

weaknesses:
- The method is sensitive to the selection of the corpus used to learn the token-token alignment.

- The pipeline is very similar to WECHSEL [1]. If I understand correctly, the method proposed in this work is a simple extension to the scenario where the source and the target languages are the same (in WICHSEL they are different). 

- WECHSEL additionally needs to align the learned fastText embeddings because the source and target embeddings are in different spaces. I guess this step is omitted in UnifyVocab because the authors assume the learned token GloVe embeddings (for tokenizer A and tokenizer B) are in the same space. However, this assumption might not hold true. Two embedding matrices learned from the same corpus can be quite different, even if they have the same vocabulary (and in your case, this does not hold true) [2]. 

- I am not sure if I agree the motivation of the paper is well-established. If a model performs well with its own tokenizer (e.g., LLama and the LLama tokenizer), why would one be interested in exchanging its tokenizer with another model's tokenizer that is intended to work on the same domain or language? I think replacing the tokenizer is mostly only meaningful when we want to have a new domain or a new language to adapt to.

[1] https://arxiv.org/abs/2112.06598  
[2] https://arxiv.org/pdf/2209.15430

questions:
$\textbf{Questions/Suggestions}$:

- I don't think Figure 3 (b) is meaningful. The authors claim that there is a negative relationship between the first-step training loss nad the BLEU. But the BLEU is very very bad, only around 2.4. For such a small BLEU, the differences between different initializations are basically negligible.

- In Table 4, does "0" in the column "#Tune (B)" without any training? In other words, does that line indicate the performance of right after replacing the tokenizer? If it is, maybe the authors can make it more clear in the caption.

- It is better to use the same color and same order in the legend of Figure 3 for better consistency. 

- There is one related paper [3] for zero-shot tokenizer transfer. They proposed ZETT where a hypernetwork is used to predict embeddings of the new tokens in the target tokenizer. The authors may consider this as a stronger baseline method.


$\textbf{Typos}$:

Line 174: "which belongs both vocabularies." -> "which belongs to overlapping vocabularies."
Line 177: "randomly chosen token from the source vocabulary." -> "a randomly chosen token from the source vocabulary."


[3] https://arxiv.org/abs/2405.07883
</content>
</thread_object>

</thread>
<thread>
creation date:
1730658895118

thread_type:
official_review

<thread_object>
creator: reviewer_qu5a

<content>
summary:
This paper tackles a vocabulary extension issue in LLMs and introduce a method called UnifyVocab to replace the vocabulary of LLM, aligning token IDs between two vocabularies. The proposed approach allows vocabularies of LLMs to get replaced based on the token-token co-occurences, enabling new vocabulary adaptation with lower costs. Experimental results show some effectiveness in (cross-lingual) knowledge transfer between models.

strengths:
- proposes vocabulary adaptation technique which will be useful in multilingual/crosslingual LLM  application
- Experimental results show some effectiveness of the proposed approach in multiple multilingual NLP tasks

weaknesses:
- There is some missing citation on vocabulary adaptation like [1]. Comparison and/or discussion would be required.  
[1] OFA: A Framework of Initializing Unseen Subword Embeddings for Efficient Large-scale Multilingual Continued Pretraining. In Proc of NAACL2024 findings

questions:
- Have you ever tried other (semantic) metrics like COMET scores instead of BLEU while evaluating the performance of alignment Matrix?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730702471100

thread_type:
official_review

<thread_object>
creator: reviewer_1xm4

<content>
summary:
This paper proposes UnifyVocab, a method to replace the vocabulary of an LLM. 
This involves using a tokenizer from another model and training new GloVe embeddings, which are aligned with cosine similarity to an existing embedding set, and then are used to replace the original embedding matrix and finetuned with the model. 

Experiments use the Pythia base model and training corpus, and experiment with replacing the vocabulary with those from Gemma, Qwen2, and LLaMa 2 & 3.  
UnifyVocab is compared to a random initialization, random permutation, and FOCUS from Dobler & de Melo (2023). 
English results are compared across 6 standard tasks, and cross-lingual transfer is compared for 12 languages (+English) on 4 standard tasks.  
Results show that the method preserves on average 98% of the original performance, and leads to improved cross-lingual transfer compared to FOCUS. 
Two-stage tuning (first finetuning the vocabulary-related parameters in the model with the rest frozen, and then fine-tuning the full model) improves performance compared to fine-tuning the full model directly. 
Token-level distillation requires less training data and generally leads to improved performance over sequence-level distillation. 

The method, though, requires ~10B tokens for training, which is a significant cost compared to past approaches applied to e.g. machine translation where separately trained embeddings may be adapted to work with a model with <20k tokens. 
Aligning embeddings with cosine similarity assumes that a) similar representation spaces are learned and so an explicit alignment step is not needed and b) the vocabularies are near-isomorphic, which are not guaranteed with the procedure used, and these assumptions are not mentioned. It would be easier to trust that the results would generalize if these were explored here and for example an explicit alignment step compared and more specific analysis about the conditions where the method is and is not successful (for example, if 6% vocabulary overlap with Gemma and Pythia makes the model much slower to converge, how similar is this to random initialization? are the cosine similarities considerably lower in this case, and/or less one-to-one mappings chosen? if something other than cosine similarity were used, how would this change?)

strengths:
Straightforward method to replace the tokenizer / vocabulary of an LLM, given sufficient data.

weaknesses:
Method is costly and does not consistently recover original model's performance. Insufficient analysis to understand the conditions where the method will succeed.

questions:
Presentation note: only the best vocabulary replacement results in the tables. In Tables 4 and 5 there are times when the original Pythia model outperforms any of the replacement methods, and so should likely be bolded instead so that this is clear.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732162579711

thread_type:
official_comment

<thread_object>
creator: area_chair_h97t

<content>
title:
Reminder: Please respond and update the score if necessary

comment:
Dear Reviewers,

Kindly ensure that you respond proactively to the authors' replies (once they are available) so we can foster a productive discussion. If necessary, please update your score accordingly. We greatly appreciate the time and effort you’ve dedicated to the review process, and your contributions are key to making this process run smoothly.

Thank you,

AC
</content>
</thread_object>

</thread>
<thread>
creation date:
1732621792945

thread_type:
public_comment

<thread_object>
creator: ~xinke_jiang1

<content>
title:
Good Work but not Complete

comment:
Hello, after reading your paper, I believe this idea is a very good piece of work, and the experiments conducted by the authors are quite thorough. However, I noticed that there are many missing details in the description, which I suspect might be due to the authors rushing to meet the deadline of ICLR. For example, the description of progressive adaptation in section 2.3 and the explanation of Figure 2 are somewhat confusing to the readers. It would be great if the authors could elaborate on these aspects to improve the clarity of the work.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732779183179

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Rebuttal Revision by Authors

comment:
- Following the suggestions of Reviewer qU5A and 8LXg, we add the results of three strong baseline methods for vocabulary adaptation.[1-3]
- Following the suggestions of Reviewer 1XM4 and 8LXg, we add the GloVe embedding alignment process using the relative representation method from Moschella et al. (2023) [4] to meet the assumption of aligning embeddings with cosine similarity. Results are reported in Table 1.
- Following the suggestions of Reviewer 1XM4, we further investigate the convergence of random initialization and the impact of different alignment matrices learned to the training convergence (Appendix B.3).
- We supplement the results of other two metrics, BLEU-1 and BertScore, to evaluate the performance of alignment matrix learned (Appendix B.4).

[1] WECHSEL: Effective initialization of subword embeddings for cross-lingual transfer of monolingual language models. https://aclanthology.org/2022.naacl-main.293

[2] FOCUS: Effective Embedding Initialization for Monolingual Specialization of Multilingual Models. https://aclanthology.org/2023.emnlp-main.829

[3] Zero-Shot Tokenizer Transfer. https://arxiv.org/abs/2405.07883

[4] Relative representations enable zero-shot latent space communication. https://arxiv.org/abs/2209.15430
</content>
</thread_object>

</thread>


Metareview:
metareview:
The paper proposes UnifyVocab, a method to replace and harmonize vocabularies in large language models (LLMs). It uses new GloVe embeddings aligned with existing ones via cosine similarity, replacing original embeddings and allowing the model to be fine-tuned. Experiments with various vocabularies show that UnifyVocab maintains 98% of original performance while improving cross-lingual transfer compared to other methods. The two-stage tuning further enhances performance, although the approach is resource-intensive, requiring about 10 billion tokens for training. The paper highlights potential improvements in embedding alignment assumptions and suggests that UnifyVocab effectively addresses vocabulary mismatches across LLMs, enabling cost-effective adaptation and improved cross-lingual knowledge transfer.

This paper received borderline scores. While the idea has merit, and reviewers agree that the method could be valuable for token embedding initialization, the overall presentation requires substantial improvement. The proposed method also has notable drawbacks, particularly concerning its efficiency and its similarities to prior work. Additional issues include unclear and poorly presented results, especially in Table 4, which reports perplexity and needs significant refinement. Moreover, the paper lacks strong motivation. Although the authors attempted to address this in their response, the paper itself needs to provide a clearer and more compelling rationale. Given the issues outlined above, I am leaning towards recommending the rejection of the paper.