The Fair Language Model Paradox

## Abstract

Abstract Large Language Models (LLMs) are widely deployed in real-world applications, yet little is known about their training dynamics at the token level. Evaluation typically relies on aggregated training loss, measured at the batch level, which overlooks subtle per-token biases arising from (i) varying token-level dynamics and (ii) structural biases introduced by hyperparameters. While weight decay is commonly used to stabilize training, we reveal that it silently introduces performance biases detectable only at the token level. In fact, we empirically show across different dataset sizes, model architectures and sizes ranging from 270 270 270 270 M to 3 3 3 3 B parameters that as weight decay increases, low-frequency tokens are disproportionately depreciated. This is particularly concerning, as these neglected low-frequency tokens represent the vast majority of the token distribution in most languages, calling for novel regularization techniques that ensure fairness across all available tokens.

## 1 Introduction

A major challenge in machine learning is designing algorithms that generalize well from training data. One of the classic methods for promoting generalization (Krogh & Hertz, ; Shalev-Shwartz & Ben-David, ) is the use of regularization techniques, such as $L_{2}$ regularization, to limit model complexity. Empirical evidence from classification problems with balanced classes shows that increasing weight decay, while effectively fitting the data and minimizing training loss, generally improves performance on unseen data. However, a recent study by Balestriero et al. () demonstrates that when training classifiers like ResNet-50 (He et al., ) on computer vision tasks such as ImageNet classification (Russakovsky et al., ), higher weight decay leads to undesired behavior, causing the model to neglect certain classes. This class-dependent effect of regularization is further amplified by imbalanced class distributions, as increasing weight decay does not result in a uniform performance decline across all classes. Instead, the model underperforms on low-probability classes while performing significantly better on more prevalent ones.

This class-dependent behavior is not unique to vision tasks. In Natural Language Processing (NLP), a shift from traditional classification settings has led to a lack of attention toward token frequency and the per-token effects of regularization. In modern large language models (LLMs) trained on text data (Brown et al., ; Radford et al., ; OpenAI et al., ; Anthropic, ; Ortiz, ), the task of predicting the next token from a large vocabulary also results in significant token frequency imbalances. For instance, as shown in Figures 3, in the IMDB dataset (Maas et al., ), 95% of the total tokens in the data are captured by the top 0.01% of tokens. This indicates that the vast majority of tokens appear infrequently, while a small set of tokens dominates, creating a substantial imbalance. Additionally, the proportion of low-frequency tokens tends to increase as the vocabulary expands. This raises a critical question:

Can regularization techniques, such as weight decay, typically used to promote generalization, ensure fairness across tokens when applied to LLMs trained on imbalanced token distributions?

Contributions. In order to study that question, we investigate the influence of weight decay on the token-level prediction performance of large language models, uncovering critical insights that are often overlooked when relying solely on aggregated performance metrics. Our study provides several key contributions:

- •
The next-token classification task suffers from severe class imbalance. Figures 3 demonstrate that the class distribution follows a heavy-tailed pattern, with the vast majority of classes being low-frequency and only a small portion being high-frequency.
- •
We trained the Apple OpenELM models with $270$M and $3$B parameters, as well as Qwen2 models with $0.5$B and $1.5$B parameters, on both the IMDB dataset and its extended version IMDB-xl, using varying levels of weight decay that yielded acceptable training losses. As observed in Figure 1, the models’ performance on low-frequency tokens significantly degrades as weight decay increases.
- •
We observe that higher-frequency tokens are consistently learned faster than low-frequency tokens across multiple random seeds, with the gap in learning speed widening as weight decay increases, suggesting that regularization may disproportionately disadvantage rare tokens.

These findings expose a critical dilemma. Practitioners often use aggressive weight decay to train LLMs—intended to stabilize training—but unintentionally and silently degrade the model’s performance on low-frequency tokens, which make up the majority of the data. While conventional wisdom advocates for increased weight decay to promote generalization, in language modeling, this strategy results in the unintended consequence of “neglecting” low-frequency tokens. This introduces harmful biases, favoring more common tokens, often reflecting the language patterns of majority groups. Our results emphasize the need for regularization techniques that explicitly address token imbalances.

## 2 Related Work

Weight Decay and Generalization. $L_{2}$ regularization was initially introduced to stabilize solutions to ill-posed problems (Tikhonov, ) and later adopted to enhance the generalization of neural networks (Krogh & Hertz, ). While many studies have linked low-norm solutions to improved generalization (Neyshabur et al., ; Golowich et al., ; Bartlett & Mendelson, ; Bartlett et al., ; Arora et al., ; Galanti et al., ; Wei & Ma, ; Li et al., ), the precise relationship between $L_{2}$ regularization and generalization remains a topic of debate. Several works (Zhang et al., ; Jiang* et al., ) argue that norm-based measures alone are insufficient to fully explain generalization in deep learning. For example, Zhang et al. () found that although weight decay can improve test accuracy, the overall effect is typically modest—around $1-2$% on ImageNet. Nonetheless, other studies have demonstrated that weight decay helps alleviate (Nakkiran et al., ; Pezeshki et al., ) the double descent phenomenon (Belkin et al., ; Nakkiran et al., ) and is critical for achieving Grokking in mathematical reasoning tasks (Power et al., ; Varma et al., ).

Weight Decay, Optimization, and Inductive Biases. Despite its modest impact on generalization, weight decay is widely employed in many state-of-the-art language models, including GPT-3 (Brown et al., ), Chinchilla (Hoffmann et al., ), and LLaMA (Touvron et al., ; ; Dubey et al., ). These models are typically trained using a “one-pass” stochastic gradient descent (SGD) regime, where the optimizer directly minimizes the population error.

