PaperlayerView exact revision on arXiv

Source-linked AI summary

Beyond Data Schedule: Fast LLM Pretraining on Data Mixtures by Loss-Reweighting

Jiazheng Li, Huaqing Zhang, Jing Xu, Tianle Cai, Jingzhao Zhang

arXiv:2610.32141v1cs.LGcs.AI

TL;DR

Multi-stage pretraining can improve specialized capabilities while forgetting earlier knowledge, and data-mixture ratios couple domain exposure with optimization scale. The paper introduces LARE, which adapts domain loss weights while preserving the sampled mixture; across five benchmarks it improves over direct mixing, including a +17.06-point GSM8K gain.

  • Problem

    Stage-based pretraining can improve a specialized capability at the cost of forgetting earlier knowledge, while mixture ratios couple domain exposure with optimization scale.

    Source: §1 · pp. 1–2

  • Method

    LARE adjusts domain loss weights from running loss statistics with a normalized inverse-square-root rule while keeping the sampled mixture fixed.

    Source: §1 · p. 3

  • Results

    LARE improves all five evaluated benchmarks over direct mixing, including +17.06 points on GSM8K.

    Source: §1 · p. 3

  • Takeaways & Limitations

    At fixed exposure, domain loss scale is an independent control that can improve joint capabilities without additional domain exposure.

    Source: §1 · p. 3, §5 · p. 10

    The study does not evaluate how reweighting may change the influence of training-data biases, harmful content, or privacy risks, and task gains are not evidence of improved safety or fairness.

    Source: Ethics statement · p. 10

Abstract

from arXiv · show

Large language models often acquire specialized capabilities through successive pretraining stages, but these gains can come at the cost of forgetting earlier knowledge. In this work, we investigate whether such trade-offs can be improved by changing how strongly each domain contributes to learning while keeping its data exposure fixed. We find that, in controlled knowledge--mathematics experiments, loss reweighting produces a single checkpoint that surpasses the best knowledge and reasoning scores attained along either stage trajectory. This reveals an independent control that data-mixture ratios alone do not capture. Building on this finding, we propose LARE (Loss-Adaptive Reweighting under fixed Exposure), which automatically adjusts domain weights from running loss statistics. Motivated by a theoretical analysis of gradient noise, its normalized inverse-square-root rule balances relative loss scales without changing the sampled mixture or requiring additional runs to search for weights. Across knowledge, mathematics, and code, \Method{} improves all five evaluated benchmarks for a 1.5B model trained on 240B tokens, averaging +7.84 pp over a matched direct mixture, including gains of +17.06 pp on GSM8K and +9.00 pp on MATH. These findings demonstrate that adjusting domain loss scales can improve joint capabilities without additional domain exposure.

1 Introduction

Multi-stage pretraining can improve a specialized capability while forgetting earlier knowledge, and direct mixtures may still leave a trade-off. The paper separates domain exposure from optimization scale, showing that loss reweighting can improve joint capability and motivating an adaptive controller.

1 Introduction

  • Successive specialized pretraining stages can improve a new capability while forgetting earlier knowledge, creating an acquisition–retention trade-off.The introduction describes these recipes as moving beyond a single stationary data distribution.Source: §1 · p. 1
  • Ratio tuning couples how much data a domain contributes with how strongly its prediction errors shape parameter updates, although these allocations need not coincide.Reweighting holds the sampled-data mixture fixed while adjusting domain loss coefficients.Source: §1 · p. 2
  • LARE derives domain loss weights from running loss statistics using a normalized inverse-square-root rule after learning-rate warmup.The rule preserves the sampled mixture and avoids auxiliary runs to search for coefficients.Source: §1 · p. 3
  • Loss reweighting at fixed exposure can exceed the best knowledge and mathematical-reasoning scores reached along either stage trajectory.The comparison holds the sampled-data mixture and token budget fixed.Source: §1 · pp. 2–3
  • Extending training with mathematics and code to 8× Chinchilla yields gains of 0.83, 17.06, 9.00, 6.71, and 5.60 points on MMLU, GSM8K, MATH, HumanEval, and MBPP, respectively, over direct mixing.Source: §1 · p. 3

2 Mixture Ratios Couple Exposure and Scale

At matched domain exposure, loss reweighting and data remixing can produce different capability outcomes because they control optimization scale and sample allocation differently. Empirical comparisons, a two-domain analysis, and matched-exposure training runs show that reweighting can expand the observed frontier and improve both knowledge and mathematical reasoning.

