DoReMi: Optimizing Data Mixtures Speeds Up Language Model Pretraining

🎯 Pitch

A tiny 280M-parameter model can find the perfect data recipe for an 8B-parameter model, slashing pretraining steps by 2.6× while improving perplexity on even downweighted domains. The secret is optimizing domain weights to minimize worst-case excess loss rather than raw loss, and this task-agnostic approach matches the performance of weights hand-tuned on downstream benchmarks.

1. Executive Summary

This paper introduces Domain Reweighting with Minimax Optimization (DoReMi), an algorithm that optimizes the mixture proportions of pretraining data domains without requiring any downstream task knowledge. DoReMi first trains a small proxy model using group distributionally robust optimization (Group DRO) over domains — where the inner maximization dynamically upweights domains with high excess loss relative to a reference model — and then takes the averaged domain weights from this process to resample a dataset for training a much larger model. Using a 280M-parameter proxy model to set domain weights for an 8B-parameter model (30× larger) on The Pile, DoReMi improves average few-shot downstream accuracy by 6.5 percentage points and reaches the baseline accuracy with 2.6× fewer training steps, while also reducing perplexity across all domains — even those it downweights. On the GLaM dataset, DoReMi attains comparable performance to domain weights that were explicitly tuned on downstream tasks, establishing that purely loss-based domain reweighting can match downstream-informed tuning only when the optimization targets worst-case excess loss rather than raw loss.

2. Context and Motivation

The Core Problem: We Don't Know How to Mix Pretraining Data Domains

The fundamental question this paper tackles is deceptively simple: if you're training a large language model on data drawn from multiple domains (Wikipedia, books, web text, code, etc.), what fraction of the training data should come from each domain? This choice — the domain weights or mixture proportions — profoundly affects the resulting model's capabilities. Every major LLM trained to date makes this choice, but the field lacks a principled, task-agnostic method for making it.

The authors frame this as a problem of data curation for pretraining (Section 1). Modern LLMs are trained on massive, heterogeneous corpora assembled from diverse sources. For example, The Pile contains 22 domains ranging from academic papers (arXiv, PubMed) to conversational text (Ubuntu IRC, OpenSubtitles) to code (GitHub). The GLaM dataset, used in training PaLM, spans filtered web pages, books, news, Wikipedia, forums, and conversations. In both cases, someone had to decide the sampling ratios. The Pile's creators used heuristic domain weights — essentially educated guesses about what proportions would produce a well-rounded model. The creators of GLaM and PaLM went further, tuning domain weights based on downstream task performance, requiring training many LMs on different weight configurations.

The gap this paper identifies is that neither approach is satisfactory. Heuristic weights are untethered from any measurable objective and are almost certainly suboptimal for model quality. Downstream-tuned weights risk overfitting to the specific evaluation tasks used for tuning and are prohibitively expensive — the paper notes this "requires training potentially thousands of LMs on different domain weights" (Section 1). Both approaches share a deeper failure: they don't answer the question of what an optimal data mixture even means in the absence of a specific downstream task.

Why This Problem Matters

The importance of domain weight selection extends far beyond academic curiosity about The Pile's default configuration. There are multiple reasons this is a high-stakes practical problem:

Training efficiency and cost. If better domain weights can achieve the same model quality with fewer training steps — the paper demonstrates a 2.6× speedup — this directly translates to reduced compute costs, energy consumption, and carbon footprint. At the scale of 8B+ parameter models, even modest improvements in training efficiency represent millions of dollars in saved compute. The authors explicitly connect to this in their broader impacts discussion (Section 6), framing DoReMi as a way to "halve the cost and energy consumption of training large language models."

The scaling imperative. As the field pushes toward ever-larger models (100B, 500B, 1T+ parameters), the cost of suboptimal data mixtures grows proportionally. A training run that processes trillions of tokens wastes enormous resources if a significant fraction of tokens comes from domains that contribute little to model capability. Worse, these wasted tokens are opportunity costs — they displace tokens from more valuable domains. Understanding how to allocate the data budget efficiently becomes more critical as models scale.

Data curation as a bottleneck. Recent work on data-centric AI has established that data quality and composition are often more important than model architecture or hyperparameters. However, data curation remains largely manual and heuristic-driven. The field has well-developed scaling laws for allocating compute between model size and data quantity (Hoffmann et al., 2022), but no analogous framework for allocating the data budget across domains within the data. DoReMi represents one of the first attempts to develop such a framework.

Fairness and representation. The authors note (Section 6) that distributionally robust optimization — the mathematical machinery underlying DoReMi — can have a "favorable impact on fairness." By optimizing for worst-case rather than average performance, DRO promotes good performance on all groups, including minority subgroups that might otherwise be underserved by a model trained to minimize average loss. Domain weights that maximize the average across all domains could systematically neglect small or "difficult" domains; DoReMi's minimax formulation explicitly guards against this.

Generalization to unseen tasks. Language models are foundation models: they are pretrained once and then deployed on many different downstream tasks, most of which are unknown at training time. Optimizing domain weights for a specific set of downstream tasks — as GLaM and PaLM did — risks producing a model that performs well on those particular tasks but generalizes poorly to others. DoReMi's key claim is that it optimizes domain weights without any knowledge of downstream tasks, using only the model's own training loss signal. If this works, it produces a more general-purpose model that isn't implicitly tuned to any particular evaluation.

Prior Approaches and Their Limitations

The paper situates itself relative to several strands of prior work, each of which falls short of providing a principled, task-agnostic solution for domain reweighting:

Heuristic domain weights (The Pile, C4). The default approach in many datasets is to assign domain weights based on intuition or simple heuristics. The Pile's domain weights were designed by the dataset creators using editorial judgment about what proportions would produce a well-rounded model. C4 (Raffel et al., 2019) uses heuristic text quality filters to clean CommonCrawl data. The fundamental limitation is obvious: human intuition about optimal data mixtures for billion-parameter models trained on hundreds of billions of tokens is unlikely to be correct. There's no feedback loop, no principled objective, and no way to verify that the chosen weights are near-optimal.

Downstream-tuned domain weights (GLaM, PaLM). GLaM (Du et al., 2021) and PaLM (Chowdhery et al., 2022) represent the state of the art in careful domain weight selection. These teams trained many models on different domain weight configurations and selected the weights that produced the best downstream task performance. This approach has several critical weaknesses that DoReMi aims to address:

Computational expense: Searching over domain weight space requires training many large models or running expensive optimization procedures (the paper cites "zero-order optimization, RL, or heuristic assumptions"). This cost limits how finely the space can be explored and makes the approach inaccessible to groups without massive compute resources.
Overfitting to specific downstream tasks: The domain weights are optimal for the specific set of tasks used during tuning. These tasks — typically question answering, reading comprehension, and similar benchmarks — represent a narrow slice of what the model will be used for in practice. A mixture optimized for TriviaQA may be suboptimal for code generation, translation, or scientific reasoning.
No generalization guarantees: The tuned weights provide no assurance about performance on domains or tasks that weren't in the tuning set. The paper's DoReMi approach, by contrast, targets worst-case loss across all domains in the pretraining data, providing a form of guaranteed coverage.

Example-level data selection (Mindermann et al., 2022; Xie et al., 2023). Recent work has explored finer-grained data selection: rather than reweighting entire domains, these methods select or filter individual examples during training. Mindermann et al. (2022) proposes selecting the top-k examples in each minibatch according to excess loss relative to a reference model — an approach that, like DoReMi, uses the concept of "learnability" as a selection criterion. Xie et al. (2023) shows that importance resampling to emphasize high-quality examples (those resembling Wikipedia and books) improves downstream performance.

DoReMi differs from these approaches in two important ways. First, it operates at the domain level rather than the example level, making it simpler to implement — the main model trains on a standard resampled dataset with no online selection logic. Second, and more fundamentally, DoReMi does not assume any target distribution for the data quality. Xie et al. (2023) explicitly favors Wikipedia-like text; this is a reasonable heuristic but encodes an assumption about what constitutes "good" data. DoReMi makes no such assumption — it lets the loss signal determine which domains deserve more weight.

Moore-Lewis selection (Moore and Lewis, 2010; Axelrod, 2017; Feng et al., 2022). Originally developed for domain adaptation in machine translation, Moore-Lewis selection chooses training examples with high cross-entropy difference between language models trained on target-domain and source-domain data. The approach requires specifying a target distribution that the selected data should resemble. DoReMi eliminates this requirement — there is no target distribution, only the goal of minimizing worst-case excess loss across the domains present in the training data. This makes DoReMi more appropriate for pretraining, where the downstream distribution is unknown and potentially very different from any single domain in the training set.

Distributionally robust optimization for language models (Oren et al., 2019). DRO-LM pioneered the application of distributionally robust optimization to language modeling, showing that training to minimize worst-case loss across domains (using a simple bigram reference model) could improve performance on minority domains without sacrificing average performance. The paper explicitly builds on this work, but with a crucial conceptual pivot: DRO-LM uses DRO to produce a robust model, while DoReMi uses DRO to produce better data. This is why DoReMi can discard the DRO-trained proxy model and keep only the domain weights — the weights are the output, not the model.

Data-centric AI in vision (Sorscher et al., 2022; Paul et al., 2021; and others). Work on data pruning and subset selection in computer vision has shown that carefully selecting training examples can beat power-law scaling and achieve the performance of much larger datasets with a fraction of the data. Methods like gradient matching, coreset selection, and uncertainty-based sampling have been extensively studied. However, these approaches typically assume that downstream data comes from the same distribution as training data, a condition that clearly fails for language model pretraining (where the training domains are a crude proxy for the vast space of possible downstream tasks). DoReMi's robust optimization approach is designed precisely for this distribution shift scenario.

How DoReMi Positions Itself

DoReMi's key positioning move is to reframe domain reweighting from a downstream optimization problem (tune weights to maximize task performance) to a pretraining optimization problem (tune weights using only the language modeling loss on training domains). This is more than a change of objective — it's a different philosophy about what makes good pretraining data:

"Instead of optimizing domain weights based on a set of downstream tasks, our approach aims to find domain weights which lead to models that perform well on all domains by minimizing the worst-case excess loss over domains" (Section 1).