As shown in (Andriushchenko et al., ), the training and validation losses remain closely aligned across different levels of weight decay. While weight decay’s effect on generalization is limited, it plays a crucial role in improving optimization. For example, both the Chinchilla paper (Hoffmann et al., ) (see Figure A7) and (Andriushchenko et al., ) (see Figure 4) demonstrate that weight decay in AdamW leads to lower training loss compared to Adam, particularly toward the end of training. Other studies (van Laarhoven, ; Zhang et al., ; Li & Arora, ; Li et al., ; Lewkowycz & Gur-Ari, ) have shown that weight decay enhances training stability by controlling the “effective learning rate” in scale-invariant neural networks. Additionally, other works (Galanti et al., ; Rangamani & Banburski-Fahey, ; Pan & Cao, ; Beneventano et al., ) reveal that weight decay contributes to various inductive biases, such as rank minimization and neural collapse, which are beneficial for network compression (Denton et al., ; Alvarez & Salzmann, ; Tukan et al., ; Yu et al., ) and downstream performance (Galanti et al., ; ).

Training with Imbalanced Classes and Minority Collapse.  Training with imbalanced classes presents a significant challenge in machine learning. Empirical studies consistently show that the weight vectors associated with the more frequent classes tend to have larger norms, which pushes the decision boundary toward the minority classes. As a result, the feature space allocated to less frequent classes shrinks, leading to a notable drop in performance (Kim & Kim, ; Kang et al., ; Cao et al., ; Ye et al., ; Liu et al., ; Kang et al., ; Balestriero et al., ). For instance, (Balestriero et al., ) showed that when training neural networks for classification of visual data, higher levels of weight decay introduce a stronger bias for the model to prioritize higher-probability classes over lower-probability ones.

To gain deeper insights into this issue, several works have investigated this phenomenon from a theoretical perspective. Fang et al. () proposed the unconstrained features model (UFM) as a simplified framework for exploring the geometric properties of the global minima in cross-entropy loss with regularization, particularly in overparameterized neural networks. In the case of a balanced dataset, they demonstrated that neural collapse (NC) occurs at any global minimizer of the loss function combined with regularization. However, in the case of class imbalance, neural networks exhibit distinct geometric patterns, and some of the NC properties no longer hold (Dang et al., ; Thrampoulidis et al., ; Hong & Ling, ; Dang et al., ). While last-layer features for samples in the same class still collapse to their respective class means (NC1), the class means and classifier weights no longer form a Simplex Equiangular Tight Frame (ETF), violating NC2 (Fang et al., ). In more extreme cases, when the imbalance becomes severe, the classifier weights for minority classes can collapse onto each other, rendering them indistinguishable from other classes (Fang et al., ). This phenomenon, referred to as “Minority Collapse,” explains the sharp decline in accuracy for minority classes in imbalanced settings.

Building on these finding, we explore whether similar behavior arises in next-token prediction tasks, where the problem can be viewed as a (very noisy) classification task, with the next token acting as the “class” for a given sequence.

## 3 Experimental Settings

Problem Setup. We focus on the task of next-token prediction in autoregressive language modeling, a self-supervised learning problem central to natural language processing. Given a sequence of tokens $x_{i}=(x_{i,1},...,x_{i,n})$, the objective is to model the conditional probability distribution $p(x_{i,t}|x_{i,<t})$ for each token position $t$. Formally, let $\mathcal{D}=\{x_{i}\}_{i=1}^{m}$ be a text corpus, where each $x_{i}=(x_{i,1},...,x_{i,n})\in\mathcal{S}=\mathcal{V}^{n}$ is a sequence of $n$ tokens, and $x_{i,t}\in\mathcal{V}$ belongs to a fixed vocabulary $\mathcal{V}$ of size $V$. We train different transformer-based models $f_{\theta}:\mathcal{S}\rightarrow[V]$ mapping from sequence to logits to minimize the regularized empirical risk:

$$ $L_{\mathcal{D}}^{\lambda}(f)=\frac{1}{m}\sum_{i=1}^{m}\left(\frac{1}{n-1}\sum_ {t=1}^{n-1}\ell(f_{\theta}(x_{i,\leq t}),x_{i,t+1})\right)+\lambda||\theta||_{ 2}^{2},$ $$

where $\ell$ denotes the cross-entropy loss, $\lambda$ is the weight decay coefficient, $||\theta||_{2}$ is the $L_{2}$ norm of the model parameters, $x_{i,\leq t}$ the sequence up to position $t$ in the $i$-th sample, and $x_{i,t+1}$ is the next token to be predicted.

Model Architecture.
For our experiments, we train multiple models, including the Apple OpenELM models with $270$M and $3$B parameters (Mehta et al., ), as well as the Qwen Qwen2 models with $0.5$B and $1.5$B parameters (Yang et al., ). The small models ($<1$B parameters) were configured with a context length of $128$ whereas the large ones with a context length of $64$. These architectures’s moderate size allows us to conduct multiple training runs with different $\lambda$ values, enabling a comprehensive exploration of weight decay’s impact on token-level dynamics across various regularization choices.

Dataset and Tokenizer.
We trained our models on the IMDB dataset (Maas et al., ), a widely used benchmark due to its balanced sentiment-labeled data. The IMDB training split contains $25000$ samples. For the scope of this study, we discarded the labels, focusing solely on the raw text data to analyze the impact of hyperparameters on token-level generation performance. Additionally, we created an extended version of the dataset, termed IMDB-xl, by incorporating all the unsupervised samples from IMDB, which increased the training set to a total of $75000$ samples. We used Byte Pair Encoding (BPE) (Gage, ; Sennrich et al., ) as our tokenization method, training a tokenizer on the IMDB dataset’s training set with a target vocabulary size of $32005$ tokens. BPE ensures that both frequent and infrequent tokens are well represented, providing a suitable basis for analyzing token-level learning. This choice allows us to examine how different tokens, particularly rare ones, are influenced by various weight decay values throughout the training process.

