Vertex sampling helps mix many language-model data domains
LinMix: Data Mixing as Linear Domain Attribution · arXiv:2610.37390
In plain words
Language models learn from text drawn from different subject areas, or domains. LinMix is a method that trains small proxy models on carefully chosen mixes of those domains, estimates which domains help a chosen measure, and uses those estimates to select a mix for training a larger model.
In tests with 288 domains, its way of choosing proxy mixes let a simple linear estimate fit about as well as the best nonlinear alternative. With 24 domains, both LinMix’s sampler and the comparison sampler supported good fits.
The research
When training text is divided into hundreds of domains, how can a limited set of small proxy models show which domains contribute to a target measure? If their training mixtures do not vary enough, the estimates used to select a larger model’s data mix may be unreliable.
LinMix varies capped proxy-model data mixtures to estimate domain contributions with a linear model and choose a target mix for pretraining.
LinMix tests whether changing the proxy mixtures can make domain effects easier to estimate as a corpus is split into more domains.
Researchers compared predicted and observed proxy-model scores: with 288 domains, vertex sampling let a linear fit match the best nonlinear alternative, while tuned Dirichlet sampling remained weak.
They then used the selected mixtures to train 4B Qwen3.5 models on 40 domains. LinMix scored higher than RegMix and Olmix on the reported coding and math tests, linking the proxy-fit result to target-model performance in that setting.
The target-model results come from one random training run per method, and training the many proxy models is costly. The proof of expected optimality applies when domain caps are equal and the prior is uniform; the paper leaves unequal caps unresolved.
What the evidence shows
- Estimate domain contributions
LinMix is a data-mixing method for language-model pretraining. It addresses a practical question: when the training corpus is divided into many domains, can a limited set of smaller proxy models reveal which domains help a target measure?
The method estimates each domain’s contribution from proxy runs, then uses those estimates to choose a capped mixture for the larger model.
- Vary the proxy mixtures
The key choice is how to form proxy mixtures. Common Dirichlet sampling draws mixtures around a prior distribution and requires setup-specific tuning.
LinMix instead fills domains to their caps in a random order until the token budget is spent, then pairs runs using the reverse order. The resulting mixtures spread across domains, giving the regression more useful contrasts to estimate their effects.
- The advantage grows at scale
Researchers compared predicted and observed proxy-model scores on benchmark training data. At 288 domains, even tuned Dirichlet mixtures produced a weak ranking agreement, while a linear fit using LinMix’s vertex sampling matched the best nonlinear alternative.
At 24 domains, all tested fits worked well with both samplers. The results therefore locate the sampling advantage in the larger domain setup.
- Target scores also favor LinMix
The researchers also trained 4B Qwen3.5 target models on 40-domain mixtures selected by LinMix, RegMix or Olmix. On MBPP, which tests coding problems, LinMix’s pass@1 score measures the fraction of problems solved by a single solution generated with greedy decoding; it was 0.444, compared with 0.304 for RegMix and 0.302 for Olmix.
On GSM8K math problems, exact match measures the fraction of answers matching the reference; LinMix scored 0.337 versus 0.135 and 0.294. These results connect the proxy comparison with stronger scores in this setup.
- Theory and evidence have bounds
A theoretical result supports the mixture design: vertex sampling is optimal in expectation when all domain caps are equal and the prior is uniform. The proof also clarifies that caps force each run to include multiple domains, limiting how much it can isolate any one domain.
The evidence supports LinMix in the tested settings, while target results use one random training run per method and the theorem’s equal-cap, uniform-prior conditions limit its reach.
Deep
- Original paper authors
- Hubert Baniecki, Przemyslaw Biecek, Martin Pawelczyk
- Submitted to arXiv
- License
- CC BY 4.0