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

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.

We argue that the low value of BLEU in Figure 3, which is the average value of BLEU-1 to BLEU-4, comes from the low BLEU-3 and BLEU-4. It is reasonable for a one-to-one mapping alignment matrix to obtain a nearly 0 value for BLEU-3 and BLEU-4. Thus we provide the results of the BLEU-1 metric and an additional semantic metric named BertScore in Figure 7. Specifically, the token ID corpus $C_{t}^{'}$, which is converted by the aligned matrix from $C_s$, is de-tokenized into a text corpus $C^{'}$ by the target Tokenize$r_t$. Then we adopt "all-mpnet-base-v2" to quantify the semantic similarity between $C^{'}$ and $C$, which is BertScore(sentence1=$C$, sentence2=$C^{'}$).

As shown in Figure 7, given different alignment matrix learned, the deviation of BLEU-1 score has increased to 8~16. Moreover, we can find a similar negative relationship between the initial training loss and BLEU-1 (slope rate=-31.186, $R^2$=0.5233) or BERTScore (slope rate=-15.737, $R^2$=0.1098).

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.

Thanks for your helpful suggestion! As you mentioned, "#Tune (B)" denotes the performance of model after initialization without any training. We have added more descriptions in the title of Table 4.

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

Thanks for your suggestion! We re-arrange the order in the legend of Figure 3 as you suggested.

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.

Thank you for reminding us! We have supplemented the results of ZeTT[4] on ${Pythia}_{1b}$ in Table 1. Our method (97.63%) recovers more performance than the other strong baseline methods like ZeTT (91.93%) given the same amount of tokens. The hyper-network trained in ZeTT costs 418.9 GPU hours for ${Pythia}_{1b}$ on a 8*A100 80GB server, while the initialization of our method only costs less than two hours for a CPU server with 128 cores to train the GloVe embedding for token-token alignment.

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

Typos in line 174 and 177

Thanks for your reminder! We have corrected the typos in the Line 174 and 177.

Thank you for your detailed review! We are thankful for your efforts during the challenging review period. We will address your concerns point by point.

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

To further investigate the impact of the corpus used, we replace the corpus with the SlimPajama[1] which is commonly used in the pre-training of the language model. Under 1B tokens amount, the corpus tokenized covers 98.83% of token IDs from the vocabulary of Pythia (49714/50304≈98.83%) and 97.21% of token IDs from the vocabulary of Gemma (248857/256000≈97.21%). Results are shown in the "w/ SlimPajama" row in Table 1, which reports a comparable results with the original settings. It further demonstrates the robustness of our method on the pre-training corpus for token embedding and alignment matrix.

[1] Daria Soboleva, Faisal Al-Khateeb, Robert Myers, Jacob R Steeves, Joel Hestness, and Nolan Dey. SlimPajama: A 627B token cleaned and deduplicated version of RedPajama. https://huggingface.co/datasets/cerebras/SlimPajama-627B.

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).

We realize that there may have been some misunderstanding regarding our method. There are two key differences between WECHSEL[2] and UnifyVocab:

1. WECHSEL needs two static word embedding for the source tokenizer and target tokenizer and an additional bilingual dictionary for alignment, while our method only requires a pre-training corpus to train the token-token alignment matrix. Besides, the training and aligning the GloVe embedding for each subword/token only cost less than 1 hour for a machine with 128 cores CPU.

2. The initialization of WECHSEL for the embedding of the language model is composed of the weighted sum of similar tokens, while UnifyVocab only re-arranges the source embedding using a one-to-one mapping function based on the token-token aligned matrix.

The pipelines for adapting new vocabulary for most of the methods are similar: first, initialize the embedding for the target tokenizer and then fine-tune the initialized model. Most of the differences between methods come from the initialization of embeddings for the target vocabulary.

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

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].

Thank you for insightful suggestions! Following Moschella et al. (2023)[3], we convert the glove embeddings into relative representations using 300 common tokens that occur in both vocabularies and conduct the left procedure of UnifyVocab to adapt the Gemma tokenizer for ${Pythia}_{1b}$. It comes to the slightly better results, which are labeled "+ Align Rep." in Table 1.

[3] Relative representations enable zero-shot latent space communication. https://arxiv.org/abs/2209.15430

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.

The importance of replacing the tokenizer lies in the fast token-level knowledge transfer from capable models, and reducing the huge cost to train a model from scratch when a much better tokenizer is found, where the special case is cross-lingual or cross domain vocabulary adaptation problem. Experiments in Section 4.1 show that token-level distillation with capable language models like LLaMA3 can significantly improve the performance of Pythi $a_{1b}$ , which is comparable with vanilla ${Pythia}_{7b}$ after only 235M tokens of token-level distillation. Moreover, UnifyVocab can also be applied to the vocabulary adaptation problem for a new language or domain (Section 4.2) and achieves better results than the traditional cross-lingual vocabulary adaptation method Focus.

Thank you for your insightful review! We are thankful for your efforts during the challenging review period. We will address your concerns point by point.

Method is costly and does not consistently recover original model's performance.

We realize that there may have been some misunderstanding regarding vocabulary adaptation methods for large language models.

