Exploiting Vocabulary Frequency Imbalance in Language Model Pre-training

Dear reviewer 5NL9,

Thank you for your time and insightful feedback. Below, we offer comprehensive responses to each of your comments.

1. "Experiments use an 85M-parameter model. Generality to ≥7B-parameter LLMs is asserted via indirect Pythia/OLMo plots."

We are currently training 450M parameter model variants and based on their training‑loss trajectories and learning dynamics, we confirm that they exhibit trends analogous to those observed in the 85M parameter models. However, due to constraints on time and GPU resources, we were unable to complete the pre-training in time. We will include this experiment in the revised manuscript.

2. "Exploiting frequency imbalance may entrench majority-language bias. Any discussions on potential mitigations could be helpful."

To address the majority language bias, one must curate corpora with more uniform language distributions achieved through oversampling underrepresented languages or undersampling majority‑language data. Additionally, developing optimization algorithms that are resilient to frequency skew such as leveraging adaptive methods like Adam in place of vanilla SGD offers a promising strategy for mitigating the effects of token‑frequency imbalance [3].

3. "The paper acknowledges higher loss on tail words but it would make the paper stronger if authors can provides more in-depth error study or harm-reduction strategy for rare-token degradation."

Exploring harm‑reduction strategies that replicate the benefits of vocabulary expansion is a compelling research direction. A prototypical example is Superbpe [2], which employs a two‑stage BPE algorithm: in the first stage (up to a threshold vocabulary size T), it operates identically to standard BPE, while in the second stage it abandons whitespace pre‑tokenization. By permitting merges across former whitespace boundaries, Superbpe directs subsequent merges toward frequent tokens, thereby preventing the introduction of new rare tokens as the vocabulary expands, prevents further increase of rare token loss.

To validate whether enlarging vocabulary size and Superbpe variants exploit the same benefits: leading to more compressible, lower complexity tokenized text, we draw on Kolmogorov complexity, which measures the minimal description length of a bitstring and thus captures its inherent compressibility [1]. For a sequence $X^{N}$ tokenized by a BPE, an asymptotic upper bound on its Kolmogorov complexity is given by

$K(X^{N}) \le NH(p) + V \log_{2} N + O(\log N)$

where $N$ is the total token count of the entire text, $H(p)$ the shannon entropy of the token distribution, and the $V \log_{2} N$ accounts for the prefix-free encoding of the token (i.e., token frequency table). Because modern language models are trained on corpora comprising billions of tokens, the $NH(p)$ dominates, yielding the practical approximation.

$K(X^{N}) \approx NH(p)$

We adopt this upper bound as a tractable proxy for data complexity, as actual Kolmogorov complexity is uncomputable. The table below reports the estimated Kolmogorov complexities for the 45.7 billion‑byte FineWeb‑Edu corpus and average downstream task performance reported by [2], comparing standard BPE at various vocabulary sizes with Superbpe variants. If no tokenizer is applied and the text is modeled as an i.i.d. source over individual bytes, the Kolmogorov complexity $K(X^{N})$ would be 45.7B bytes. These results demonstrate that both vocabulary expansion and the Superbpe variants leverage the same benefit: they reduce the complexity of the tokenized text by segmenting common character sequences into single tokens, thereby facilitating the model to capture and learn frequent patterns in the training data more easily.

Table A

Table B

Concluding Remarks

Thank you for highlighting the importance of scalability, fairness, and rare-token robustness. Your comments prompted us to (i) extend our experiments to 450 M-parameter models now in training, confirming that the key trends observed with 85 M parameters persist at larger scales; (ii) add a focused discussion of corpus-balancing and adaptive-optimizer strategies that can mitigate majority-language bias; and (iii) deepen our error analysis on tail words, contrasting vocabulary-expansion with SuperBPE as complementary harm-reduction approaches. We believe these additions substantially strengthen the paper’s claims and practical relevance.

We appreciate your insights and remain eager to pursue any further analyses or experiments that would make the study more useful to the community.

References

[1] “The No Free Lunch Theorem, Kolmogorov Complexity, and the Role of Inductive Biases in Machine Learning” ICML 2024

[2] “SuperBPE: Space Travel for Language Models” COLM 2025

[3] “Heavy-Tailed Class Imbalance and Why Adam Outperforms Gradient Descent on Language Models” Neurips 2024

Dear reviewer Md1b,

We are grateful for your time and thoughtful feedback. In the sections that follow, we provide detailed responses to your comments.

