Everything in Moderation: Per-Domain Coverage Optima and Alignment-Resistant Domain Gaps in Multi-Domain Mid-Training

AI summary

Overview

Research area: Large language model training pipelines — specifically data-mixture design during the mid-training stage and whether downstream alignment (SFT/RL) can compensate for imbalanced per-domain coverage.

Technical level: Intermediate. The paper is readable with a working understanding of the standard pretraining → mid-training → SFT → RL pipeline and basic inferential statistics (permutation tests, bootstrap CIs, Welch t-tests). Methods are conceptually simple; the statistical caveats require care.

Scope: A single-setting controlled study on Qwen3-8B-Base (with a 4B replication) using five semantically rule-disjoint KOR-Bench reasoning domains, varying per-domain mid-training token coverage across a five-domain simplex while holding all subsequent training stages fixed.

What This Paper Is About

In multi-stage LLM training, the per-domain composition of mid-training data is usually determined by whatever data happens to be available, not by design. This paper asks two questions: is there a "right" amount of each domain to include, and once an imbalance is baked in at mid-training, can a later supervised fine-tuning and RL pass undo it? The authors vary domain coverage systematically across a constrained allocation space and measure whether the resulting accuracy gaps persist after alignment.

Key Contributions

1. A systematic coverage sweep over a simplex. The authors train mid-training-only checkpoints under 24 controlled five-domain coverage configurations (plus six allocations withheld from fitting), spanning severely imbalanced to balanced to counter-skewed mixtures, each at five random seeds — a granular map of how per-domain token share relates to per-domain accuracy.

2. A quantified negative result on alignment gap-closure. Using a fixed-budget compensatory SFT pass applied to all 24 checkpoints (reweighting toward under-covered domains), the authors show the mid-training-induced pairwise gaps essentially survive, and that an equal-budget uniform-SFT control behaves almost identically — ruling out the compensation formula itself as the cause.

3. Evidence for interior per-domain coverage optima. No domain is best served by its largest or smallest assigned share; the moderate band (10–40%) is best for all five domains, with a calibrated permutation test for quadratic interiority at P ≈ 0.010 and fitted peaks between 9.9% and 35.1%.

4. Explicit documentation of the design's confounds. The paper carefully separates "mixture-level marginal" statements from causal claims, reports a simplex-constrained saddle-point classification, and includes a FineWeb-Edu-only control showing that the zero-coverage collapse is entangled with generic distributional drift from any continued pretraining.

Main Findings

• Interior optima, not monotonic returns. Every domain's own mid-training share is non-monotonically related to its accuracy: accuracy rises with a domain's share, peaks at a moderate value, then falls. The moderate band (10–40%) beats both low and high bands for all five domains without any curve fitting. Fitted split-Gaussian peaks at 8B land at approximately Cipher 9.9%, Logic 15.0%, Operation 15.9%, Counterfactual 26.0%, and Puzzle 35.1%. Six held-out allocations reproduce the curve shapes out-of-sample (closely for four domains, only approximately for Operation), but the specific peak locations are not validated on held-out data.

• Alignment resists gap closure. Compensatory SFT raised absolute accuracy in 116 of 120 configuration–domain cells (mean +4.32 pp) yet bridged 0/240 domain pairs at a 5 pp difference threshold and only 30/240 at a 10% ratio threshold. A uniform-SFT control matched this almost exactly (gains in all 120 cells, mean +4.20 pp, 0/240 and 32/240). A permutation null that redistributes the same observed gain magnitudes at random would bridge 13.8 ± 3.3 pairs at 5 pp and 77.9 ± 8.5 at 10% (P < 0.001), implying the gains are placed in a systematically gap-preserving way.

• Closure is possible but costly. Sharpening the SFT policy concentration closes more pairs (0/60 → 12/60 at the 5 pp metric) but roughly halves the mean gain (+4.34 → +2.26 pp), indicating a tension between gap closure and average accuracy under a fixed budget.