2.1 Reweighting Expands the Remix Frontier

  • Remixing changes domain sample proportions and gradient-estimation sample counts, whereas reweighting keeps sample allocation fixed and scales domain gradient estimates.The effective-rate identity is exact for SGD but only an abstraction for optimizers with nonlinear gradient transformations.Source: §2.1 · p. 3, §2 · p. 3
  • At fixed 7:1 or 3:1 exposure, the corresponding 1:2 reweighted run Pareto-dominates its direct-mixture anchor; larger weights further improve reasoning but can trade away knowledge accuracy.The figure compares MMLU and GSM8K after 30B sampled tokens.Source: §2.1 · pp. 3–4, Figure 2 · p. 4

2.2 Theory: Why Reweighting Can Outperform Remixing

  • In the two-domain linear regression analysis, separately tuned domain rates can achieve losses that no remix with a shared learning rate can Pareto-dominate.The result is illustrated at a fixed equal mixture and in the model’s remix Pareto frontier.Source: §2.2 · p. 5, Figure 3 · p. 5

2.3 Same Domain Exposure, Different Outcomes

  • With equal token budgets, reversing stage order reverses which capability is forgotten; direct mixing moderates but does not eliminate the trade-off.The knowledge-then-mathematics run ends at 29.45 MMLU and 36.62 GSM8K, while the reverse order ends at 31.78 and 0.76.Source: §2.3 · pp. 5–6
  • 32.27 MMLU and 38.59 GSM8K: at the same 7:1 exposure, 1:2 loss weighting exceeds direct mixing and the best scores along either stage trajectory.The direct-mixture scores are 31.40 MMLU and 27.67 GSM8K; the stage-trajectory maxima are 31.85 and 37.91.Source: §2.3 · p. 6, Figure 4 · p. 6

3 Lare: Loss-Adaptive Reweighting under Fixed Exposure

The section motivates separating domain exposure from optimization scale, derives an inverse-square-root weighting rule from gradient-noise analysis, and introduces Lare to adapt weights online. In a controlled run, Lare selects a useful loss scale and improves the final direct-mixture endpoint.

3 Lare: Loss-Adaptive Reweighting under Fixed Exposure

  • Fixed-weight grid search requires a separate pretraining run per candidate, and the best setting may not transfer across mixtures, budgets, or optimizers.Lare is introduced as an online alternative that infers and adapts relative weights from training signals.Source: §3 · p. 6

3.1 From Domain Loss to Optimization Scale

  • For the stylized Gaussian model, the root-mean-square stochastic-gradient norm scales as Θ(√(dℓ⋆ρ)) as label-flip probability ρ decreases.The analysis motivates the exponent rather than claiming an exact language-model scaling law.Source: §3.1 · p. 7

3.2 The Adaptive Reweighting Rule

  • After learning-rate warmup, Lare normalizes inverse-square-root weights from smoothed domain losses and counts, keeping the mixture-averaged coefficient at one.The rule changes relative domain loss scales without changing the sampled mixture.Source: §3.2 · p. 7

3.3 Validating the Reweighting Dynamics

  • At the final checkpoint, Lare improves direct 1:1 weighting by 0.84 MMLU and 12.13 GSM8K points, with an endpoint competitive with fixed 1:2.The comparison holds the model, sampling distribution, optimizer, schedule, and token budget constant; only the loss-weighting rule varies.Source: §3.3 · p. 7

4 Experiments: Ablations and Scaling up

The experiments compare optimizers and evaluate loss reweighting across training budgets and domain compositions. Lare improves all matched benchmarks in the reported 4× and 8× settings, while ActGD provides the strongest direct-mixture optimizer result.

4.1 Experimental Setup

  • Experiments use a shared 1.5B GPT-2-style decoder-only Transformer and knowledge, mathematics, and code domains, with 30B, 120B, and 240B-token budgets.Reported cells are generally single runs evaluated at the final checkpoint.Source: §4.1 · pp. 7–8

4.2 Ablating over Optimizers

  • ActGD obtains the highest MMLU and GSM8K scores, beating Muon by 0.42 and 0.91 points and AdamW by 0.52 and 6.67 points, respectively.It applies one update family and shared hyperparameter set across transformer blocks, embeddings, and the output head.Source: §4.2 · p. 8, Table 1 · p. 9

4.4 Further Analyses

  • Further analyses examine robustness to data-mixture choices, reweighting across optimizers and fixed-weight baselines, and scaling across model sizes at corresponding Chinchilla budgets..Source: §4.4 · p. 9