1. "Why conduct experiments only with an untied embedding model? Provide additional experiments with a tied embedding setting."

We adopt untied input and output embeddings because contemporary language models typically do not share these parameters, and this design has been employed by several pre-training interpretability studies [1,2,3] to facilitate clearer analysis. Nevertheless, in response to the reviewer’s comment, we conducted additional experiments with a tied embedding architecture. The table below presents the average frequent word loss and cross-entropy loss on Fineweb-edu. Tied embedding shows a similar tendency to untied variants.

Table A

2. "No results on downstream benchmarks such as MMLU"

Thank you for raising this point. We evaluated our models on MMLU, but their scores clustered near the random-guess baseline of 25 %, providing essentially no discriminatory power at our current compute scale. Instead, as shown in the table below, we broadened our evaluation to lighter benchmarks—OpenBookQA, Lambada (openai), WinoGrande, and BLiMP— due to the scarcity of benchmarks that allow fair downstream comparisons for small language models [3]. The table below shows that the larger vocabulary size improves downstream performance on these benchmarks. These suites have proven far more sensitive to the effects under study. If there is a particular additional benchmark you feel would shed further light on our claims, we would be happy to incorporate it in a revised version.

Table B

3. "Write Sentence concisely + add uniform distribution line in figure 1a"

We will shorten our sentences and, in response to your suggestion, add a reference line for the uniform distribution in Figure 1a and expand the y‑axis range to 0.5–1 in both Figures 1b and 1c. Thank you for your feedback.

4. "Why is Global Cross-Entropy Loss defined through the Average Per-Vocabulary Loss and Total Loss + definition issue"

Global cross‑entropy loss is the cross‑entropy computed over the entire training or validation corpus. Because cross‑entropy quantifies the average uncertainty in next‑token prediction, it can be expressed as a weighted sum of each vocabulary’s mean loss, with weights proportional to their relative frequencies. In the revised manuscript, we will clarify that t denotes a sentence drawn from the N documents and v represents a specific vocabulary entry in the tokenizer. Thank you for your feedback.

5. "Increasing the model size would lower the whole token loss, including rare tokens?"

In response to the reviewer’s query, we conducted additional experiments demonstrating that scaling model capacity disproportionately benefits the prediction of high‑frequency tokens relative to low‑frequency tokens. Consequently, the reduction in frequent‑token loss emerges as the principal contributor to the overall decrease in cross‑entropy loss. Because BPE vocabulary IDs do not reflect token frequency, we cannot directly distinguish common and rare tokens in the Pythia and OLMo‑2 training corpora. Instead, we utilize the Google Web Trillion Word Corpus frequency list, which records counts for 333,333 unique words appearing at least 10,000 times in a 1.0249 × 10^12 words. We classify the top 2,500 entries as “frequent” words and the bottom 10,000 as “rare” words. As shown in the table below, although cross‑entropy loss declines for both groups with increasing model size, the reduction for frequent tokens is markedly greater. We therefore conclude that the dominant effect of model scaling is the lowering of frequent‑token loss, which in turn drives the overall cross‑entropy reduction.

Table C

Table D

Concluding Remarks

Thank you for your close reading and constructive criticism. Your comments have prompted several concrete improvements that we believe make the paper clearer and more convincing. To address your main concerns, we will:

1. Tied-embedding validation – replicate all core experiments with tied embeddings to verify that the vocabulary-scaling trends hold.
2. Broader downstream tests – add evaluations on OpenBookQA, LAMBADA, WinoGrande, and BLiMP so that readers can observe consistent gains from larger vocabularies, even at our small-model scale.
3. Clearer writing and figures – shorten long sentences, add a uniform-distribution line to Figure 1 for easier comparison, and explain the global-loss formula in plain language.
4. Scaling analysis – extend our study from 1B to 7B parameters (Pythia, OLMo-2) to show that increased model size primarily reduces frequent-token loss, driving the overall cross-entropy drop.

We believe these revisions directly address the issues you raised and strengthen the paper’s generality and clarity. We remain eager to incorporate any further suggestions that could make the work even more convincing.

References

[1] “Paloma: A Benchmark for Evaluating Language Model Fit” Neurips 2024

[2] “Pythia: A Suite for Analyzing Large Language Models Across Training and Scaling” ICML 2023

[3] “Polypythias: Stability and Outliers Across Fifty Language Model Pre-Traning Runs” ICLR 2025

