The Lexical Imprint of Pre-training Data

AI summary

Overview

Research area: Natural Language Processing — pre-training data analysis, zero-shot evaluation, and data selection for large language models.

Technical level: Advanced. The paper is built on information-theoretic arguments (cross-entropy and KL divergence over word distributions) and requires controlled from-scratch pre-training runs across multiple model scales and corpora.

Scope (one sentence): The paper proposes word-level unigram cross-entropy (UCE) as a tokenizer-agnostic, model-free measure of lexical alignment between a pre-training corpus and a benchmark, and shows that this single statistic tracks zero-shot accuracy at the corpus level, can be used to select training documents at the document level, and can route prompts to specialized drafters at the prompt level.

What This Paper Is About

Zero-shot benchmark scores are usually read as measurements of model capability, but they also depend on how closely the pre-training corpus resembles the benchmark's wording. This paper asks how much of zero-shot accuracy can be explained purely by word-level distributional alignment between a pre-training corpus and a benchmark target.

To answer this cleanly, the authors run controlled from-scratch pre-training that varies only the corpus (C4, DCLM, FineWeb-Edu, OpenWebText) while holding the training recipe fixed, and measure UCE between each corpus and each benchmark. They then push the same statistic from description into control: using it to select training documents, and using it to route prompts to corpus-specialized drafters during speculative decoding.

Key Contributions

1. A model-free alignment measure (UCE). Word-level unigram cross-entropy between a reference corpus distribution P and a benchmark target distribution Q, defined as UCE(Q,P) = −Σ_{w∈V_b} Q(w) log P(w) = H(Q) + D_KL(Q‖P). It is tokenizer-agnostic, computationally efficient, and avoids the segmentation confounds of subword token-level cross-entropy.

2. A controlled corpus-level sweep. Ten zero-shot benchmarks, four pre-training corpora, and three model scales (400M, 1.33B, 3.36B) trained from scratch with a fixed recipe, showing that lower UCE consistently accompanies higher zero-shot accuracy on every benchmark tested.

3. The UCE-filter for document-level control. A one-shot, row-local selection rule that analytically computes each training document's marginal contribution ΔH_r to corpus–benchmark cross-entropy and discards the rows whose removal most lowers it. It is tested against a uniform random control and a reproduction of DSIR at a fixed token budget, plus a contamination audit.

4. Prompt-level routing as a third instantiation. The same KL functional applied per prompt routes each prompt to the specialist drafter whose corpus is most lexically aligned with it, yielding throughput gains over a generalist drafter in speculative decoding.

Main Findings

• Corpus-level alignment matches corpus rankings. For a 3.36B model trained on 60B tokens, UCE predicted the accuracy ordering of C4, DCLM, and FineWeb-Edu on all eight benchmarks shown in Figure 2. The "higher-quality" DCLM and FineWeb-Edu led with the lowest UCE on six of them (FineWeb-Edu on ARC Easy/Challenge, SciQ, OpenBookQA; DCLM on SocialIQA and SWAG), while C4 had lower UCE where it scored better (HellaSwag and PIQA for both, plus SocialIQA and SWAG for FineWeb-Edu).

• Corpora differ in mass allocation, not vocabulary coverage. All four corpora attain near-saturated coverage of benchmark vocabulary at these scales. The differentiator is how probability mass is spread: the higher a corpus scores, the smaller the share of benchmark word mass falling on its rare-word tail.

• The relationship survives random subsetting. Across 16 runs (five 8.5B-token subsets each from C4, DCLM, FineWeb-Edu, and one from OpenWebText, each training a 400M model), the inverse relationship held with R² = 0.928 on ARC Easy and 0.805 on PIQA. Within a corpus, subsets had nearly identical UCE and differed little in accuracy, consistent with run-to-run noise.

• Selecting for lower UCE lowers measured UCE. On the targeted 7-task suite (CSR7), UCE-filter subsets reached 10.92 bits from DCLM (versus 11.17 for both Full and Random, and 11.08 for DSIR) and 10.93 bits from FineWeb-Edu (versus 11.32 for Full and Random, 11.20 for DSIR). Targeted UCE of Random and DSIR varied by less than 0.001 bits across selection seeds. Out-of-target UCE values stayed close to the random control for both selectors.