The reference to "excess loss" — the gap between the model being trained and a pretrained reference model — is the technical linchpin. Excess loss captures learnability: domains where the reference model achieves low loss (the data is predictable, not just noise) but the current model still has high loss (the model hasn't learned it yet) are the ones that should receive more weight. This elegantly handles both ends of the difficulty spectrum. Trivially easy domains (where the current model already matches the reference) don't need more samples. Inherently noisy domains (where even the reference model has high loss) also don't benefit from more samples — there's nothing to learn. The domains that matter are the ones with headroom for improvement.

The paper also positions DoReMi relative to the broader scaling picture established by Hoffmann et al. (2022). While Chinchilla scaling laws optimize the allocation of compute between model size and training tokens, DoReMi operates within the training tokens, optimizing their composition across domains. The two approaches are complementary: one could first use DoReMi to determine domain weights, then apply compute-optimal scaling to determine the appropriate model size and total token budget. The paper doesn't make this connection explicitly, but it's a natural extension of the framework.

Finally, DoReMi positions itself as enabling a small-to-large transfer paradigm. The entire domain weight optimization runs on models that are 30× smaller than the final model, using only 8% of the main model's training compute. This is a practical breakthrough: it means that data mixture optimization, which previously required training many large models, can now be done cheaply with small models. The paper demonstrates that the optimized weights transfer effectively across a 30× scale gap, suggesting that the "optimal" data mixture for a given corpus may be a property of the data itself rather than the model architecture or scale — a finding with implications for how we think about pretraining data curation more broadly.

3. Technical Approach

3.1 Reader Orientation

This is a data-centric optimization paper that builds a system for automatically determining how much weight each data domain should receive during language model pretraining. The core idea is that by training a small "proxy" model using distributionally robust optimization over domains—where the optimizer dynamically upweights domains with high excess loss relative to a reference model—the resulting domain weights can be transferred to train a much larger model more efficiently, without ever needing to know what downstream tasks the model will be evaluated on.

3.2 Big-Picture Architecture (Diagram in Words)

The DoReMi system has five major components arranged in a three-step pipeline:

Reference Model Training (Step 1): A small language model (280M parameters) is trained on the default or reference domain weights. This model captures a baseline level of "learnability" for each domain—it tells us which domains are inherently predictable (low loss) and which are noisy (high loss).
Group DRO Proxy Training with Domain Weight Updates (Step 2): A second small model of identical architecture (the proxy model) is trained using Group Distributionally Robust Optimization. At each training step, the system computes the per-domain excess loss (proxy loss minus reference loss, clipped at zero), updates domain weights via exponentiated gradient ascent to upweight domains with high excess loss, then updates the proxy model using these new weights to reweight the training objective. Crucially, the domain weights themselves—averaged over training—are the output of this step, not the proxy model.
Main Model Training with Resampling (Step 3): The averaged domain weights from Step 2 define a new sampling distribution. A new dataset is created by resampling from the original domains according to these optimized weights, and a large model (8B parameters, 30× larger than the proxy) is trained on this resampled dataset using standard training procedures.

Information flows as follows: reference domain weights → Step 1 produces a reference model → Step 2 uses the reference model's per-token losses and a freshly-initialized proxy model to iteratively update domain weights via Group DRO → the averaged domain weights are extracted from the DRO training trajectory → Step 3 uses these weights to resample the training data → the main model trains on this reweighted data.

3.3 Roadmap for the Deep Dive

First, the formal definition of domain weights and the data distribution they induce, since this notation underpins everything else.
Second, the reference model (Step 1)—how it's trained, what role it serves, and why a reference model is necessary rather than using raw loss.
Third, the core DRO objective function (Equation 1) and the excess loss concept, because this is the mathematical heart of DoReMi and the mechanism by which domain weights are optimized without downstream task knowledge.
Fourth, the Group DRO training algorithm (Algorithm 1) in full operational detail—the interleaved updates, the token-level loss aggregation, the smoothing parameter, and the final averaging that produces the optimized domain weights.
Fifth, the main model training with resampled data (Step 3) and how it differs from the proxy model's reweighted training.
Sixth, iterated DoReMi—the extension that runs multiple rounds of the process to refine domain weights, along with its convergence criterion.
Seventh, the design choices and their justifications—why excess loss rather than raw loss, why Group DRO's online learning optimizer rather than online subset selection, why domain-level rather than example-level reweighting, why the reference and proxy models share the same architecture, and why the smoothing parameter exists.

3.4 Detailed, Sentence-Based Technical Breakdown

Domain Weights and the Data Distribution

Before explaining DoReMi's mechanisms, the paper establishes precise notation for what domain weights are and how they define a training distribution. Given $k$ domains (e.g., Wikipedia, GitHub, arXiv), each domain $D_i$ contains a set of preprocessed text examples. A domain weight vector $\alpha \in \Delta^k$ specifies a probability distribution over the $k$ domains—formally, $\alpha$ lies in the $k$ -dimensional probability simplex $\Delta^k$ , meaning each entry $\alpha_i \geq 0$ and $\sum_{i=1}^k \alpha_i = 1$ .

The domain weights induce a training distribution $P_\alpha$ through hierarchical sampling:

$P_\alpha = \sum_{i=1}^k \alpha_i \cdot \text{unif}(D_i)$

where $\text{unif}(D) = \frac{1}{|D|} \sum_{x \in D} \delta_x$ is the uniform distribution over examples in domain $D$ , and $\delta_x$ is a Dirac delta at example $x$ (a point mass that assigns probability 1 to $x$ and 0 to everything else).

What it computes: $P_\alpha$ is the marginal distribution over training examples after first sampling a domain $i$ with probability $\alpha_i$ , then sampling an example uniformly from that domain's pool. Multiplying $\alpha_i$ by $\frac{1}{|D_i|}$ means that giving domain $i$ weight $\alpha_i$ means a specific example from that domain has sampling probability $\alpha_i / |D_i|$ .

Why this form: This hierarchical sampling structure is the standard way pretraining datasets are constructed—you decide domain-level proportions first, then draw examples uniformly within each domain. The uniform-within-domain assumption encodes that all examples within a domain are equally valuable a priori, with only the domain-level mixing weights being optimized. This is a coarse-grained control mechanism compared to example-level selection, but it's computationally simpler and sufficient to capture the main effects.

A subtle but important detail: the domain weights $\alpha_i$ are based on token count in this paper's experiments. This means that two domains assigned the same weight $\alpha_i$ may contribute different numbers of examples if their example lengths differ. The paper notes this depends on the tokenizer (Appendix C), and all experiments chunk data into fixed-length 1024-token examples with a SentencePiece tokenizer (vocabulary size 256k), which standardizes the weight-to-token relationship across domains.

Step 1: Training the Reference Model

DoReMi's first step trains a small reference model $p_{\text{ref}}$ using standard next-token language modeling on data sampled from reference domain weights $\alpha_{\text{ref}}$ . The reference model serves as a baseline of learnability for each domain: its per-token negative log-likelihood $\ell_{\text{ref}}(x) = -\log p_{\text{ref}}(x)$ tells us how predictable each token is given an adequate model trained on the full data mixture.

Training procedure. The reference model is trained for $T$ steps with batch size $b$ , using the same training hyperparameters as the main model (Appendix C). For The Pile experiments, the reference domain weights $\alpha_{\text{ref}}$ are the default Pile domain weights—the heuristically chosen proportions from Gao et al. (2020). For the GLaM dataset, $\alpha_{\text{ref}}$ is uniform across the 8 domains. The reference model has 280M parameters in the main experiments, matching the proxy model's architecture exactly.

What the reference model provides. After training, $p_{\text{ref}}$ can be queried to get $\ell_{\text{ref},j}(x) = -\log p_{\text{ref}}(x_j | x_1, \dots, x_{j-1})$ for each token $j$ in any example $x$ . These per-token losses serve as the baseline in the excess loss computation (Step 2). A domain where $\ell_{\text{ref}}$ is low is one where the reference model finds the text predictable—it has learned statistical regularities in that domain's language. A domain where $\ell_{\text{ref}}$ remains high is inherently high-entropy or noisy; even a well-trained model cannot predict the next token confidently.

Why a reference model rather than raw loss? The paper explicitly contrasts this with DRO-LM (Oren et al., 2019), which used a simple bigram reference model. Using a full Transformer reference model of the same architecture as the proxy ensures the losses are on a comparable scale. More fundamentally, using the reference loss rather than raw loss transforms the optimization from "find domains that are hard for the current model" into "find domains where the current model has the most room to improve relative to what's achievable." This is the critical distinction between excess loss and raw loss, and the paper's ablation (Section 4, Figure 6 right) shows that optimizing raw loss—either preferring hardest domains ( $\ell_\theta(x)$ ) or easiest domains ( $-\ell_{\text{ref}}(x)$ )—fails to replicate DoReMi's gains.

The reference model can be small (280M) even when the final model is large (8B), because its role is purely analytical—it's never deployed or used for generation, only for computing baseline losses during proxy training. The reference model consumes approximately 4% of the main model's training FLOPs (half of the 8% total for reference + proxy).

The Core DRO Objective and Excess Loss

The mathematical heart of DoReMi is the distributionally robust optimization objective in Equation 1. This objective defines what the proxy model optimizes during Step 2, and—through the domain weight updates—indirectly determines the final domain weights.

$\min_{\theta} \max_{\alpha \in \Delta^k} L(\theta, \alpha) := \sum_{i=1}^k \alpha_i \cdot \left[ \frac{1}{\sum_{x \in D_i} |x|} \sum_{x \in D_i} \ell_\theta(x) - \ell_{\text{ref}}(x) \right] \quad (1)$

where:

$\theta$ are the proxy model parameters (weights of the Transformer being trained in Step 2).
$\alpha \in \Delta^k$ is a domain weight vector in the $k$ -dimensional probability simplex.
$D_i$ is the set of preprocessed examples in domain $i$ .
$|x|$ is the number of tokens in example $x$ .
$\ell_\theta(x) = -\log p_\theta(x)$ is the negative log-likelihood (next-token language modeling loss) of example $x$ under the current proxy model.
$\ell_{\text{ref}}(x) = -\log p_{\text{ref}}(x)$ is the negative log-likelihood of example $x$ under the previously trained reference model (frozen during Step 2).
$\ell_\theta(x) - \ell_{\text{ref}}(x)$ is the excess loss for example $x$ —how much worse the proxy model is than the reference model on this example.
The bracketed term $\frac{1}{\sum_{x \in D_i} |x|} \sum_{x \in D_i} (\dots)$ computes the per-token average excess loss for domain $i$ .

What the objective computes: The inner maximization $\max_{\alpha \in \Delta^k}$ finds the domain weights that put all probability mass on the single domain with the highest per-token average excess loss. The outer minimization $\min_\theta$ then optimizes the proxy model parameters to minimize this worst-case excess loss. This is a minimax formulation: the proxy model is trained to perform well specifically on the domain where it currently has the largest gap relative to the reference model.

The summation over domains with $\alpha_i$ multiplying each domain's average excess loss is a convex combination. When the inner maximization is solved exactly, $\alpha$ becomes a one-hot vector pointing to the worst domain. In practice, the Group DRO optimizer maintains a soft $\alpha$ that is updated incrementally and smoothed, so multiple domains can receive non-zero weight simultaneously.

Why excess loss rather than raw loss. If the objective used raw proxy loss $\ell_\theta(x)$ instead of excess loss, the worst-case domain would always be the one with the highest entropy (most unpredictable tokens), regardless of whether that domain is actually learnable. A domain of random noise would dominate the optimization, and the proxy model would waste capacity trying to fit noise. By subtracting the reference loss, domains where the reference model already has high loss (noisy/unlearnable domains) have small or even negative excess loss, effectively removing them from contention. Similarly, domains where the reference model has very low loss and the proxy model has also achieved low loss (already-learned domains) have near-zero excess loss.

The excess loss specifically highlights domains with headroom: where the proxy model is currently underperforming relative to what is achievable. This is the definition of "learnable but not yet learned"—the exact property that makes a domain a good candidate for additional training emphasis.

Why per-token averaging. The normalization by $\sum_{x \in D_i} |x|$ (total tokens in the domain) ensures that the domain-level excess loss is a per-token average rather than a per-example sum. Without this normalization, domains with longer average examples would receive disproportionate weight simply because they have more total tokens. The per-token normalization makes the domain weights dependent on the rate of excess loss rather than the volume of excess loss.

The Group DRO Training Algorithm (Algorithm 1)

Algorithm 1 provides the complete pseudocode for Step 2—the iterative process that trains the proxy model while simultaneously updating domain weights. The algorithm interleaves two types of updates at each training step: an exponentiated gradient ascent update on the domain weights $\alpha_t$ , and a standard optimizer update on the proxy model parameters $\theta_t$ .

Initialization. The algorithm takes as input the domain datasets $D_1, \dots, D_k$ , the number of training steps $T$ , batch size $b$ , domain weight step size $\eta$ (set to $\eta = 1$ in all experiments), and a smoothing parameter $c$ (set to $c = 1e-3$ or 0.001). The proxy model weights $\theta_0$ are randomly initialized (standard Transformer initialization). The domain weights $\alpha_0$ are initialized to uniform: $\alpha_0 = \frac{1}{k} \mathbf{1}$ , meaning every domain starts with equal weight $1/k$ .

Minibatch sampling (Line 4 of Algorithm 1). At each step $t$ , a minibatch $B = \{x_1, \dots, x_j\}$ of size $b$ is sampled from the uniform domain distribution $P_u$ where $u = \frac{1}{k} \mathbf{1}$ . This is a key design choice: the minibatch is always sampled uniformly across domains, not according to the current domain weights $\alpha_{t-1}$ or the reference weights $\alpha_{\text{ref}}$ . This ensures that every domain contributes data to every minibatch (in expectation), allowing the algorithm to compute per-domain excess losses even for domains with very small $\alpha_{t-1}$ .

Per-domain excess loss computation (Lines 6-7 of Algorithm 1). For each domain $i$ , the algorithm computes a scalar $\lambda_t[i]$ representing the current per-token excess loss on domain $i$ :

$\lambda_t[i] \leftarrow \frac{1}{\sum_{x \in B \cap D_i} |x|} \sum_{x \in B \cap D_i} \sum_{j=1}^{|x|} \max\{\ell_{\theta_{t-1},j}(x) - \ell_{\text{ref},j}(x), 0\}$

where:

$B \cap D_i$ is the subset of the current minibatch that belongs to domain $i$ .
$|x|$ is the token length of example $x$ (bounded by the maximum sequence length $L = 1024$ ).
$\ell_{\theta_{t-1},j}(x) = -\log p_{\theta_{t-1}}(x_j | x_1, \dots, x_{j-1})$ is the proxy model's per-token loss at position $j$ .
$\ell_{\text{ref},j}(x) = -\log p_{\text{ref}}(x_j | x_1, \dots, x_{j-1})$ is the reference model's per-token loss at the same position.
$\max\{\dots, 0\}$ clips the per-token excess loss at zero, discarding tokens where the proxy model actually outperforms the reference model.

What $\lambda_t[i]$ computes: the average per-token clipped excess loss for domain $i$ within the current minibatch. If domain $i$ has no examples in the minibatch, $\lambda_t[i]$ is effectively zero (or undefined—the paper's implementation presumably handles this gracefully). This is a stochastic estimate of the full-domain excess loss, computed on-the-fly without requiring a separate evaluation pass.

Why clipping at zero. The Group DRO optimizer requires non-negative losses because the exponentiated gradient update (next step) uses the loss directly in the exponent. If excess loss could be negative, domains where the proxy model outperforms the reference would receive decreased weight—but this would be counterproductive, since the reference model is frozen and represents a fixed capability baseline. The clipping ensures that only domains where the proxy is worse than the reference affect the weight updates; domains where the proxy has caught up or surpassed the reference are ignored.

Why token-level computation and aggregation. Computing excess loss at the per-token level rather than per-example ensures that long examples don't dominate the domain's average. Each token contributes equally to $\lambda_t[i]$ , regardless of which example it belongs to. This is consistent with how language modeling loss is typically computed—as an average over tokens—and aligns with the per-token normalization in the overall objective.

Domain weight update (Lines 8-9 of Algorithm 1). Given the per-domain excess losses $\lambda_t$ , the domain weights are updated via exponentiated gradient ascent:

$\alpha'_t \leftarrow \alpha_{t-1} \exp(\eta \lambda_t)$

where the exponential is applied entrywise: $\alpha'_t[i] = \alpha_{t-1}[i] \cdot \exp(\eta \cdot \lambda_t[i])$ . Since $\eta = 1$ , this is simply multiplying each domain's current weight by $\exp(\lambda_t[i])$ , which is at least 1 (since $\lambda_t[i] \geq 0$ ). Domains with high excess loss get their weight multiplied by a larger factor, while domains with zero excess loss keep their weight unchanged.

After the multiplicative update, the weights are renormalized and smoothed:

$\alpha_t \leftarrow (1 - c) \frac{\alpha'_t}{\sum_{i=1}^k \alpha'_t[i]} + c u$

where $u = \frac{1}{k} \mathbf{1}$ is the uniform distribution and $c = 1e-3$ is the smoothing parameter.

What the renormalization and smoothing compute: The first term $(1-c) \frac{\alpha'_t}{\sum_i \alpha'_t[i]}$ projects the multiplicatively updated weights back onto the probability simplex, ensuring they sum to 1. The second term $c u$ adds a small uniform component, ensuring that no domain's weight ever drops below $c/k \approx 4.5 \times 10^{-5}$ for The Pile's 22 domains. This prevents any domain from being completely zeroed out—the proxy model always maintains some exposure to every domain.

Why smoothing is necessary. Without smoothing ( $c = 0$ ), once a domain's weight decays to zero, it can never recover because it will never appear in the weighted objective, and thus its excess loss will never be computed again (or would be computed from zero examples). The smoothing term $c u$ acts as an epsilon-greedy exploration mechanism, ensuring the proxy model continues to see a small amount of data from all domains. This allows domains to "come back" if they become the worst-case domain later in training. The specific value $c = 1e-3$ was used in all experiments without extensive tuning—the paper states they "did not extensively tune these hyperparameters."

Why exponentiated gradient ascent. This update rule is the standard algorithm for online learning over the simplex—it's the instance of the Hedge/Freund-Schapire algorithm applied to the domain weight vector. The key property is that the cumulative product of exponentials (equivalently, the sum of losses in log-space) determines the final weights. This makes the weight update proportional to the exponential of cumulative excess loss, which rapidly amplifies differences between domains.

Proxy model update (Line 10 of Algorithm 1). After updating the domain weights, the proxy model parameters $\theta$ are updated for the objective $L(\theta_{t-1}, \alpha_t)$ using a standard optimizer (Adafactor in all experiments). This means the proxy model's training loss on each example is weighted by the domain weight of that example's domain. Concretely, if domain $i$ has current weight $\alpha_t[i] = 0.5$ , then every example in the minibatch from domain $i$ has its loss contribution scaled by 0.5 (relative to the uniform-minibatch baseline).

This update is computationally inexpensive compared to alternatives because it's just a loss reweighting—all examples in the minibatch still get forward and backward passes, but their gradients are scaled by different amounts. The alternative approach used in DRO-LM (Oren et al., 2019)—selecting a worst-case subset of each minibatch—would require running forward passes on a larger minibatch (e.g., double the size) and then selecting only a subset for the backward pass, which wastes forward-pass computation.

Final domain weight extraction (Return statement). After $T$ steps, DoReMi returns the average of the domain weights over the entire training trajectory:

$\bar{\alpha} = \frac{1}{T} \sum_{t=1}^T \alpha_t$

This is the standard approach from Nemirovski et al. (2009) for converting an online learning algorithm's trajectory into a single solution. The averaging has a crucial smoothing effect: early steps where weights are volatile contribute less to the final average as $T$ grows, while the converged/late-stage weights dominate. This makes $\bar{\alpha}$ more stable than using the final iterate $\alpha_T$ .

Why average over the trajectory rather than using the final weights. The domain weights can oscillate during training—domains alternately become the worst-case, their weight increases, the proxy model improves on them, their excess loss drops, and another domain takes over. Averaging smooths out these oscillations and captures the time-integrated importance of each domain. This is analogous to how Polyak averaging produces better parameter estimates in stochastic optimization.

Hyperparameter settings. The paper reports minimal tuning of the DRO-specific hyperparameters: domain weight step size $\eta = 1$ , smoothing parameter $c = 1e-3$ , uniform minibatch sampling, and Adafactor as the base optimizer. The number of training steps $T$ for the proxy model matches the main model's training steps (200k for The Pile, 300k for GLaM), and the batch size $b = 512$ also matches. The proxy model architecture is identical to the reference model (280M parameters, 12 layers, 12 attention heads, model dimension 768, hidden dimension 3072).

Step 3: Training the Main Model with Resampled Data

Once the optimized domain weights $\bar{\alpha}$ are obtained from Step 2, they define a new training distribution $P_{\bar{\alpha}}$ . Step 3 trains the large main model on data sampled from this distribution.

Resampling procedure. The paper explicitly states that the main model is trained by resampling the dataset with domain weights $\bar{\alpha}$ . This means a new training dataset is created by drawing examples according to $P_{\bar{\alpha}}$ : first sample a domain $i \sim \bar{\alpha}$ , then sample an example uniformly from $D_i$ . This resampled dataset is then used for standard training with the same hyperparameters as the baseline (200k steps, batch size 512, maximum sequence length 1024, Adafactor optimizer, learning rate 1e-3 with exponential decay to 1e-4, weight decay 1e-2, gradient clipping to norm 1, linear warmup for 6% of steps).

Resampling vs. reweighting. The paper draws an important distinction between the proxy model's training (which uses loss reweighting—multiplying each example's loss contribution by its domain's current $\alpha_t[i]$ ) and the main model's training (which uses data resampling—physically creating a new dataset with different domain proportions). This distinction is motivated by practical considerations: the main model is large (8B parameters) and benefits from a standard, unmodified training loop. Loss reweighting would require modifying the training code to accept per-example weights, which may interact poorly with optimizers, learning rate schedules, and other infrastructure assumptions. Resampling produces a standard dataset that can be fed into any training pipeline.

More subtly, resampling and reweighting are not equivalent in expectation for optimizers with momentum or adaptive learning rates (like Adafactor). With reweighting, the effective learning rate varies by domain (high-weight domains receive larger gradient contributions), which can destabilize training. With resampling, every example in the dataset has equal weight—the domain proportions are baked into the data rather than the loss—so the optimizer sees a consistent effective learning rate across all examples.

Why the proxy model underperforms the main model. Table 3b reveals an important phenomenon: when the proxy model and main model are the same size, the proxy model (trained with loss reweighting via DRO) typically underperforms the main model (trained with resampling on the averaged weights). At 280M, the proxy model beats the baseline on 19/22 domains; at 1B, the proxy model underperforms the baseline on all 22 domains, while the 1B main model trained on DoReMi weights beats the baseline on 19/22 domains. The paper hypothesizes that the mismatch between training paradigms (reweighting vs. resampling) explains this gap, and suggests that "a resampling-based Group DRO optimizer may improve DoReMi for larger proxy models."

This finding is crucial for understanding DoReMi's robustness: the proxy model itself can be relatively poor quality (as at 1B), but the domain weights it produces still transfer effectively to the main model. The domain weights encode which data to emphasize, not how good the proxy model is—and apparently the signal about domain importance survives even when the proxy model's training is suboptimal.

Compute budget breakdown. Training the reference model (Step 1) and the proxy model (Step 2) each consume roughly 4% of the main model's training FLOPs, for a total of approximately 8% overhead. For the 8B model trained for 200k steps at batch size 512, the main training processes $200,000 \times 512 \times 1024 \approx 1.05 \times 10^{11}$ tokens. The two 280M models process the same number of tokens each (since steps and batch size are matched), but their per-token FLOPs are roughly $280M / 8B \approx 3.5\%$ of the main model's, so $2 \times 3.5\% \approx 7\%$ , consistent with the paper's 8% figure.

Iterated DoReMi

The paper extends DoReMi by running it for multiple rounds, using the optimized domain weights from one round as the reference domain weights for the next round. This is called iterated DoReMi.

Procedure. After running DoReMi once (Steps 1–3) to obtain $\bar{\alpha}$ from the initial reference weights $\alpha_{\text{ref}}$ , the process is repeated with $\alpha_{\text{ref}} \leftarrow \bar{\alpha}$ . For the second round, a new reference model is trained on $P_{\bar{\alpha}}$ , a new proxy model is trained using this reference model, and new domain weights $\bar{\alpha}^{(2)}$ are extracted. This cycle continues.

Convergence criterion. Iteration stops when the maximum absolute change in any domain weight between rounds falls below a threshold:

$\|\bar{\alpha} - \alpha_{\text{ref}}\|_\infty < 1e-3$

where $\|\cdot\|_\infty$ is the infinity norm (maximum absolute entry). This means no domain's weight changes by more than 0.001 between successive iterations. On the GLaM dataset, convergence occurred after 3 rounds. The paper does not report the number of rounds to convergence on The Pile.

Why iteration helps. The first round of DoReMi uses a reference model trained on the default domain weights. The reference model's per-token losses therefore reflect what is learnable given the default data mixture. But the default mixture may be suboptimal—some domains may be undersampled, causing the reference model to underperform on them and making them appear artificially "hard" (high reference loss). This would incorrectly signal that these domains have low excess loss (headroom is small because reference loss is already high), causing DoReMi to downweight them.

In the second round, the reference model is trained on the optimized weights from round 1, which samples more heavily from domains that were identified as "learnable but not yet learned." This reference model should have lower loss on those domains, revealing new headroom—potentially shifting the optimal weights further. The iteration allows the domain weight optimization to account for the fact that the reference model's own quality depends on the data mixture it was trained on.

Empirical behavior on GLaM. Table 2 shows the domain weight evolution across rounds. Round 1 substantially upweights filtered webpages (from 0.125 uniform to 0.44) and downweights Wikipedia (from 0.125 to 0.09). Round 2 further adjusts: filtered webpages increase to 0.51, conversations nearly double (0.10 → 0.22), books increase (0.11 → 0.17), while Wikipedia, forums, and news decrease further. Round 3 is almost identical to Round 2 (all weights within 0.01), satisfying the convergence criterion.

Notably, the Round 2 weights broadly match the downstream-tuned weights: filtered webpages (0.51 vs. 0.42), conversations (0.22 vs. 0.27), books (0.17 vs. 0.20), Wikipedia (0.05 vs. 0.06), forums (0.04 vs. 0.02), news (0.02 vs. 0.02). DoReMi recovers similar domain weights without any downstream task data.

Design Choices and Their Justifications

Excess loss over raw loss. The paper's ablation (Figure 6, right) demonstrates that optimizing raw proxy loss ( $\ell_\theta$ , "hardest domains") or negative reference loss ( $-\ell_{\text{ref}}$ , "easiest domains") does not match DoReMi's downstream accuracy gains. The excess loss combines both signals: it identifies domains that are learnable (low reference loss, so they're not pure noise) but not yet learned (high proxy loss, so there's room to improve). This is the operational definition of "worth training on" that DoReMi operationalizes.

Group DRO optimizer over online subset selection. DRO-LM (Oren et al., 2019) selects a worst-case subset of each minibatch for gradient updates, requiring oversampling (e.g., sampling a minibatch of size 2b, computing forward passes on all 2b examples, then backpropagating on the b worst). DoReMi's Group DRO optimizer instead reweights the loss—all examples contribute to the gradient, just with different weights. The paper argues this is "more expensive" for online subselection since it requires forward passes on discarded examples, while Group DRO uses all examples. However, the reweighting approach means the proxy model sees a different loss scale per example, which becomes problematic at larger model sizes (Table 3b).

Domain-level over example-level reweighting. Operating at the domain level makes DoReMi practical for large-scale pretraining: the output is a small set of $k$ numbers (22 for The Pile, 8 for GLaM) that can be used to resample the dataset once, after which standard training proceeds without modification. Example-level reweighting would require either online selection during main model training (as in Mindermann et al., 2022) or storing importance weights for billions of examples. The domain-level approach also naturally handles the pretraining/downstream distribution shift via robust optimization—the DRO objective explicitly optimizes for performance on all domains, providing worst-case guarantees that example-level methods matching a target distribution do not.

Reference and proxy models share architecture. Using identical architectures ensures the losses are directly comparable. If the reference model had a different capacity or architecture, its per-token losses would not reflect what the proxy model could achieve—a weaker reference model would make domains appear artificially "learnable" (low reference loss because the model is underspecified), while a stronger reference would make them appear artificially "unlearnable" (high reference loss because the model is overfitting). Same-size, same-architecture models provide a fair baseline.

The paper explicitly contrasts this with DRO-LM's use of a bigram reference model, noting that a full Transformer reference "ensure[s] that the losses are on a similar scale." The bigram model would have much higher loss on any domain requiring long-range context, making excess loss unreliable as a learnability signal.

Smoothing parameter $c = 1e-3$ . This small uniform component ensures no domain is ever completely removed from the proxy model's training distribution. Without it, domains that are initially downweighted might never recover because the proxy model never sees enough data from them to accurately estimate their excess loss. The smoothing acts as epsilon-greedy exploration in the domain weight space. The specific value $1e-3$ adds 0.1% uniform weight to all domains, which is negligible for highly-weighted domains (e.g., 0.5 weight becomes 0.4995 + 0.000045 ≈ 0.4995) but prevents zero-weight domains from vanishing entirely.

Minibatch sampled uniformly regardless of $\alpha_t$ . Algorithm 1 always samples minibatches from the uniform domain distribution $P_u$ , even when the current domain weights are highly skewed. This ensures that every domain contributes examples to every minibatch in expectation, providing the information needed to update all entries of $\lambda_t$ . If minibatches were sampled according to $\alpha_t$ , domains with small weights would appear rarely, making their $\lambda_t[i]$ estimates noisy or undefined. The uniform sampling ensures that the domain weight update is based on current (although noisy) information about all domains, which is important for the online learning dynamics to correctly identify when a previously low-weight domain has become the new worst case.

Why 280M proxy model size. The main experiments use 280M proxy/reference models for an 8B main model. Figure 6 (left) shows that increasing proxy model size from 70M to 280M improves downstream accuracy at 8B, but increasing to 1B does not continue the trend—it's slightly worse than 280M. The paper hypothesizes this is because the Group DRO optimizer works less well at larger scales (as seen in Table 3b). The 280M size hits a sweet spot: large enough to produce meaningful loss estimates and domain weight signals, but small enough that the DRO training remains stable and the total overhead is only 8% of main model training.

4. Key Insights and Innovations

Innovation 1: The Pivot from "Train a Robust Model" to "Extract the Data Recipe"

The most conceptually distinctive move in DoReMi is not the DRO mathematics itself—which follows established Group DRO formulations (Oren et al., 2019; Sagawa et al., 2020)—but the repurposing of DRO from a model-training objective into a data-discovery mechanism. In the standard DRO-LM framework, the entire point of doing distributionally robust optimization is to produce a model that performs well under worst-case domain shifts. The trained robust model is the output. DoReMi discards the DRO-trained proxy model entirely after Step 2 and keeps only the domain weight trajectory extracted from the optimization process.

This is a fundamentally different stance toward the optimization. In DoReMi, the proxy model is not a product—it is an instrument for measuring which domains the training process would need to emphasize in order to close the gap with a reference model. The DRO machinery, rather than solving the problem of how to train robustly, becomes a cheap simulation of what data mixture a much larger model would benefit from most. The actual robust training happens implicitly through the resampled dataset in Step 3, where the large model trains with standard ERM but on better-proportioned data. This is a conceptual reframing of the relationship between robust optimization and data curation: DRO figures out where the leverage is, and a separate, larger training run exploits that leverage.

Prior work treated data selection and robust optimization as separate activities. Mindermann et al. (2022) used excess loss for online example selection during training, but selection and model training were intertwined in a single process. Xie et al. (2023) used importance resampling with a target distribution defined by heuristics (preferring Wikipedia-like text). DoReMi is the first to show that the internal dynamics of a DRO training run—specifically, the sequence of domain weights produced by the online learning updates—contains actionable information about optimal data proportions that transfers across a 30× model scale gap. This transfer property is the key empirical finding that validates the pivot: if the domain weights didn't generalize from 280M to 8B models, discarding the proxy model would be pointless. The fact that they do (Figure 4: perplexity improvement on all 22 domains at 8B, Table 3a: 22/22 domains beating baseline) is evidence that the extracted weights capture something fundamental about the data rather than something idiosyncratic to the small model's training dynamics.

The significance of this pivot is primarily conceptual and economic. Conceptually, it separates the discovery of good data mixtures from the execution of training on them, establishing data mixture selection as a distinct optimization problem solvable with cheap proxy runs. Economically, it makes data mixture optimization accessible: training two 280M models (8% of main model FLOPs) replaces "training potentially thousands of LMs on different domain weights" that the GLaM/PaLM downstream-tuning approach required. This is not a small efficiency gain—it changes the problem from being computationally prohibitive to being a modest overhead.

Innovation 2: Excess Loss as a Task-Agnostic Learnability Metric That Handles Both Ends of the Difficulty Spectrum

The excess loss $\ell_\theta(x) - \ell_{\text{ref}}(x)$ is not a new quantity—it appears in Oren et al. (2019) (with a bigram reference) and is closely related to the cross-entropy difference in Moore-Lewis selection (Moore and Lewis, 2010; Axelrod, 2017) and the "learnability" criterion in Mindermann et al. (2022). What makes DoReMi's use distinctive is the dual role it plays in a minimax objective and the specific way it resolves the tension between upweighting hard domains and upweighting learnable ones.

Prior data selection methods based on loss magnitude face an inherent ambiguity: a domain can have high loss either because it is difficult but learnable (the model has room to improve) or because it is intrinsically noisy (no model can predict it well). Raw proxy loss $\ell_\theta(x)$ cannot distinguish these cases—it would upweight both equally, wasting capacity on noise. Conversely, methods based on a reference model's loss alone ( $\ell_{\text{ref}}(x)$ ) would upweight domains that are already well-predicted by the reference, effectively emphasizing easy domains that don't need more samples.

The excess loss resolves this ambiguity by subtracting the reference loss, not by thresholding or heuristics, but as an algebraic operation embedded in the DRO objective. Domains where the reference model already has high loss (noisy/unlearnable domains) have small or even negative excess loss—they drop out of the maximization. Domains where the proxy has caught up to the reference (low excess loss) also drop out. What remains are domains where the reference has low loss (the data is learnable) but the proxy still has high loss (the model hasn't learned it yet). This is the operational definition of "worth training on more."

The paper demonstrates the necessity of this dual signal through a clean ablation (Section 4, Figure 6 right, Table 7). Optimizing for the hardest domains alone (raw proxy loss $\ell_\theta$ ) produces domain weights that actually worsen perplexity compared to the baseline—worst-case log-perplexity degrades from 2.39 to 2.66, and not a single domain improves (Table 7). Optimizing for the easiest domains ( $-\ell_{\text{ref}}$ ) is even worse (worst-case log-perplexity 4.27). Neither component alone works; both together, in the form of excess loss, improve all 22 domains. This ablation transforms the excess loss from "one reasonable choice among many" into a necessary condition for DoReMi's gains—the components are individually harmful, and only their combination in the specific algebraic form of a difference works.

The theoretical contribution here is a principled justification for why baselining matters that goes beyond what prior work provided. Oren et al. (2019) used a bigram baseline and observed it helped, but didn't isolate the mechanism. The DoReMi paper's simple unigram example (Appendix D) provides a toy model showing that domains with the highest and lowest entropy can both be downweighted without harm, and that excess loss naturally identifies the intermediate-entropy domains that benefit most from additional samples. The example is minimal (3 domains, 3 tokens, Bayesian estimator with Dirichlet prior) but it crystalizes the intuition: the parameter error for low-entropy domains decreases rapidly with samples (they need few), the error for high-entropy domains is minimized at zero samples (they're already at the prior, which matches the ground truth), so optimal allocation shifts samples to intermediate domains. DoReMi's domain weights on The Pile (Table 1) exhibit this pattern: arXiv (low entropy, formulaic) drops from 10.5% to 0.4%, Books3 (intermediate) drops only modestly, and Pile-CC (high diversity, intermediate learnability) surges from 11.2% to 60.6%.

This is fundamentally a diagnostic contribution rather than a new algorithm—it gives the field a way to think about data quality in terms of learnability rather than superficial features like "cleanliness" or "resemblance to Wikipedia." The excess loss provides an operational answer to the question "which data is worth training on?" that requires no human judgment, no downstream tasks, and no assumptions about what good data looks like.

Innovation 3: The Empirical Demonstration That Domain Weights Generalize Across a 30× Model Scale Gap

The paper's most striking empirical result is not that DoReMi improves training—it's that domain weights optimized using a 280M-parameter proxy model transfer essentially intact to training an 8B-parameter model and produce consistent improvements. This is a non-obvious empirical discovery rather than a theoretical prediction, and it has significant implications for how the field thinks about data mixture optimization.

Before this work, there was no evidence that optimal data proportions for small models would generalize to large ones. Several plausible arguments cut against transferability: larger models have greater capacity to absorb information from high-entropy domains, potentially changing which domains are most valuable; the loss landscape and optimization dynamics differ at scale; and the interaction between domain weights and model size might be non-separable—what's optimal at 280M could be suboptimal at 8B. The paper's evidence contradicts these concerns in the specific regime tested.

The key evidence appears in Table 3a and Figure 6 (left). When varying the proxy model size from 70M to 280M to 1B and evaluating the resulting 8B main model, DoReMi (280M → 8B) produces the best results (worst-case log-perplexity 1.46, average 1.40), but even DoReMi (70M → 8B)—using a proxy model more than 100× smaller than the main model—still improves all 22 domains over the 8B baseline (Table 4). The 1B proxy model's domain weights continue to improve 22/22 domains at 8B despite the 1B proxy model itself being relatively poor quality (Table 3b: it underperforms the 1B baseline on all domains). Remarkably, iterated DoReMi on the GLaM dataset finds domain weights (Table 2) that match the pattern of weights manually tuned on downstream tasks—without using any downstream data. This suggests that optimal domain weights may be more a property of the data composition than of the specific model architecture or training recipe.

The conceptual implication is that data mixture optimization can be decoupled from model scale optimization, at least up to the 30× factor tested. This means the scaling problem can be factored: first determine good domain proportions with a cheap proxy run, then separately determine the model size and token budget using existing scaling laws (Hoffmann et al., 2022). The two optimizations don't need to be entangled in a single expensive search. This is a practical insight that lowers the barrier to data-centric optimization—a team that can only afford to train a 280M model can still optimize the data mixture for a much larger training run they plan to execute later.

The limits of this transferability are an open question. The paper shows success at 30× (280M → 8B) but doesn't test the limits—would 280M weights still transfer to a 70B or 500B model? The failure of the 1B proxy model to outperform the 280M proxy in the 8B results (Figure 6 left) hints at a non-monotonic relationship, and the paper's hypothesis that Group DRO's loss reweighting becomes problematic at larger proxy scales (Section 4) suggests the transferability may be bounded by optimizer dynamics rather than an inherent property of domain weights. This is a productive ambiguity—it defines a clear research question about how far and under what conditions data mixture information propagates across model scales.

The significance of this finding is primarily practical: it validates the core premise that makes DoReMi economically viable. If domain weights had to be optimized at the target model scale, DoReMi's approach would be no cheaper than the downstream-tuning methods it aims to replace. The 30× transfer factor turns an intractable optimization into an 8% overhead.

Innovation 4: The Identification of "No-Tradeoff" Domain Reweighting as an Empirical Phenomenon Distinct from Domain-Level Overfitting

One of the paper's most counterintuitive results is that DoReMi improves perplexity on every single domain—even domains it substantially downweights (Figure 4, Table 4). For example, PubMed Central drops from 10.7% weight in the baseline to 0.5% under DoReMi (280M), yet its log-perplexity improves from 1.60 to 1.46. arXiv drops from 10.5% to 0.4%, and log-perplexity improves from 1.64 to 1.38. StackExchange drops from 9.3% to 1.5%, and log-perplexity improves from 1.61 to 1.39. This is not the expected behavior under a naive tradeoff model where reweighting toward one domain necessarily degrades others.

The standard concern with domain reweighting—and indeed with any form of data curation—is that emphasizing certain domains will cause the model to overfit to their characteristics at the expense of performance on other domains. This is precisely the tradeoff that DRO is designed to mitigate (by optimizing worst-case rather than average loss), but DoReMi's result goes beyond mitigation: the worst-case domain also improves. The 8B baseline's worst-case log-perplexity was 1.71 (Ubuntu IRC), and DoReMi (280M → 8B) reduces it to 1.46 (Table 3a). This means DoReMi is not merely redistributing quality—it is producing a Pareto improvement.

The paper offers a specific hypothesis for why this can happen (Section 3.2): domains at the extremes of the learnability spectrum consume disproportionate samples under uniform or heuristic weights without providing proportional benefit. Very-low-entropy domains (like arXiv, where the language is highly formulaic and repetitive) require few samples to learn—the model saturates quickly. Very-high-entropy/noisy domains (perhaps OpenSubtitles) are, at the token level, close to a uniform prior—the model's random initialization already approximates them, and learning their specific statistics transfers poorly to other domains. The intermediate-entropy domains—where language is structured but varied, like the diverse web text in Pile-CC—benefit most from additional samples, and positive transfer from learning these domains well improves performance on all domains, including the extremes.

This hypothesis is supported by the simple theoretical model in Appendix D, where a 3-domain, 3-token setup with a Dirichlet-Multinomial Bayesian estimator shows that the lowest-entropy domain (predictability 1.0) and the highest-entropy domain (uniform distribution) should both be downweighted in favor of the intermediate domain (entropy 0.46), and that doing so reduces parameter error for all three domains. The simulation with DoReMi on this toy example recovers the predicted weights exactly (0.39, 0.61, 0.0 after rounding).

The significance here is conceptual and practical. Conceptually, it changes the framing of data mixture optimization from "which domains do we want the model to specialize in?" to "which domains provide the most generalizable learning signal?"—two questions with very different answers. Practically, it means domain reweighting can be a strictly beneficial intervention rather than a zero-sum negotiation between domain-specific performance goals. This removes a major objection to automated data curation—the fear that optimizing for one metric will silently degrade another—and suggests that the current heuristic weights in widely-used datasets like The Pile are not just suboptimal but are leaving easy Pareto improvements on the table.

5. Experimental Analysis

Evaluation Methodology

Dataset. The paper uses two datasets: The Pile (Gao et al., 2020), an 800GB publicly available text corpus with 22 domains (listed in Table 1), and the GLaM dataset (Du et al., 2021), which contains text from 8 domains (Table 2) and was also used in training PaLM (Chowdhery et al., 2022). For The Pile, the test set is held-out validation data from each domain, used for perplexity evaluation; for downstream evaluation, the paper uses the generative one-shot tasks from GPT-3 (Brown et al., 2020): TriviaQA, NaturalQuestions, WebQuestions, SQuADv2, and LAMBADA, with standard exact-match accuracy. The GLaM dataset results are reported on the same five downstream tasks.
Base model(s). All models are decoder-only Transformers (Vaswani et al., 2017) trained from scratch with the standard next-token language modeling loss. The proxy and reference models use 280M parameters (12 layers, 12 attention heads, model dimension 768, hidden dimension 3072; full architecture table in Appendix Table 9). The main model uses 8B parameters (32 layers, 32 attention heads, model dimension 4096, hidden dimension 24576). Ablation experiments cover scales from 70M to 1B parameters (Appendix Table 9). All models use a SentencePiece tokenizer with 256k vocabulary and are trained on examples chunked to length 1024 tokens. The paper states the 280M model is "representative of the capabilities of many contemporary LLMs" though this claim is specific to the PaLM 2-derived architecture family.
Metrics. The primary metrics are: (1) per-domain log-perplexity on held-out validation data, measured as negative log-likelihood per token; (2) average one-shot downstream accuracy (exact match) averaged across the 5 generative tasks (TriviaQA, NaturalQuestions, WebQuestions, SQuADv2, LAMBADA); (3) worst-case log-perplexity across the 22 domains (the maximum per-domain perplexity); (4) number of domains where the model beats the baseline (out of 22). For the FLOPs-matched comparisons, the paper also reports the training steps needed to reach baseline accuracy (e.g., "2.6× faster").
Baselines. The paper compares against: (1) Baseline domain weights: for The Pile, the heuristically-chosen default domain weights from Gao et al. (2020), computed by counting examples per domain and multiplying by the number of epochs specified in the original Pile paper (Appendix C); for GLaM, uniform domain weights across the 8 domains. (2) Downstream-tuned domain weights (GLaM only): the domain weights from Du et al. (2021) that were optimized based on downstream task performance — treated as an "oracle" comparison since these tasks overlap with the evaluation set. (3) Hardest domains ablation: running DoReMi's DRO optimizer but replacing excess loss with raw proxy model loss ℓ_θ(x), which prefers domains that are hardest for the proxy model. (4) Easiest domains ablation: running the DRO optimizer with negative reference loss -ℓ_ref(x), which prefers domains easiest for the reference model. (5) For the scale analysis, the paper compares DoReMi (X → X) where proxy and main model share size X ∈ {280M, 510M, 760M, 1B}, against a baseline model of size X trained with standard domain weights.
Generation budget / compute accounting. Compute is equalized by fixing the number of training tokens processed: 200k steps at batch size 512 with maximum sequence length 1024 for The Pile (~1.05 × 10^11 tokens), and 300k steps for GLaM. The domain weight optimization (Steps 1–2) consumes additional compute: training the 280M reference and proxy models each costs approximately 4% of the main 8B model's training FLOPs, for a total overhead of ~8%. All FLOPs come from standard forward and backward passes. The paper measures training speedup in terms of "steps to reach baseline accuracy" — for example, DoReMi (280M → 8B) reaches the 8B baseline's accuracy within 75k steps versus 200k for the baseline, a 2.6× speedup.
Cross-validation / statistical protocol. The paper does not employ cross-validation or statistical significance testing. Results are reported as point estimates from single training runs. For the iterated DoReMi convergence criterion (GLaM dataset), the paper defines convergence as maximum absolute change in any domain weight ‖ᾱ - α_ref‖_∞ < 1e-3 and reports that this takes 3 rounds. The downstream evaluation uses the standard test sets for the 5 generative tasks with exact-match scoring, following the evaluation protocol from Brown et al. (2020). The per-domain perplexity is computed on held-out validation data from each domain of The Pile; the paper does not specify the size of these validation splits.

Main Quantitative Results

Perplexity Improvements on The Pile (8B Main Model)

The central perplexity result appears in Figure 4 and Table 4. DoReMi (280M → 8B) reduces per-domain log-perplexity on all 22 domains of The Pile compared to the 8B baseline trained with default domain weights. The worst-case log-perplexity across domains drops from 1.71 (baseline) to 1.46, while the average log-perplexity drops from 1.64 to 1.40 (Table 3a). This improvement is not driven by uniformly upweighting a few domains at the expense of others — even heavily downweighted domains improve. For example (Table 4): PubMed Central drops from 10.7% weight to 0.46% weight under DoReMi (280M), yet its log-perplexity improves from 1.60 to 1.46; arXiv drops from 10.5% to 0.36% and improves from 1.64 to 1.38; StackExchange drops from 9.3% to 1.5% and improves from 1.61 to 1.39.

Varying proxy model scale at fixed 8B main model (Table 3a, Table 4). DoReMi (70M → 8B) already improves all 22 domains (worst-case log-perplexity 1.63, average 1.53). DoReMi (150M → 8B) further improves (worst-case 1.56, average 1.52). DoReMi (280M → 8B) achieves the best results (worst-case 1.46, average 1.40). DoReMi (1B → 8B) shows slightly worse results than 280M (worst-case 1.58, average 1.54), though still improving all 22 domains over baseline. This non-monotonic relationship — the 1B proxy model produces domain weights that transfer less effectively than the 280M proxy model's weights — is attributed by the authors to the Group DRO optimizer working less well at larger proxy scales, a hypothesis supported by the proxy model quality data (Table 3b, discussed below under "Proxy vs. main model quality").

Downstream Accuracy Improvements on The Pile (8B Main Model)

The headline downstream result (Figure 2, Figure 3 left, Table 5): DoReMi (280M → 8B) improves average one-shot downstream accuracy by 6.5 percentage points over the 8B baseline (26.56% vs. 20.03% average across the 5 tasks) and achieves the baseline accuracy within 75k training steps — 2.6× faster than the baseline's 200k steps.

Per-task breakdown (Table 5): LAMBADA improves from 20.10% to 29.19% (+9.09 points), NaturalQuestions from 4.35% to 7.73% (+3.38), SQuADv2 from 44.43% to 51.89% (+7.46), TriviaQA from 24.55% to 34.86% (+10.31), and WebQuestions from 6.74% to 9.15% (+2.41). All 5 tasks improve, with TriviaQA and LAMBADA showing the largest absolute gains.

Varying proxy model scale at fixed 8B main model (Figure 6 left, Table 5). Average downstream accuracy varies with proxy model size: DoReMi (70M → 8B) achieves 21.11% (vs. baseline 20.03%, +1.08 points), DoReMi (150M → 8B) achieves 22.51% (+2.48), DoReMi (280M → 8B) achieves 26.56% (+6.53), and DoReMi (1B → 8B) achieves 22.35% (+2.32). The non-monotonicity mirrors the perplexity results: 280M proxy yields the best 8B main model, with 1B proxy performing worse than 280M and roughly equivalent to 150M. The paper notes that even the suboptimal DoReMi (1B → 8B) still achieves the baseline accuracy almost 2× faster (Figure 6 left: the 1B curve crosses the baseline accuracy around 100k steps).

Speedup quantification (Figure 2, Figure 3 left). The 2.6× speedup is measured as: baseline reaches its final accuracy of ~20% at 200k steps; DoReMi (280M → 8B) reaches the same ~20% accuracy at approximately 75k steps. The DoReMi curve continues improving beyond this point, ending at 26.56% at 200k steps. This means DoReMi provides both a speedup (reaching baseline quality faster) and an absolute improvement (reaching higher final quality).

Results on the GLaM Dataset (8B Main Model)

Figure 3 (right) shows the GLaM results. Round 1 of DoReMi achieves comparable downstream accuracy to the uniform baseline (~33% vs. ~32% at 300k steps), while Round 2 of iterated DoReMi achieves comparable accuracy to downstream-tuned domain weights (~34.5% for both at 300k steps). This is the key "oracle-matching" result: DoReMi, using no downstream task information, produces domain weights that perform as well as weights explicitly optimized on the downstream evaluation tasks.

The GLaM domain weight evolution across rounds (Table 2) converges quickly. Round 1 shifts weights substantially from uniform (0.125 each): filtered webpages increase to 0.44, conversations to 0.10, books to 0.11, while Wikipedia drops to 0.09, forums to 0.16, news to 0.10. Round 2 further adjusts: webpages to 0.51, conversations nearly double to 0.22, books to 0.17, Wikipedia drops to 0.05, forums to 0.04, news to 0.02. Round 3 is essentially identical to Round 2 (all changes < 0.01), satisfying the convergence criterion. The downstream-tuned weights (from Du et al., 2021) show a similar pattern: webpages 0.42, conversations 0.27, books 0.20, Wikipedia 0.06, forums 0.02, news 0.02. The two remaining GLaM domains (not shown in the main table body — "Other" categories) presumably account for the remaining ~0.02 probability mass.

Domain Weights Analysis (The Pile and GLaM)

Table 1 presents the full DoReMi (280M) domain weights for The Pile's 22 domains versus the baseline weights. The most dramatic changes (difference column): Pile-CC increases by +0.4936 (from 0.1121 to 0.6057 — more than half the total probability mass on a single domain), YouTubeSubtitles +0.0460, PhilPapers +0.0247, HackerNews +0.0059. The largest decreases: PubMed Central −0.1025 (from 0.1071 to 0.0046), ArXiv −0.1016 (from 0.1052 to 0.0036), StackExchange −0.0776 (from 0.0929 to 0.0153), PubMed Abstracts −0.0732 (from 0.0845 to 0.0113), Books3 −0.0452 (from 0.0676 to 0.0224), Wikipedia (en) −0.0220 (from 0.0919 to 0.0699). Several small domains receive moderate upweighting: Enron Emails, EuroParl, Ubuntu IRC, BookCorpus2, NIH ExPorter all increase by 0.001–0.006.

A notable pattern: academic and reference domains are heavily downweighted (ArXiv, PubMed Central, PubMed Abstracts, StackExchange, FreeLaw, USPTO Backgrounds all drop to < 0.02), while diverse web text is dramatically upweighted (Pile-CC from 0.11 to 0.61). The paper notes that a 1B proxy model produces qualitatively different weights: OpenWebText2 becomes the most upweighted domain (0.3289) rather than Pile-CC (0.1199), though many other domains show similar relative changes (Appendix Table 8). This suggests "there may be multiple possible local minima in the domain weight space."

The GLaM domain weights (Table 2) show a less extreme redistribution — the dataset has only 8 domains versus 22 for The Pile, providing less room for dramatic reallocation. The general pattern converges toward the downstream-tuned weights within 2–3 rounds.

Results Across Scales: DoReMi with Matched Proxy and Main Model Sizes

Section 4 and Figure 5 present experiments where the proxy/reference models have the same size as the main model trained with DoReMi domain weights (DoReMi X → X for X ∈ {280M, 510M, 760M, 1B}). These experiments are "for scientific purposes since this does not save compute in practice" — they analyze DoReMi's behavior without the proxy-to-main scale gap.

Downstream accuracy (Figure 5). DoReMi consistently improves average one-shot downstream accuracy across all four scales. The absolute gap between DoReMi and baseline at 200k steps ranges from ~1.5–3 percentage points across scales, with the gap narrowing slightly at larger scales (the 510M point appears to have a smaller gap than 280M or 760M, though the paper does not quantify this precisely). The paper reports that DoReMi "achieves the baseline accuracy 4× faster on average across scales" — meaning the DoReMi curves reach the baseline's final accuracy at roughly 50k steps rather than 200k.

Perplexity (Table 6). DoReMi improves the worst-case log-perplexity at all scales except 510M, where it is equal to baseline (2.14). At 280M: worst-case drops from 2.39 to 2.19, average from 2.32 to 2.13, with 22/22 domains beating baseline. At 1B: worst-case drops from 1.94 to 1.92, average from 1.87 to 1.83, with 19/22 domains beating baseline. The number of domains beating baseline is 22/22 (280M), 15/22 (510M), 17/22 (760M), 19/22 (1B). The paper notes that "DoReMi improves the worst-case perplexity on all scales and improves 18 of 22 individual domain perplexities on average across scales."

GLaM results across scales (Appendix A, Figure 7). DoReMi is "comparable or better than both the baseline (uniform) domain weights and downstream-tuned domain weights" across all four scales. Interestingly, at 280M, Round 2 weights achieve slightly worse downstream accuracy than Round 1 weights when used to train 280M models, but transfer better to training 8B models. This is the only evidence in the paper of domain weights having scale-dependent transfer properties, and it's presented as an observation without detailed analysis.

Proxy Model vs. Main Model Quality

Table 3b presents a critical comparison: when proxy and main models have the same size, how does the proxy model's perplexity compare to the main model trained on the proxy's domain weights? The proxy model systematically underperforms the main model, and this gap worsens at larger scales.

At 280M: the proxy model beats the baseline on 19/22 domains (average log-perplexity 2.27 vs. baseline 2.32), but underperforms the DoReMi main model trained on its weights (2.27 vs. 2.13, with the main model beating baseline on 22/22 domains). At 1B: the proxy model underperforms the baseline on all 22 domains (average log-perplexity 2.02 vs. baseline 1.87), while the DoReMi main model trained on the proxy's domain weights beats the baseline on 19/22 domains (average 1.83 vs. baseline 1.87).

This "proxy degradation" phenomenon is important for understanding DoReMi's robustness. Even when the DRO-trained proxy model itself is of poor quality (1B proxy worse than baseline on every domain), the domain weights extracted from its training trajectory still produce a main model that significantly improves over baseline (19/22 domains beating baseline, worst-case log-perplexity 1.92 vs. 1.94, average 1.83 vs. 1.87). The paper hypothesizes that the mismatch between proxy model training (loss reweighting) and main model training (data resampling) explains the proxy model's poor performance, and suggests a "resampling-based Group DRO optimizer may improve DoReMi for larger proxy models."

Ablation Studies and Robustness Checks

Excess loss components ablation (Figure 6 right, Table 7). The paper ablates the two components of the excess loss by running DoReMi with only the proxy loss ℓ_θ(x) (preferring "hardest domains") or only negative reference loss −ℓ_ref(x) (preferring "easiest domains"). At 280M scale with matched proxy/main models: Hardest produces worst-case log-perplexity 2.66 (vs. baseline 2.39, a significant degradation), average 2.62 (vs. 2.32), and 0/22 domains beating baseline. Easiest produces worst-case 4.27, average 4.18, and 0/22 domains beating baseline. Neither ablation comes close to DoReMi's performance (worst-case 2.19, average 2.13, 22/22 domains beating baseline). Downstream accuracy (Figure 6 right) shows the same pattern: the DoReMi curve is substantially above both the Hardest and Easiest curves throughout training, and neither ablation outperforms the baseline. This demonstrates that both components of excess loss are necessary — raw loss alone (whether of the proxy or reference model) is worse than useless.

Proxy model scale variation for 8B main model (Figure 6 left, Table 3a, Table 4). Discussed in the main results above. Key finding: increasing proxy size from 70M to 280M monotonically improves 8B main model performance, but 1B proxy underperforms 280M. All proxy sizes (including 70M) improve all 22 domains over the 8B baseline (Table 4). All proxy sizes achieve the baseline accuracy at least 2× faster at 8B scale.

Matched proxy/main model scale sweep (Figure 5, Table 6). Discussed in the main results above. Key finding: DoReMi improves downstream accuracy and perplexity consistently from 280M to 1B when proxy and main models share the same size, though the improvement is not monotonic across scales and is less pronounced at 510M.

Iterated DoReMi rounds (Figure 3 right, Table 2). On the GLaM dataset, Round 2 domain weights are nearly identical to Round 3 (convergence by the ‖·‖_∞ < 1e-3 criterion), and Round 2 achieves comparable performance to downstream-tuned weights. The paper does not report the effect of iterated DoReMi on The Pile, nor the number of rounds to convergence on The Pile. This is a notable omission — the per-domain perplexity improvements in Figure 4 and Table 4 are from a single round of DoReMi, and it's unknown whether iteration would further improve The Pile results.

Reference domain weight sensitivity. The paper inadvertently provides evidence on sensitivity through the GLaM experiments, where the reference weights are uniform (vs. The Pile's heuristic default). DoReMi converges from uniform initialization to domain weights matching downstream-tuned proportions (Table 2) within 3 rounds, suggesting robustness to the choice of initial reference weights. However, no systematic sweep of reference weight configurations is performed on either dataset.

Domain weight trajectory analysis (Appendix B, Figure 8). The paper shows exponential moving average (smoothing parameter 0.99) of domain weights during DoReMi runs for both 280M and 1B proxy models. In both cases, domain weights change rapidly in the first ~50k steps and then stabilize. The paper notes this suggests: (1) smaller compute budgets may require drastically different domain weights (since early-training domain weights differ substantially from the final average), and (2) "we may be able to save compute by extrapolating the domain weights after 50k steps." No extrapolation method is developed or tested.

Domain weights at different proxy scales (Appendix Table 8). DoReMi (280M) puts most weight on Pile-CC (0.6057); DoReMi (1B) puts most weight on OpenWebText2 (0.3289). Other domains show generally similar relative patterns (e.g., arXiv, PubMed Central, and StackExchange are among the most heavily downweighted in both), but the specific web domain that receives the "lion's share" differs. This is presented as evidence of "multiple possible local minima in the domain weight space" without further investigation.

Proxy model quality analysis (Table 3b). Discussed above. The proxy model's quality relative to baseline degrades with scale (19/22 domains beating baseline at 280M, 0/22 at 1B), yet domain weights remain useful. This is an important robustness finding: DoReMi does not require a high-quality proxy model to produce useful domain weights.

Critical Assessment

Does DoReMi Actually Improve Language Model Training Efficiency?

The paper's central claim is that DoReMi "speeds up language model pretraining" — specifically, that it improves downstream accuracy by 6.5pp and achieves baseline accuracy 2.6× faster on The Pile (Section 3.2, Figure 2). The experiments support this claim with important caveats about what is being measured.

The 2.6× speedup is computed as: baseline 8B model reaches ~20% average accuracy at 200k steps; DoReMi (280M → 8B) reaches ~20% at 75k steps. This is a valid measurement of training efficiency — the DoReMi model reaches the same quality level faster. However, the 8% overhead for domain weight optimization (training the 280M reference and proxy models) is not included in this speedup calculation. The speedup is measured purely in terms of the main model's training steps. When accounting for overhead, the total compute for DoReMi (280M → 8B) to reach baseline accuracy is the proxy/reference compute plus main model compute: roughly (0.08 × 200k) + 75k = 91k step-equivalents, which is still a ~2.2× speedup over the baseline's 200k. The paper never does this combined accounting explicitly, but the overhead is small enough at 8% that the speedup remains substantial.

More importantly, the speedup measurement is relative to a specific baseline: the heuristic domain weights from The Pile. We don't know how much of the 2.6× speedup comes from DoReMi specifically versus simply from any rebalancing away from The Pile's default proportions. The paper does not compare against simpler reweighting heuristics (e.g., uniform weights, or weights proportional to some measure of domain quality, or weights from a single round of perplexity-based reweighting without the DRO machinery). The improvement is measured against one specific (and possibly poor) baseline, not against a range of alternatives.

Does DoReMi Improve Perplexity on All Domains Without Tradeoffs?

Figure 4 and Table 4 show that DoReMi (280M → 8B) reduces per-domain log-perplexity on all 22 domains of The Pile, including heavily downweighted domains like ArXiv and PubMed Central. This is a genuine and striking result that the experiments robustly support. The finding is replicated across all proxy model scales (70M, 150M, 280M, 1B all improve all 22 domains at 8B — Table 4) and across matched proxy/main model sizes (DoReMi X → X improves 15–22 of 22 domains depending on scale — Table 6). The result holds on both The Pile and GLaM datasets.

However, the paper's explanation for why this happens — the hypothesis that lowest- and highest-entropy domains can be downweighted without harm — is supported only by the toy theoretical model in Appendix D, not by any analysis of the actual entropy characteristics of The Pile's domains. The paper does not measure the entropy of each domain, does not verify that downweighted domains are indeed at the extremes of the entropy distribution, and does not demonstrate positive transfer from upweighted domains to downweighted ones. The "no-tradeoff" phenomenon is empirically demonstrated but mechanistically unexplained beyond the suggestive toy example.

A related concern: the perplexity improvements are measured on held-out data from the same domains that DoReMi optimizes over. It's possible that the improvements reflect better fit to the domain-specific validation distributions (which come from the same sources as the training data) rather than genuine improvements in language modeling capability that would generalize to new domains not in the training set. The downstream task results partially address this — the 5 generative tasks use different data sources than The Pile's domains — but a broader out-of-distribution evaluation would strengthen the claim.

Do DoReMi Domain Weights Transfer Across the 30× Model Scale Gap?

The paper's most practically significant claim is that domain weights optimized with a 280M proxy model transfer effectively to an 8B model. The evidence supports this transfer but the experimental design only tests one specific scale gap, and the non-monotonicity with proxy scale raises questions.

The supporting evidence: DoReMi (280M → 8B) improves all 22 domains (Table 4) and downstream accuracy by 6.5pp (Table 5). The improvements are larger than DoReMi (70M → 8B) and DoReMi (150M → 8B), suggesting meaningful benefit from the larger proxy. However, DoReMi (1B → 8B) performs worse than DoReMi (280M → 8B) despite using a larger proxy model closer to the main model scale. The paper attributes this to Group DRO optimizer degradation at larger scales (supported by Table 3b data). This means the transfer result is not simply "larger proxy → better transfer" — there is an optimal proxy scale, and the paper does not establish how to find it without running experiments at multiple scales.

A deeper unaddressed question: does transfer work in the opposite direction or across different scale gaps? Would 8B-optimized domain weights work for a 280M model? Would 280M weights transfer to a 70B or 500B model? The paper argues the domain weights capture "something fundamental about the data" (Section 4, Section 6) but provides evidence for only one transfer direction (small to large) at one specific ratio (30×). The finding that Round 1 GLaM weights outperform Round 2 weights at 280M but the reverse holds at 8B (Appendix A, Figure 7) suggests the transfer relationship may be more complex than the paper's narrative implies.

Does DoReMi Match Downstream-Tuned Domain Weights Without Downstream Data?

On the GLaM dataset, iterated DoReMi achieves downstream accuracy comparable to the domain weights from Du et al. (2021) that were explicitly tuned on downstream tasks. This result is genuinely impressive but has important limitations in scope and strength of evidence.

First, the result is on a single dataset (GLaM, 8 domains) and the effect size is small. Figure 3 (right) shows the Round 2 DoReMi curve essentially overlapping with the downstream-tuned curve, but the absolute accuracy difference between uniform baseline and downstream-tuned is only ~2–3 percentage points (from ~32% for baseline to ~34.5% for downstream-tuned/DoReMi R2). On The Pile (22 domains), the gap between baseline and DoReMi is 6.5pp — much larger — but no downstream-tuned weights exist for comparison (since tuning on 22 domains would be prohibitively expensive, which is the whole motivation for DoReMi).

Second, the "downstream-tuned" weights from Du et al. (2021) are the only oracle comparison. The paper doesn't establish whether these weights are actually near-optimal — they represent one particular tuning procedure applied to one particular set of downstream tasks. DoReMi matching them is evidence of DoReMi's quality, but doesn't prove DoReMi is as good as task-tuned weights in general.

Third, and most critically, the downstream tasks used for evaluation overlap with the tasks used to tune the GLaM weights. This means the oracle comparison is at its strongest possible advantage — the downstream-tuned weights were literally optimized for these exact tasks. Yet DoReMi matches them. This is strong evidence that DoReMi's loss-based optimization captures something genuinely correlated with downstream task performance, without overfitting to any particular task.

Weaknesses, Missing Baselines, and Missing Experiments

Single training run per configuration. The paper reports results from what appear to be single training runs (no error bars, no standard deviations, no multiple seeds). At 8B scale, training multiple runs would be expensive, but at the 280M scale where most ablations are performed, running 3–5 seeds would have been feasible and would provide information about the stability of the domain weights and downstream improvements. Without variance estimates, we cannot assess whether the 6.5pp improvement is reliably distinguishable from noise, or whether the non-monotonicity in proxy model scale is a real effect or within-run variance.

No comparison to simpler reweighting methods. The paper compares DoReMi against The Pile's heuristic weights and GLaM's uniform weights. Missing are comparisons against: uniform weights on The Pile (how much of the improvement is just from not using The Pile's specific heuristics?), perplexity-proportional weights (sample more from high-perplexity domains), inverse-perplexity weights (sample more from low-perplexity domains), or a single-iteration approach that adjusts weights based on the reference model's per-domain perplexities without DRO. These simpler baselines would help isolate what the DRO machinery specifically contributes beyond the general idea of "adjust domain weights using loss information."

No out-of-distribution or robustness evaluation. All evaluations are in-distribution: perplexity on held-out data from the same domains present in training, downstream tasks that are standard and widely reported. The paper does not evaluate on distribution shifts (e.g., domains not in The Pile), adversarial examples, or tasks measuring robustness, fairness, or calibration. This is particularly relevant given DoReMi's DRO heritage — distributionally robust optimization is designed for robustness to domain shifts, but DoReMi's robustness properties are never tested. The paper's "Broader impacts" section mentions fairness benefits of DRO but provides no empirical evidence.

The GLaM domain weight tuning procedure is not described in detail. The paper cites Du et al. (2021) for the downstream-tuned weights but does not reproduce the tuning methodology. The reader cannot assess whether the comparison is truly against an optimized set of weights or against one particular (potentially suboptimal) tuned configuration. If the GLaM weights were tuned with limited search budget, DoReMi matching them may be less impressive than if they were extensively optimized.

No analysis of what makes a domain "upweightable" vs. "downweightable." The paper presents the domain weights (Table 1) and the resulting perplexity improvements (Table 4), but never analyzes why specific domains are upweighted or downweighted beyond the general excess-loss mechanism. A deeper analysis correlating domain weight changes with measurable properties of domains (average token entropy, vocabulary diversity, syntactic complexity, factual density, temporal coverage, etc.) would provide insight into what DoReMi is actually selecting for.

No experiment on example-level or finer-grained reweighting. The paper mentions that "using fine-grained domains could improve the gains from DoReMi" (Section 6) and notes DoReMi is more effective on The Pile (22 domains) than GLaM (8 domains), but no experiments test finer-grained domain definitions. Clustering-based domain discovery or sub-domain splitting would test whether the benefits scale with domain granularity.

Limited model diversity. All experiments use the same Transformer architecture family with the same tokenizer. There is no evidence about whether DoReMi's domain weights transfer across architectures (e.g., encoder-decoder vs. decoder-only, different tokenizers, different context lengths). Given the paper's claim that domain weights capture "something fundamental about the data," testing across architectures would be a natural validation.

The 8% overhead figure is optimistic for practical deployment. The 8% figure assumes the proxy model is trained for the same number of steps as the main model. Appendix Figure 8 suggests domain weights stabilize after ~50k steps (25% of the full 200k run), which could reduce the overhead to ~2% if domain weights can be reliably extrapolated. But no extrapolation method is tested, and the paper acknowledges that "smaller compute budgets may require drastically different domain weights" (Appendix B), which would undermine early stopping.

6. Limitations and Trade-offs

The Proxy Model Scale Sweet Spot Is Empirically Discovered, Not Predicted

The entire DoReMi pipeline rests on the assumption that domain weights optimized with a small proxy model will transfer effectively to a much larger main model. The paper demonstrates this transfer works at one specific scale ratio (280M proxy → 8B main model, a ~30× gap), but it also discovers that the relationship between proxy model size and transfer quality is non-monotonic and not predictable a priori.

The evidence: DoReMi (70M → 8B) improves all 22 domains (Table 3a), DoReMi (150M → 8B) improves further, DoReMi (280M → 8B) achieves the best results (worst-case log-perplexity 1.46, downstream +6.5pp), but DoReMi (1B → 8B) performs worse than the 280M proxy variant (worst-case 1.58, downstream +2.32pp — Figure 6 left, Tables 3a and 5). The 1B proxy model is, in absolute terms, closer in scale to the 8B main model than the 280M proxy is, yet it produces inferior domain weights for training the 8B model.

The consequence: a practitioner cannot simply choose the largest proxy model they can afford and expect the best results. There is an optimal proxy scale somewhere below the main model scale, and this optimum depends on factors the paper does not fully characterize. The paper hypothesizes that the Group DRO optimizer degrades at larger proxy scales because loss reweighting (how the proxy model trains) diverges from data resampling (how the main model trains), and Table 3b provides supporting evidence: the 1B proxy model underperforms the 1B baseline on all 22 domains (average log-perplexity 2.02 vs. 1.87), while the 280M proxy beats its baseline on 19/22 domains (2.27 vs. 2.32). But this hypothesis is post-hoc — the paper provides no method for predicting the optimal proxy scale from, say, the ratio of proxy to main model parameters or the number of domains.

Mitigation status: The paper does not address this beyond the empirical observation. The authors suggest that "a resampling-based Group DRO optimizer may improve DoReMi for larger proxy models" (Section 4), but this is future work. No scaling law or heuristic is provided for selecting the proxy model size.

The 8% Overhead Figure Excludes Difficulty Estimation Costs That Would Be Required in Practice

The paper reports that DoReMi's domain weight optimization consumes approximately 8% of the main model's training FLOPs — training two 280M models (reference and proxy) for the same number of steps as the 8B main model. This figure is used to frame DoReMi as a lightweight preprocessing step. However, this accounting only covers the cost of Steps 1–2 of the algorithm. It does not include the cost of determining what the reference domain weights should be, nor any cost for evaluating whether the optimized weights are actually better — both of which would be necessary in a realistic deployment where the baseline domain weights are unknown or untrusted.

More concretely: in the paper's experiments on The Pile, the reference domain weights are the default Pile weights from Gao et al. (2020). In a new pretraining scenario — say, a company assembling a corpus from dozens of internal and external sources — there are no pre-existing default weights. The practitioner must either (a) guess at reasonable reference weights, train the reference model, run DoReMi, and hope the result is better but have no way to verify without training the main model, or (b) run multiple DoReMi experiments with different reference weight configurations to find a good one, multiplying the overhead. The paper's GLaM experiment uses uniform reference weights, which is a natural starting point but may be far from optimal — Round 1 DoReMi on GLaM achieves only baseline-level downstream accuracy (Figure 3 right), with Round 2 needed to match downstream-tuned weights. This means the practitioner may need to budget for multiple rounds of iteration (each round trains a new reference and proxy model, adding ~8% overhead per round).

The consequence: the 8% figure is a lower bound that assumes a known-good reference configuration. The true overhead in an unfamiliar data regime could be 2–3× higher if iteration or reference-weight search is needed. More fundamentally, there is no mechanism in DoReMi to verify that the optimized weights are actually better than the reference weights without training the main model — the proxy model's own perplexity is demonstrably unreliable as a quality signal (Table 3b: the 1B proxy model is terrible yet produces useful weights).

Mitigation status: The paper does not account for these costs in its headline efficiency numbers. The iterated DoReMi experiments on GLaM (Section 3.2) implicitly demonstrate that multiple rounds may be necessary, but the overhead of iteration is not quantified relative to the gains. The paper acknowledges in Section 6 that "most of the variation in the domain weights during a DoReMi run seems to occur in the beginning of training" (Appendix Figure 8) and suggests extrapolation from early stopping as future work, but this addresses the cost of one DoReMi run, not the cost of discovering a good reference configuration.

Domain Weights Are Transferred Across Scales Without Any Understanding of When Transfer Breaks

The core practical value proposition of DoReMi — that domain weights optimized cheaply on a small model can be used to train a much larger model — requires that the optimized weights are scale-invariant across the specific gap being bridged. The paper demonstrates transfer across a 30× gap (280M → 8B), but provides no evidence at larger gaps and some evidence that the transfer relationship is fragile.

The most concerning evidence is in Appendix A (Figure 7): on the GLaM dataset, Round 2 domain weights achieve slightly worse downstream accuracy than Round 1 weights when used to train a 280M main model (DoReMi 280M → 280M), but Round 2 weights outperform Round 1 when transferred to an 8B model (DoReMi 280M → 8B). This means the ranking of domain weight configurations reverses depending on the target model scale. If this inversion occurs at other scale gaps, then domain weights optimized with a 280M proxy could in principle be optimal for 8B but suboptimal for other scales — and the practitioner would have no way to know without running the large-scale experiment.

The consequence: the paper's claim that domain weights capture "something fundamental about the data" (Section 4) is too strong given the evidence. The transfer property is empirically demonstrated at one specific ratio but may not be universal. A team that runs DoReMi with a 280M proxy intending to train a 70B model (a 250× gap, vs. the 30× tested) is operating outside the paper's evidence envelope. The optimal domain weights might shift — domains that are noise-like at 280M may contain learnable structure at 70B, or vice versa. The paper provides no theoretical framework for predicting when transfer will succeed, no scaling law relating proxy scale to transfer quality, and no upper bound on the safe transfer ratio.

Mitigation status: The paper acknowledges this as an open question in Section 6: "Understanding why the domain weights can be transferred across scales and the limits of how far these domain weights transfer are important questions to answer in future work." But the paper's own abstract and introduction present the transfer property as a fait accompli ("DoReMi... first trains a small proxy model... to produce domain weights... We then resample a dataset with these domain weights and train a larger, full-sized model") without caveats about the tested scale range.

The Method Has Only Been Validated on Two English-Language Text Corpora with a Single Model Architecture

All experiments in the paper use decoder-only Transformer models from a single architecture family, trained on English-language text from two specific corpora (The Pile and GLaM). The downstream evaluation consists of five question-answering and language understanding tasks, all in English, all from the GPT-3 evaluation suite (Brown et al., 2020). This is a narrow evaluation footprint given the breadth of LLM applications and the paper's claim to provide a general data mixture optimization method.

Several plausible generalization gaps are untested and unaddressed:

Multilingual corpora: DoReMi's excess loss mechanism fundamentally depends on the reference model providing a meaningful per-token learnability baseline. In a multilingual setting where some languages are low-resource and the reference model's loss is uniformly high due to data scarcity rather than inherent noise, DoReMi might systematically downweight low-resource languages — the exact opposite of what fairness and inclusivity considerations would demand. The DRO framework's theoretical guarantee is worst-case performance over domains, but this guarantee is only as good as the reference model's ability to distinguish "learnable but not yet learned" from "inherently difficult due to insufficient data."
Different architectures: The paper uses standard decoder-only Transformers. It is unknown whether domain weights optimized with this architecture would transfer to encoder-decoder models (T5-style), mixture-of-experts architectures, or non-Transformer architectures. The interaction between model architecture and optimal data mixture is unexplored.
Non-text modalities and multimodal data: DoReMi is formulated for text domains; extension to images, audio, or multimodal corpora would require defining domain-level excess losses in potentially incompatible loss spaces (cross-entropy for text, reconstruction loss for images, contrastive loss for multimodal).

The consequence: the paper's framing as a general method for "language model pretraining" overstates the breadth of validation. The method may work differently — or not at all — in settings where domain characteristics differ from those in The Pile and GLaM, or where the model architecture differs from the tested Transformer variants.

Mitigation status: Not addressed. The paper does not discuss the potential for corpus-specific or architecture-specific behavior. The closest the paper comes to acknowledging scope limitations is the observation that "DoReMi is more effective on The Pile (22 domains) than the GLaM dataset (8 domains)" (Section 6), which is a within-scope observation about domain granularity rather than an out-of-scope acknowledgment about corpus diversity.

The "No-Tradeoff" Perplexity Result Lacks Mechanistic Explanation and Out-of-Distribution Validation

The paper's most visually compelling result — DoReMi improves perplexity on all 22 domains of The Pile, even heavily downweighted ones (Figure 4, Table 4) — is presented as evidence that domain reweighting can be a Pareto improvement rather than a zero-sum redistribution. The paper offers a hypothesis (lowest-entropy and highest-entropy domains can be downweighted without harm, while intermediate-entropy domains provide positive transfer) and a toy model (Appendix D, 3 domains, 3 tokens, Dirichlet-Multinomial) to illustrate the mechanism. However, the hypothesis is never tested with measurements on the actual The Pile domains, and the "no-tradeoff" result may partially reflect the fact that the validation data comes from the same sources as the training data rather than genuine improvements in language modeling capability.

The key unexamined question: when a heavily downweighted domain like ArXiv (weight drops from 10.5% to 0.4%) shows improved perplexity (from 1.64 to 1.38), is this because the model genuinely understands academic language better (due to positive transfer from the upweighted web text), or because the model has overfit less to the idiosyncrasies of ArXiv's validation split? In the latter case, the improved perplexity would not reflect better language modeling of academic text in general — it would be an artifact of the evaluation being correlated with the training data in ways that regularize out when the domain is downweighted.

The downstream task results provide partial evidence of genuine improvement (all 5 tasks improve, Table 5), but the downstream tasks are not matched to the downweighted domains — TriviaQA draws from Wikipedia and web sources, not from ArXiv or PubMed Central. We cannot verify that performance on academic-domain tasks (e.g., scientific question answering) improved, or that the perplexity gains on downweighted domains generalize beyond their specific validation splits.

The consequence: the "no-tradeoff" result, while genuine within the paper's evaluation framework, may overstate DoReMi's ability to improve all domains simultaneously. Without out-of-distribution evaluation (e.g., new academic text not from the ArXiv or PubMed sources in The Pile), we cannot distinguish between genuine cross-domain transfer and artifact reduction. This matters for practitioners: if the improved perplexity on downweighted domains is partly an artifact of reduced overfitting to domain-specific quirks, then deploying the model on truly new text from those domains might not show the same improvements.

Mitigation status: The paper does not address this concern. The toy model in Appendix D is suggestive but not validated against empirical measurements. No analysis of domain entropy, cross-domain similarity, or transfer patterns is provided. The paper's abstract states that "DoReMi improves perplexity across all domains" without qualifying the in-distribution nature of the evaluation.

The Method Provides No Guidance on Choosing the Number or Definition of Domains

DoReMi treats the set of domains as given — it takes D_1, ..., D_k as input and optimizes weights over them. But the choice of how to partition the data into domains is itself a consequential design decision that DoReMi provides no guidance on, and that interacts with the method's effectiveness in ways the paper observes but does not resolve.

The evidence: the paper notes that "DoReMi is more effective on The Pile (22 domains) than the GLaM dataset (8 domains)" (Section 6). On The Pile, DoReMi improves downstream accuracy by 6.5pp over baseline and reduces perplexity on all 22 domains. On GLaM, the downstream improvement over uniform baseline is ~2–3pp (Figure 3 right), and Round 1 of DoReMi achieves only baseline-level performance. The authors interpret this as evidence that "using fine-grained domains could improve the gains from DoReMi," which implies that the 22-domain partition may still be suboptimally coarse — further splitting some domains could yield additional gains.

But this creates a practical dilemma: how many domains is enough? If 8 domains produce modest gains and 22 domains produce larger gains, would 100 domains (clustered automatically or split along sub-domain lines) produce even larger gains? The DRO objective in Equation 1 treats each domain as an independent group whose worst-case performance is protected. With very fine-grained domains, the maximization over α in the DRO objective could place all weight on a tiny set of worst-case examples — a form of over-pessimism that the paper acknowledges: "When domains are very fine-grained, it will be important to control the pessimism of DRO (e.g., DRO can put all the weight on a small set of worst-case examples)" (Section 6). Yet no experiments test this boundary or provide heuristics for domain granularity selection.

The consequence: a practitioner implementing DoReMi on a new corpus must decide how to partition the data, but the paper provides no methodology for making this choice beyond the implicit suggestion that more domains are better (up to some unknown limit where DRO pessimism becomes harmful). This is a significant practical gap, because the domain definition can be arbitrary — by source (as in The Pile), by topic, by writing style, by language, by time period — and different partitionings may lead to different optimized weights with different downstream effects.

Mitigation status: The paper flags this as an open direction in Section 6, suggesting "automatically finding fine-grained domains (e.g., via clustering as in DRO-LM (Oren et al., 2019)) and reweighting the data at an example level," and acknowledges the DRO pessimism risk. But no experiments or analyses guide practitioners on how to choose k or how to define domains for a new corpus.

7. Implications and Future Directions

How This Work Changes the Landscape

DoReMi shifts the conversation around language model pretraining from a focus on model scaling toward a recognition that data mixture optimization is an equally consequential—and separately tractable—design dimension. Before this work, the dominant framework for allocating pretraining compute was the Chinchilla scaling laws (Hoffmann et al., 2022): given a FLOPs budget, how should one trade off model size against training tokens? DoReMi introduces an orthogonal question that sits inside the token budget: given a fixed number of training tokens drawn from a multi-domain corpus, what fraction should come from each domain? This is not a refinement of scaling laws—it is an entirely new axis of optimization that was previously handled by heuristics (The Pile) or expensive downstream tuning (GLaM, PaLM). The paper demonstrates that this axis matters enormously: simply changing the sampling probabilities over domains yields a 2.6× training speedup and a 6.5 percentage point downstream accuracy improvement at 8B scale on The Pile, with only 8% compute overhead for the optimization.

This is best characterized as a methodological reframing with significant practical impact, not a paradigm shift. The individual components—DRO, excess loss, online domain weight updates—all existed prior to this work (Oren et al., 2019; Sagawa et al., 2020; Nemirovski et al., 2009). What is new is the repurposing of these tools from model training to data discovery, and the empirical demonstration that domain weights extracted from a small proxy model's DRO training trajectory transfer across a 30× model scale gap to improve a much larger model. This conceptual pivot—"we are not training a robust model; we are using DRO as a measurement instrument for data quality"—is the paper's lasting intellectual contribution.

The work also resolves a latent tension in the data curation literature. Prior work split along two lines: heuristic filtering (Raffel et al., 2019; Gao et al., 2020) that relies on human judgment about what constitutes "good" data, and downstream-informed selection (Du et al., 2021; Chowdhery et al., 2022) that optimizes data composition for specific evaluation tasks. These approaches sit at opposite ends of a spectrum—one task-agnostic but untethered from measurable quality, the other quality-aware but potentially overfit to specific benchmarks. DoReMi demonstrates that a third point on this spectrum exists: loss-based, task-agnostic optimization that correlates with downstream performance. The excess loss ℓ_θ(x) − ℓ_ref(x) provides a measurable signal about which data is "learnable but not yet learned" without requiring any downstream task labels, human judgments about data quality, or assumptions about what good text looks like. This is a genuine conceptual advance—it operationalizes a definition of data value that is both principled (rooted in the minimax theory of DRO) and practical (computable with small models).

Several research directions become more attractive in light of this work:

Data mixture optimization as a standard preprocessing step. The paper's 8% overhead figure makes domain reweighting cheap enough to be incorporated into any large-scale training pipeline. This shifts the default from "use heuristic weights and hope" toward "run DoReMi first, then train." The paper provides a concrete recipe (two small model trainings, no hyperparameter tuning of the DRO-specific parameters) that lowers the barrier to adoption.
Scaling laws for data composition. Just as Hoffmann et al. (2022) established that optimal model size and token count follow power-law relationships with compute budget, DoReMi opens the question of how optimal domain weights scale with model size, training duration, and total data volume. The paper's finding that domain weight trajectories stabilize after ~50k steps (Appendix Figure 8) hints at regularities that could be captured by scaling laws.
Verifier and data quality signals beyond loss. DoReMi uses excess loss as its sole signal for data value. The paper's success with this simple metric invites exploration of richer signals: gradient alignment across domains, mutual information between domains and downstream capabilities, or learned difficulty estimators that don't require training a reference model.
Understanding positive transfer. The "no-tradeoff" result—perplexity improves on all domains, even heavily downweighted ones—is empirically demonstrated but mechanistically unexplained beyond the toy model in Appendix D. This invites research into when and why training on one domain improves performance on another, which would enable more targeted data mixture optimization.

Conversely, some research directions become less attractive. The paper's evidence that DoReMi matches downstream-tuned domain weights on GLaM (Figure 3 right, Table 2) suggests that expensive, task-specific data mixture tuning may be unnecessary for general-purpose pretraining—the loss signal alone captures much of what downstream tuning provides. Similarly, the failure of the "hardest domains" and "easiest domains" ablations (Figure 6 right, Table 7) to improve over baseline suggests that naive loss-based reweighting without the excess loss formulation is worse than doing nothing, which should discourage simple heuristics like "train more on high-loss domains."

Follow-Up Research This Work Enables

Stress-testing the scale transfer limit: At what proxy-to-main ratio does DoReMi break? The paper demonstrates transfer across a 30× gap (280M → 8B) but the non-monotonic relationship—1B proxy underperforms 280M proxy at 8B (Figure 6 left, Table 3a)—indicates there is an optimal proxy scale rather than a monotonic "bigger proxy is better" relationship. A systematic study should run DoReMi with proxy models at 150M, 280M, 500M, 1B, and 3B to train a fixed 8B main model, measuring both domain weight stability (do the weights converge to the same distribution across proxy scales?) and downstream accuracy. The critical test: does the optimal proxy scale depend on the main model size, the number of domains, or the total training budget? A negative result—finding that optimal proxy scale is unpredictable and that poor proxy choices silently degrade main model quality—would significantly limit DoReMi's practical applicability and motivate the development of a resampling-based Group DRO optimizer that the paper hypothesizes would improve larger proxy models (Section 4).

Does DoReMi's "no-tradeoff" perplexity improvement generalize to out-of-distribution evaluation within the same domain? The paper's most striking result—perplexity improves on all 22 domains even when domain weights drop by 10× or more (Figure 4, Table 4)—is measured on held-out validation data drawn from the same sources as the training data. This leaves open the possibility that the improvement partly reflects reduced overfitting to domain-specific validation quirks rather than genuine gains in domain capability. A strong follow-up would evaluate the 8B baseline and DoReMi (280M → 8B) models on out-of-distribution text from the same domain categories but different sources: new ArXiv papers published after The Pile's cutoff date, GitHub repositories from different organizations, PubMed abstracts from journals not in The Pile's training data, and so on. If perplexity improvements persist on genuinely new data from downweighted domains, it would validate the positive transfer hypothesis and strengthen the case for domain reweighting as a Pareto improvement. If improvements vanish or reverse, it would reveal that DoReMi's benefits are partly an artifact of in-distribution evaluation and motivate work on regularized reweighting that preserves out-of-distribution robustness.

Extending DoReMi to example-level reweighting with pessimism control. The paper notes that "using fine-grained domains could improve the gains from DoReMi" and that "when domains are very fine-grained, it will be important to control the pessimism of DRO" (Section 6). The natural extension is to treat each example (or each cluster of similar examples) as its own domain, running DoReMi at the example level. This requires solving the pessimism problem: with k = millions of examples, the DRO maximization will place all weight on a tiny fraction of worst-case examples that may be outliers. A concrete experiment: cluster The Pile's 22 domains into N sub-domains using embedding-based clustering (varying N from 22 to 100 to 1000), run DoReMi on each granularity level with a pessimism penalty (e.g., KL divergence from uniform) on the domain weights, and measure whether finer granularity with controlled pessimism outperforms the 22-domain DoReMi baseline. The KL penalty strength would need to be tuned; the paper's smoothing parameter c = 1e-3 is a crude form of this, and more principled approaches from the DRO literature (e.g., χ²-divergence balls, CVaR-style constraints) could yield better tradeoffs.

Verifying domain weight transfer across architectures and tokenizers. All experiments use a single architecture family (decoder-only Transformers) with a SentencePiece tokenizer (256k vocabulary). The paper's claim that domain weights capture "something fundamental about the data" (Section 4, Section 6) predicts that weights should transfer across architectures. This is testable: run DoReMi (280M) on The Pile with the standard Transformer architecture, then train an 8B model with the resulting domain weights but using a different architecture—an encoder-decoder model (T5-style), a model with a different tokenizer (e.g., byte-level BPE), or a model with a different context length. If the domain weights transfer successfully (perplexity improves on all 22 domains relative to the new architecture's baseline with default Pile weights), it would validate that DoReMi captures data properties rather than model-specific artifacts. If transfer fails, it would reveal an important scope limitation and suggest that domain weights need to be optimized for each architecture family separately.

Developing cheap difficulty estimators to replace the reference model. Training a reference model costs ~4% of the main model's FLOPs. The paper's trajectory analysis (Appendix Figure 8) shows domain weights stabilize after ~50k steps, and the authors suggest extrapolation as a way to reduce costs. A more ambitious direction: can we predict DoReMi's domain weights without training a reference model at all? One approach: use a small, pretrained language model (e.g., a 125M-parameter model from an existing model family) as a frozen reference, running DoReMi's proxy training against this off-the-shelf reference. If the resulting domain weights approximate those from a trained reference model, the overhead drops to ~4% (proxy training only). Another approach: train a lightweight classifier to predict domain-level excess loss from simple corpus statistics—vocabulary size, token entropy, sequence length distribution, n-gram diversity—without any model training. A concrete experiment: extract these features for each of The Pile's 22 domains, train a small regressor to predict the excess loss that DoReMi would assign, and check whether the regressor's recommended domain weights yield similar downstream improvements to full DoReMi. Success here would reduce the overhead from 8% to near zero.

Evaluating DoReMi in a multilingual, code-switching, or low-resource setting. The paper's evaluation is entirely on English text. The excess loss mechanism is potentially problematic for multilingual corpora: low-resource languages will have high reference loss because the reference model cannot learn them well from limited data, causing DoReMi to interpret them as "unlearnable noise" and downweight them—exactly the opposite of what fairness principles would dictate. A strong stress-test would apply DoReMi to a multilingual corpus (e.g., mC4, Wikipedia in 50+ languages) and measure whether low-resource languages are systematically downweighted, and whether this degrades downstream performance in those languages. If DoReMi does downweight low-resource languages, it would motivate the development of fairness-aware variants that constrain minimum domain weights or use language-specific reference models to avoid penalizing data scarcity as unlearnability. The paper's broader impacts discussion (Section 6) acknowledges that DRO "can have a favorable impact on fairness" but provides no empirical evidence; this experiment would test that claim directly.

Practical Applications and Downstream Use Cases

Cost-efficient pretraining for mid-sized organizations training 1B–10B parameter models. The paper's headline result—2.6× training speedup on The Pile—translates directly to reduced compute costs. For an organization training an 8B-parameter model from scratch on a multi-domain corpus, running DoReMi with two 280M models (8% overhead, ~$2,000–5,000 in cloud compute at current prices) before the main training run can reduce the main training cost by more than half (reaching baseline quality at ~75k steps instead of 200k) or improve final quality by 6.5 percentage points for the same cost. The paper's robustness to proxy model scale (DoReMi 70M → 8B still achieves nearly 2× speedup; Table 3a, Figure 6 left) means the overhead can be reduced further for organizations with tighter budgets. The key practical requirement is that the corpus must be partitionable into coarse domains with meaningful differences—the paper shows gains on 22-domain and 8-domain corpora, suggesting the approach works for any heterogeneous pretraining dataset.

Data mixture optimization for domain-specific model pretraining. The paper evaluates on general-purpose downstream tasks, but DoReMi's mechanism—identifying domains with high excess loss and upweighting them—should be particularly valuable for domain-specific pretraining where the goal is strong performance on a target domain (e.g., biomedical text, legal documents, code). In this setting, the practitioner can include the target domain as one of the k domains in DoReMi, and the DRO objective will naturally upweight it if it has high headroom relative to the reference model. A concrete deployment: a legal-tech company pretraining an LM on a corpus mixing general web text, case law, statutes, and legal textbooks. Running DoReMi would automatically discover the optimal mixing weights without manual tuning and without overfitting to specific downstream legal tasks (which might not exist as labeled benchmarks). The GLaM result—DoReMi matching downstream-tuned weights (Figure 3 right)—suggests this approach would be competitive with task-specific tuning while being cheaper and generalizing better to unseen legal tasks.

Data curation for continual pretraining and model updates. Large LMs are increasingly updated with new data over time (continual pretraining) rather than trained once from scratch. DoReMi provides a principled method for determining which new data domains to emphasize during these updates: run DoReMi with the existing model serving as the reference model (no separate reference training needed—the deployed model already provides ℓ_ref), train a small proxy model on the new data mixture using DRO, and extract domain weights for the update. The excess loss would naturally upweight new domains where the existing model has high loss relative to what is achievable—exactly the data that provides the most new information. This reuses the core insight that excess loss measures "learnable but not yet learned" but in a continual learning context where the reference model is the previous checkpoint. The paper's trajectory analysis (Appendix Figure 8) suggests domain weights stabilize quickly, making this potentially cheap enough to run before each model update cycle.

Informing data acquisition strategies. The paper's domain weight analysis on The Pile (Table 1) provides a concrete signal about which types of data are most valuable for improving language model quality: diverse web text (Pile-CC) receives 60.6% of the weight under DoReMi (280M), while academic reference domains (ArXiv, PubMed Central) receive less than 0.5% despite being 10.5% and 10.7% of the baseline weights, respectively. For an organization deciding how to allocate a data acquisition budget—should we license more books, crawl more academic papers, or collect more web text?—running DoReMi on their existing corpus would provide an empirical answer grounded in the model's own learning dynamics rather than intuition. The caveat is that DoReMi's recommendations are relative to the existing corpus composition; the optimal weights might shift if entirely new types of data were added. But for incremental data acquisition decisions within an existing domain taxonomy, DoReMi provides a quantitative, loss-based alternative to guessing.