Hyperparameters and Training.
We conducted experiments across a range of weight decay values $\lambda\in(0.0,2.0)$. To account for the stochastic nature of training, each configuration was run with the same 5 different random seeds, with results averaged and presented with confidence intervals. The models were trained using the AdamW optimizer (Loshchilov & Hutter, ) to decouple the weight decay from gradient updates, with a learning rate of $5$e$-5$ and a cosine decay schedule. A warm-up period of $10$ training steps was used for stabilization during early training. We trained each LLM for a total of $10000$ steps, evaluating token-level metrics every $100$ steps to understand how different tokens are learned and represented under varying weight decay values. For small models (less than $1$B parameters), we used a batch size of $64$ per device on a single NVIDIA A100 $32$GB GPU, whereas for large models, we used a batch size of $16$ with a gradient accumulation step of $4$ on a single NVIDIA A100 $80$GB GPU. Additionally, mixed precision training (fp16) was employed to optimize GPU memory usage, reduce computational load, and accelerate convergence.

Evaluation Metrics. Our evaluation used token-level metrics to assess how individual tokens were generated under varying weight decay configurations. Unlike traditional metrics like perplexity, we focused on whether specific tokens were under-represented during inference, revealing how certain configurations induce bias in token representation, even when aggregate performance appears stable. The average training loss, computed with cross-entropy loss over all tokens in a batch, is given by:

$$ $\mathcal{L}_{\text{avg}}=-\frac{1}{BC}\sum_{b=1}^{B}\sum_{c=1}^{C}\sum_{v=1}^{ V}y_{b,c,v}\log p_{b,c,v},$ $$

where $y_{b,c,v}$ and $p_{b,c,v}$ are the ground truth and predicted probabilities, respectively. This method implicitly favors frequent tokens, as they appear more often, leading the model to prioritize them over low-frequency tokens. To counter this imbalance, we also used a token-balanced training loss, ensuring each token contributes equally, regardless of frequency. First, we compute the cross-entropy loss for each token $\ell_{b,c}=-\sum_{v=1}^{V}y_{b,c,v}\log p_{b,c,v}$. We then average these losses by token type, yielding the final token-balanced loss:

$$ $\mathcal{L}_{\text{tok-bal}}=\frac{1}{V}\sum_{v=1}^{V}\frac{1}{|\{(b,c):y_{b,c ,v}=1\}|}\sum_{\{(b,c):y_{b,c,v}=1\}}\ell_{b,c}.$ $$

This approach ensures that low-frequency tokens contribute equally to the loss, leading to a more balanced optimization process, though it may challenge the model’s handling of rare tokens, especially under strong regularization. The computation of per-token metrics is summarized in Algorithm 1.

Per-Token Learning Speed. To quantify how quickly the model learns individual tokens during training, we introduce the per-token learning speed metric, which measures how rapidly the model minimizes the cross-entropy loss for each token. Specifically, we compute the area under the curve (AUC) of the token’s normalized loss trajectory over time, where a smaller AUC indicates faster learning. For each token, we normalize its cross-entropy losses across all training steps to the range $[0,1]$, ensuring comparability between tokens with different loss scales. Let $\ell_{t}$ represent the cross-entropy loss for a token at training step $t$, and let $\mathcal{L}=\{\ell_{1},\ell_{2},\dots,\ell_{T}\}$ be the set of losses over $T$ steps. We define the normalized loss at each step as: $\tilde{\ell_{t}}=\frac{\ell_{t}-\min(\mathcal{L})}{\max(\mathcal{L})-\min(
\mathcal{L})}$. The area under the normalized loss curve is calculated as: $\text{AUC}(\mathcal{L})=\int_{0}^{T}\tilde{\ell_{t}}dt$. The learning speed $S$ is then defined as: $S=1-\frac{\text{AUC}(\mathcal{L})}{T}$. This ensures that $S$ ranges from 0 to 1, with higher values indicating faster learning. If the range of losses is zero (i.e., $\max(\mathcal{L})=\min(\mathcal{L})$), we set $S$ to zero, as no learning occurs.

## 4 Empirical Results

### 4.1 Textual Data is Highly Imbalanced

As a preliminary step in our study, we investigated how imbalanced textual data actually is. For this purpose we conducted several statistical evaluations of the IMDB dataset in order to verify that indeed the vast majority of tokens are of low-frequency and a small minority of the tokens appear very often in the data as predicted by the Zipf law.

Token Frequency Percentiles. Figure 3(a) highlights the extreme token frequency imbalance in the IMDB dataset, tokenized using a BPE tokenizer with a vocabulary size of $32005$. The vast majority of the token frequency mass is concentrated in a very small portion of high-frequency tokens. Specifically, 95% of the total tokens in the data is captured by the top 0.01% of tokens, which demonstrates the steep distribution of token frequencies, where very few tokens dominate. To illustrate the calculation: suppose we have a dataset with 100 samples, each consisting of 10 tokens. If a set of 1% of the tokens in the vocabulary appear 800 times across those 1000 tokens, we say that 80% of the total tokens in the data is captured by the top 1% of tokens. Figure 3(b) further emphasizes this imbalance, showing how the proportion of low-frequency tokens (those below the $95$th percentile in the data) grows as the vocabulary size increases. This suggests that token imbalance is inherent to language and further amplified as the vocabulary size expands with larger tokenizers.

Figure: Figure 3: Comparison of token frequency distribution and the ratio of low-frequency tokens across varying vocabulary sizes for the IMDB dataset. The left plot shows the token frequency distribution with cumulative frequency thresholds (50%, 80%, and 95%) marked. The right plot illustrates how the ratio of tokens below the 95th percentile increases with vocabulary size, converging to $\approx 0.85$.
Refer to caption: x5.png

