Upsample or Upweight? Balanced Training on Heavily Imbalanced Datasets (2025)

Abstract: Data abundance across different domains exhibits a long-tailed distribution: few domains have abundant data, while most face data scarcity. Our work focuses on a multilingual setting, where available data is heavily skewed toward high-resource languages, creating significant imbalances in training data sizes across languages. This disparity challenges training language models to perform uniformly well in all languages. Two common strategies to address this issue are upsampling low-resource languages (Temperature Sampling) and upweighting their loss functions (Scalarization). These methods are often assumed to be equivalent, but this equivalence has not been rigorously established, prompting our investigation.Through theoretical and empirical analysis, we identify when these two methods are equivalent and when they diverge. We prove that they are equivalent under full gradient descent but differ under stochastic gradient descent due to differences in gradient variance. Specifically, Temperature Sampling exhibits lower variance in gradient estimation, leading to faster convergence but a higher risk of overfitting. Based on these insights, we propose Cooldown, a strategy that starts by heavily upsampling low-resource languages to accelerate convergence and gradually reduces the upsampling to prevent overfitting—achieving the best of both worlds. Our method competes effectively with existing data re-weighting techniques while offering computational efficiency.

## introduction information on the internet ranges from common knowledge, such as famous landmarks, to rare details, such as local folklore and specialized scientific theories. data availability across different domains is often long-tailed @xcite @xcite , where very few domains have abundant data. however, the standard language model training objective treats each training instance equally, putting no emphasis on domains that suffer from data scarcity. this work focuses on language modeling on a natural divide of domains with a heavy mismatch: different languages in multilingual language modeling. multilingual language models are often trained on corpora with an overwhelming amount of english and other high-resource languages (hrls) and tiny amounts of data for low-resource languages (lrls) @xcite @xcite . for example, the multilingual c4 corpus @xcite contains 2733 billion english tokens but only 1 billion swahili tokens. uniformly sampling from the combined dataset would result in the language model optimized heavily towards performance on hrls (e.g. english), sacrificing performance on lrls (e.g. swahili). two methods are often employed to address domain mismatches: scalarization and temperature sampling. scalarization adjusts the losses for individual domains by re-weighting them under uniform sampling @xcite . in this case, we assign a larger weight to lrls to emphasize their importance. temperature sampling, weights each training instance uniformly and handles the mismatch by over-sampling lrls and/or down-sampling hrls @xcite @xcite . intuitively, scalarization modifies the loss while temperature sampling modifies the dataset. scalarization and temperature sampling are widely regarded as equivalent. @xcite denotes "we follow convention and implement scalarization via proportional sampling". @xcite and @xcite implement sampling probabilities by multiplying losses with per-domain renormalized weights. the underlying assumption is that temperature sampling and scalarization are equivalent, and we can use them interchangeably. however, to the best of our knowledge, this equivalence has not been rigorously established. we closely investigate this assumed equivalency in theory ( §3.1). specifically, we prove that although they are equivalent in full gradient descent (theorem 1), temperature sampling induces lower variance in the context of stochastic gradient descent (theorem 2). moreover, the variance induced by scalarization increases as the approximated temperature increases or the domain distribution's skewness increases (theorem 3). based on our theoretical results and connecting to the literature on lower variance between stochastic gradients accelerates convergence @xcite , we make the following hypothesis: hypothesis 1. temperature sampling converges much faster than scalarization at higher temperatures or on heavily imbalanced domain distributions. we empirically verify our hypothesis ( §3.2) and find that temperature sampling does converge faster but is more prone to overfitting. we identify that the temperature controls the speed of convergence and hence can be used as a control knob to adjust the convergence speed. we thus propose cooldown: to use a large temperature initially for fast convergence, then decrease the tempera-ture to prevent overfitting to the lrls. figure 1 illustrates the effectiveness of cooldown , which significantly accelerates convergence on the lrl due to a high temperature (aggressive upsampling of lrls) at the beginning of training and reduces overfitting to the lrl due to the lowering the temperature during training. to sum up, our contribution is two-fold: we show the effectiveness of cooldown in multilingual settings. ## cooldown: balanced training for heavily imbalanced datasets based on the theoretical analysis in §3.1 and the empirical results in §3.2, we conclude that temperature sampling with a large temperature converges faster than scalarization with equivalent weights but is more prone to overfitting on the lrl when using a large temperature. we thus hypothesize that we can employ a large temperature during the beginning of training to speed up convergence and then decrease the temperature to prevent overfitting on low-resource directions. our proposed method: we design a simple temperature scheduling method: cooldown which starts with a high temperature (aggressive upsampling of lrls) and then lowers the temperature to τ = 1 (proportional sampling) at a fixed iteration to prevent overfitting. we describe our experiment setup at §4.1, report results at §4.2, and discuss our findings at §4.3. ## conclusion we examined two common balancing methods for multi-domain language modeling with data imbalances: scalarization and temperature sampling. although both yield the same loss, temperature sampling converges faster but risks overfitting. to mitigate this, we propose cooldown, a variant that adjusts temperatures during training to maintain fast convergence while reducing overfitting.