Data Mixing and Domain Scheduling in Large Language Models: How DoReMi, RegMix, and Multi-Stage Annealing Shape Pre-Training Dynamics

In large language model pre-training, data composition is as consequential as parameter count and compute budget. While early foundation models relied on raw natural frequencies or manual heuristic filtering to construct training corpora, empirical scaling laws have shown that arbitrary domain ratios cause severe compute inefficiencies. Over-sampling redundant text yields diminishing returns, while under-sampling dense domains such as code and mathematics stalls reasoning performance.

Optimizing the pre-training data mixture requires balancing two interconnected mechanisms: determining the static equilibrium of domain proportions and scheduling dynamic shifts in data composition across the training lifecycle. Modern pre-training workflows have transitioned from intuition-driven heuristics to formal optimization frameworks, including Group Distributionally Robust Optimization, regression modeling on proxy runs, and multi-stage annealing curricula.

The Mathematical Formulation of Data Mixing

Pre-training corpora comprise $K$ heterogeneous domains $\mathcal{D} = \{D_1, D_2, \dots, D_K\}$, representing distinct sources such as filtered web text (e.g., Common Crawl), curated literature, code repositories, conversational transcripts, and academic publications. Let $\boldsymbol{\alpha} = (\alpha_1, \alpha_2, \dots, \alpha_K)$ represent the domain mixing weights residing on the probability simplex:

$$\Delta^K = \left\{ \boldsymbol{\alpha} \in \mathbb{R}^K : \sum_{k=1}^K \alpha_k = 1, \quad \alpha_k \ge 0 \right\}$$

When training a model with parameters $\theta$, the expected objective is a convex combination of individual domain cross-entropy losses:

$$\mathcal{L}(\theta, \boldsymbol{\alpha}) = \sum_{k=1}^K \alpha_k \mathcal{L}_k(\theta) = \sum_{k=1}^K \alpha_k \mathbb{E}_{x \sim D_k} \left[ -\sum_{i=1}^{|x|} \log P_\theta(x_i \mid x_{<i}) \right]$$

In naive uniform or natural frequency sampling, the mixture weights $\boldsymbol{\alpha}$ reflect the raw token availability in storage. However, natural distributions are dominated by low-entropy, repetitive web text. Training under raw proportions forces gradient steps toward high-frequency patterns, creating severe sample inefficiency and starving information-dense domains.

Domain Reweighting via Minimax Optimization: DoReMi

To systematically discover domain weights without requiring manual downstream task annotations, DoReMi (Domain Reweighting with Minimax Optimization) reformulates data mixing through group distributionally robust optimization (Group DRO).

The Excess Loss Metric

Optimizing purely for worst-case domain loss $\max_k \mathcal{L}_k(\theta)$ fails because different data domains exhibit inherently disparate baseline entropies. An intrinsically noisy domain (such as raw transcripts) maintains high loss regardless of training progress, whereas structured code achieves lower cross-entropy early.

DoReMi resolves this disparity by measuring excess loss relative to a reference model $\theta_{\text{ref}}$ trained with equal or default domain weights:

$$e_k(\theta) = \mathcal{L}_k(\theta) - \mathcal{L}_k(\theta_{\text{ref}})$$

The excess loss $e_k(\theta)$ isolates how much a specific domain is lagging behind its theoretical capacity under the current training trajectory.

Proxy Optimization Dynamics

DoReMi trains a small proxy model (for example, 280M parameters) alongside the reference model. During proxy training, domain weights $\boldsymbol{\alpha}$ are treated as an adversarial player updated via exponentiated gradient ascent:

$$\alpha_k^{(t+1)} = \frac{\alpha_k^{(t)} \exp\left(\eta \cdot e_k(\theta^{(t)})\right)}{\sum_{j=1}^K \alpha_j^{(t)} \exp\left(\eta \cdot e_j(\theta^{(t)})\right)}$$

