Generalization v.s. Memorization: Tracing Language Models' Capabilities Back to Pretraining Data

1 INTRODUCTION

• Prior studies focused on verbatim recall, although exact reproduction of long text is uncommon for high-level capabilities and short outputs such as knowledge answers.Earlier work also examined small models or dataset subsets and often required expensive retraining to estimate counterfactual memorization.Source: §1 · pp. 1–2
• The paper defines distributional memorization as correlation between LLM output and pretraining-data distributions, and distributional generalization as their divergence.The authors use the shorter terms memorization and generalization throughout the paper.Source: §1 · p. 2
• The proposed task-gram language model counts co-occurrences of semantically related input-output n-gram pairs to model task-relevant pretraining distributions and compare them with LLM predictions.The authors describe it as modeling long-range task-relevant dependencies, unlike classical n-gram models that capture local lexical dependencies.Source: §1 · p. 2
• Experiments use Pythia models pretrained on the Pile across translation, factual QA, world-knowledge questions, and math reasoning.The task datasets are WMT, TriviaQA, MMLU, and GSM8K, respectively.Source: §1 · p. 2
• TriviaQA shows the strongest memorization, while MMLU is weaker and WMT insignificant; with larger models, memorization matters more for TriviaQA and generalization more for MMLU, WMT, and GSM8K.The authors relate this pattern to factual recall, reasoning, and transferable translation skills.Source: §1 · p. 3
• Gradient-based estimates likewise find the greatest pretraining-document influence on TriviaQA, followed by MMLU and then WMT.Documents containing task-gram pairs exert greater influence than documents containing only individual n-grams.Source: §1 · p. 3

2 METHOD

• Semantically similar n-grams mined from supervised task inputs and outputs form task-specific pairs that characterize the task.The method uses cosine similarity between n-gram embeddings to measure semantic closeness, with a task-specific threshold hyperparameter.Source: §2 · p. 3
• A task-gram language model estimates an output n-gram's probability given an input n-gram from their co-occurrence counts in the same pretraining-corpus document.When the input n-gram appears in unseen text, this conditional probability estimates the chance that the corresponding output n-gram appears.Source: §2 · p. 3
• For test examples, the method compares task-gram and LLM probabilities for matching input-output n-gram pairs, scoring each output n-gram in the prompted input context.The LLM probability is formed from the probabilities of the output n-gram's tokens.Source: §2 · p. 4
• Distributional memorization is the Spearman correlation between task-gram and LLM probabilities over test-set n-gram pairs; distributional generalization is defined as its opposite.Greater distributional correlation implies less distributional generalization, and vice versa.Source: §2 · p. 4
• The authors report that task-gram probabilities capture task-relevant distributions and explain LLM behavior better than the ∞-gram language model.They describe distributional memorization as how well task-gram probabilities predict LLM probabilities.Source: §2 · p. 4

Models and pretraining corpus

• Pythia models range from 13M to 12B parameters and are trained on the 207B-token Pile; the study also includes 1B and 7B OLMo models trained on Dolma.The OLMo models use a 3T-token pretraining corpus.Source: §3 · p. 4

Downstream tasks

• The four task types are machine translation, factual question answering, world knowledge questions, and reasoning.The corresponding datasets include WMT-09, TriviaQA, MMLU, and GSM8K.Source: §3 · pp. 4–5
• WMT-09 supplies 2.5K test examples, and its European-language coverage was chosen because those languages are more prominent in the Pile.The WMT test set includes six languages: Hungarian, Czech, German, Italian, Spanish, and French.Source: §3 · pp. 4–5
• TriviaQA has a 10K test set, while MMLU covers 57 subjects and includes questions requiring logical or mathematical reasoning.GSM8K uses grade-school math word problems with step-by-step solutions, requiring a generated reasoning step before the final answer.Source: §3 · pp. 4–5

Analysis overview

• Figure 3 presents task performance against Pile n-gram-pair counts for WMT, TriviaQA, MMLU, and GSM8K.Figure 2 depicts the expected distributional memorization and generalization trends across task types.Source: Figure 2 · p. 4, Figure 3 · p. 5

Searching pretraining data at scale

• WIMBD counts n-gram-pair co-occurrences, while ∞-gram counts individual n-gram occurrences and probabilities in the Pile.WIMBD is used for low-frequency pairs because ∞-gram's approximate counts can miss them.Source: §3 · p. 5

Performance and n-gram pair counts