Figure: Figure 4: Impact of Weight Decay on Cross-Entropy. Average training loss (blue) and class-balanced loss (orange) increase with weight decay. The class-balanced loss is more sensitive due to its focus on low-frequency tokens.
Refer to caption: x7.png

**Table 1: Impact of weight decay ($\lambda$) on model performance metrics (mean$\pm$std).**
| Weight Decay $\lambda$ | 0.0 | 0.1 | 0.3 | 0.5 | 1.0 |
| --- | --- | --- | --- | --- | --- |
| Training Loss | 0.051$\pm$0.000 | 0.052$\pm$0.000 | 0.054$\pm$0.000 | 0.056$\pm$0.000 | 0.068$\pm$0.001 |
| Per-Token Loss | 0.066$\pm$0.001 | 0.070$\pm$0.001 | 0.078$\pm$0.001 | 0.090$\pm$0.002 | 0.163$\pm$0.004 |
| Per-Token PPL | 1.069$\pm$0.003 | 1.073$\pm$0.003 | 1.081$\pm$0.003 | 1.095$\pm$0.003 | 1.177$\pm$0.003 |
| Per-Token Accuracy (%) | 98.798$\pm$0.031 | 98.781$\pm$0.029 | 98.778$\pm$0.032 | 98.759$\pm$0.028 | 98.714$\pm$0.029 |

Figure: Figure 6: (a) The relationship between token frequency (log-scale) and per-token cross entropy (symlog-scale) across frequency bins, representing powers of $3$. Lower-frequency tokens exhibit significantly higher cross entropy, indicating weaker learning. (b) The relationship between token frequency and accuracy across frequency bins. Higher-frequency tokens achieve higher accuracy, while lower-frequency tokens demonstrate more variability in performance. Both plots are $\lambda=1.0$.
Refer to caption: extracted/5929542/figures/impair/imdb-wd10-per-token-loss.png

Figure: Figure 7: We compare the (a) per-token cross-entropy loss and (b) learning speed across token frequencies when training an LLM with weight decay $\lambda=0.1$ and $\lambda=1.0$, on the IMDB and IMDB-xl datasets using a BPE tokenizer with a vocabulary size of $32005$. As weight decay increases, the model disproportionately disregards low-frequency tokens, which make up the vast majority of tokens in language datasets. Low-frequency tokens suffer from higher cross-entropy loss and reduced learning speed, while high-frequency tokens remain largely unaffected. Critically, the degradation of low-frequency token performance happens silently, as the average training loss, monitored by practitioners, remains largely unchanged across different levels of weight decay.
Refer to caption: x8.png

### 4.2 Performance Impairment in Low-Frequency Tokens

Average vs. Token-Balanced Training Loss. To highlight the effect of weight decay on low-frequency tokens, we compared the average training loss with the token-balanced training loss. Figure 5 presents both metrics at the end of training for models trained with varying weight decay values ($\lambda\in\{0.0,0.1,0.3,0.5,1.0\}$). As shown, the token-balanced loss rises sharply with increasing weight decay, while the average training loss shows only a slight increase as $\lambda$ moves from $0.0$ to $1.0$. This discrepancy occurs because weight decay disproportionately affects low-frequency tokens, which are given equal importance in the token-balanced loss. In contrast, the average training loss places much less emphasis on low-frequency tokens, making the effects of weight decay less noticeable. For exact values, including per-token perplexity and training accuracy, see Table 1.

Per-Token Performance vs. Weight Decay. Beyond comparing the average and token-balanced training losses, we also investigated how increasing weight decay influences performance on low- and high-frequency tokens separately. In this experiment, we examine how increasing weight decay affects the per-task loss function for tokens of different frequencies. To this end, we trained multiple versions of the same model (e.g., Apple OpenELM $270$M) with varying degrees of weight decay on a given dataset (e.g., IMDB) and compared the average loss function for low-frequency and high-frequency tokens. As shown in Figure 1, the performance on high-frequency tokens remains largely unaffected by the increase in weight decay, in contrast to the loss for low-frequency tokens, which increases significantly with higher weight decay. Here, low-frequency tokens are those that appear between $3^{0}$ and $3^{1}$ times as the next token in the data, while the highest-frequency tokens appear between $3^{9}$ and $3^{10}$ times as the next token.

Per-Token Performance vs. Frequency. As a next step, we would like to visualize the distribution of per-token losses and accuracies across tokens of varying frequencies. In Figure 6, tokens are grouped into bins, where the $i$th bin contains tokens with frequencies between $3^{i-1}$ and $3^{i}$. For each token, we plotted in Figure 6 (a) its average per-token cross-entropy and in Figure 6 (b) the per-token accuracy, averaging over all sequences where that token is predicted and over all random seeds. As shown, model performance improves significantly and consistently as token frequency increases, with high-frequency tokens benefiting the most. In Figure 5, we replicated this analysis for per-token learning speed, which also monotonically improves as token frequency increases, following the same pattern observed for loss and accuracy.

High-Frequency Tokens are Learned Faster. In Figure 7 (b), tokens are divided into bins, selected in the same way as before, to compare the average token learning speed (as defined earlier) and its standard deviation for two models: one trained without weight decay ($\lambda=0$, blue) and the other with weight decay ($\lambda=1$, orange). The percentage of tokens in each bin is also shown, with the majority concentrated in the low-frequency bins. As observed, less frequent tokens are learned more slowly by both models. Notably, the gap in learning speed between the models widens for lower-frequency tokens but diminishes for higher-frequency tokens. This indicates that increasing weight decay disproportionately deprioritizes low-frequency tokens.

## 5 Theoretical Discussion