where $\eta > 0$ represents the domain weight step size. Simultaneously, the proxy model parameters $\theta$ are updated via gradient descent to minimize $\sum_k \alpha_k^{(t)} \mathcal{L}_k(\theta)$.

                      DoReMi Optimization Loop
                      
 +-----------------------------------------------------------------+
 |                                                                 |
 |   [Domain Corpus D_k] ---> Mini-Batch Sampling ~ \alpha_k       |
 |                                    |                            |
 |                                    v                            |
 |                         [Proxy Model \theta]                    |
 |                                    |                            |
 |                                    v                            |
 |                        Compute Domain Loss L_k                  |
 |                                    |                            |
 |    [Reference \theta_ref] ---> Excess Loss e_k = L_k - L_ref    |
 |                                    |                            |
 |                                    v                            |
 |                       Exponentiated Gradient:                   |
 |                \alpha_k <- \alpha_k * exp(\eta * e_k)           |
 |                                    |                            |
 +------------------------------------+----------------------------+
                                      | (Output static \alpha*)
                                      v
                      [Full-Scale Foundation Model]

Empirical evaluations show that DoReMi shifts sampling budgets away from low-entropy, highly predictable domains toward medium-entropy web distributions and technical corpora. The resulting static weights allow a full-scale model (such as an 8B parameter model) to reach baseline downstream accuracy with up to 2.6x fewer training steps.

Data Mixture as Regression: RegMix

While DoReMi optimizes domain weights dynamically via adversarial tracking, RegMix treats data mixture determination as an explicit parametric regression problem.

Proxy Mixture Exploration

RegMix explores the mixture space by sampling diverse data proportions from a Dirichlet distribution:

$$\boldsymbol{\alpha}^{(m)} \sim \text{Dirichlet}(\boldsymbol{\beta}), \quad m = 1, \dots, M$$

For each candidate mixture $\boldsymbol{\alpha}^{(m)}$, an ultra-lightweight proxy model (ranging from 1M to 60M parameters) is trained for a bounded token budget (e.g., 1 billion tokens).

Performance Fitting and Closed-Form Allocation

After training $M$ proxy runs, validation losses $\mathcal{L}_{\text{val}, k}^{(m)}$ across all domains are recorded. RegMix fits a regression model predicting validation performance as a function of mixture weights:

$$\hat{\mathcal{L}}_k(\boldsymbol{\alpha}) = w_0 + \sum_{j=1}^K w_j \log(\alpha_j) + \epsilon$$

Once the performance response surface is estimated, an optimization problem is solved to find the mixture $\boldsymbol{\alpha}^*$ that minimizes aggregate target validation loss:

$$\boldsymbol{\alpha}^* = \arg\min_{\boldsymbol{\alpha} \in \Delta^K} \sum_{k=1}^K c_k \hat{\mathcal{L}}_k(\boldsymbol{\alpha})$$

where $c_k$ denotes the target weighting for domain $k$.

Because neural scaling laws demonstrate consistent relative ordering across model sizes, the optimal data mixture identified using 1M-parameter proxies correlates strongly with the performance of 1B to 7B parameter foundation models. RegMix consumes less than 2% of total pre-training compute to discover mixtures that match or exceed human-designed recipes.

Multi-Stage Curriculum and Dynamic Data Annealing

Static domain weighting assumes that the optimal token distribution remains constant throughout training. However, research into learning dynamics demonstrates that models benefit from structured curricula where domain ratios evolve over time.

       Timeline of Multi-Stage Data Scheduling & LR Cooldown
       
Learning
Rate \eta
  ^
  |      +--------------------------------+
  |     /                                  \
  |    /   STABLE PRE-TRAINING PHASE        \   ANNEALING / COOLDOWN
  |   /    (Broad Generalist Data Mix)       \  (High-Signal Specialized Mix)
  |  /                                        \
  +--+-----------------------------------------+--------------------> Steps
  0                                          T_decay                T_max
  
  [Data Mix Proportions]
  * Stable Phase (0 -> T_decay):
    - Web text (Common Crawl / FineWeb): 70% - 80%
    - Code & Tech: 15%
    - Books & Academic: 10%
    
  * Annealing Phase (T_decay -> T_max):
    - High-Quality Reasoning & Math: 35% - 40%
    - Curated Synthetic Instruction Data: 25%
    - High-Grade Code: 25%
    - General Web: 10% - 15%

Warmup-Stable-Decay Schedules

Under traditional cosine learning rate schedules, the learning rate constantly decays, entangling optimizer dynamics with data arrival. In modern architectures (such as MiniCPM and Llama 3), practitioners utilize Warmup-Stable-Decay (WSD) schedules.