4.3 Scaling Across Training Budgets and Domains

  • At 8× Chinchilla with K:M:C, Lare improves all five benchmarks over direct mixing by a mean of 7.84 points, with larger gains than at 4× on every benchmark.The comparison matches model, sampled mixture, optimizer, and token budget.Source: §4.3 · p. 9, Table 2 · p. 10

5 Conclusion

The conclusion treats domain exposure and optimization scale as separate pretraining controls and reports gains from reweighting across domains and budgets. The paper also identifies unassessed safety-related risks and situates loss reweighting alongside prior data-mixture approaches.

5 Conclusion

  • At 8× Chinchilla, Lare improves all five evaluated benchmarks by 7.84 points on average over matched direct mixing, supporting separate treatment of data allocation and loss scaling.The conclusion states that exposure determines what the model sees, while reweighting determines how strongly each domain shapes learning.Source: §5 · p. 10

Ethics statement

  • The study does not evaluate how reweighting changes the influence of training-data biases, harmful content, or privacy risks, and task gains do not establish improved safety or fairness.The authors call for appropriate safety evaluation before deployment.Source: Ethics statement · p. 10

Reproducibility statement

  • The authors state that code, dataset, and model are provided in supplementary materials and an anonymous GitHub repository, with reproduction steps in the README..Source: Reproducibility statement · p. 10

A Related Work

  • Prior data-mixture methods change what the model sees, whereas loss reweighting changes the gradient contribution of already-sampled data.The related-work discussion describes loss reweighting as complementary to sampling-based approaches.Source: A Related Work · p. 17

B.1 Proof of the Loss--Gradient-Noise Scaling

The proof analyzes gradient covariance at the optimum of a Gaussian logistic-regression model with independently flipped labels. It establishes loss-dependent noise scaling and motivates inverse-square-root reweighting, while limiting the argument to unconstrained classifier-scale optimization.

B.1 Proof of the Loss--Gradient-Noise Scaling

  • In the Gaussian model, gradient-noise components scale differently along and orthogonal to the true decision direction, yielding a root-mean-square gradient norm proportional to the square root of optimal loss.The proof evaluates the covariance at the population-risk minimizer under independently flipped labels.Source: §B.1 · pp. 18–19
  • The low-loss limit requires optimizing classifier scale; with unit-norm weights, even separable Gaussian halfspace data retain strictly positive logistic loss.The theorem is a mechanism result for unconstrained scale optimization, not for unit-normalized classifiers.Source: §B.1 · p. 19
  • The analysis motivates weighting proportional to inverse square-root loss when effective dimensions and hidden constants are comparable.The theorem’s variance and standard-deviation scaling distinction underlies the proposed weighting exponent.Source: §B.1 · p. 19
  • Lare substitutes an online smoothed loss for the unavailable optimum loss, a practical modeling choice rather than a consequence of the asymptotic theorem..Source: §B.1 · p. 19

B.2 Proof of Proposition 1

The proof shows that each domain’s optimized loss decreases strictly with its mixture share, while independent domain learning rates can produce a point that shared-rate remixing cannot match in both losses. A numerical instance illustrates the resulting frontier improvement.

B.2 Proof of Proposition 1

  • Each domain’s optimized loss has an attained positive-rate minimum and decreases strictly as its mixture share increases.The proof applies this single-domain result to each domain’s loss.Source: §B.2 · pp. 21–22
  • With separate per-domain optimal rates, the reweighting point has at least one lower loss coordinate than any alternative remix using a shared rate.The proof establishes this separation for remixes with a different mixture allocation.Source: §B.2 · pp. 20, 22

B.3 Numerical instantiation of Figure 3

  • In the numerical instance, reweighting improves both domain losses at a fixed 1:1 mixture, while no mixture with a shared rate can match both losses.The instance uses two orthogonal coordinate blocks with different curvature scales, making the domains’ optimal rates differ.Source: §B.3 · p. 22
  • The figure compares per-domain optimal rates with a shared rate minimizing the mixture-weighted loss, and overlays their attainable frontiers.Losses are computed exactly by iterating the second-moment recursion, without Monte Carlo.Source: §B.3 · pp. 22–23

C Data Sources and Processing

The appendix documents the knowledge, mathematics, and code corpora, their processing and masking, how sampled-token mixtures relate to realized loss shares, and the experiment settings.

C.1 Sources and Filtering

  • The prepared corpora comprise full-text web documents for knowledge, sourced problems and solutions for mathematics, and code-oriented prompts, responses, and source files.The mathematics pool is filtered for benchmark overlap and deduplicated; the code pool is also filtered and deduplicated.Source: C Data Sources and Processing · p. 24, §C.1 · p. 24