• Selection raises accuracy mainly on the targeted suite. In a 2×2 grid of source (DCLM, FineWeb-Edu) and scale (1.34B, 3.33B) trained on 70B tokens, the UCE-filter beat the random control on the Targeted average in every cell. Examples: DCLM at 1.34B, 62.49 ± 0.26 versus 58.85 ± 0.56; DCLM at 3.33B, 66.26 ± 0.20 versus 62.95 ± 0.03; FineWeb-Edu at 1.34B, 61.85 ± 0.15 versus 59.12 ± 0.41; FineWeb-Edu at 3.33B, 64.64 ± 0.21 versus 62.80 ± 0.09. Out-of-target averages stayed close to the random control, and DSIR fell below the random control on the out-of-target average in every cell.

• The selection gain is not explained by contamination. A word-level 13-gram audit found that nine of the 12 benchmarks had under 0.3% of items matched in every corpus sample; BoolQ, SciQ, and MMLU peaked on FineWeb-Edu. On SWAG, SocialIQA, and LAMBADA (at most 0.01% of items matched on C4, DCLM, or FineWeb-Edu), lower-UCE corpora still scored better. After removing flagged DCLM documents (0.13% of rows) and retraining at 1.34B, the UCE-filter kept 3.40 pp of its 4.13 pp CSR7 advantage; dropping flagged evaluation items moved results by at most 0.05 pp per cell.

• The signal isolates a gain at production scale. In a 30B-parameter Stage-2 case study on a 2T-token production mixture, replacing only the web component with its UCE-filtered subset and training on 200B tokens reached MMLU 74.28 and MMLU-Pro 48.81, against 75.61 / 48.64 for the full production mixture at ~2T tokens and 71.17 / 44.46 for the Stage-1 endpoint. The authors describe this as isolation of the gain rather than improvement over a curated mix.

• Prompt-level routing improves decoding throughput. With DeepSeek-R1-Distill-Llama-8B as verifier, UCE-based routing among three specialist drafters (OpenThoughts-114k math, code, science) gave a 55.3% average speedup and τ = 1.550, versus 50.0% and 1.534 for the generalist drafter — a 5.3 pp average gain. Speedups rose on every benchmark, including MMLU (55.3% versus 45.7%) and MedQA (49.6% versus 45.7%), which no specialist targets.

• The statistic has clear boundaries. UCE does not order grammar (BLiMP), arithmetic (MathQA), or non-English (Global-PIQA) benchmarks, and it is blind to syntax, multi-token composition, and factual co-occurrence by construction.

Methodology in Plain English

The authors summarize any text collection — a pre-training corpus or a benchmark — by its word frequencies. Words are extracted with Unicode NFKC normalization, punctuation stripping, and whitespace splitting, without lowercasing. They then compute the cross-entropy of the benchmark's word distribution under the corpus's word distribution. Because the benchmark's own entropy is a constant for a given benchmark, this quantity is essentially the KL divergence between the benchmark and the corpus; lower means the corpus puts probability mass on the words the benchmark actually uses.

To test whether this number predicts behavior, they train small Transformers (400M, 1.33B, 3.36B parameters) from scratch on fixed token budgets (8.5B, 26B, 60B, paired with model size at or near the Chinchilla ratio of 20 tokens per parameter) using identical hyperparameters, varying only which corpus the tokens came from. OpenWebText is only about 9B tokens, so it is evaluated at the 8.5B scale only. Zero-shot accuracy is measured on a 12-benchmark pool, with the corpus-level study using the 10 benchmarks other than CommonsenseQA and BoolQ.

For document-level control, the UCE-filter computes the change in corpus–benchmark cross-entropy caused by removing each individual document. Because removing a document only alters counts for words that appear in it, this change can be computed analytically in one linear pass instead of re-estimating the full distribution. The 20% of documents whose removal most increases cross-entropy (i.e., those most aligned with the target) are retained. The comparison uses equal-token random subsets and a reproduction of DSIR, with three training seeds per cell.