1. Previous works for the vocabulary adaptation of LLM cost a similar or more amount of tokens, e.g., 65.5B for GPT-2 (WECHSEL)[1], 12.8B for XLM-R (Focus)[2]. In this work, the cost of 10B tokens mainly comes from the 2M tokens batch size, which follows the pre-training setting of Pythia, and training steps are only 5k. However, as shown in Table 6, the amount of tokens required can be reduced to 2B tokens by decreasing the batch size to 1M and training steps to 2k, which recovers the average 95.66% performance of the vanilla model.

2. The performance of the original model is hard to recover with a limited token amount for vocabulary adaptation comparing the one of pre-training (10B << 300B token amount in pre-training). The results of Table 2 in ZeTT[3], which replaces the Mistral tokenizer(32.0k) to the one of GPT-2 (50.3k), are further demonstrated that the phenomenon of performance loss during replacing another tokenizer. We supplement the results of the other three vocabulary adaptation methods on Pythia. As shown in Table 1, our method (97.63%) recovers more performance than the other strong baseline methods like ZeTT (91.93%) given the same amount of tokens. The hyper-network trained in ZeTT costs 418.9 GPU hours for $Pythia_{1b}$ on a 8*A100 80GB server, while the initialization of our method only requires less than two hours for a cpu server with 128 cores to train the GloVe embedding for token-token alignment.

[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

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.

Thank you for insightful review! To meet these assumptions, we follow the relative representation alignment method from Moschella et al. (2023)[4]. Specifically, we convert the GloVe embeddings into relative representations using 300 common tokens that occur in both two vocabularies, and conduct the left procedures of UnifyVocab to adapt the Gemma tokenizer for Pythia. It comes to the slightly better results, which are denoted "+ Align" in Table 1.

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

More specific analyses. 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?

As shown in Figure 7(a), we plot the training loss of random initialization under replacing different tokenizers, which shows a similar phenomenon in Figure 4(a) where the convergence of Gemma are slower than the one of llama3 and qwen2.

We argue that the difference in converge rate may come from the different initial training losses for the one-to-one mappings chosen in our method. To evaluate this hypothesis and imitate the cases of worse methods other than cosine similarity are adopted, we randomly shuffled the learned alignment matrix for Qwen2 [40%, 60%, 80%]. The initial training loss increases from 5.35 to 11.06(randomly shuffle 80%) when replacing the Qwen2 tokenizer for $Pythia_{1b}$. As shown in Figure 7(b), the convergence of randomly shuffle 80% token-token alignment is approaching the ones of Gemma in Figure 4(a).

Presentation note in Table 4 and 5

Thanks for your helpful suggestions! We have bold the better performance from vanilla model in Table 4 and 5 as suggested.

Dear Reviewer 1XM4,

As the discussion phase is approaching the end, we sincerely hope you could find some time to review our response. We hope to fully address your concerns.

We understand that your time is valuable and you may be busy with other things. However, your insights would be extremely valuable for improving our work.

We greatly appreciate your consideration.

Best,

The Authors

Thank you for your insightful review! We are thankful for your efforts during the challenging review period. We will address your concerns point by point.

Missing vocabulary adaptation baselines

Thanks for your reminder! OFA[1] and WECHSEL[2] require additional embeddings for the source language and target language to compose the parameters of tokens in the target tokenizer. Our method reuses the parameter of the most similar source token from the aligned matrix learned by the tokenized corpus. We have supplemented the results of the following three strong baselines[1-3] on ${Pythia}_{1b}$ in Table 1. Our method (97.63%) recovers more performance than the other strong baseline methods like ZeTT (91.93%) given the same amount of tokens.

[1] OFA: A Framework of Initializing Unseen Subword Embeddings for Efficient Large-scale Multilingual Continued Pretraining. https://aclanthology.org/2024.findings-naacl.68

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

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

More semantic metrics to evaluate the performance of alignment Matrix

Thanks for your suggestion! We investigate another semantic metric BertScore for COMET requires the source input. Specifically, the token ID corpus $C_{t}^{'}$, which is converted by the aligned matrix from $C_s$, is de-tokenized into a text corpus $C^{'}$ by the target Tokenize$r_t$. Then we adopt "all-mpnet-base-v2" to quantify the semantic similarity between $C^{'}$ and $C$, which is BertScore(sentence1=$C$, sentence2=$C^{'}$). [ Test corpus $C$ --(tokenized by the source Tokenize$r_s$) → $C_s$ --(converted by align matrix $M_{s\to t}$ using the most similar target token) → $C_{t}^{'}$ --(de-tokenized by the target Tokenize$r_t$) → $C^{'}$ ]. As shown in Figure 8(b), we find that the BertScore($C$, $C^{'}$) is negatively proportional (slope rate=-15.737, $R^2$=0.1098) with the initial training loss. In other words, the semantic metric BertScore can also be used to the evaluation of alignment matrix.

Dear Reviewer qU5A,

As the discussion phase is approaching the end, we sincerely hope you could find some time to review our response. We hope to fully address your concerns.

We understand that your time is valuable and you may be busy with other things. However, your insights would be extremely valuable for improving our work.

We greatly appreciate your consideration.

Best,

The Authors

Thank you for your appreciation to our work! We supplement more descriptions, experiments and analyses to support the claims and effectiveness of our method. If you have any other confusion, please feel free to raise more comments :)