1. Warmup Phase (0% to 2% steps): Learning rate increases linearly from 0 to $\eta_{\text{max}}$.
2. Stable Phase (2% to 85-90% steps): Learning rate is held constant at $\eta_{\text{max}}$. The model is exposed to a high-volume, diverse foundation mixture (predominantly high-coverage web tokens) to build linguistic representation and broad world knowledge.
3. Decay / Annealing Phase (Last 10% to 15% steps): Learning rate decays rapidly toward zero (via cosine or linear annealing). Concurrently, the data mixture shifts drastically toward high-quality, high-signal tokens.

The Annealing Data Mix

During learning rate cooldown, the model undergoes rapid parameter consolidation. Injecting dense reasoning data during this stage provides dramatic performance gains:

• Mathematical Proofs and Formal Logic: Upweighted from 3% to 20-30% of the active batch.
• Synthetically Verified Code: Filtered through unit test execution.
• Synthetic Multi-Step Rationales: High-quality generated reasoning traces.

Studies by Blakeney et al. (2024) confirm that introducing high-quality data during the decay phase yields up to 28x greater token efficiency compared to distributing the same high-quality data uniformly across the entire pre-training run.

Engineering Implementation: Sampling and Sharding

Operationalizing dynamic data mixtures at scale introduces significant distributed infrastructure requirements.

Online Multinomial Sampling vs. Offline Pre-Sharding

| Architecture | Mechanism | Trade-Offs | | :--- | :--- | :--- | | Offline Pre-Sharding | Tokens are mixed and concatenated into fixed shards before training starts. | Maximum I/O throughput; zero runtime scheduling overhead; inflexible if mixture ratios must change mid-run. | | Online Multinomial Sampling | Independent data loaders maintain per-domain buffer queues; batches are dynamically mixed using weighted index samplers. | Complete flexibility to alter $\boldsymbol{\alpha}(t)$ on the fly; requires distributed synchronization to avoid data loader worker starvation. |

In online sampling pipelines, each training step draws domain indices $k \sim \text{Multinomial}(\boldsymbol{\alpha}(t))$. Worker processes fetch pre-tokenized chunks from separate memory-mapped files or network streams.

Epoch Budgets and Overfitting Controls

High-quality domains (such as high-tier math corpora or niche textbooks) often contain fewer tokens than the total training budget. If a 10-billion-token math dataset is sampled at 20% in a 15-trillion-token run, it will be repeated 300 times.

Excessive repetition leads to memorization and validation loss divergence. Pre-training guidelines enforce strict epoch limits:

• Web data: 1 to 2 epochs maximum.
• High-quality code and math: 4 to 8 epochs maximum.
• Synthetic reasoning: 2 to 4 epochs maximum.

When a domain reaches its repetition limit, its weight in $\boldsymbol{\alpha}(t)$ must be downscaled to prevent overfitting, redistributing the remaining sampling probability to adjacent high-quality sources.

Summary of Core Principles

1. Natural Frequencies Are Suboptimal: Raw internet text over-indexes on repetitive, low-information tokens.
2. Excess Loss Guides Allocation: Minimax approaches like DoReMi equalize learning progress across domains by targeting excess loss relative to baseline reference models.
3. Mixtures Scale Across Model Sizes: Relative domain utility identified on ultra-small proxy models (1M to 60M parameters via RegMix) transfers reliably to multi-billion parameter architectures.
4. Decoupled Data Curricula: Shifting the data mixture toward dense reasoning and synthetic verification during the learning rate cooldown phase extracts significantly higher capability per token than static mixing.

Sources

• DoReMi: Optimizing Data Mixtures Speeds Up Language Model Pretraining (Xie et al., 2023)
• RegMix: Data Mixture as Regression for Language Model Pre-training (Liu et al., 2024 / ICLR 2025)
• The Llama 3 Herd of Models (Dubey et al., 2024)
• MiniCPM: Unveiling the Potential of Small Language Models with Scalable Training Strategies (Hu et al., 2024)
• Scaling Laws and Compute-Optimal Training Beyond Fixed Training Durations (Blakeney et al., 2024)