For inference, drafters are trained on separate domain corpora; each incoming prompt is converted to a word distribution and matched to the drafter with the lowest cross-entropy against it, which reduces to a sparse dot product at negligible latency. Finally, they audit all four corpora with a 13-gram index to check whether benchmark text in the training data could explain the results.

Why This Matters

The paper reframes a piece of zero-shot benchmark variance as a measurable property of pre-training data rather than of model capability, and it supplies a cheap, model-free probe for it.

Impact on research:

• Offers a controlled, from-scratch complement to post-hoc analyses on released checkpoints, where corpus identity is entangled with training recipe, scale, and training history.
• Gives a diagnostic that separates accuracy attributable to corpus–benchmark lexical overlap from accuracy reflecting generalization beyond it.
• Connects to the compression view of language modeling by grounding the argument in the fact that likelihood training minimizes KL divergence to the training distribution.

Real-world applications:

• Pre-training data curation: choosing or mixing web corpora based on alignment with intended evaluation suites.
• Domain adaptation: selecting a subset of a large crawl that is lexically aligned with a target domain at a fixed token budget.
• Inference cost reduction: routing prompts to domain-specialized drafters in speculative decoding, which in this study raised average decoding speedup from 50.0% to 55.3% over the verifier alone.
• Benchmark auditing: flagging evaluations whose scores may be substantially driven by lexical overlap with the pre-training corpus.

Industry relevance: The 30B Stage-2 case study speaks directly to production pre-training, where the practical question is how much of a curated mixture's gain can be recovered with far fewer tokens — here roughly 10× fewer (200B versus ~2T).

Future Directions

• Paraphrase robustness. The authors propose a paraphrase protocol to test how far the relationship depends on exact lexical realization rather than meaning.
• Restricting UCE to benchmark-characteristic vocabulary. A lexical decomposition that limits UCE to a characteristic vocabulary is proposed as a follow-up to the current whole-vocabulary measure.
• Post-training persistence. Whether the lexical imprint survives SFT and RLHF is unresolved; the paper states this requires a controlled study designed for those objectives, since they differ in data volume, training objective, and what is measured.
• Broader data properties and languages. A general account of generalization would need designs that vary other pre-training data properties alongside the lexical one, and whitespace tokenization limits UCE's direct applicability to non-segmented scripts (e.g., Chinese, Japanese) and highly agglutinative morphologies.

Target Audience

Researchers and engineers working on pre-training data curation, data selection, and benchmark methodology for large language models. It is most useful to readers comfortable with cross-entropy and KL divergence who want a quantitative, corpus-side account of why benchmark scores move with training data — and to practitioners deciding how to allocate a fixed token budget across candidate corpora. Readers looking for a standalone data-selection method should note the authors' explicit framing of the UCE-filter as a controlled probe rather than a general-purpose selector.

Authors’ abstract

How much of zero-shot benchmark accuracy is captured by word-level alignment between a pre-training corpus and the benchmark target? We study this question through word-level unigram cross-entropy (UCE), a tokenizer-agnostic, model-free measure of lexical alignment between a reference corpus and a target benchmark. Likelihood training shapes the model's predictive distribution from the pre-training corpus, and zero-shot scores under likelihood-based evaluation depend in part on the lexical alignment between that corpus and the benchmark target. Cleanly tracing how the lexical imprint of pre-training data scales with corpus--benchmark alignment requires controlled from-scratch pre-training that varies corpus identity while holding the training recipe fixed. Across this controlled sweep, spanning $10$ zero-shot benchmarks, $4$ pre-training corpora, and $3$ model scales, we find that corpus--benchmark lexical alignment, UCE, consistently tracks zero-shot accuracy on every benchmark, with corpora more aligned with the benchmark (lower UCE) achieving higher accuracy. The same signal extends from measurement to control at finer scopes. At the document level, strengthening the imprint by selecting training documents to maximize lexical alignment with the benchmark produces the accuracy gains predicted by the mechanism. At the prompt level, matching the imprint per prompt by routing each prompt to the drafter whose domain corpus is most lexically aligned with it produces measurable throughput gains over a generalist-drafter baseline. Taken together, pre-training corpora leave a measurable lexical imprint on base-model zero-shot behavior, and UCE provides a model-free probe of lexical alignment between reference and target data.