C.2 Training Records and Loss Masking

  • Mathematics and code examples are serialized as prompt–response records, while knowledge uses full-text documents; masks determine which targets contribute to loss.Mathematics responses retain worked solutions when available and end with a normalized boxed answer; code records follow source-specific parsing rules.Source: §C.2 · pp. 24–25

C.3 Mixture Ratios and Realized Loss-Token Shares

  • Mixture ratios count sampled tokens before masking, so realized valid-target shares can differ; a nominal 1:1 K:M mixture has 43.58% mathematics loss targets.Controlled comparisons reuse the same frozen stream, including sample order and loss masks.Source: §C.3 · p. 25, Table 4 · p. 25

D.1 Detailed Hyperparameters

  • The 1.5B experiments use a 4096-token context, global batches of 128 sequences, and a single learning-rate schedule spanning any multi-stage run.The 1×, 4×, and 8× budgets correspond to approximately 30B, 120B, and 240B sampled tokens.Source: §D.1 · p. 26
  • After warmup, Lare derives normalized inverse-square-root loss weights from domain loss EMAs and valid-target count shares.During warmup, all loss coefficients remain one; controlled comparisons use matched streams and aggregate per-domain statistics across data-parallel ranks.Source: §D.1 · p. 26

D.2 Matched-Exposure Results

Under matched 30B-token exposure, fixed 1:2 reweighting improves MMLU and GSM8K over direct mixing across three optimizers. The Muon reweighted endpoint also exceeds the best scores along either stage trajectory.

D.2 Matched-Exposure Results

  • The comparison holds model configuration, optimizer assignment, learning-rate schedule, token budget, and domain-token allocation fixed across stage and stationary recipes.Each run receives 26.25B knowledge tokens and 3.75B mathematics tokens.Source: §D.2 · p. 27
  • Fixed 1:2 reweighting improves both MMLU and GSM8K over direct mixing under all three optimizers; the Muon endpoint exceeds the best scores along either stage trajectory.Table 7 compares matched-exposure recipes across optimizers.Source: §D.2 · p. 27, Table 7 · p. 27
  • Across all three optimizers, the two-stage orders retain different capabilities, while fixed 1:2 reweighting improves both MMLU and GSM8K over direct mixing under every optimizer.The Muon reweighted endpoint exceeds the best scores along either stage trajectory; under AdamW, Math →Knowledge retains a higher MMLU score.Source: §D.2 · p. 27

D.3 Detailed Data-Mixture Results

Across tested data mixtures, Lare improves reasoning and, when code is included, knowledge, mathematics, and code benchmarks. The detailed results show the trade-off at the reasoning-heavy K:M boundary and gains across every tested code mixture.

D.3 Detailed Data-Mixture Results

  • At 1× Chinchilla, GSM8K improves by 4.85–7.43 points across K:M ratios 7:1, 3:1, and 1:1, while MMLU changes by +0.17, −0.03, and −0.94 points, respectively.The largest MMLU cost occurs at the 1:1 reasoning-heavy boundary.Source: §D.3 · p. 27, Figure 6 · p. 27, Table 8 · p. 28
  • With code added, Lare improves MMLU by 0.04–0.85 points, GSM8K by 7.51–11.07 points, and MBPP by 2.00–4.00 points at every tested ratio.The tested mixtures are K:M:C ratios 7:0.5:0.5, 3:0.5:0.5, and 1:0.5:0.5.Source: Figure 6 · p. 27, §D.3 · p. 28, Table 8 · p. 28
  • At the 1:1 K:M boundary, mathematics contributes approximately 43% of realized loss tokens but 68% of the weighted objective at the final checkpoint.The passage describes this split as consistent with the larger MMLU cost.Source: §D.3 · p. 28

D.4 Additional Optimizer Results

Additional optimizer tests find that Lare’s GSM8K gains persist across optimizer–ratio settings, while MMLU shifts depend on the sampled ratio. The appendices also specify evaluation protocols, longer-budget trajectory plots, and a joint model–token scaling setup.

D.4 Additional Optimizer Results

  • Across all nine optimizer–ratio settings, Lare improves GSM8K; the six added pairs gain 5.24–12.13 points, while MMLU shifts vary by ratio.MMLU changes are positive at 7:1, nearly unchanged at 3:1, and more negative at 1:1.Source: Table 9 · p. 28, §D.4 · pp. 28–29
  • At 7:1, Lare improves both benchmarks over direct mixing under every optimizer; at Muon 3:1, GSM8K rises 11.52 points with a 0.11-point MMLU decrease.The online endpoint is bracketed by fixed-weight endpoints in several settings, but under AdamW fixed 1:3 is higher by 0.24 MMLU and 2.20 GSM8K points.Source: §D.4 · p. 29, Table 10 · p. 29