It is well known that the class frequency is positively correlated with the norm of the top layer classifier of the given class (Kang et al., ; Huang et al., ; Kim & Kim, ). For instance, Dang et al. () considered a variant of unconstrained features model (UFM) (Fang et al., ), in which the features are constrained to be non-negative, motivated by the fact that features are usually the output of ReLU activations in many common architectures. We use their framework to analyze the influence of the token frequency and the weight decay on the per-token loss function. Formally, suppose we have a set of possible tokens $\mathcal{V}$ (one-hot encodings of the numbers in $[V]$) and a dataset $\mathcal{D}=\cup^{V}_{k=1}\{x_{k,i}\}^{n_{k}}_{i=1}$ of sequences $x_{k,i}=(x_{k,i,1},\dots,x_{k,i,n})$ of tokens with the next token being $k\in[V]$ (whose one-hot encoding is $y_{k}$). Within the model proposed in (Dang et al., ), for each sequence $x_{k,i}$ we learn an unconstrained feature representation $h_{k,i}\in\mathbb{R}^{d}$ together a linear layer $W\in\mathbb{R}^{V\times d}$ using the Cross-Entropy loss:

$$ $\displaystyle\min_{W,H}$ $\displaystyle\frac{1}{N}\sum_{k=1}^{V}\sum_{i=1}^{n_{k}}\ell_{\textnormal{CE}} (Wh_{k,i},y_{k})+\frac{\lambda_{W}}{2}\|W\|_{F}^{2}+\lambda_{H}\|H\|_{F}^{2},$ (1) s.t. $\displaystyle H\geq 0,~{}\lambda_{W}>0,~{}\lambda_{H}>0,$ $$

where $\ell_{\textnormal{CE}}(z,y_{k})=-\log\left({\textstyle\frac{\exp(z_{k})}{\sum^
{V}_{i=1}\exp(z_{i})}}\right)$, $H:=[h_{1,1},\ldots,h_{1,n_{1}},h_{2,1},\ldots,h_{V,n_{V}}]\in\mathbb{R}^{d
\times N}$ (where $N=\sum^{V}_{k=1}n_{k}$) are the learned representations for each sample $(k,i)$ and $H\geq 0$ denotes entry-wise non-negativity. In addition, $W=[w_{1},w_{2},\ldots,w_{V}]^{\top}\in\mathbb{R}^{V\times d}$ to be the last-layer weight matrix, with $w_{k}\in\mathbb{R}^{d}$ being the $k$-th row of $W$.

While the above formulation does not exactly match the practice (since $H$ is a matrix of free parameters instead of the outputs of a neural network), this abstraction can help understand how the per-token frequency $n_{k}$ and the regularization parameters $\lambda_{W}$ and $\lambda_{H}$ influence the per-token loss and classifier. For instance, according to Theorem 4.1 and Proposition 4.3 in (Dang et al., ), if $d\geq V$ then any global minimum of Eq. 1 satisfies:

- (a)
Within-class feature collapse: $\forall~{}k\in[V],~{}i\neq j\in[n_{k}]:~{}h_{k,i}=h_{k,j}=\mu_{k}$.
- (b)
Class-mean orthogonality: $\forall~{}k\neq l:~{}\mu_{k}^{\top}\mu_{l}=0$.
- (c)
Class-mean norm: $\|\mu_{k}\|^{2}=\sqrt{\frac{\lambda_{W}(V-1)}{\lambda_{H}Vn_{k}}}M_{k}$.
- (d)
Weight norms: $\|w_{k}\|^{2}=\sqrt{{\textstyle\frac{\lambda_{H}}{\lambda_{W}V^{3}(V-1)}}}
\left((V-1)^{2}\sqrt{n_{k}}M_{k}+\sum^{V}_{j=1}\sqrt{n_{j}}M_{j}\right)$.

Here, $M_{k}=\left[\log\left((V-1)\left({\textstyle\frac{\sqrt{n_{k}}}{N\sqrt{\frac{V
-1}{V}\lambda_{W}\lambda_{H}}}}-1\right)\right)\right]_{\diamond}$, where the function $[x]_{\diamond}$ returns $x$ if $x$ is defined and is positive and 0 0 otherwise.

The above series of observations imply that $\mu_{k}=0$ if and only if $M_{k}=0$. As a result (summarized in Corollary 4.6 in (Dang et al., )), if $n_{k}\leq\lambda_{W}\lambda_{H}N^{2}{\textstyle\frac{V-1}{V}}$ (which implies $M_{k}=0$) the model avoids learning the $k$th token. In particular, the number of tokens that are neglected by the model increased when increasing the level of weight decay.

###### Proposition 5.1 .

Suppose $d\geq V$, then any global minimizer $(W,H)$ of the problem obeys $\ell_{\textnormal{CE}}(Wh_{k,i},y_{k})~{}=~{}\log\left(\sum^{V}_{j=1}\exp\left
({\textstyle\frac{M_{j}}{V^{2}}}\right)\right)-M_{k}$.

We observe that the loss function $\ell_{ij}=\ell_{\textnormal{CE}}(Wh_{k,i},y_{k})$ can be decomposed into two components: the first part, $\ell^{\prime}$, is independent of both $k$ and $i$, while the second part depends on $-M_{k}$. Consequently, tokens $k$ associated with larger values of $M_{k}$ incur a smaller per-token loss. Since $M_{k}$ increases with $n_{k}$, it follows that the per-token loss is smaller for tokens of higher frequency. This aligns with the results in Fig. 7 (a), where the loss function decreases monotonically for higher-frequency tokens when training with weight decay.