• The study checks for contamination by requiring no large n-gram overlap between the Pile and test data at n = 8 and n = 14.The check follows the method used in the GPT-3 technical report.Source: §4 · p. 5
• A test example's estimated pretraining-corpus probability is based on whether any task-gram-table n-gram pair appears in that example.WMT performance is measured with BLEU, TriviaQA and MMLU with accuracy, and GSM8K with BERTScore precision for generated versus reference chain-of-thoughts.Source: §4 · pp. 5–6
• TriviaQA has more n-gram-pair counts than WMT, MMLU, and GSM8K in the Figure 3 comparison.Figure 3 plots task performance against per-example pair counts; WMT aggregates counts across languages, while the other tasks use bins.Source: §4 · p. 5
• For models below 410M, WMT and TriviaQA show near-zero performance regardless of n-gram pair count, while MMLU reaches below 10% around 150 pairs, significantly below random guessing.Test examples in this interval contain more reasoning or math problems, which appear harder for Pythia models.Source: §4 · p. 6
• Above 410M parameters, performance increases with task-related n-gram-pair counts across all tasks, with a stronger trend as model size grows.GSM8K's strongest increasing trend occurs for Pythia 2.8B; larger models show a less significant trend.Source: §4 · p. 6

• WMT shows no statistically significant distributional memorization, while larger models generate more novel n-gram pairs.The authors interpret the novelty trend as greater distributional generalization and hypothesize transfer of translation skills between languages.Source: §5 · pp. 6–7
• TriviaQA has significant memorization, strongest for n = 3, while its n = 5 and ∞-gram measures are below 0.25.The n = 3 measure exceeds 0.35; memorization also increases with model size.Source: §5 · p. 7
• For MMLU, knowledge-intensive tasks show stronger memorization than reasoning-intensive tasks, most clearly at n = 3; the gap is similar at n = 5.Memorization decreases with model size on both task groups; the authors suggest larger models may need to adjust recall probabilities for specialized, less common knowledge.Source: §5 · pp. 7–8
• GSM8K shows no significant memorization for Pythia or OLMo, while distributional generalization increases with model size in both families.The analysis uses normalized Kendall tau ranking distance for GSM8K because memorization is not significant.Source: Figure 4 · p. 7, §5 · p. 8
• The task-gram LM yields memorization measures at least as high as the ∞-gram LM on TriviaQA and MMLU, which the authors take as evidence of better distribution modeling.The ∞-gram LM predicts token probabilities using the longest available corpus prefix, with backoff when a longer n-gram is absent.Source: §5 · pp. 6, 8
• Larger models generally show decreased memorization and increased generalization, though the pattern differs across tasks.Source: §5 · p. 8

6 INFLUENCE OF n-GRAM DATA THROUGH PRETRAINING

• The estimated influence of a pretraining document is based on the dot product between its pretraining-loss gradient and the test-loss gradient, accumulated through model checkpoints.The test loss evaluates the model’s probability for an output n-gram conditioned on the task example; the document loss evaluates that n-gram in its pretraining context.Source: §6 · p. 8
• Figure 5 plots average training influence against Pythia model size for WMT, TriviaQA, and MMLU, contrasting documents containing n-gram pairs with documents containing only the output n-gram.Green lines denote documents containing the pairs; blue lines denote documents containing only the output n-gram.Source: Figure 5 · p. 9, §6 · p. 9
• Average influence is computed over test examples and randomly retrieved pretraining documents, with R = 50 because full-gradient computation is expensive.The authors describe this as complementary causal evidence, not an analysis covering the full pretraining corpus.Source: §6 · p. 9

7 PRACTICAL IMPLICATIONS: PROMPT OPTIMIZATION

• The authors propose rewriting instructions to resemble pretraining data for memorization, or to be less similar to it for generalization.This prompt-design suggestion is based on the hypothesis that prompt distributions strongly affect LLM generation distributions.Source: §7 · pp. 9–10
• The prompt optimizer uses GPT4o to rewrite a task prompt iteratively, providing average pretraining-corpus n-gram counts as feedback to maximize or minimize.The paper reports zero-shot testing accuracy for TriviaQA and GSM8K using Pythia and OLMo models.Source: §7 · p. 10
• TriviaQA benefits from prompts more similar to pretraining data, whereas GSM8K benefits from prompts less similar to it.The reported optimized prompts are not significantly different in length.Source: §7 · p. 10

8 RELATED WORK

• Prior studies of LLM capabilities often use synthetic experiments or small-scale analyses, while some training-data dependency methods rely on only a fraction of pretraining data.WIMBD supports efficient n-gram retrieval at scale; other cited work examines n-gram novelty, syntactic templates, or local text dependencies.Source: §8 · p. 10
• The paper analyzes zero-shot capabilities using WIMBD and ∞-gram frameworks to explore pretraining data at full scale.Source: §8 · p. 10
• Counterfactual memorization measures performance differences when training examples are included or excluded, but retraining requirements make this approach infeasible for large-scale LLM analysis.The paper proposes distributional memorization based on n-gram counts as a scalable alternative.Source: §8 · p. 10