D.5 Evaluation Settings

  • MMLU uses 0-shot acc_norm across 57 subjects; GSM8K and MATH use greedy completions, while HumanEval and MBPP score one completion against benchmark tests.GSM8K uses its main test split and MATH uses MATH-500; MBPP uses three fixed demonstrations.Source: §D.5 · pp. 29–30, Table 11 · p. 29

D.6 Training Trajectories at Longer Budgets

  • Figure 7 plots matched K:M:C = 7:0.5:0.5 runs with direct mixing in blue-gray and Lare in red; lighter and darker curves end at 120B and 240B tokens.Curves are centered 11-checkpoint moving averages.Source: Figure 7 · p. 30, §D.6 · p. 30

E.1 Transferred Setup

  • The transfer study compares direct mixing, Lare, and two stage schedules at K:M = 7:1 across 124M, 350M, and 774M models at their nominal 1× Chinchilla token budgets.Within each size, direct mixing and Lare differ only in loss reweighting; stage schedules preserve total domain exposure.Source: §E.1 · p. 31

E.2 Replication Results

Replication results show a growing mathematical-reasoning advantage under joint model–token scaling, with little MMLU change at the largest allocation. The optimizer appendix describes ActGD’s activation-based update geometry and row-wise direction.

E.2 Replication Results

  • At 774M/15.73B tokens, Lare gains +7.36 GSM8K and +0.02 MMLU points over direct mixing; stage endpoints remain order-sensitive, and results are single seed.The 124M result remains at the low-accuracy GSM8K floor.Source: Table 12 · p. 31, Figure 8 · p. 32
  • Math validation loss falls 5.70%, 6.73%, and 6.55% across allocations, while model size and token budget co-vary, so parameter-count effects are not isolated.The downstream GSM8K benefit becomes visible only after the model leaves the low-accuracy regime.Source: §E.2 · p. 31

F Activation Gradient Descent

  • ActGD maps batch gradients to parameter updates, whereas Lare weights domain contributions before gradient formation; ActGD targets layer-activation changes rather than Euclidean parameter displacement.The activation-matching construction is local to one optimization step and does not assume equivalence with model merging.Source: F Activation Gradient Descent · p. 33, §F.1 · p. 33

F.1 From Activation Matching to a Whitened Update

  • Whitening expresses updates in activation geometry, making the constraint isotropic and preventing covariance-spectrum tails from setting update scale.The unconstrained minimizer maps back to a covariance-adjusted parameter update.Source: §F.1 · pp. 33–34

F.2 The Row-Wise (2, ∞) Direction

  • The adopted (2, ∞) direction normalizes each output row independently in activation metric, preserving within-row gradient structure without coupling rows through one matrix-wide factor.This gives output rows comparable activation-metric steps despite differences in raw gradient norms or input variances.Source: §F.2 · p. 34, Table 13 · p. 34

F.3 Stateful Optimizer Used in Our Experiments

ActGD uses smoothed gradient and covariance estimates to construct updates, with a shared configuration across model matrices. Its implementation reuses and shards activation statistics, but adds computation and memory costs that token-matched comparisons do not establish as equal to alternatives.

F.3 Stateful Optimizer Used in Our Experiments

  • ActGD smooths batch gradients and input covariances with separate EMAs, then uses a stabilized Newton–Schulz inverse approximation for its row-wise direction.With exact running estimates and inverse approximation, the practical equations reduce to the row-wise activation-space direction.Source: §F.3 · pp. 34–35
  • The main experiments use one ActGD configuration for embeddings, the language-model head, and transformer-block matrices.These parameters share the row-wise update family and hyperparameter values.Source: §F.3 · p. 35, Table 14 · p. 35

F.4 Parameter Coverage and Distributed Implementation

  • For linear layers, ActGD estimates covariance from activations, reuses query/key/value statistics, and handles embeddings through token-frequency row scaling rather than a vocabulary-sized covariance matrix.Distributed training averages activation statistics across data-parallel ranks and assigns inverse and update computation to owner ranks.Source: §F.4 · p. 35
  • Covariance formation and approximate inversion dominate ActGD’s extra computation; subsampling, covariance reuse, caching, and sharding reduce overhead.Comparisons match sampled tokens, not wall-clock time or peak memory, so they establish optimization quality and sample efficiency rather than equal throughput or memory.Source: §F.4 · pp. 35–36

Related papers

Loading 2610.32141v1…