• Zero coverage causes a transient collapse. Driving a domain to zero mid-training coverage collapses its mid-training-only accuracy even where the base prior was high — Counterfactual dropped 83.6% → 45.6%, Operation 60.4% → 33.2%. The tested recipe partly reverses this: Operation returns to +5.2 ± 0.3 pp over its SFT-only baseline by mid+SFT; Counterfactual only turns positive (+1.2 ± 3.4 pp) by the RL row and is not statistically resolved. A small residual deficit persists on Logic (−2.4 pp) and Puzzle (−2.0 pp).

• The collapse is not purely coverage starvation. A FineWeb-Edu-only control — no KOR-Bench data at all — dropped Counterfactual further (83.6% → 29.2%, −54.4 pp) than the zero-KOR-coverage condition (−38.0 pp). The authors therefore frame the interference effect as co-mingled with generic distributional drift from any continued pretraining.

• The best allocation's advantage is descriptive, not significant. An exploratory θ* allocation (Cipher 9.7%, Operation 11.8%, Logic 15.8%, Counterfactual 30.4%, Puzzle 32.3%) achieved the largest full-pipeline gain of the three trained end-to-end (+4.36 pp vs. SFT+RL alone, compared to +0.80 pp for balanced and +0.64 pp for imbalanced). Welch t-tests give p ≈ 0.052 and p ≈ 0.030, neither surviving Bonferroni correction. The gain is domain-concentrated: Logic +16.7 pp and Counterfactual +6.4 pp, while Cipher loses −5.0 pp.

• The five optima are marginals, not a jointly optimal mixture. Because coverage proportions sum to 100%, each domain's share cannot vary independently. Fitting a simplex-aware quadratic response surface in isometric log-ratio coordinates reads as a saddle in every domain — meaning the interior data neither establish a jointly optimal mixture nor exclude a shallow one.

• External benchmarks give only limited consistency. On nine configurations, Spearman rank correlations between KOR-Bench average and external benchmarks were positive for CounterBench (ρ = +0.67) and ZebraLogic Grid (ρ = +0.53), near zero for ProofWriter (ρ = +0.34), and negative for ZebraLogic MC (ρ = −0.12). The external corpus draws 34.8% of tokens from the ProofWriter rule family, so none of these is a clean zero-exposure transfer test.

Methodology in Plain English

The authors start from a public base model with no instruction tuning and run a fixed sequence of stages: continued language modeling on a synthetic reasoning corpus (mid-training), supervised fine-tuning on problem–response pairs (SFT), and a short sequence-level RL pass with binary correctness rewards (GSPO). What they vary is only the proportion of mid-training tokens given to each of five reasoning domains — ciphers, custom mathematical operations, formal logic, constraint puzzles, and counterfactual reasoning — all drawn from the KOR-Bench rule definitions but with training instances freshly synthesized by a deterministic solver, never reused from the evaluation set.

The five proportions must sum to 100%, so the design lives on a four-dimensional simplex. They pick 24 points spanning that space — some heavily skewed toward one domain, some balanced, some counter-skewed — train each at five seeds, and measure KOR-Bench accuracy, then fit per-domain coverage–accuracy curves. To test gap closure, they apply to every checkpoint a compensatory SFT pass that reweights SFT data toward under-covered domains under a fixed token budget, and compare it to an equal-budget uniform SFT pass. They report pairwise closure counts over all 240 domain pairs and build a permutation null to ask whether the observed gains close as many gaps as random reallocation would. Six additional allocations were kept out of the curve-fitting and used purely for out-of-sample validation. A FineWeb-Edu-only baseline served as a control for generic drift.

Why This Matters

Research impact: The paper reframes mid-training data composition from an engineering afterthought into a design variable with measurable, persistent consequences. It provides a rare quantified negative result about alignment's limits — a fixed-budget SFT pass cannot readily undo coverage choices — while being unusually careful about what it does not show. It also extends prior work on pretraining exposure determining RL generalization (Zhang et al., 2025) from a synthetic setting to a naturally imbalanced multi-domain corpus, and offers a simplex-aware response-surface framing that other mixture-optimization researchers can adopt.

Real-world applications:

• Data-curation planning for domain-adapted models. Teams fine-tuning a base model for legal, medical, or financial reasoning can treat coverage allocation as an explicit tuning knob rather than a byproduct of available corpora, and can audit whether their mixture sits in the moderate band.