Now, suppose we set $\lambda_{W}=\lambda_{H}=\lambda$. The derivative of the per-token loss $\ell_{k,i}$ with respect to $\lambda$ is given by: $\frac{\partial\ell_{k,i}}{\partial\lambda}=\frac{\partial\ell^{\prime}}{
\partial\lambda}+\left(\frac{1}{\lambda}+\frac{N\sqrt{V-1}}{\sqrt{n_{k}V}-
\lambda N\sqrt{V-1}}\right)\cdot\mathbb{I}\left[n_{k}>\lambda^{2}N^{2}\frac{V-
1}{V}\right]$. We note that the first term of the derivative, $\frac{\partial\ell^{\prime}}{\partial\lambda}$, is independent of $n_{k}$, while the second term is monotonically decreasing with respect to $n_{k}$, provided that $n_{k}>\lambda^{2}N^{2}\frac{V-1}{V}$ ($k$ is a non-neglected token). Therefore, the derivative of the loss for non-neglected tokens $k$ is higher for smaller values of $n_{k}$. In particular, the loss for low-frequency tokens grows at a faster rate compared to high-frequency tokens when increasing $\lambda$. This can be observed in Fig. 7 (a) and Fig. 5, where, although the losses generally increase across all types of tokens as a function of $\lambda$, they increase more significantly for low-frequency tokens.

As a final note, although one might think that the observations made in Sec. 4 could be due to poor training, these theoretical results emphasize that they are actually caused by a fundamental issue when training next-token predictors. The results above apply to the global minima of the objective, indicating that they pertain to situations where the training was in fact optimal.

###### Proposition 5.1 .

## 6 Conclusion

We investigated the impact of weight decay on token-level learning dynamics in large language models. Our findings reveal critical insights into how weight decay affects the learning process of individual tokens—effects that are hidden when relying solely on aggregated metrics. We demonstrated that increasing weight decay disproportionately harms the performance of low-frequency tokens, even when the overall average loss remains largely unchanged. Additionally, we observed that higher-frequency tokens are generally learned faster than their low-frequency counterparts. This interplay between token frequency, performance, and regularization highlights the nuanced effects of training techniques on different parts of the model’s vocabulary.

These results expose a significant oversight in current LLM training practices. Weight decay is commonly employed to reduce overfitting and enhance optimization. While this seems beneficial at first—due to improved convergence and stability in overall loss metrics—our analysis uncovers a hidden pitfall: weight decay can severely compromise the model’s ability to handle low-frequency tokens. Crucially, this degradation goes unnoticed when only aggregated loss metrics are considered. This discrepancy between aggregated performance and token-specific learning underscores the need for fine-grained, token-level evaluations. Without such assessments, models risk sacrificing performance on rare or specialized vocabulary, potentially limiting their effectiveness in domains that require precise handling of low-frequency terms.

As illustrated in Figure 3, the imbalance becomes more pronounced as vocabulary sizes increase. This is particularly concerning as vocabulary sizes are continuously expanded to improve model performance and broaden capabilities (Takase et al., ; Tao et al., ; Toraman et al., ). For example, while LLaMA-v1 (Touvron et al., ) and LLaMA-v2 Touvron et al. () used a vocabulary of $32000$ tokens, LLaMA-v3 (Dubey et al., ) expanded this to $128256$ tokens, Qwen2 (Yang et al., ) further extended it to $151936$ tokens, and Gemma-2 (Team et al., ) increased the size to $256128$ tokens.

Broader Impact. Our work enhances understanding of how weight decay affects token-level performance in LLMs while remaining undetected by aggregated metrics. This is crucial for ensuring fairness and reliability in real-world applications. While our societal impact is indirect, it underscores the need for more granular evaluation metrics to detect and mitigate potential biases on low-frequency tokens. These findings contribute to the development of more equitable and robust AI systems.

## Appendix A Proof of Proposition 5.1

See 5.1

###### Proof.

By combining (a-d) in Theorem 4.1 in (Dang et al., ), we have:

$$ $\displaystyle w^{\top}_{k}h_{k,i}{}$ $\displaystyle=~{}w^{\top}_{k}\mu_{k}$ $\displaystyle=~{}\sqrt{\frac{\lambda_{H}}{\lambda_{W}V(V-1)}}\left(V\sqrt{n_{k }}\mu^{\top}_{k}\mu_{k}-\sum_{m=1}^{V}\sqrt{n_{m}}\mu^{\top}_{m}\mu_{k}\right)$ $\displaystyle=~{}\sqrt{\frac{\lambda_{H}}{\lambda_{W}V(V-1)}}\left(V\sqrt{n_{k }}\sqrt{\frac{V-1}{V}\frac{\lambda_{W}}{\lambda_{H}}\frac{1}{n_{k}}}M_{k}\right)$ $\displaystyle=~{}M_{k}$ $$

Similarly, we can show that $w^{\top}_{j}h_{k,i}=\frac{M_{j}}{V^{2}}$. Hence, the loss function is equal to:

$$ $\ell_{\textnormal{CE}}(Wh_{k,i},y_{k})~{}=~{}-\log\left({\textstyle\frac{\exp( w^{\top}_{k}h_{k,i})}{\sum^{V}_{j=1}\exp(w^{\top}_{j}h_{k,i})}}\right)~{}=~{} \log\left(\sum^{V}_{j=1}\exp\left({\textstyle\frac{M_{j}}{V^{2}}}\right)\right )-M_{k}.$ $$

∎

###### Proof.

## Appendix B Per-Token Metrics

This appendix presents Algorithm 1, which details our method for computing per-token metrics used throughout this paper. This algorithm is central to our analysis, enabling the calculation of fine-grained token-level performance measures that underpin our study’s findings.

Figure: Algorithm 1 Token-level Metric Computation Algorithm

## Appendix C IMDB Token Distribution

As an additional evaluation, we examine the token frequency distribution of the IMDB dataset, as shown in Figure 8. Using a BPE tokenizer with a 32005-token vocabulary, we illustrate the distribution of high-frequency and low-frequency tokens. The histogram reveals a striking imbalance typical of natural language datasets: a few tokens occur extremely frequently, while most unique tokens appear rarely. This long-tail distribution is a fundamental characteristic of language data.

