Overcoming Vocabulary Mismatch: Vocabulary-agnostic Teacher Guided Language Modeling

Thank you for the reviewer’s helpful feedback. We are pleased that the reviewer recognized the motivation and analysis of our work. We provide our response as below.

Clarification of the “Vocabulary Mismatch” Notion. (Claims And Evidence)

We appreciate the chance to clarify our terminology and the motivation behind it.

• Tokenizer Differences Originate from "Vocabulary" Choices: The issue of "tokens extracted by different models" fundamentally stems from the fact that these models use different tokenizers. When building a tokenizer, the primary goal is to construct a vocabulary. Since different model series typically start with tokenizers based on distinct vocabularies, we believe that the fundamental reason for the observed mismatch lies in the vocabulary itself. Thus, we refer to this problem as a vocabulary mismatch.

• "Vocabulary" Mismatch Leads to "Token" Sequence and Logit Mismatches: The consequences of vocabulary mismatch manifest as differences in token sequences and divergences in logit distributions. Since the token sequence is determined by the tokenizer’s vocabulary, we treat “token sequence mismatch” as a downstream result of vocabulary mismatch. Moreover, in transformer-based models, the vocabulary size directly determines the dimensionality of the final logit layer. Therefore, “logit distribution divergence” is more precisely attributed to differences in vocabulary, rather than just token-level differences.

• Terminological Consistency with Prior Work: A large body of prior work [1,2,3] has used terms in a way that implicitly refers to mismatches in vocabulary. In this context, using the term vocabulary mismatch aligns with established usage in the field. The suggestion that it is misleading may in fact run counter to prevailing terminology practices in related literature.

[1] Cui et al., Multi-level optimal transport for universal cross-tokenizer knowledge distillation on language models. In AAAI 2025.

[2] Boizard et al., Towards cross-tokenizer distillation: the universal logit distillation loss for LLMs. TMLR, 2025.

[3] Xu et al., Bridging the gap between different vocabularies for llm ensemble. In NAACL 2024.

Highlighting Our Contribution on Vocabulary Mismatch in Pretraining. (Weaknesses)

We would like to highlight that addressing vocabulary mismatch in the pretraining stage is one of the main contributions of our work. While knowledge distillation during pretraining is a well-established technique, in practice, vocabulary mismatch often limits the use of diverse teacher models. To our best knowledge, this aspect remains underexplored in the pretraining literature, and we believe this is the first work to directly address vocabulary mismatch in this context. VocagnoLM offers a simple yet effective solution to this problem, with practical benefits as follows:

• Practicality of Token-Level Lexical Alignment

  • While our token-level lexical alignment may not constitute a novel theoretical contribution, we argue that this is precisely what makes it a practically effective solution to the vocabulary mismatch problem. Our proposed method performs lexical alignment based on character offsets, which theoretically ensures that every student token must be contained within one or more teacher tokens. In practice, as shown in Figure 6, we observe 100% Intersection-of-String (IoS), demonstrating that two token sequences can be efficiently and accurately aligned using only a simple approach.
  • Furthermore, this practical nature allows for both online mapping during training and offline mapping during preprocessing. This flexibility is particularly advantageous in large-scale pretraining scenarios, where optimization at the training stage is crucial.
  • We highlight the simplicity and efficiency of our mapping function, which supports various weighting options in the aggregation step. While we only explore a simple aggregation function in this paper, we leave the exploration of more advanced variants as future work.

• Advantages of Loss-Based Teacher Guidance for Logit Distribution Divergence

  • Even after resolving the token sequence mismatch, the issue of logit distribution divergence remains. Traditional KL divergence cannot adequately address this, and existing distance-based methods such as ULD, can incur additional information loss during alignment.
  • In contrast, our approach is based on the teacher’s loss, which allows us to circumvent the problem arising from the dimensional mismatch between vocabularies.

In summary, we decompose the issues caused by vocabulary mismatch into two sub-problems (token sequence mismatch / logit distribution divergence) and design separate solutions for each. We demonstrate that applying these techniques in the pretraining stage leads to effective improvement.

We sincerely appreciate the reviewer’s insightful comments and valuable suggestions. We are glad that the reviewer acknowledged the importance of the problem and appreciated our comprehensive analysis. Below, we provide our detailed responses to the points raised.