9 CONCLUSION

• Across four tasks evaluated with Pythia, LLMs tend to memorize more on simpler, knowledge-intensive tasks and generalize more on harder, reasoning-intensive tasks.The task-gram language model is presented as a scalable approach to modeling task-relevant pretraining-data distributions.Source: §9 · p. 10

A LIMITATIONS

• Pythia and the Pile are somewhat outdated, while limited training data and checkpoints for newer open models constrain scaling studies.WIMBD also needs improved search and retrieval for larger corpora such as Dolma, at 3T tokens.Source: A LIMITATIONS · p. 14
• Task-relevant n-gram pair quality is highly sensitive to filtering, and better methods could improve the analysis.The paper identifies improved filtering as future work.Source: A LIMITATIONS · p. 14

B ETHICS STATEMENT

• Understanding how LLMs use pretraining data may help researchers diagnose and mitigate bias, robustness, and unexpected-behavior issues.The ethics statement links the extended memorization definition to improved interpretability.Source: B ETHICS STATEMENT · p. 14
• The findings raise privacy and security concerns about possible training-data recall and underscore the need for rigorous data handling and anonymization.The statement also connects this awareness to policy and data-use practices.Source: B ETHICS STATEMENT · p. 14
• Understanding model capabilities and limitations may inform business and policymaker deployment decisions, with the stated aim of more reliable and trustworthy AI.The ethics statement frames this as a possible societal and economic implication.Source: B ETHICS STATEMENT · p. 14

C RELATED WORK DETAILS

• Prior studies link pretraining structure and distributions to reasoning, in-context learning, and mathematical abilities, including aggregation of reasoning paths and local structure.The cited work also discusses latent-variable inference and parallel structures in training data.Source: C RELATED WORK DETAILS · p. 14
• Other scalable approaches search corpus n-grams, but local-context distribution models may not capture complex LLM capabilities.The discussion mentions embedding-based estimation over around 0.3% of pretraining data and retrieval systems for larger-scale n-gram search.Source: C RELATED WORK DETAILS · p. 15
• The paper aims to analyze the origins of zero-shot capabilities through full searches across the pretraining corpus using WIMBD and the ∞-gram framework.Source: C RELATED WORK DETAILS · p. 15
• Unlike verbatim-recall accounts and counterfactual measures, this work proposes distributional memorization based on n-gram counts for large-scale LLM analysis.The paper motivates the change by noting that counterfactual analysis requires retraining an LLM from scratch for each data point.Source: C RELATED WORK DETAILS · p. 15

D EXPERIMENT DETAILS

• The dataset choices balance task types, and the authors report that GSM8K results align with their earlier findings.The appendix also captions Figure 7 as comparing GSM8K accuracy with n-gram pair count across OLMo model sizes in Dolma.Source: D EXPERIMENT DETAILS · p. 15, Figure 7 · p. 17
• The study uses LASER for translation and E5 for other tasks, with lower cosine-similarity thresholds for larger n-gram pairs.Figure 6 reports that the TriviaQA memorization trend does not change across γ thresholds 0.7, 0.75, and 0.8.Source: D EXPERIMENT DETAILS · pp. 15–16, Figure 6 · p. 16
• Task-gram construction filters supervised input-output n-gram pairs by embedding cosine similarity, then retains pairs with nonzero pretraining-corpus counts.WMT uses Europarl and LASER embeddings; TriviaQA uses its training set, while MMLU and GSM8K use test sets for mining.Source: D EXPERIMENT DETAILS · pp. 15–16

E PROMPT OPTIMIZATION DETAILS

• Prompt optimization targets task clarity while adjusting a memorization score, with lower scores encouraging novel outputs and higher scores encouraging alignment with learned knowledge.The appendix includes a table captioned as prompts optimized for memorization and generalization on TriviaQA and GSM8K.Source: Table 2 · p. 17, E PROMPT OPTIMIZATION DETAILS · p. 18

F n-GRAM PAIR EXAMPLES

• The paper presents representative mined-pair examples from translation and question-answering, varying model sizes and n-gram lengths.Examples are captioned for translation at 2-gram and 4-gram lengths, MMLU at bigram and 5-gram lengths, and TriviaQA at trigram length.Source: F n-GRAM PAIR EXAMPLES · p. 18, Figure 8 · p. 19, Figure 9 · p. 19

G MODEL GENERATIONS VS. n-GRAM FREQUENCY IN PRETRAINING

• As n-gram frequency increased, generated model pairs were more aligned with pairs found in pretraining.The examples span translation and question-answering tasks, with different model sizes and n-gram lengths.Source: G MODEL GENERATIONS VS. n-GRAM… · p. 18, Figure 12 · p. 21