A study of sequence weighting at scale

A study of sequence weighting across in-house and open-weight language models found non-monotonic scaling behavior: as models grow from small to medium scale they shift from learning general patterns independent of data weight to learning data-specific patterns proportional to those weights, then at large scale learn all patterns independent of data weight again. The research, which assigns random sequence weights to each training sequence and measures loss reduction, parallels the MAI-Thinking-1 technical report (section 2.5.2), where the quality ordering of a code-heavy versus a STEM-heavy training mix reversed as model size increased. The authors characterize data mixing as prone to such "aberrant" scaling behaviors that cannot be predicted by extrapolating small-scale results.

TL;DR: We study the scaling laws of data weighting across in-house and open-weight LMs, finding non-monotonic behavior across scales. We vary the weight assigned to sequences during training and measure how strongly the model’s loss reduction on a sequence depends on the sequence’s weight. Taken together, our results are consistent with a general trend: as models transition from small to medium scale, they transition from learning general patterns independent of data weight to learning data-specific patterns proportional to the data weights. As models then transition from medium to large scale they are able to learn all patterns present in the data, once again independent of data weight. When training neural networks, we spend a lot of time thinking about precisely what we want our models to learn. How much do we care about our models being generally intelligent across all coding tasks versus being good at OCaml specifically? How useful is data from an older market regime compared to data collected recently? Generally, how much do we want to upweight high-quality data versus allowing our models to learn from everything? These questions are ultimately data mixing questions, questions about how and where to spend our models’ representational capacity and our FLOPs. Further complicating these questions is that in practice, many of the experiments we would like to run are prohibitively expensive. It’s intractable to run a dense multi-dimensional grid search over hyperparameters even at medium scales, let alone at the largest scales where we may only be able to train one model for a given task . The standard approach to this is to fit hyperparameters like data mix weights via a scaling law approach, in which we fit hyperparameters at small scales and then extrapolate those hyperparameters to larger scales the canonical example of this being Chinchilla scaling https://arxiv.org/abs/2203.15556 . A core challenge for scaling laws is that for the extrapolation to hold, large-scale models must have behaviors that are either a invariant with scale e.g., MuP-style results https://arxiv.org/abs/2203.03466 in which many hyperparameter optima remain stable across scales or b predictable via extrapolating changes in behavior at small scale e.g., Kaplan-style results https://arxiv.org/abs/2001.08361 in which loss decreases predictably with scale . Unfortunately, not all behaviors have these characteristics: some behaviors are emergent at scale and are not predictable via scaling laws. For the rest of this post, we’ll call such non-predictable scaling behaviors aberrant . Our internal experiments have shown that data mixing is a setting particularly prone to aberrant scaling behaviors. This parallels findings in the MAI-Thinking-1 technical report section 2.5.2 https://microsoft.ai/pdf/mai-thinking-1.pdf , where the relative ordering in quality of training with a code-heavy and a STEM-heavy mix reversed as model size increased. We therefore set out to precisely characterize the extent to which language models actually learn patterns in the data proportional to the weight accorded to that data in a data mix, and how that behavior changes across scales. Data mix experiments can conflate the effect of the data’s weight in a corpus with other important but different considerations. Most notably, results from data mix experiments are heavily affected by variation in data quality across different sources and by differences in the uniqueness of marginal tokens in any given data source. To isolate the effects of mix weight alone, in this post we specifically analyze data weighting , a variant of data mixing in which each individual sequence in our training dataset receives a different weight in our loss. We measure the extent to which a model has learned a training sequence with a given weight by first assigning random sequence weights to each sequence in the dataset, training a model on the dataset, then re-evaluating it on the same training dataset. We then define a metric, the effective sequence weight exponent denoted by that measures what power of a sequence’s weight is closest to proportional to the model’s expected loss decrease on that sequence e.g., says that the model’s expected loss decrease on sequences is directly proportional to their weight, while says that the model’s expected loss decrease on sequences is the same for all possible sequence weights . Previous work makes mixed predictions for how scales: Byrd & Lipton 2019 https://arxiv.org/abs/1812.03372 and the general phenomenon of large-scale interpolation https://arxiv.org/abs/1611.03530 predict that larger models should asymptote towards , while Li et al. 2026 https://arxiv.org/abs/2608.14071 find that at fixed tokens-per-parameter, optimal token repetition count for valuable domains mildly increases with model size. We ask: is the effective sequence weight exponent a smooth power law with scale, or is it aberrant? We evaluate the effects of sequence weighting on an internal text-based benchmark. To test across model families and scales, we evaluate two families of in-house-pretrained LLMs ranging from tens of millions to hundreds of billions of parameters and one open-weight model family Qwen 2.5 https://arxiv.org/abs/2412.15115 ranging from 500M to 72B parameters. We find a number of aberrant scaling behaviors reasoning about data mixing across scales is complicated and some striking trends. Taken together, our results are consistent with a non-monotonic rise-then-fall in the effective sequence weight exponent: small-scale models fit small effective sequence weight exponents, learning patterns across the entire dataset independent of data weight; medium-scale models fit larger effective sequence weight exponents, learning patterns in data proportional to their data weight. Large-scale models once again fit small effective sequence weight exponents, learning all patterns present in the data regardless of weight. Epoching shifts the effective sequence weight peak towards smaller models. Aberrant scaling laws are inconvenient for our model training methodologies, but identifying them lets us avoid bad extrapolations. In our science we are vigilant for such aberrant behaviors, and when we discover one, we redesign our scaling experiments and evaluations such that they become predictable.<sup 1 fn:1 </sup Our experiments suggest a number of possible remedies, from evaluating and extrapolating data-mix results only from sufficiently large models, to adjusting training weights to compensate for the observed .<sup 2 fn:2 </sup And we keep on measuring how far off our models are from our predictions to keep on catching the next place that our scaling laws break down, and to allow us to train larger, better models with even more confidence. Methodology Dataset. We evaluate the effects of sequence weighting on an internal text benchmark. We weight sequences in this dataset using sequence weights drawn log-uniform between 0.01 and 10. We train for 3 epochs and evaluate after each epoch. Models. We evaluate three different model families on this dataset: - JS-dense , a family of in-house pretrained dense LMs comprising 9 models - JS-sparse , a family of in-house pretrained MoE LMs comprising 8 models - Qwen 2.5 https://arxiv.org/abs/2412.15115 , an open-weight family of dense LMs ranging from 500M parameters to 72B parameters. The JS-dense and JS-sparse models range from tens of millions to hundreds of billions of parameters. Note that in all contexts, held-out performance is increasing with scale for these models. Measuring the effect of sequence weighting. We first train our model on the training dataset with each sequence weighted corresponding to , re-evaluate the model on its training dataset, and for each sequence measure the loss reduction on that sequence achieved from training. We then fit a such that appropriately normalized the loss reduction on a given sequence with weight is best explained by . We make this fit more precise in an attached note /a-study-of-sequence-weighting-at-scale/sequence-weighting-note.pdf , but to better understand what different values mean, let’s consider a few scenarios: - : the expected loss reduction on sequences is not affected by their weight i.e., we don’t reduce the loss on high-weight sequences any more than we do on low-weight sequences . - : higher weight sequences have their loss reduced more than lower-weight sequences, but the reduction scales sublinearly according to exponent when , the loss reduction of a sequence is proportional to its train weight . - : higher weight sequences have loss reduced more than lower weight sequences in a more concentrated way than implied by the sequence weights an example: if the model memorized for the sequence with highest weight and randomly sampled a prediction for all other rows, we would see Hyperparameters. We tune hyperparameters to minimize loss against a validation set.<sup 3 fn:3 </sup Because the largest in-house models are too large to economically tune hyperparameters on, we set hyperparameters via a Chinchilla-style power law scaling law, fitting optimal hyperparameters as a function of scale. Results and Analysis Here are the interactive 3D surfaces for each model family: Within a given model family and number of epochs, often increases at small scales particularly through tens of billions of parameters and decreases at larger scales as we near the hundreds of billions of parameters regime for reference, the largest pictured JS-dense rung has fewer than 100B parameters . Training for more epochs shifts the peak towards smaller model scales. Between model families, we find little relationship between the actual values of – the smallest model in JS-sparse is substantially larger than the smallest model in JS-dense, yet has a smaller ; the largest model in JS-sparse is also substantially larger than the largest model in Qwen 2.5, yet has a larger . Why does sometimes rise then fall? We hypothesize that this is because models learn different patterns present in the data at different rates. Some patterns generalize well across all sequences for example, understanding English , while others are idiosyncratic to smaller groups of sequences. In this dataset, there is more loss reduction possible due to idiosyncratic patterns than there is to general patterns. We therefore hypothesize that smaller models use their limited capacity to learn general patterns and cannot represent idiosyncratic patterns, medium-sized models learn both the general and the most important idiosyncratic patterns by sequence weight , and large models have the capacity to learn all idiosyncratic patterns with all weights. Epoching allows the model to learn more of both the general and idiosyncratic patterns. More broadly, we find that the behavior of sequence weighting at small scales does not cleanly predict behavior at the largest scales. This nicely encapsulates an all-too-common result; larger models can have behaviors that differ in kind from those exhibited by smaller models. Accordingly, our scaling research requires spending as much if not more time and care trying to understand why we’re off our predicted curves as it does fitting the curves themselves. If this kind of thing is interesting to you, consider applying https://www.janestreet.com/join-jane-street/machine-learning/ . You’ll join a close-knit group of brilliant, supportive colleagues, harnessing tens of thousands of GPUs, petabytes of training data, and the agility and resources to invest in the best ideas. 1. A classic example of such a fix to an aberrant scaling law is in multiple-choice question answering https://arxiv.org/abs/2304.15004 : while small-scale models often hover around the noise floor with sampled responses until a discrete jump when the model learns to follow instructions, measuring continuous metrics like bits-per-byte on answers often reveals smooth, non-aberrant scaling. ↩ fnref:1 2. In particular, if effective sequence weights scale as , then achieving target loss reduction weights might require training with damped weights e.g. . We have not tested this hypothesis. ↩ fnref:2 3. We apply heavy regularization in all cases such that the resulting model approximately maximizes performance on a held-out set; while we could likely tune all but the very smallest of these architectures to memorize the entire input data in the first pass, this isn’t representative of the scenario that we care about. ↩ fnref:3