Thank you for your response! As my weaknesses have been addressed, I raised my score to borderline accept. The main reason why I have not gone higher is that the conclusion of the paper does not change the current guidance for building tokenizers and models. It does indeed help use explain the current relationship between frequent and rare tokens.

To the authors, I highly suggest including many of these experiments in the paper and working on cleaning up the notation and figures.

Dear reviewer 5AqE,

We sincerely thank you for your time and feedback. Below, we address your comments in detail:

1. "Missing key methodological details (BPE) and too narrow y range exaggerate trivial fluctuations in figure."

We only use BPE with GPT-2 pre-tokenization rule. We will increase the y ranges in a revised manuscript. Thank you for your feedback.

2. "metric name is confusing + how you measure word-level average per vocab loss?"

Thank you for the suggestion. We will standardize our terminology as follows: word type for unique word forms, word for individual occurrences of a word type, vocabulary type for subword units in the tokenizer, and token for individual occurrences of those vocabulary types. We will also revise Equation 2 as Average per-word loss to reflect these definitions. Your interpretation of the word-level per-vocabulary loss calculation is correct: in the updated manuscript, Equation 2—renamed Average per-word loss—sums the cross-entropy losses of each word’s constituent tokens, and Figure 2 reports the mean of these per-word losses across the 2,500 most frequent words in FineWeb-Edu.

3."Why embednorm does not reflect token frequency imbalance during training? + please explain this clearly in your manuscript"

Throughout training, both input and output embeddings reflect token-frequency imbalances, so frequent tokens develop larger embedding norms and thus produce higher logits. Under standard cross-entropy training, the gradient for the output embedding row of a target token $t$ is $\frac{\partial \mathcal{L}}{\partial E_{\mathrm{out},t}} = (p_t -1)h$ and $\frac{\partial \mathcal{L}}{\partial E_{\mathrm{out},j}} = p_j h\ (j\neq t)$, respectively. $p_t$ is very small at initialization, the gradient of the target token is approximately $-h$, whereas non-target rows get $p_j h$ closer to 0. Consequently, each occurrence of a frequent token pulls its embedding by roughly $\lVert h\rVert$ along the hidden state direction, causing its norm to grow in proportion to its count, while rare tokens accrue only negligible norm increases. Although $\ell_{2}$ weight decay can slow this norm inflation for frequent tokens, its uniform shrinkage causes rare-token embedding norm, which receives negligible positive gradients, to shrink near zero [6]. By constraining a fixed unit $\ell_{2}$ norm on every embedding row on every step during training, we prevent embedding norms from encoding token-frequency imbalance during training. We have included this explanation in Appendix D, but will present this more clearly in the rest of the paper. Thank you for your guidance.

4."What artifacts specifically will be released?"

We will release a 24K-196K tokenizer trained on Fineweb-edu (10B gpt2 tokens sample), openwebtext, and tokenized text, respectively. We will also release base and embednorm models trained on fineweb-edu (5 different seeds from 24K-196K vocab size).

5."cite Magnusson et al 2024 in section 3.5 properly. what’s section 3.5 novelty against Magnusson et al 2024?"

We appreciate your correction and will incorporate your feedback. While [1] first identified that global loss is dominated by common tokens, our study further demonstrates that increasing the vocabulary further amplifies this dominance. Our work explains that a larger vocabulary reduces the complexity of the tokenized text and makes it more compressible by merging frequent non-iid character patterns into single tokens and boosting their relative frequency. This increased compressibility enables the model to learn and predict these frequent patterns more easily (even under an unigram assumption), which in turn lowers the cross‑entropy loss.

6. "Why exactly bigger vocabularies increase certainty on the top tokens?"

Thank you for this insightful question. Expanding the vocabulary increases the relative frequency of common tokens as the total token count of the training corpus reduces. Even if frequent character strings are segmented as a single token (Figure 1b), expanding the vocabulary size further decreases the overall token count, which keeps increasing the relative frequency of common tokens. This amplified token‑frequency imbalance facilitates the language model’s ability to learn and predict these frequent tokens, as they occur more often with larger vocabularies. We provided a detailed analysis in Appendix A.

7. "line 34: what is meant by “marginal ones” and mistakes on “them”?"

The marginal probabilities represent each token’s relative frequency, and “mistakes on them” indicates mistakes on rare tokens. Because the conditional probability of a rare token is very low, mispredicting such tokens incurs a large penalty.

8. "Explain clearly what is the definition of segmentation efficiency. What do you want to show in figure 1c?"

