Pretraining Data Mixtures: Designing the Token Distribution

This witness is a toy unigram language model. It is faithful to the central mixture-and-loss equations because the learned distribution is exactly the curated mixture distribution, and validation loss is exact cross-entropy under held-out toy distributions. It is not a transformer simulation.

from math import log
from collections import defaultdict

VOCAB = ["story", "fact", "proof", "equation", "def", "loop", "bug", "clickbait"]
DOCS = [
    ("web_a",  "web",  0.65, "w1", False, {"story": 8, "fact": 5, "clickbait": 2}),
    ("web_b",  "web",  0.35, "w2", False, {"story": 7, "clickbait": 7, "fact": 1}),
    ("web_c",  "web",  0.70, "w1", False, {"story": 8, "fact": 5, "clickbait": 2}),
    ("edu_a",  "edu",  0.92, "e1", False, {"proof": 5, "equation": 5, "fact": 3}),
    ("edu_b",  "edu",  0.58, "e2", True,  {"proof": 6, "equation": 5, "fact": 1}),
    ("code_a", "code", 0.86, "c1", False, {"def": 5, "loop": 5, "bug": 1}),
    ("code_b", "code", 0.40, "c2", False, {"def": 2, "bug": 7, "loop": 1}),
]

EVALS = {
    "general": {"story": 0.45, "fact": 0.45, "clickbait": 0.10},
    "math": {"proof": 0.45, "equation": 0.45, "fact": 0.10},
    "code": {"def": 0.45, "loop": 0.45, "bug": 0.10},
}

def curate(thresholds, dedup=True):
    kept, seen = [], set()
    for doc_id, src, quality, group, leak, counts in DOCS:
        if quality < thresholds.get(src, 0.0):
            continue
        if dedup and group in seen:
            continue
        seen.add(group)
        kept.append((doc_id, src, quality, group, leak, counts))
    return kept

def normalize(counts, eps=1e-9):
    total = sum(counts.values()) + eps * len(VOCAB)
    return {tok: (counts.get(tok, 0.0) + eps) / total for tok in VOCAB}

def source_stats(kept):
    counts_by_src = defaultdict(lambda: defaultdict(float))
    tokens_by_src = defaultdict(float)
    leaks_by_src = defaultdict(float)
    for _, src, _, _, leak, counts in kept:
        n = sum(counts.values())
        tokens_by_src[src] += n
        if leak:
            leaks_by_src[src] += n
        for tok, n in counts.items():
            counts_by_src[src][tok] += n
    return counts_by_src, tokens_by_src, leaks_by_src

def normalize_weights(weights, active_sources):
    active_sources = list(active_sources)
    if not active_sources:
        return {}
    total = sum(weights.get(src, 0.0) for src in active_sources)
    if total <= 0:
        return {src: 1.0 / len(active_sources) for src in active_sources}
    return {src: weights.get(src, 0.0) / total for src in active_sources}

def mix(source_distributions, weights):
    return {
        tok: sum(weights.get(src, 0.0) * dist[tok] for src, dist in source_distributions.items())
        for tok in VOCAB
    }

def evaluate(thresholds, early_mix, late_mix=None, late_fraction=0.0, dedup=True):
    kept = curate(thresholds, dedup=dedup)
    counts_by_src, tokens_by_src, leaks_by_src = source_stats(kept)
    source_distributions = {src: normalize(counts) for src, counts in counts_by_src.items()}
    active = list(source_distributions.keys())
    w_early = normalize_weights(early_mix, active)
    q_early = mix(source_distributions, w_early)

    if late_mix and late_fraction > 0:
        w_late = normalize_weights(late_mix, active)
        q_late = mix(source_distributions, w_late)
        q = {tok: (1 - late_fraction) * q_early[tok] + late_fraction * q_late[tok] for tok in VOCAB}
        avg_w = {src: (1 - late_fraction) * w_early.get(src, 0.0) + late_fraction * w_late.get(src, 0.0) for src in active}
    else:
        q = q_early
        avg_w = w_early

    losses = {
        name: -sum(p * log(max(q.get(tok, 0.0), 1e-12)) for tok, p in p_eval.items())
        for name, p_eval in EVALS.items()
    }
    contamination = sum(avg_w[src] * leaks_by_src[src] / tokens_by_src[src] for src in active)
    repetition = {src: 180 * avg_w[src] / tokens_by_src[src] for src in active}

    return {
        "kept_docs": [doc[0] for doc in kept],
        "losses": {k: round(v, 3) for k, v in losses.items()},
        "contamination_mass": round(contamination, 3),
        "repetition_pressure": {k: round(v, 2) for k, v in repetition.items()},
    }

baseline = evaluate(
    thresholds={"web": 0.0, "edu": 0.0, "code": 0.0},
    early_mix={"web": 0.70, "edu": 0.15, "code": 0.15},
    dedup=False,
)

curated_late = evaluate(
    thresholds={"web": 0.55, "edu": 0.60, "code": 0.60},
    early_mix={"web": 0.50, "edu": 0.25, "code": 0.25},
    late_mix={"web": 0.35, "edu": 0.35, "code": 0.30},
    late_fraction=0.20,
    dedup=True,
)

print("baseline:", baseline)
print("curated + late:", curated_late)

assert curated_late["losses"]["math"] < baseline["losses"]["math"]
assert curated_late["losses"]["code"] < baseline["losses"]["code"]
assert curated_late["losses"]["general"] > baseline["losses"]["general"]
assert curated_late["contamination_mass"] < baseline["contamination_mass"]