mixle.task.data_mixture module¶
DoReMi-style data-mixture optimization: domain weights as a bandit/DOE problem (roadmap F8).
A pretraining corpus is normally split into named domains (web text, code, books, …) mixed by a hand-picked weight vector. DoReMi (Xie et al.) instead treats the weight vector itself as something to OPTIMIZE: run many small, cheap proxy trainings at different mixtures, score each on held-out loss, and search the simplex for the mixture that generalizes best – then apply that learned mixture to the real, much larger run. This module is the small, honest version of that loop, reusing mixle’s own optimization machinery rather than inventing new search code:
SyntheticDomain– a synthetic “domain”: a name plus a stand-in data-generating distribution (a fixed periodic token pattern with configurable noise, or pure noise for an unlearnable domain).
proxy_run_score()– one proxy run: build a token stream from a mixture of domains, train a real (tiny)mixle.models.language_model.LMfor a handful of steps, and return the mean held-out NLL across domains (lower is better). This is the “small-run proxy” DoReMi scores mixtures with.
optimize_mixture()– the DoReMi search loop: repeated proxy runs scored viaproxy_run_score(), with candidate mixtures proposed bymixle.task.bandit’sThompsonGaussian(discrete arms on a simplex-lattice design frommixle.doe.mixture) ormixle.doe.optimizer’sBayesianOptimizer(continuous search over a softmax-reparameterized simplex). No new optimizer machinery – both paths are the same modules F5/I1/D5 already reuse this session.
estimate_near_duplicate_rate()– a minimal, honest corpus dedup/quality receipt: a MinHash estimate of the fraction of documents with a near-duplicate elsewhere in the corpus.
F5 (scaling-law fits) integration point: F5’s fitted scaling laws could extrapolate a proxy run’s
held-out loss at this scale to a prediction at the real target scale, letting the search compare
mixtures by extrapolated real-scale loss instead of raw proxy-scale loss. F5’s branch was not reachable
from this worktree at the time F8 was built, so that extrapolation is not wired in here – the natural
integration point is inside proxy_run_score(), mapping its returned proxy-scale loss through a
fitted mixle.task.<f5-module> law before it reaches the optimizer. The search loop itself
(optimize_mixture()) does not need to change: it only requires a scalar score per mixture, however
that score is produced.
- class SyntheticDomain(name, vocab, period=8, noise_p=0.0, pattern_seed=0)[source]
Bases:
objectOne synthetic “domain”: a fixed periodic token pattern, optionally corrupted by noise.
pattern_seedfixes a length-periodsequence of token ids (drawn once, from0..vocab)) that repeats forever – the domain’s learnable structure. Each sampled token then has independent probabilitynoise_pof being replaced by a uniform-random token, sonoise_p=0is a perfectly learnable domain andnoise_p=1(orperiod=None) is pure, irreducible noise: no amount of training data lowers a model’s achievable loss on it belowlog(vocab). Distinct(period, pattern_seed, noise_p)triples give genuinely different data-generating distributions, standing in for e.g. “web text” vs “code” vs “books” without needing real corpora.
- estimate_near_duplicate_rate(corpus, *, shingle_size=5, num_hashes=64, threshold=0.8, seed=0)[source]
Estimate the fraction of documents in
corpusthat have a near-duplicate elsewhere in it.A minimal, honest MinHash quality/dedup receipt: each document is reduced to its set of word-
shingle_sizeshingles, each shingle set to anum_hashes-entry MinHash signature (an unbiased estimator of Jaccard similarity), and two documents are called near-duplicates when their signatures agree on at leastthresholdof their entries. Returns|{documents with >= 1 near-duplicate partner}| / |corpus|.O(n^2)in the corpus size – fine for the receipt-sized corpora this is meant for, not a production LSH dedup pipeline.
- optimize_mixture(domains, proxy_steps, budget, *, method='bandit', proxy_kwargs=None, seed=0)[source]
Learn domain mixture weights via repeated short proxy runs (DoReMi-style search).
budgetproxy runs (eachproxy_run_score()atproxy_stepsgradient steps) are used to search the mixture-weight simplex.method="bandit"(default) discretizes the simplex into a lattice of candidate mixtures (mixle.doe.mixture.simplex_lattice) and searches them withmixle.task.bandit.ThompsonGaussian(reward = negative held-out loss);method="doe"searches continuously viamixle.doe.optimizer.BayesianOptimizerover a softmax-reparameterized simplex. Returns the learned weight vector (one entry per domain, summing to 1).
- proxy_run_score(mixture_weights, domains, proxy_steps, *, batch_size=16, d_model=16, n_layer=1, n_head=2, block=8, lr=3.0e-3, eval_tokens=512, seed=0, eval_seed=999_000, return_detail=False)[source]
Run one short proxy training and return the mean held-out NLL across
domains(lower is better).Builds a training token stream by drawing
mixture_weights[i]-proportional tokens from each domain (concatenated; the number of tokens is chosen so training runs roughlyproxy_stepsgradient steps atbatch_size), trains a real (tiny)mixle.models.language_model.LMon it for one epoch, then scores held-out NLL oneval_tokensfresh tokens from EACH domain (independent of the mixture) and returns the unweighted mean across domains – the DoReMi objective is generalizing to every domain, not just the ones the mixture over-samples.return_detail=Truealso returns the per-domain NLL dict, keyed by domain name.seedcontrols the training-data draw (and so varies across repeated proxy runs, e.g. insideoptimize_mixture()’s search loop);eval_seedcontrols the held-out draw and is fixed by default so different mixtures proposed during a search are scored against the SAME held-out set – comparing candidate mixtures on a moving eval target would swamp the (often small) between-mixture signal in eval-sampling noise.