Distinction between "Teacher-generated Knowledge" and "Teacher-Guided Language Modeling" (Exp Designs or Analyses #1)

We would first like to clarify that the scope of this work specifically targets the pretraining stage. Unlike (sequence-level) knowledge distillation [1] approaches trained on generated tokens in the fine-tuning stage, it is difficult to define the “generated tokens” of the teacher model during the pretraining stage. In the extreme case, a teacher model could reproduce the same input corpus or generate an entirely new 15B corpus using the teacher model’s internal knowledge.

Besides, training the student model on new knowledge generated by the teacher model differs from our objective of teacher-guided language modeling within a given corpus. However, while such an approach goes beyond the scope of this work, we appreciate the suggestion and believe it as a valuable direction worth exploring further.

[1] Kim et al., Sequence-Level Knowledge Distillation, In EMNLP 2016.

Comparison with the standard KD using same tokenizer teacher (Exp Designs or Analyses #2, Weakness)

We appreciate the reviewer’s suggestion for further experiments. We present an additional experimental result on 2B tokens that offers relevant insight. Specifically, we compare our method against the standard KD approach using two teacher models of comparable performance: MetaMath-Llemma-7B (which shares the same tokenizer as the student model) and Mistral-ProXMath-7B (which uses a different tokenizer).

As shown in the table below, our method performs better even when using a teacher model with a different vocabulary, despite both two teacher models achieving the same performance on the AVG(w/o SAT) metric.

We also clarify that our Token-level Lexical Alignment is a deterministic and straightforward solution based on character offset. While simple, it ensures 100% overlap (Figure 6), and plays a critical role in enabling reliable teacher guidance across different tokenizers. We believe that our alignment mechanism contributes to the effectiveness of our method even when compared to standard KD with the same tokenizer, as demonstrated by the results in the table.

Clarification on Computational Overhead. (Exp Designs or Analyses #3)

To report the additional computational overhead, we decompose into two components compared to standard KD, Token-level Lexical Alignment and Loss-based Guidance.

• First, the loss-based guidance shares most of its computational operations with standard KLD. When measuring 1M tokens, KLD and loss-based guidance require approximately 15.36 TFLOPs, showing no significant difference in computation cost.
• On the other hand, Token-level Lexical Alignment is a CPU-bound operation, so we report latency instead of FLOPs.. When measuring on a 2048 token sequence using the Mistral-ProXMath teacher model, the alignment step takes approximately 0.047 seconds, averaged on 1000 repeated runs. While this mapping process introduces a small amount of overhead, it can be amortized during preprocessing time. Especially in pretraining stage, corpus is typically packed to the maximum sequence length, so the mapping process can be efficiently performed during preprocessing time.

We hope this addresses the reviewer’s concern regarding computational overhead. We’ll include this point in the final version.

Response for "Essential References Not Discussed"

Thank you for the new suggestion. Minitron (Muralidharan et al., 2024) has already been cited in the introduction (Line 48-49). We also agree that Pretraining Exploration is a valuable reference that highlights the importance of knowledge distillation during pretraining. In the final version, we will revise the first paragraph of the introduction to emphasize on the importance of KD-based pretraining.

We appreciate the reviewer's thoughtful feedback and constructive suggestion. We address the points raised below.

Impact of top-k threshold (Selected Ratio) on scaling trends. (Claims and Evidence #1, #2)

• As discussed in Appendix B and Figure 7a, we explore the effect of the top-k threshold using a 2B-token subset of the corpus, following the experimental setup of Lin et al [2]. Based on this analysis, we adopt a top-k threshold of 40% ( 60% Masking Ratio ) in our study (L253).
• Consistent with prior work [1,2], we observe that selecting more tokens from a fixed corpus does not necessarily lead to better performance. Importantly, as shown in Figure 7a, the optimal threshold may vary depending on the choice of teacher model. Taking into account both our empirical results in Appendix B and findings from previous study [1,2], we chose the 40% threshold.

Insights from established findings, and highlight our representative tasks. (Claims and Evidence #3)

• As mentioned in our limitations section, we agree that VocAgnoLM is designed to be broadly applicable across different models and datasets.. Due to our limited pretraining resources, we focused our efforts on the moderately scaled 15B OpenWebMath corpus. As reviewer bj8e acknowledged, although our work focuses on math/reasoning tasks, we made every effort to evaluate the model through a diverse set of representative tasks to ensure broad applicability within the domain.
• As a complementary perspective, we would like to share an intuition that motivated our design choice. Several prior works [1,2,3] have provided promising evidence that selective token/dataset training may generalize beyond narrow domains, not only on math domain corpus, but also on general-domain corpora such as ThePile, SlimPajama, StarCoderData. Motivated by these findings, our study aims to validate whether the proposed method can effectively follow the teacher model’s expertise within a domain, using a moderately scaled math-domain web corpus.
• We hope that these prior findings, along with our own results (e.g., Figure 5, which shows performance improvement using an instruction-tuned teacher model) can help estimate VocagnoLM’s potential for general-domain performance. We consider extending our experiments to broader general-domain corpora a valuable direction for future work.

[1] Mindermann et al., Prioritized Training on Points that are learnable, Worth Learning, and Not Yet Learnt, In ICML 2022.

[2] Lin et al., RHO-1: Not All Tokens Are What You Need, In NeurIPS 2024.

[3] Xie et al., Data Selection for Language Models via Importance Resampling, In NeurIPS 2023.

Comparison with scaled CPT baseline. (Experimental Designs and Analyses)

Thank you for the insightful suggestion. Following your suggestion, we conducted an additional experiment on 2B tokens.

In VocAgnoLM, we apply a top 40% threshold based on teacher guidance. To match the signal strength of important tokens, we rescaled learning rate by a factor of lr / 0.4 = 2.5x.

As shown in the table below, our method, using various teacher models, outperforms CPT (lr /= 0.4). Although CPT delivers a similar amount of signal strength by increasing the learning rate, it also amplifies the effect of unimportant tokens. These results support the effectiveness of our teacher guidance, and allow us to estimate the impact of various teacher models in terms of CPT-equivalent training scale.

Difference with MiniPLM. (Questions for Authors)

Thank you for pointing this out. We noticed that the order of the last two sentences in Section 7.2 may have been mistakenly switched. MiniPLM enables cross-vocabulary distillation through an offline KD strategy. However, our method supports both online and offline KD, offering more flexibility in the choice of pretraining strategies. Additionally, while MiniPLM performs distillation at the instance-level, our approach focuses on token-level, which marks a key difference. We will revise the sentence accordingly and describe this difference in the final version.