Albalak et al. (2023) — ARXIV-2312.02406
Key takeaways
- The problem is that mixing proportions are fixed before training and so cannot adapt to training dynamics, while data selection methods that could adapt are too slow to run online
- ODM combines the two: group data into domains as mixing does, but choose the proportions online as selection would, formulated as a multi-armed bandit over the domains
- The reward is the per-domain training loss, computed on the batches the run is already taking — “fortuitously requires minimal overhead”. The argument for it is information-theoretic: perplexity measures expected information gain, so sample most from the domain with most left to learn
- The algorithm is Exp3. Its policy is a Gibbs distribution over the
importance-weighted rewards, mixed with a uniform distribution for
exploration at rate
ℰₜ:πₜ(Dᵢ) = (1 − K·ℰₜ)·exp(ℰₜ₋₁·R̂ᵢ)/Σⱼexp(ℰₜ₋₁·R̂ⱼ) + ℰₜ - Results: 1B parameters, 50B tokens, 22 domains of The Pile. Reaches the next best method’s final perplexity in 19% fewer iterations, and improves 5-shot MMLU by 1.9% relative accuracy, at negligible wall-clock cost
Standing in the anthology
The source of SOTA-103.
Read in full for #114; the reading is NOTE-002.
Four practices had been drawn from the replaced takeaways and none of them described this method. One survives, restated: SOTA-103 had the loop right and the signal wrong — validation performance rather than training loss — which is the difference the paper’s efficiency claim rests on entirely.