Figure: Figure 8: Token frequency distribution of the IMDB dataset Maas et al. () using a BPE tokenizer (vocab size $32005$). y-axis: number of unique tokens per frequency bin. x-axis: token frequency bins (logarithmic scale). Rightmost bin shows 3 tokens (the, punctuation) appearing $3^{9}$ to $3^{10}$ times each. This highlights high-frequency tokens (the, punctuation) and rare tokens ("pilgr", "cron") in lower bins, illustrating the long-tail distribution typical in natural language datasets.
Refer to caption: x16.png

## Appendix D Class Imbalance in Vision

In order to show that the behavior reported in the main text is not specific to LLMs but holds also for other types of domains, we experimented with image classification. We present an analysis of how dataset balance and weight decay affect the performance of a ResNet9 model (He et al., ) trained for CIFAR10 (Krizhevsky & Hinton, ) classification with imbalanced classes. This highlights the interplay between data distribution, regularization, and model performance. To obtain the imbalanced CIFAR10 dataset, we target a number of samples per class of [91, 142, 222, 347, 541, 845, 1317, 2055, 3205, 5000]. We randomly subsampled the corresponding classes in order to obtain those number of images per class.

Figure: Figure 9: Performance comparison of ResNet9 on balanced (left) and imbalanced (right) CIFAR10 datasets. The x-axis represents the 10 CIFAR10 classes, while the y-axis shows the test accuracy for each class. Different lines correspond to various weight decay values used during training. In the balanced dataset, performance is relatively uniform across classes, with overall homogenous degradation at high weight decay. The imbalanced dataset reveals a clear trend where less frequent classes perform worse, with this effect exacerbated by stronger weight decay values.
Refer to caption: x17.png

## Peer Reviews

confidence: 3
contribution: 3
initial_rating: 5
presentation: 3
questions: Did the authors consider alternative regularization methods beyond weight decay during their experiments?

How would the findings differ if tested on corpora with varied linguistic features, such as highly inflected languages or low-resource languages?

Would the proposed metrics, such as per-token learning speed, generalize effectively to larger and more diverse datasets?
soundness: 3
strengths: The paper brings forward a nuanced perspective on weight decay, highlighting an often-overlooked effect on low-frequency tokens in LLMs. This is particularly timely given the widespread use of weight decay without token-level monitoring.


The study uses multiple models with varying architectures and sizes across different datasets, demonstrating the robustness of the findings.
summary: The paper titled The Fair Language Model Paradox presents an investigation into token-level biases in Large Language Models (LLMs) induced by weight decay, a common regularization method. The authors explore how weight decay affects low-frequency tokens disproportionately, leading to performance degradation in these tokens even as aggregated training loss metrics remain stable. This study reveals the hidden biases against low-frequency tokens, calling for more equitable regularization techniques to ensure fairness across the token distribution.
weaknesses: The use of only the IMDB dataset (including an extended version) raises concerns about the generalizability of the results across other types of text data. Testing on a more varied set of corpora (e.g., diverse languages or topics) would strengthen the claims about low-frequency token bias.

The paper’s theoretical discussion on the link between token frequency, regularization, and loss functions feels dense and somewhat disjointed from the empirical findings. A clearer integration of these theoretical insights into the experimental results would enhance the readability and cohesion.

The broader impact section is sparse, particularly given the potential implications of token biases in LLMs on marginalized dialects or low-resource languages. The authors could deepen their exploration of societal impacts to underscore the relevance of their findings.

---

confidence: 4
contribution: 2
initial_rating: 6
presentation: 3
questions: I don't have questions.
soundness: 3
strengths: 1. Innovative thinking on model weight decay for unbalanced class distribution data: the article proposes that the increased weight decay of large language models leads to model underperformance on low-frequency tokens and significantly better performance on high-frequency tokens, which can lead to model bias and unfairness. It triggers further thinking in the field of NLP on the contradiction between model generalization performance and model bias under long-tailed data, and focuses the attention of large language models on token-level performance. 
2. informative textual analysis and experimental validation: the article experimentally validates the average model performance, the impact of token-level performance under weight decay and own word frequency, and theoretically analyzes why the loss function of high-frequency tokens monotonically decreases when trained with weight decay, rigorously arguing the point of view from the theory and experiments
3. challenges to existing practices: the paper challenges the weight decay technique commonly used in current LLM training practices, and emphasizes the need to develop new regularization techniques to ensure the fairness of all tokens.
4. concise language and clear logic: the paper is concise and logical, and the experimental results are well organized to help readers clearly understand its research contributions.
summary: The article proposes that the increased weight decay of large language models leads to model underperformance on low-frequency tokens and significantly better performance on high-frequency tokens, which can lead to model bias and unfairness. It triggers further thinking in the field of NLP on the contradiction between model generalization performance and model bias under long-tailed data, and focuses the attention of large language models on token-level performance.
weaknesses: 1.dataset limitation: although the paper uses the IMDB dataset for experiments, the dataset is limited in types and domains, and may not be able to fully represent the model's performance in diverse tasks and domains.
2.Lack of different regularization comparison experiments: the paper lacks comparison experiments for the effects of different regularization techniques, for example, comparison with other types of regularization methods (e.g., dropout, data augmentation, etc.), which can make the experimental results more convincing.
3.There's still room to explore: although the article puts forward the contradiction between the fairness at the token level and model generalization under the existing regularization techniques, it does not put forward a proven solution, which is regrettable.

---

