M i L e Loss: a New Loss for Mitigating the Bias of Learning Difficulties in Generative Language Models (2024)

Abstract: Generative language models are usually pre-trained on large text corpus via predicting the next token (i.e., sub-word/word/phrase) given the previous ones. Recent works have demonstrated the impressive performance of large generative language models on downstream tasks. However, existing generative language models generally neglect an inherent challenge in text corpus during training, i.e., the imbalance between frequent tokens and infrequent ones. It can lead a language model to be dominated by common and easy-to-learn tokens, thereby overlooking the infrequent and difficult-to-learn ones. To alleviate that, we propose a **MiLe Loss** function for **mi**tigating the bias of **le**arning difficulties with tokens. During training, it can dynamically assess the learning difficulty of a to-be-learned token, according to the information entropy of the corresponding predicted probability distribution over the vocabulary. Then it scales the training loss adaptively, trying to lead the model to focus more on the difficult-to-learn tokens. On the Pile dataset, we train generative language models at different scales of 468M, 1.2B, and 6.7B parameters. Experiments reveal that models incorporating the proposed MiLe Loss can gain consistent performance improvement on downstream benchmarks.

## introduction generative language models like gpt-3 @xcite are generally pretrained on extensive textual data, in the manner of predicting the next token given the previous ones for each training text. recently, large generative language models have been exhibiting impressive performance on various downstream natural language tasks, like dialogue system, classification, sequence labeling, etc. @xcite @xcite et al., 2022), and attracting much attention from both academia and industry. however, previous works have overlooked an inherent issue in natural language corpus that might affect the pretraining of a language model, i.e., frequent tokens far outnumber infrequent ones. actually, zipf's law @xcite highlights the inherent imbalance of tokens in natural language datasets, i.e., a few frequent tokens would dominate a dataset while many infrequent ones only form a minor portion. for instance, 50% of the brown corpus @xcite , which comprises over a million tokens, is covered by only the top 135 most frequent tokens. the imbalance of tokens is essentially a class imbalance problem. we argue that infrequent tokens are difficult to learn due to their fewer occurrences, in contrast to the frequent ones that can be learned adequately @xcite . to confirm that, we utilize the remarkable language model llama @xcite with 6.7b parameters on the pile @xcite validation set and perform a detailed perplexity (ppl) analysis at the token level. it's worth noting that a higher perplexity is indicative of a token's higher learning difficulty. in our analysis, all tokens are grouped into three frequency buckets: high, medium, and low, based on their counts in the whole pile dataset 1 . here, we calculate the frequency of each token and sort them in descending order of frequency. then, we categorize the top tokens that cover 80% of the target sequence：i like playing basketball input sequence: i like playing ___ figure 1 : an example where predicting the next token is more like a multi-label classification problem. it is a straightforward idea to use the notable focal loss @xcite from the field of object detection as an alternative to the prevalent cross-entropy loss for the next token prediction. this modification aims to intensify the language model's focus on the infrequent and difficult-to-learn tokens. focal loss is a dynamically scaled version of the cross-entropy loss, where the scaling factor decreases as the predicted probability w.r.t the ground-truth token increases. specifically, focal loss decreases the weights of the easy-to-learn tokens, as their predicted probabilities are higher, and meanwhile increases the weights of the difficultto-learn ones, as their predicted probabilities are lower. in that way, it compels the language model to pay more attention to difficult-to-learn tokens. nevertheless, focal loss @xcite only takes into account the probability w.r.t the groundtruth token when assessing its learning difficulty, and is intuitively designed for the multi-class classification problem where an object is only associated with a single class label. indeed, in language model pretraining, when predicting the next token given the previous ones, there might exist multiple valid tokens besides the ground-truth one. this makes predicting the next token more like a multi-label classification problem, where an object can be associated with multiple class labels @xcite . for example, as shown in figure 1 , given the previous tokens "i like playing ", there are multiple valid next tokens, like "basketball", "football", "golf", etc. suppose the target training token sequence is "i like playing basketball". as the valid tokens would divide up almost the total probability (i.e., 1.0), the groundtruth token "basketball" would be given a smaller probability (e.g., 0.18). then focal loss would treat "basketball" for the position as a difficult-tolearn token. however, as all the other valid tokens are also correct for the position in the view of language modeling, only allowing "basketball" to be predicted is unsuitable. thus, the learning difficulty assessed by focal loss for "basketball" is imperfect in such a multi-label classification case. in this paper, we propose a new loss function termed mile loss to better enable a language model to pay more attention to the difficult-to-learn tokens in such multi-label classification cases. we observe that when a next target token is easy-tolearn, the minor valid tokens would divide up almost the total probability while others are associated with very low probabilities, resulting in a low information entropy of the predicted probability distribution over the vocabulary. on the contrary, if a next token is difficult-to-learn, the predicted probability distribution would be more uniform, resulting in a higher information entropy. therefore, instead of relying on the single probability of the ground-truth token as focal loss, the proposed mile loss uses the information entropy of the predicted probability distribution for assessing learning difficulties, which can better handle cases with multiple valid tokens. then, tokens exhibiting high-entropy, possibly being difficult-to-learn, will be assigned increased weights during language model pretraining. to validate the effectiveness of the proposed mile loss, we train three different-sized models on the pile dataset @xcite . experimental results indicate that mile loss steadily outperforms focal loss and cross-entropy loss on downstream benchmarks. our contributions can be summarized as follows. ## experiments we train three generative language models of different capacities, i.e., 468m, 1.2b, and 6.7b parameters, on the open-source pile dataset @xcite as @xcite @xcite , and make comparisons among different loss functions. ## analyses we conduct further experiments to provide more insightful analyses on the proposed mile loss. ## conclusions in this paper, we present our observation of the bias of learning difficulties among tokens during language model pretraining, mainly caused by the inherent token imbalance in textual training data. we initially introduce focal loss as an attempt to mitigate the bias of learning difficulties. however, we find that considering the single probability of the ground-truth next token for assessing its learning difficulty is unreasonable, especially in cases with multiple valid next tokens. to tackle that, we propose mile loss, which assesses the learning difficulty of a token by taking into account the global information entropy of the predicted probability distribution over the vocabulary. extensive experiments demonstrate that, compared with both cross-entropy loss and focal loss, the proposed mile loss achieves superior performance for various downstream tasks in zero-shot and few-shot learning settings.