cs.LGOct 23, 2026

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

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

Organizations: Tsinghua University · ByteDance Seed

Abstract

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.

Explore similar work

CardsList
  1. Rethinking Data Mixing from the Perspective of Large Language Models

    Apr 9, 2026Yuanjian Xu, Tianze Sun, Changwei Xu +7LLM TrainingImportance Weighting

  2. Data Mixing for Large Language Models Pretraining: A Survey and Outlook

    Mar 25, 2026Zhuo Chen, Yuxuan Miao, Supryadi +1Data Mixture OptimizationLanguage Model Pretraining