• Mixture-optimization tooling. Existing data-mixture methods (DoReMi, RegMix) optimize for aggregate perplexity; this work suggests a per-domain accuracy objective with interior optima, motivating tooling that optimizes predicted accuracy rather than relative gain.

• Alignment budget allocation. Organizations spending compute on SFT and RL data reweighting now have evidence that reweighting alone may preserve existing gaps, informing where to invest alignment effort.

• Diagnostic benchmarking. The zero-coverage collapse and the FineWeb-Edu control give practitioners a diagnostic recipe for detecting when a domain's regression is due to coverage starvation versus generic drift from continued pretraining.

Industry relevance: Any organization running a multi-stage training pipeline on domain-specific data — where the mid-training corpus is assembled from multiple heterogeneous sources — faces exactly the allocation decision this paper studies. The finding that moderate coverage tends to outperform both extremes, and that later alignment cannot cleanly repair imbalance, has direct implications for compute budgeting and for how teams should think about the sequencing of data decisions across training stages.

Future Directions

• Disentangling the simplex confound. The authors explicitly call for a simplex-aware joint response-surface model with enough allocations to determine whether the stationary point is truly a saddle or a shallow flat region. Until then, the five per-domain optima remain marginal statements rather than coordinates of one achievable mixture.

• Validating the fitted peaks on held-out allocations. The curve shapes generalize, but the peak locations do not. A design that withholds allocations specifically near predicted optima — or perturbs around them — would test whether the fitted 9.9–35.1% peaks are real or artifacts of the 24-point sweep.

• Testing whether other alignment policy families close gaps. The paper tested one compensation formula, one budget, and a short RL leg. Larger budgets, different RL objectives, preference-based methods, or multi-turn RLHF could plausibly behave differently, and the authors frame the claim as specific to the tested family.

• Separating gradient dominance from distributional drift. The FineWeb-Edu control shows the zero-coverage collapse is entangled with generic drift. Mechanistic work — measuring representations, gradient conflict, or parameter overlap during mid-training — could determine which of the two competing hypotheses (gradient dominance versus drift) actually drives the observation.

Target Audience

This paper is most valuable to LLM training practitioners and data-curation engineers who assemble multi-domain corpora for continued pretraining or mid-training, and to researchers studying data-mixture optimization and the limits of post-training alignment. It is also useful for teams designing domain-adapted models who need to justify allocation decisions with evidence rather than intuition. Readers without a background in multi-stage LLM training or basic statistics will find the statistical framing dense, but the core findings are stated in accessible terms.

Authors’ abstract

Mid-training, the stage between pre-training and alignment, is where a model's per-domain data composition is typically set by data availability rather than principled design. We ask what that decision buys, and whether a later alignment pass can undo it. In a controlled logical-reasoning setting (Qwen3-8B-Base, with a 4B replication; five semantically rule-disjoint KOR-Bench domains) we train 30 allocations spanning the five-domain simplex, 24 sweep configurations plus six withheld from the fit, at five seeds each. Three findings emerge. First, every domain has an interior coverage optimum: the moderate band ($10\%$-$40\%$) is best for all five domains, and a calibrated permutation test for quadratic interiority gives $P\approx0.010$; the fitted mid-training-only curves, with 8B peaks between $9.9\%$ and $35.1\%$, reproduce for curve shape but not peak location. Second, the gaps survive a fixed-budget alignment pass: compensatory SFT raises 116/120 cells (mean $+4.32\%$) yet bridges $0/240$ pairs at a $5\%$ threshold and $30/240$ at a $10\%$ ratio, an equal-budget uniform control behaves almost identically, and a permutation null would bridge $13.8\pm3.3$ and $77.9\pm8.5$ pairs ($P&lt;0.001$). Third, zero coverage collapses mid-training-only accuracy, though a FineWeb-Edu-only control shows the collapse is commingled with generic drift. An exploratory $θ^*$ allocation attains the largest full-pipeline gain ($+4.36\%$ vs. $+0.80\%$/$+0.64\%$\,pp) but is marginal under Welch test.