Data Mixing › Code Implementation
Language AI HandbookAll chapters Data Mixing › Code ImplementationOne section · Sections in this chapter6 / 9 ↓
Reading tools Definitions & inline links
Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.
Article links
Make inline references clickable
Let us work through a simulation that captures the key ideas in data mixing. We will implement a simplified mixing framework that demonstrates domain proportion effects, quality weighting, and the proxy model evaluation paradigm.
This simulation uses synthetic capability scores rather than real model training, which makes it fast enough to run interactively. The core dynamics, including the oversampling ceiling, the advantage of curated recipes over proportional sampling, and the DoReMi weight update, are accurately represented. The synthetic model captures the most important qualitative behaviors while abstracting away the engineering complexity of actual large-scale training.
Setting Up Domain Profiles
We start by defining synthetic domain profiles that represent a realistic pretraining corpus.
import numpy as np
# Define domain profiles with realistic characteristics
# Each domain has: corpus size (tokens), base quality score, capability contributions
domains = {
"web_common_crawl": {
"size_billions": 3000, # Very large pool
"quality": 0.4, # Mixed quality
"coding": 0.05,
"reasoning": 0.3,
"facts": 0.4,
"language": 0.5,
},
"web_filtered": {
"size_billions": 500,
"quality": 0.7,
"coding": 0.1,
"reasoning": 0.5,
"facts": 0.6,
"language": 0.7,
},
"books": {
"size_billions": 80,
"quality": 0.9,
"coding": 0.05,
"reasoning": 0.7,
"facts": 0.7,
"language": 0.95,
},
"wikipedia": {
"size_billions": 20,
"quality": 0.95,
"coding": 0.1,
"reasoning": 0.6,
"facts": 0.95,
"language": 0.85,
},
"code_github": {
"size_billions": 300,
"quality": 0.75,
"coding": 0.98,
"reasoning": 0.6,
"facts": 0.2,
"language": 0.3,
},
"scientific_papers": {
"size_billions": 60,
"quality": 0.88,
"coding": 0.2,
"reasoning": 0.85,
"facts": 0.9,
"language": 0.75,
},
"reddit_filtered": {
"size_billions": 150,
"quality": 0.55,
"coding": 0.15,
"reasoning": 0.35,
"facts": 0.3,
"language": 0.6,
},
}
total_corpus = sum(d["size_billions"] for d in domains.values())
proportional_weights = {
k: v["size_billions"] / total_corpus for k, v in domains.items()
}Total corpus size: 4,110B tokens Proportional weights (sampling uniformly by size): web_common_crawl : 0.730 (73.0%) web_filtered : 0.122 (12.2%) books : 0.019 (1.9%) wikipedia : 0.005 (0.5%) code_github : 0.073 (7.3%) scientific_papers : 0.015 (1.5%) reddit_filtered : 0.036 (3.6%)
When sampling proportionally to corpus size, the 3 trillion token CommonCrawl pool dominates the mixture at over 72% by default. Every other domain, including books and Wikipedia, receives a tiny fraction. If training used these natural proportions, the model would rarely encounter the highest-quality content in the corpus. Books would account for less than 2% of all training tokens, meaning the model spends the vast majority of its budget learning patterns from content of much lower average quality.
Simulating Capability Scores
Now we simulate how different mixing recipes affect model capabilities using a simplified proxy model. The simulation captures the most important dynamics: capability improves with exposure to relevant domains, quality modulates how much each exposure is worth, and oversampling has diminishing returns after a threshold.
def simulate_capability(mixing_weights, noise_std=0.02, rng=None):
"""
Simulate capability scores under a given mixing recipe.
This is a toy model: real capability depends on many more factors.
We capture the core dynamic: each domain contributes in proportion to its
sampling weight, quality, and relevance to the capability being measured.
"""
if rng is None:
rng = np.random
capability_scores = {
"coding": 0.0,
"reasoning": 0.0,
"facts": 0.0,
"language": 0.0,
}
for domain, weight in mixing_weights.items():
quality_adjusted_weight = weight * domains[domain]["quality"]
for cap in capability_scores:
capability_scores[cap] += (
quality_adjusted_weight * domains[domain][cap]
)
# Map weighted contributions into a plausible proxy-score range, then add
# small run-to-run variation.
capability_scores = {
k: float(np.clip(0.45 + 0.90 * v + rng.normal(0, noise_std), 0, 1))
for k, v in capability_scores.items()
}
return capability_scores
def overall_score(caps):
"""Weighted aggregate capability score."""
weights = {
"coding": 0.25,
"reasoning": 0.30,
"facts": 0.25,
"language": 0.20,
}
return sum(caps[k] * weights[k] for k in caps)Comparing Mixing Recipes
Let us compare three mixing strategies: proportional (sampling by corpus size), curated (a practitioner's informed guess), and quality-weighted (sampling proportionally to domain quality scores).
# Strategy 1: Proportional (baseline)
recipe_proportional = proportional_weights.copy()
# Strategy 2: Curated (typical LLM recipe, broadly similar to Llama 2)
recipe_curated = {
"web_common_crawl": 0.10,
"web_filtered": 0.35,
"books": 0.15,
"wikipedia": 0.05,
"code_github": 0.20,
"scientific_papers": 0.10,
"reddit_filtered": 0.05,
}
# Strategy 3: Quality-weighted (boost quality, penalize noise)
quality_scores = np.array([domains[d]["quality"] for d in domains])
quality_weights = quality_scores / quality_scores.sum()
recipe_quality = dict(zip(domains.keys(), quality_weights))
# Evaluate each recipe
recipe_rng = np.random.RandomState(0)
runs_per_recipe = 20
results = {}
for recipe_name, recipe in [
("Proportional", recipe_proportional),
("Curated", recipe_curated),
("Quality-Weighted", recipe_quality),
]:
scores = [
simulate_capability(recipe, rng=recipe_rng)
for _ in range(runs_per_recipe)
]
aggregate = [overall_score(s) for s in scores]
cap_means = {
k: np.mean([s[k] for s in scores])
for k in ["coding", "reasoning", "facts", "language"]
}
results[recipe_name] = {
"mean": np.mean(aggregate),
"std": np.std(aggregate),
"caps": cap_means,
}Recipe Comparison (proxy model simulation): Recipe Overall Coding Reasoning Facts Language ----------------------------------------------------------------- Proportional 0.616+-0.011 0.528 0.621 0.639 0.692 Curated 0.800+-0.011 0.638 0.847 0.827 0.898 Quality-Weighted 0.841+-0.011 0.614 0.885 0.920 0.957
The proportional recipe underperforms on all capability dimensions because the model spends most of its training budget on lower-quality, lower-diversity CommonCrawl text. The curated recipe shows the biggest gain in coding because it deliberately includes much more code. The quality-weighted recipe leads on reasoning and factual accuracy because it favors books, Wikipedia, and scientific papers. Notice that no single recipe dominates all capability dimensions: curated mixing wins on coding, quality weighting wins on reasoning and facts, and both beat proportional on language quality.
Simulating the Oversampling Ceiling
We now look at what happens when we push oversampling to its limits for the books domain. This experiment varies the weight assigned to books from 1% to 40% and measures how overall capability and language quality respond as the oversampling ratio increases.
# Vary how much we oversample books while keeping other weights proportional
book_weights = [0.01, 0.05, 0.10, 0.15, 0.20, 0.30, 0.40]
oversample_results = []
oversample_rng = np.random.RandomState(1)
for book_weight in book_weights:
# Distribute remaining weight proportionally among other domains
remaining = 1.0 - book_weight
other_domains = {
k: v for k, v in proportional_weights.items() if k != "books"
}
other_sum = sum(other_domains.values())
recipe = {k: (v / other_sum) * remaining for k, v in other_domains.items()}
recipe["books"] = book_weight
# Calculate actual oversample ratio for books domain
# (books has 80B tokens, we're training on 1000B total)
oversample_ratio = (book_weight * 1000) / domains["books"]["size_billions"]
scores = [
simulate_capability(recipe, rng=oversample_rng) for _ in range(30)
]
# A compact toy penalty represents memorization displacing generalization
# once the same book tokens have been repeated more than about three times.
repeat_penalty = 0.015 * max(oversample_ratio - 3.0, 0.0) ** 2
mean_score = np.mean([overall_score(s) for s in scores]) - repeat_penalty
language_score = np.mean([s["language"] for s in scores]) - repeat_penalty
oversample_results.append(
{
"book_weight": book_weight,
"oversample_ratio": oversample_ratio,
"mean_score": mean_score,
"language_score": language_score,
}
)Effect of oversampling books:
Book Weight Oversample Ratio Overall Score Language Score
--------------------------------------------------------------
1% 0.1x 0.615 0.683
5% 0.6x 0.630 0.705
10% 1.2x 0.643 0.729
15% 1.9x 0.659 0.763
20% 2.5x 0.676 0.792 <-- sweet spot
30% 3.8x 0.696 0.834 <-- sweet spot
40% 5.0x 0.676 0.834The results confirm the oversampling ceiling: language scores improve as you oversample books up to around 3 to 4 times, but gains plateau and may slightly reverse at very high ratios. The mechanism is that repeated exposure to the same book tokens eventually causes the model to memorize rather than learn generalizable patterns. The gradient signal from well-memorized text approaches zero, and those training steps contribute nothing to the model's ability to generalize.
Implementing a Simple DoReMi Step
We simulate one step of the DoReMi weight update to illustrate the algorithm's behavior. The key thing to observe is how the update naturally directs weight toward domains where the proxy model is struggling.
def doremi_weight_update(domain_weights, proxy_losses, ref_losses, eta=0.1):
"""
One step of the DoReMi exponentiated gradient update.
Increases weight for domains where proxy model struggles (high excess loss),
decreases weight for domains where proxy model has caught up.
"""
weights = np.array([domain_weights[d] for d in domains])
proxy_l = np.array([proxy_losses[d] for d in domains])
ref_l = np.array([ref_losses[d] for d in domains])
# Excess loss: positive means proxy is worse than reference
excess_loss = proxy_l - ref_l
# Exponentiated gradient update
updated_weights = weights * np.exp(eta * excess_loss)
# Project back to simplex (normalize)
updated_weights = updated_weights / updated_weights.sum()
return dict(zip(domains.keys(), updated_weights))
# Simulate a reference model with uniform weights trained for some steps
# Simulate proxy model losses: code and reasoning are harder domains
loss_rng = np.random.RandomState(7)
ref_losses = {
"web_common_crawl": 2.1 + loss_rng.normal(0, 0.05),
"web_filtered": 2.0 + loss_rng.normal(0, 0.05),
"books": 2.4 + loss_rng.normal(0, 0.05),
"wikipedia": 2.2 + loss_rng.normal(0, 0.05),
"code_github": 2.8 + loss_rng.normal(0, 0.05),
"scientific_papers": 2.6 + loss_rng.normal(0, 0.05),
"reddit_filtered": 1.9 + loss_rng.normal(0, 0.05),
}
# Proxy model with uniform weights is better on easy domains, worse on hard ones
proxy_losses = {
"web_common_crawl": ref_losses["web_common_crawl"] - 0.20,
"web_filtered": ref_losses["web_filtered"] - 0.15,
"books": ref_losses["books"] + 0.10, # Struggling
"wikipedia": ref_losses["wikipedia"] - 0.05,
"code_github": ref_losses["code_github"] + 0.25, # Struggling most
"scientific_papers": ref_losses["scientific_papers"] + 0.12, # Struggling
"reddit_filtered": ref_losses["reddit_filtered"] - 0.18,
}
# Initial uniform weights
initial_weights = {d: 1.0 / len(domains) for d in domains}
# Run several DoReMi update steps
weights = initial_weights.copy()
weight_history = [weights.copy()]
for step in range(20):
weights = doremi_weight_update(weights, proxy_losses, ref_losses, eta=0.15)
weight_history.append(weights.copy())
final_weights = weight_history[-1]DoReMi weight evolution (uniform start -> adaptive end): Domain Initial Final Change ---------------------------------------------------------- web_common_crawl 0.143 0.073 -0.070 web_filtered 0.143 0.085 -0.058 books 0.143 0.179 +0.036 wikipedia 0.143 0.114 -0.029 code_github 0.143 0.281 +0.138 scientific_papers 0.143 0.190 +0.047 reddit_filtered 0.143 0.077 -0.065
The DoReMi update naturally increases weights for code and scientific papers (where the proxy model struggles relative to the reference), while reducing weights for Reddit and common web crawl (where the proxy model performs well). This mirrors what human practitioners would do manually when they look at per-domain loss curves: boost the hard domains, suppress the easy ones. The difference is that DoReMi does this automatically and continuously throughout training, rather than requiring a human to review curves and adjust manually.
Key Parameters
The key parameters for a data mixing configuration are:
- domain_weights: The probability vector specifying how frequently each domain is sampled. Must sum to 1. This is the primary lever you control.
- training_budget: The total number of tokens to train on. Combined with domain weights, this determines the effective oversampling ratio for each domain.
- quality_threshold: If applying quality filtering before mixing, the minimum quality score for a document to be included. Raising this threshold reduces corpus size but increases average quality.
- eta: The learning rate for DoReMi-style adaptive mixing. Higher values cause more aggressive weight shifts; lower values produce smoother, more conservative updates.
- minimum_floor: A lower bound on any domain's weight, making sure no domain is completely excluded from training. Typical values are in the range of 1 to 5%.