confidence: 4
contribution: 2
initial_rating: 3
presentation: 2
questions: Page 4, line 221 mentions that the experiments were conducted on an A100 32GB GPU, but Nvidia A100 does not have a 32GB version. It is suspected that this should be Nvidia V100 instead.
soundness: 2
strengths: 1. The paper presents a novel perspective for analyzing the performance of large language models. The author observed the difference in the learning high-frequency and low-frequency tokens, and identifies the cause of the differences, namely, the weight decay regularization technique. The experimental results demonstrate a significant correlation between weight decay and the loss of low-frequency tokens.

2. In addition to empirical conclusions, the authors also provide a theoretical disscussion on the impact of weight decay on per-token loss for different token frequencies.
summary: The authors analyze the relationship between weight decay and token-level loss in large language models. As a regularization technique that stabilizes model training, weight decay is widely applied in the training of large language models. Through the training of models ranging from 270M to 3B,  the authors have discovered that as weight decay increases, the ability of the model to learn low-frequency tokens deteriorates, which is reflected by an increasing loss for these low-frequency tokens. Additionally, the gap between the losses for low-frequency and high-frequency tokens also grows larger. This phenomenon suggests that there is a need to develop new regularization techniques to avoid this issue.
weaknesses: 1. The experiments in this paper use the IMDB corpus for model training. However, this corpus is biased and differs significantly from mainstream pre-training corpora. Consequently, it may not adequately reflect potential issues in mainstream large language model training.

2.  The experiments in this paper are based on training sequences of lengths 128 and 64, which are somewhat too short for large language model (LLM) training. For instance, in Figure 2, the tokenized tokens using the llama3 tokenizer already consists of 92 tokens, which appears to be relatively short text even in common pre-training corpora. A widely recognized viewpoint is that the context window has a significant impact on the per-token loss of LLMs, and longer context windows can help the model learn better. Mainstream models typically use a length of around 8192, and there is a considerable gap between this window length and the lengths used in the authors' experiments. Consequently, whether these conclusions can be generalized to mainstream large language models remains to be further validated.

3. This experiment compared the impact of weight decay \(\lambda\) ranging from 0.0 to 2.0 on the model. From Figure 1, Figure 4, and Table 1, it can be observed that starting from \(\lambda = 0.3\), there is a noticeable change in the per-token loss for low-frequency tokens. However, most current LLMs set the weight decay \(\lambda\) to 0.1, which, as shown in Table 1, has a negligible impact on the model. Therefore, whether the issue raised by the authors is universal remains to be further examined.

---

confidence: 3
contribution: 2
initial_rating: 5
presentation: 3
questions: My main point of confusion revolves around whether low-frequency tokens and high-frequency tokens should be treated equally in large language models. Could you provide some practical examples to illustrate this?
soundness: 3
strengths: 1.This represents a novel finding, as prior work has typically focused on aggregated training loss measured at the batch level, neglecting the detailed dynamics of individual tokens.
2.The experiments demonstrated that the models' performance on low-frequency tokens significantly deteriorates as weight decay increases.
summary: Inspired by a discovery in the computer vision field regarding unbalanced sample classification tasks, where regularization methods are more effective for larger classes, the authors investigated the impact of weight decay—a common regularization technique—on token-level learning dynamics in large language models. The experimental results demonstrated that as weight decay increases, the performance of low-frequency tokens is disproportionately affected, while high-frequency tokens are learned faster than their low-frequency counterparts. This is a novel finding, as previous work typically relies on aggregated training loss measured at the batch level, overlooking token-specific dynamics.
weaknesses: 1.The paper does not offer specific insights on the implementation of regularization techniques that ensure fairness across all tokens. While the impact of weight decay on low-frequency tokens is highlighted, there is no detailed discussion on how to address this imbalance or propose alternative regularization methods that might mitigate the disproportionate effect on low-frequency tokens, ensuring more equitable performance across the entire vocabulary.
2.There is a lack of experiments or other evidence demonstrating the necessity of treating low-frequency tokens with the same level of importance as high-frequency tokens. Furthermore, it remains unclear whether this approach could lead to other issues, such as affecting the stability of model training or overall model performance. Addressing these concerns would be essential to understanding the broader implications of implementing equal importance for all tokens in language models.

Meta-review: The paper investigates the impact of weight decay, a common regularization technique, on token-level learning dynamics in LLM. The key finding is that as weight decay increases, the performance of low-frequency tokens is disproportionately degraded, while high-frequency tokens are learned faster. This contradicts the typical evaluation of LLMs based on aggregated training loss, which fails to capture these token-level biases. The authors demonstrate this effect empirically across different model sizes, architectures, and dataset sizes, highlighting the ubiquity of this issue in LLM training.

*Strengths：*
- The paper brings forward a novel and important perspective on the hidden biases against low-frequency tokens introduced by weight decay, a widely used regularization method in LLM training.
- The study uses multiple models and datasets, demonstrating the robustness of the findings.
- The paper challenges the current practices in LLM training and calls for the development of more equitable regularization techniques.

*Weaknesses*
- Generalizability: The use of only the IMDB dataset, even in an extended version, raises significant concerns about the generalizability of the results across other types of text data and domains. The reviewers felt that testing on a more diverse set of corpora, including different languages and topics, would be necessary to strengthen the claims about low-frequency token bias.
- Incomplete theoretical integration: The theoretical discussion on the link between token frequency, regularization, and loss functions feels disjointed from the experimental findings. The reviewers suggested that a clearer integration of these theoretical insights into the empirical results would enhance the overall coherence and impact of the paper.
- Lack of proposed solutions: While the paper highlights an important issue regarding the disproportionate impact of weight decay on low-frequency tokens, it does not provide any concrete proposals or alternative regularization methods to address this problem. The reviewers felt that the absence of such solutions limits the practical significance of the work.

Given these concerns, the current version of the paper does not meet the bar for acceptance at ICLR. I encourage the authors to address the reviewers' feedback and consider resubmitting the work to a future ICLR conference or another suitable venue.

Decision: Reject