We apologize for the previous ambiguity in our definition of segmentation efficiency. Segmentation efficiency gauges the average token count per word, that is, how many tokens the tokenizer needs, on average, to encode each of the 2,500 most frequent words. In Figure 1c, we aimed to illustrate that the frequent words are highly overlapping across datasets by showing that the overlap between the top N words of fineweb-edu and openwebtext remains consistently high. In response to the reviewer’s suggestion, we will expand the y‑axis range for greater clarity. Thank you for your advice.

9. "What frequent words have you used in figure 2?"

We use frequent 2,500 words in the fineweb-edu.

10. "define frequent words coverage in the caption."

We define “frequent‑words coverage” as the proportion of the model’s total loss attributable to the tokens comprising the 2,500 most frequent words in the FineWeb‑Edu corpus.

11. "Maybe just plot the loss on infrequent words in figure 2b?"

In Figure 2, we will present the average per-vocab loss for infrequent tokens for each tokenizer instead of rare words, since rare words vary across datasets and most of them are not represented by a single token.

12. "What is the shaded area showing here? STD over all strings?"

Shaded area showing 5 different runs on different random seeds

13. "Why use a box plot and why average per vocab loss is so wide here?"

We employ box plots to depict the variability in loss for the 2,500 most frequent words: the shaded box spans one standard deviation above and below the mean, and the central black line corresponds exactly to the average per‑word loss reported in Figure 2a. To enhance interpretability, we will substitute these box plots with a line graph.

14. "Frequent word loss of embednorm decreases as vocab scales up? Hard to see in box plots"

For clarity, the following table reports the exact average frequent word loss for both the baseline model and the embedding‑norm variant.

Table A

15. "why you use a common crawl dump? Isn’t it enough to show common token overlap between Fineweb-edu and downstream benchmarks?"

We want to verify on an independent validation set that expanding the vocabulary leads to reduced cross‑entropy loss by enabling the model to better capture high‑frequency words. Moreover, empirical findings in [2] report a strong negative Pearson correlation (–0.93) between CC‑Main‑2023‑40 cross‑entropy loss and reasoning benchmark performance, further motivating its use as a predictive validation set.

16. "again the caption refers to loss on infrequent tokens but doesn’t show it"

We will revise Figure 5 to match the presentation style of Figure 2. Following your recommendation, Figure 5a will illustrate the average loss for infrequent tokens.

17. "Do you have any explanation on why larger models should lower loss more on common tokens? "

It is also an intriguing research direction to understand why larger models lower loss on common tokens. We hypothesize that this effect derives from simplicity bias [3,4,5]: the inductive bias of deep networks to favor lower complexity solutions, and the bigger the model, the stronger the bias. [3] demonstrates that larger GPT‑3 models produce output text with lower complexity, which mirrors that of natural data. Moreover, [4] shows that simplicity bias increases with depth, and [5] demonstrates that the representation of models converges as model size grows. Given a training corpus that is both sufficiently large and diverse, simplicity biases enable larger models to more effectively capture the shared, high‑frequency patterns across natural text, thereby achieving lower cross‑entropy loss.

Concluding Remarks

Thank you again for your thoughtful and detailed feedback. Your comments have helped us clarify key methodological details, tighten our terminology, and strengthen both our analyses and visualizations. Should any further questions arise or additional experiments be desired, we would be delighted to address them. Your insights have greatly improved our paper, and we look forward to submitting a stronger revision.

References

[1] “Paloma: A Benchmark for Evaluating Language Model Fit” Neurips 2024 (benchmark)

[2] “compression represents intelligence Linearly” COLM 2024

[3] “The No Free Lunch Theorem, Kolmogorov Complexity, and the Role of Inductive Biases in Machine Learning” ICML 2024

[4] “The Low-Rank Simplicity Bias in Deep Networks” TMLR 2023

[5] “The Platonic Representation Hypothesis” ICML 2024

[6] “Fishing for Magikarp: Automatically Detecting Under-trained Tokens in Large Language Models” EMNLP 2024

Dear reviewer 8nWv,

We sincerely appreciate your time and effort. We respond to your comment in what follows:

1. "Altering vocabulary axis invariably changes the embedding and prediction-layer parameters, thus it may enhance model expressivity."

In response to the reviewer’s comment, we conducted a controlled experiment comparing a 24K vocabulary model (122 M total parameters) with a 49K vocabulary variant of comparable size (124 M parameters) to ascertain whether the benefits of vocabulary enlargement arise primarily from increased embedding capacity. To isolate this effect, we reduced the hidden dimension from 768 to 648 and the feed‑forward dimension from 2048 to 1728, while maintaining identical model depth which known to have a dominant impact on performance [2,3]. As shown in the table below, the 49K model incurs higher cross‑entropy loss and delivers inferior results on downstream benchmarks (ARC‑E, HellaSWAG, PIQA, and SCIQ) compared to its 24K counterpart, indicating that simply expanding the embedding layer contributes far less to model expressivity than is often assumed.

Table A

Additionally, we compare ce-loss of base model and tied embedding variants. Were increased embedding capacity the primary driver of loss reduction, the tied‑embedding model would be expected to incur substantially higher loss. However, as shown in the table below, the difference in cross‑entropy loss between these two configurations is negligible.

Table B

2."6e-4 is an optimal learning rate for every experiments (24K-196K)?"

We selected a learning rate of 6e-4 because it was employed for GPT‑3 125M models [5]. To address the reviewer’s question, we conducted a learning‑rate sweep using four rates (2.4e-3, 1.2e-3, 1.5e-4, 7.5e-5). The table below reports the resulting cross‑entropy losses and word-level average per vocab loss (we’ll rename this as average per word loss) of frequent 2500 words in fineweb-edu across different vocabulary sizes and learning rates, using the same validation dataset described in Section 3.1. Although 1.2e-3 looks more optimal learning rate, same loss curve observed across different vocabulary size.

Table C

Table D

3."Constraining the embeddings would also significantly increase rare token loss."

To address the reviewer’s concern, we conducted an additional experiment to evaluate the effect of embedding‑norm constraints on cross‑entropy loss for both frequent and infrequent tokens. Specifically, we computed the average per‑token loss for the 2,500 most frequent and 2,500 least frequent tokens using 5B character subset of fineweb-edu (our validation dataset, see section 3.1). Gap quantifies the loss difference between the base models and embedding-normalized models. While the norm constraint does raise loss on rare tokens, we find that the primary driver of the overall increase in cross‑entropy loss is the elevated loss on frequent tokens. This effect occurs because constraining the embedding norm prevents the model from encoding token‑frequency imbalance during training. We provided a detailed explanation in Appendix D. We are grateful to the reviewer for prompting this clarification.

Table E

Table F

Table G

Table H

4."“word” and “token” are used interchangeable,which as the authors also point out those are not the same things."

We apologize for having used “word” and “token” interchangeably throughout the manuscript. In the revised version, we will distinguish and apply these terms consistently. Thank you for your feedback.

5."Why use JSD for measuring heavy-tailedness of distribution?"

We use the Jensen–Shannon divergence from a uniform distribution to quantify token frequency imbalance and thereby evaluate the compressibility (or complexity) of the tokenized text alongside Shannon entropy. This work focuses exclusively on token‑frequency imbalance, not on the distribution’s tail shape; references to “heavy‑tailedness” were unintended. In particular, we will remove the phrase “leading to a heavier long‑tail distribution” from Figure 1’s caption. Thank you for drawing our attention to this issue.

6."How section 5 (discussion) connects with rest of paper"

Section 5 addresses the apparent discrepancy with prior work [4], which finds that increased token‑frequency imbalance degrades machine translation performance. We wanted to explain what leads to different conclusion in section 5. In machine translation, source and target vocabularies have minimal overlap—terms frequent in English rarely appear in French and vice versa—whereas in monolingual pre‑training, the training corpus and evaluation benchmarks share extensive lexical overlap, including many common high‑frequency tokens. We wanted to point out that in machine translation, gains from reducing frequent token loss in the source language during training do not necessarily translate into improved performance in the target language and this fundamental contrast accounts for the somewhat different conclusions of [4] and our study.

Concluding Remarks

Your feedback has been instrumental in this process, and we sincerely extend our gratitude for your invaluable insights. Should you have any inquiries or require clarifications about our rebuttal, please don't hesitate to reach out. We are eager to address any concerns and elucidate potential ambiguities in greater depth.

Reference

[1] “Over-Tokenized Transformer: Vocabulary is Generally Worth Scaling” ICML 2025

[2] “SCALE EFFICIENTLY: INSIGHTS FROM PRE-TRAINING AND FINE-TUNING TRANSFORMERS” ICLR 2022

[3] “The Impact of Depth and Width on Transformer Language Model Generalization” Arxiv

[4] “Tokenization and the Noiseless Channel” ACL 2023

[5] “Language Models are Few-Shot Learners”