GENERALIZATION V.S. MEMORIZATION: TRACING LANGUAGE MODELS’ CAPABILITIES BACK TO PRETRAINING DATA

🎯 Pitch

Larger models do not gain abilities uniformly—performance on factual QA improves through memorization of training data, while translation and math reasoning improve via generalization, producing novel outputs. By measuring the correlation between model output probabilities and pre-training data frequency using a novel task-gram language model, the study reveals a stark capability divide: knowledge-intensive tasks rely on rote recall, but reasoning-intensive tasks do not.

1. Executive Summary

This paper introduces distributional memorization, an extended concept that measures the correlation between an LLM’s output probabilities and the frequency of task-relevant patterns in its pretraining data, alongside a novel task-gram language model for capturing those frequencies by counting co-occurrences of semantically related n-gram pairs from task inputs and outputs (e.g., translation phrase pairs, question-answer n-gram pairs). Using the Pythia model family trained on the Pile dataset, the authors evaluate four tasks—machine translation (WMT), factual question answering (TriviaQA), world knowledge (MMLU), and math reasoning (GSM8K)—finding that factual QA exhibits the strongest memorization effect with significant Spearman correlations (>0.35 for 3-gram task-gram LM), while translation and reasoning show no significant memorization and instead exhibit increased distributional generalization (the divergence between LLM output distribution and pretraining data distribution) as model size grows. The headline finding is that larger models improve on knowledge-intensive tasks primarily through increased memorization—TriviaQA’s distributional memorization rises with model size—while harder reasoning tasks achieve gains through greater generalization, establishing that memorization drives capability improvements only when tasks rely on factual recall from training data.

2. Context and Motivation

The Core Problem: We Don't Know Whether LLM Capabilities Come from Memorization or Genuine Generalization

The fundamental question this paper tackles is deceptively simple yet has profound implications: when a large language model performs impressively on a task—translating a sentence, answering a trivia question, solving a math problem—is it genuinely understanding and reasoning, or is it simply reproducing patterns it saw during pretraining? This question matters because the answer determines whether we should expect LLMs to be robust on genuinely novel inputs, whether they can be trusted in high-stakes settings, and what kind of progress we can expect from simply scaling up models and data.

The paper frames this as a tension between two competing hypotheses about the source of LLM capabilities. On one side, the memorization hypothesis suggests that LLMs function primarily as vast pattern stores, retrieving and recombining text fragments from their training data. Under this view, impressive task performance reflects the density of task-relevant patterns in the pretraining corpus rather than any deeper capability. On the other side, the generalization hypothesis suggests that LLMs learn transferable skills—linguistic rules, reasoning strategies, cultural knowledge structures—that can be applied to inputs that differ substantially from any specific training example.

This tension is not merely academic. It has direct consequences for:

Deployment safety: If LLMs predominantly memorize, they may fail catastrophically on inputs that deviate from their training distribution, making them unreliable for safety-critical applications like medical diagnosis or legal reasoning.
Privacy and data rights: Strong memorization implies that training data can be extracted from models, raising concerns about the inadvertent inclusion of personal information, copyrighted material, or toxic content in pretraining corpora.
Scaling expectations: If performance gains from larger models primarily reflect increased memorization capacity, then scaling may asymptotically hit fundamental limits—you can only memorize so much, and at some point the model must genuinely generalize to improve. If, instead, larger models improve through better generalization, scaling may continue to yield benefits on novel tasks.
Scientific understanding of intelligence: The memorization-versus-generalization question in LLMs connects to longstanding debates in cognitive science and machine learning about the nature of intelligence itself—whether it is fundamentally about pattern retrieval or about abstract reasoning.

The Gap: Existing Memorization Definitions Are Too Narrow

Prior to this work, the study of memorization in LLMs was dominated by a particular operational definition: verbatim recall. That is, a model "memorizes" a piece of text if it can reproduce that text exactly when prompted with a prefix. This definition has generated valuable findings—Carlini et al. (2022) demonstrated that exact copy-and-paste behaviors are more prevalent in larger models, Zhang et al. (2023) showed that LLMs can memorize rare long-tail data including private information, and Jiang et al. (2024) studied how test set contamination manifests as verbatim reproduction.

However, the paper argues—and this is the key gap it identifies—that verbatim recall is too restrictive to capture the kind of memorization that matters for task performance. Consider the following scenario: a pretraining corpus contains a document that says "Marie Curie discovered radium and polonium, winning Nobel Prizes in both Physics and Chemistry." A test question asks "What elements did Marie Curie discover?" The model answers "radium and polonium." There is no verbatim overlap between the training sentence and the test output, yet the model is clearly retrieving factual knowledge from pretraining rather than reasoning from first principles. By a verbatim recall definition, this would not count as memorization—yet it is exactly the kind of memorization that underpins factual question answering.

More generally, the paper identifies several limitations with the verbatim recall framing:

1. Exact reproduction is rare for high-level capabilities. When an LLM performs translation, reasoning, or produces short factual answers, it rarely outputs text that appears verbatim in pretraining data. The pretraining corpus might contain a document discussing Marie Curie's discoveries in French, and the test question asks about them in English. A verbatim search would find no match, but the model's ability to answer depends on having encountered the factual content during training.

2. The relationship between pretraining data and task output is often long-range and structural, not local and lexical. Consider a translation example where the source sentence contains "economic growth" and the target contains "croissance économique." These phrases might co-occur in the same document (a bilingual parliamentary proceeding) but separated by paragraphs—or even in different aligned segments within the same document. A standard n-gram model looking at local windows of 3–5 words would never capture this dependency. The paper argues that task-relevant supervision signals exist as correlated n-gram pairs across potentially large distances within documents, and capturing these requires modeling co-occurrence, not contiguity.

3. Counterfactual memorization approaches don't scale. Several works (Feldman, 2020; Feldman and Zhang, 2020; Zhang et al., 2023) have defined memorization via counterfactuals: a training example is "memorized" if removing it from the training set changes the model's prediction on that example. This is theoretically elegant but practically infeasible for large-scale LLMs because it requires retraining the model from scratch to analyze a single data point. The paper explicitly notes this constraint:

"this definition is impractical for large-scale analysis of LLMs as it requires retraining an LLM from scratch to analyze one data point."

4. Existing pretraining data analysis tools can't capture task-relevant patterns. The paper acknowledges recent advances in infrastructure for searching massive pretraining corpora—specifically WIMBD (Elazar et al., 2024) for efficient n-gram retrieval and the ∞-gram framework (Liu et al., 2024) for building unbounded n-gram language models. However, these tools are designed to capture local contextual dependencies: the ∞-gram model predicts each token based on the longest matching prefix in the training corpus, which models contiguity within a single sequence. This cannot capture the long-range, cross-sequence patterns that characterize task-relevant pretraining data. As the paper puts it:

"such distribution only models local contextual dependency, which might not be useful for understanding complex LLM capabilities."

The Conflation Problem: Why Prior Work Reached Conflicting Conclusions

A subtle but important motivation running through this paper is what we might call the conflation problem in prior memorization research. Studies that defined memorization narrowly (via verbatim recall) naturally found less evidence for it, particularly when examining complex capabilities. Studies that defined it more broadly (via counterfactual influence) could detect more nuanced forms of memorization but couldn't scale to realistic settings. The field lacked a definition that was simultaneously:

Broad enough to capture the practical forms of memorization that matter for task performance (retrieval of facts, reuse of translation patterns, recognition of problem structures).
Scalable enough to apply to full-sized LLMs with trillion-token pretraining corpora without requiring model retraining.
Precise enough to distinguish between different types of tasks and quantify how the memorization-generalization balance shifts with model scale.

This paper's core motivation is to fill that gap: to provide a framework that can quantify the degree to which any task capability originates from pretraining data patterns, at the scale of real-world LLMs, without requiring retraining.

How This Paper Positions Itself

The paper positions its contribution along several axes:

1. A new definitional framework for memorization that captures task-relevant dependency on pretraining data. Rather than asking "did the model reproduce this exact string?", the paper asks "does the model's probability distribution over outputs—the chance it assigns to different possible answers—correlate with the frequency of those answers (or their components) in the pretraining data?" This is distributional memorization: a measure of how much the LLM's behavior can be predicted by a language model constructed directly from task-relevant n-gram frequencies in the pretraining corpus.

The crucial shift is from a binary notion (memorized vs. not memorized) to a continuous one (how strongly does pretraining data frequency predict LLM output probability?). This allows the paper to measure degrees of memorization across different tasks, model sizes, and n-gram types, rather than simply classifying examples.

2. A scalable method for modeling task-relevant pretraining data distributions. The paper introduces the task-gram table and task-gram language model as practical tools for capturing the long-range, cross-sequence patterns that standard n-gram models miss. A task-gram table is built by: (a) extracting all candidate n-grams from task inputs and their corresponding outputs, (b) filtering for semantically similar pairs using embedding cosine similarity, and (c) counting how often these pairs co-occur within the same documents in the pretraining corpus. This produces an empirical conditional probability $P(\text{output n-gram} \mid \text{input n-gram})$ that captures task-relevant supervision signals from the pretraining data, even when the input and output n-grams appear far apart in the original document.

This is explicitly inspired by phrase tables from statistical machine translation (Passban et al., 2016), which captured translation pairs from parallel corpora. The paper generalizes this idea to arbitrary tasks: any task that can be characterized by input-output n-gram pairs can have its pretraining data distribution modeled via a task-gram LM. For TriviaQA, the task-gram pairs are question-answer n-gram co-occurrences; for MMLU, they are question stem-answer choice co-occurrences; for GSM8K, problem text-solution step co-occurrences.

3. An empirical taxonomy of when memorization versus generalization drives capability improvements. Rather than arguing that LLMs are "just memorizing" or "truly understanding," the paper takes the more nuanced position that the answer depends on the task and the model scale. It predicts—and then empirically demonstrates—that:

Knowledge-intensive tasks with high-frequency factual content (TriviaQA) show strong and increasing memorization as models grow larger. Better performance comes from better recall of training data patterns.
Hard reasoning tasks (GSM8K, reasoning-intensive MMLU) show decreasing memorization and increasing generalization as models grow larger. Better performance comes from the ability to combine and transform training patterns in novel ways, not from closer adherence to training distributions.
Translation occupies an intermediate position: it shows no significant memorization effect at any model size (the LLM output distribution is not predicted by pretraining frequency), consistent with translation being a learned skill that transfers across language pairs rather than a retrieval task.

This taxonomy moves the debate from a binary "memorization vs. generalization" framing to a more nuanced understanding of when each is the primary mechanism, which has direct implications for model design, deployment, and evaluation strategies.

4. Anchoring correlation findings with causal evidence. The paper goes beyond mere correlation between pretraining frequency and LLM behavior by providing gradient-based influence analysis (following Pruthi et al., 2020). This estimates how much individual pretraining documents affect test example predictions through the training process, using dot products between training loss gradients and test loss gradients accumulated across model checkpoints. The finding that documents containing task-gram pairs exert systematically more influence on test predictions than documents containing only individual output n-grams provides causal evidence—not just correlational—that the n-gram pairs identified by the task-gram table are genuinely important for model behavior.

This combination of distribution-level correlation analysis with instance-level gradient influence analysis represents a multi-pronged approach to tracing capabilities back to pretraining data, providing converging evidence that strengthens the central claims.

Connection to Broader Debates

The paper explicitly situates itself within several ongoing conversations in the field:

The scaling laws debate. The finding that memorization increases with model size for some tasks but decreases for others provides nuance to the observation by Carlini et al. (2022) that verbatim recall is more common in larger models. It suggests that "larger models memorize more" is task-dependent—larger models do encode more of their training distribution, but they also develop stronger generalization capabilities that can override memorization for tasks that require novel combinations of knowledge.

The emergence debate. Recent work has debated whether LLM capabilities "emerge" discontinuously with scale or develop gradually (Srivastava et al., 2024). The paper's framework offers a way to investigate this: if a capability emerges at a certain scale, does it emerge because the model suddenly has enough capacity to memorize the relevant training patterns, or because it develops a generalization mechanism that didn't exist at smaller scales? The distributional memorization metric provides a tool for answering this question.

The training data attribution problem. A growing body of work seeks to attribute model outputs to specific training examples, both for interpretability and for legal/compliance reasons (e.g., copyright, data licensing). The paper's task-gram approach offers a way to identify which pretraining documents contain task-relevant patterns, which could be extended to attributing specific model outputs to specific training documents—a capability the gradient influence analysis in Section 6 begins to explore.

3. Technical Approach

3.1 Reader Orientation

This paper builds a scalable analysis framework — not a new model or training method, but a set of tools and metrics for understanding why existing LLMs perform as they do on different tasks. The core system is a pipeline that takes a supervised task dataset and a pretraining corpus, then quantifies the degree to which the LLM's predictions are explained by patterns in the training data (memorization) versus deviations from those patterns (generalization). The problem it solves is the inability to distinguish these two sources of capability at the scale of real-world LLMs without expensive retraining.

3.2 Big-Picture Architecture (Diagram in Words)

The analysis pipeline has five major components connected in sequence:

Task Data — a supervised dataset of input-output pairs for a specific task (e.g., English sentences paired with French translations, trivia questions paired with short answers). This defines what "correct behavior" looks like for the task.
Task-Gram Table Constructor — takes the task data and produces a filtered set of semantically related n-gram pairs (input n-gram, output n-gram) that characterize the task. Uses embedding cosine similarity to determine which n-gram pairs are meaningfully related.
Pretraining Corpus Search Engine — using WIMBD (Elazar et al., 2024) and ∞-gram (Liu et al., 2024), searches the full pretraining corpus to count how often each n-gram pair co-occurs within the same document, and how often individual n-grams appear. This produces the empirical frequency data.
Task-Gram Language Model — converts the co-occurrence counts into a conditional probability distribution: for each input n-gram, what is the probability of seeing a particular output n-gram in the same document? This is a simple empirical LM built from the task-gram table frequencies.
Memorization Evaluator — takes the task-gram LM probabilities and the LLM's predicted probabilities on test examples, computes their Spearman rank correlation, and reports this as the distributional memorization score. Statistical significance tests separate real effects from noise.

Information flows left to right: task data → task-gram table → pretraining corpus search → task-gram LM → correlation with LLM predictions → memorization score. A separate gradient-influence analysis (Section 6) provides complementary causal evidence by estimating how much individual pretraining documents influenced test predictions during training.

3.3 Roadmap for the Deep Dive

First, the task-gram table construction (Definition 1) — because this is the novel representational innovation that makes everything else possible, and understanding it requires knowing what n-gram pairs are, how they're filtered, and why cosine similarity is used.
Second, the task-gram language model (Definition 2) — because this converts raw n-gram pair counts into a probability distribution that can be compared with LLM predictions, and its conditional probability formulation is the foundation for the memorization metric.
Third, the LLM probability extraction procedure (Equation 2) — because comparing two probability distributions requires defining what "LLM probability" means for a given output n-gram, and the autoregressive chain rule formulation has subtle implications.
Fourth, the distributional memorization metric itself (Definition 3 and Equation 3) — because this is the paper's primary analytical tool, and understanding why Spearman correlation (rather than Pearson, KL divergence, or other metrics) is chosen clarifies what the metric does and doesn't measure.
Fifth, the alternative ∞-gram comparison — because the paper explicitly contrasts its task-gram LM against traditional n-gram LMs to demonstrate that capturing long-range, cross-sequence dependencies matters.
Sixth, key design choices and hyperparameters — because the sensitivity of the task-gram table to embedding model choice, cosine similarity threshold, and n-gram size determines what the analysis can and cannot detect.

3.4 Detailed, Sentence-Based Technical Breakdown

This is primarily a measurement and analysis paper whose core idea is that the degree to which LLM capabilities originate from pretraining data can be quantified by correlating LLM output probabilities with an n-gram language model constructed from task-relevant, cross-sequence co-occurrence patterns in the pretraining corpus.

The foundational component of the analysis pipeline is the task-gram table, which is a filtered set of semantically related n-gram pairs extracted from a supervised task dataset. The purpose of this table is to characterize what kinds of input-output relationships define a particular task, so that the pretraining corpus can be searched for evidence of those relationships. The table is formally defined in Definition 1.

Step 1: Extract all candidate n-gram pairs from task data.

Given a task dataset DT = {(xi, yi)}i consisting of input-output pairs (e.g., English source sentence xi and French target sentence yi for translation, or question xi and answer yi for QA), the paper first extracts all possible n-grams from both sides. For each example (x, y), let Gn(x) denote the set of all contiguous n-grams (sequences of n tokens) in the input x, and Gn(y) denote all n-grams in the output y. The set of all possible input-output n-gram pairs across the entire dataset is:

$A_n(T) = \bigcup_i \left[ G_n(x_i) \times G_n(y_i) \right]$

where An(T) is the set of all candidate n-gram pairs for task T using n-gram size n, Gn(x_i) is the set of all n-grams in the i-th input example, Gn(y_i) is the set of all n-grams in the i-th output example, and × denotes the Cartesian product (all combinations of one input n-gram and one output n-gram from the same example). The union runs over all examples in the task dataset.

What it computes: for each training example, it takes every possible n-word chunk from the input, every possible n-word chunk from the output, and pairs them together, then collects all such pairs across the entire dataset. For a single translation example with 10 words on the source side and 8 words on the target side, using 3-grams, this produces roughly (10-3+1) × (8-3+1) = 8 × 6 = 48 candidate pairs from that one example alone. Across a dataset of thousands of examples, the total number of candidate pairs is enormous.

Why this form: pairing every input n-gram with every output n-gram from the same example captures all possible associations that a learner might extract — including incorrect or spurious ones. The subsequent filtering step (cosine similarity) is what separates the meaningful associations from the noise. The Cartesian product is deliberately exhaustive because the paper wants to capture a wide range of potential task-relevant patterns; filtering after extraction is more conservative than pre-filtering and potentially missing important relationships.

Step 2: Filter by semantic similarity using cosine similarity threshold.

The raw Cartesian product An(T) contains mostly meaningless pairs — "the cat sat" on the input side paired with "le chat était" on the output side might be meaningful, but "the cat sat" paired with a random function word from the output is not. To filter for semantically related pairs, the paper uses a threshold on the cosine similarity between n-gram embeddings. The formal definition is:

$H_n(T) = \{(s_j^x, s_j^y) \mid \cos(E(s_j^x), E(s_j^y)) > \gamma_T, s^x \neq s^y, (s_j^x, s_j^y) \in A_n(T)\}$

where H_n(T) is the task-gram table for task T with n-gram size n, s_j^x and s_j^y are specific n-gram strings from the input and output respectively, E(·) is a pretrained text embedding model that maps any text string to a dense vector, cos(·, ·) is the cosine similarity between two embedding vectors (ranging from -1 to 1), γ_T is a task-specific threshold hyperparameter between 0 and 1, and the constraint s^x ≠ s^y excludes pairs where the input and output n-grams are identical (which would capture simple repetition rather than task-specific relationships).

What it computes: for each candidate pair, both n-grams are passed through a frozen embedding model to get dense vector representations. The cosine similarity between these vectors is computed. If the similarity exceeds the threshold γ_T, the pair is included in the task-gram table; otherwise, it is discarded. The constraint s^x ≠ s^y removes trivial pairs where the same string appears on both sides (e.g., proper names that appear unchanged in both source and target).

Why this form: cosine similarity over pretrained embeddings captures semantic relatedness rather than surface-level lexical overlap. This is crucial because task-relevant relationships are often paraphrastic or cross-lingual — "economic growth" and "croissance économique" share no tokens but should be treated as related, while "economic growth" and "economic decline" share tokens but are semantically opposite. The embedding-based filter captures the former and rejects the latter. Identity pairs are excluded because they don't provide information about the task — they just reflect that some strings appear unchanged in both input and output, which is a property of the data format (e.g., numbers in math problems) rather than a task-relevant pattern.

Embedding model choices. The paper uses different embedding models for different tasks, selected based on the nature of the task:

For machine translation (WMT), it uses LASER embeddings (Schwenk and Douze, 2017), which are language-agnostic sentence representations specifically designed to map sentences in different languages to similar vectors when they have the same meaning. This is appropriate because translation pairs are cross-lingual by definition, and a general-purpose English embedding model would fail to recognize that French and English sentences are semantically equivalent.
For all other tasks (TriviaQA, MMLU, GSM8K), it uses E5 embeddings (Wang et al., 2022), which are general-purpose multilingual sentence representations trained with contrastive learning on a diverse range of tasks. E5 is chosen because these tasks are primarily English and involve finding question-answer or problem-solution relationships rather than cross-lingual equivalence.

Cosine similarity thresholds. The paper specifies task-specific thresholds that decrease as n-gram size increases:

For WMT (translation): the thresholds are 0.85 for 2-grams, 0.8 for 3-grams, 0.75 for 4-grams, and 0.7 for 5-grams. The threshold decreases with n because longer n-grams are inherently more specific and less likely to randomly match, so requiring the same similarity threshold would be too strict.
For TriviaQA and MMLU: the thresholds are 0.75 for 3-gram pairs and 0.65 for 5-gram pairs. The absolute values are lower than for WMT because these tasks involve question-answer relationships rather than translation equivalence, which is a looser semantic relationship.

Data sources for task-gram mining. The supervised data used to construct the task-gram table varies by task:

For WMT, the paper uses the Europarl corpus (Koehn, 2005), which consists of approximately 20 million parallel sentences extracted from the proceedings of the European Parliament. This is chosen because it is a large, high-quality parallel corpus that provides rich phrase-pair statistics for translation. It is separate from the WMT-09 test set used for evaluation.
For TriviaQA, the paper uses its training set to mine n-gram pairs. This contains question-answer pairs that characterize the task.
For MMLU, since the training set is very small (100–500 examples per subtask, for 57 subtasks), the paper mines n-gram pairs directly from the test set. This is a practical compromise driven by data availability, but it raises a methodological concern: the n-gram pairs that define the task are extracted from the same data that will be used for evaluation, potentially creating a circularity where the task-gram table is tailored to the test examples. The paper does not discuss this potential issue.
For GSM8K, the paper mines n-gram pairs from the 1K test set directly, similar to MMLU.

Final step: filter by occurrence in pretraining data. After constructing the initial task-gram table via embedding similarity filtering, the paper performs a final filtering step: it searches the pretraining corpus for each n-gram pair and retains only those pairs that have nonzero co-occurrence count in the corpus. This is a practical step — pairs that appear in the supervised data but never co-occur in pretraining cannot contribute to the memorization analysis, and including them would artificially depress the correlation by adding zero-probability entries that the LLM couldn't possibly have memorized. The paper states that this reduces the size of the task-gram table to the set of pairs that are actually attested in the pretraining data, making the subsequent language model more meaningful.

The Task-Gram Language Model: From n-Gram Co-occurrence Counts to Conditional Probabilities

Once the task-gram table H_n(T) is constructed, the next step is to build a language model from it. This model estimates the probability of seeing a particular output n-gram s^y given that a particular input n-gram s^x appears in the same document, based purely on empirical frequencies in the pretraining corpus. This is formalized in Definition 2.

Counting functions. The paper defines two counting operations on the pretraining corpus D:

C((s^x, s^y), D) — the number of documents in the pretraining corpus where the n-gram pair (s^x, s^y) co-occurs, meaning both n-grams appear somewhere within the same document (not necessarily adjacent or even close).
C(s^x, D) — the number of documents in the pretraining corpus where the input n-gram s^x appears at least once, regardless of whether s^y also appears.

Conditional probability definition. The task-gram language model computes the conditional probability of an output n-gram given an input n-gram as the ratio of these counts:

$P_{n,D}(s^y \mid s^x) = \frac{C((s^x, s^y), D)}{C(s^x, D)}$

where P_{n,D}(s^y|s^x) is the estimated probability that output n-gram s^y appears in a document given that input n-gram s^x appears in the same document, C((s^x, s^y), D) is the co-occurrence count of the pair, and C(s^x, D) is the occurrence count of the input n-gram alone. This probability is defined only for pairs (s^x, s^y) that are in the task-gram table H_n(T).

What it computes: for a specific pair of semantically related n-grams, it divides the number of pretraining documents containing both by the number containing the input n-gram. If s^x appears in 1000 documents and s^y appears alongside it in 50 of those documents, then P_{n,D}(s^y|s^x) = 50/1000 = 0.05. This is a straightforward maximum-likelihood estimate of the conditional probability — it asks: "in the pretraining data, when you encounter this input pattern, how often do you also encounter this output pattern in the same document?"

Why this form: the ratio estimator is the empirical frequency, which is the simplest possible conditional probability model. No smoothing, no backoff, no neural parameterization — this is deliberately kept simple because the goal is to measure the raw statistical signal in the pretraining data, not to build an accurate language model. Any smoothing would obscure the relationship between pretraining frequency and LLM behavior that the paper is trying to measure. The document-level co-occurrence (rather than sentence-level or adjacent-window) is chosen because pretraining data is not aligned at the sentence level — a document might discuss a concept on one page and mention the related fact on another page. Using document-level co-occurrence captures these long-range dependencies that would be missed by local window approaches.

Critical distinction from standard n-gram LMs. A standard n-gram language model estimates P(token_i | token_{i-1}, ..., token_{i-n+1}) — the probability of the next token given the immediately preceding n-1 tokens, within a single contiguous sequence. The task-gram LM estimates P(output_n-gram | input_n-gram) — the probability of an output text fragment given an input text fragment, where the two fragments appear anywhere within the same document, potentially separated by thousands of tokens. This is a fundamentally different conditioning structure: standard n-gram LMs model local sequential dependencies, while the task-gram LM models long-range, cross-sequence, task-relevant associations.

Relationship to phrase tables in statistical machine translation. The paper explicitly draws inspiration from phrase tables used in statistical machine translation (Passban et al., 2016), which store translation pairs and their probabilities extracted from parallel corpora. Just as a phrase table captures that "economic growth" in English corresponds to "croissance économique" in French with some estimated probability, the task-gram table captures that an input n-gram corresponds to an output n-gram based on document co-occurrence in a (potentially non-parallel) pretraining corpus. The key difference is that phrase tables are extracted from explicitly aligned parallel data, while task-gram tables are extracted from unaligned, monolingual or mixed-language documents using embedding similarity as a proxy for alignment.

Search infrastructure. To actually compute these counts over a 207B-token corpus like the Pile, the paper uses two complementary search systems:

WIMBD (What's In My Big Data?, Elazar et al., 2024) is used for accurately counting the co-occurrence of n-gram pairs. WIMBD is chosen for this task because co-occurrence frequencies are typically low (many pairs appear in only a handful of documents), and the approximate counting in ∞-gram sometimes fails to capture these low-frequency events. WIMBD provides exact counts by maintaining an index of all n-grams in the corpus and their document locations.
∞-gram (Liu et al., 2024) is used for counting the occurrence of single n-grams, which appear much more frequently in the corpus, making approximation errors less consequential. ∞-gram uses a prefix-based efficient retrieval system that can handle queries about n-grams of effectively unbounded length.

Extracting LLM Probabilities for Output n-Grams

To compare the task-gram LM's probabilities with the LLM's behavior, the paper needs to define what "the LLM's probability" means for a given output n-gram s^y conditioned on an input containing s^x. This is done via the chain rule of probability applied to the LLM's autoregressive token predictions.

LLM probability definition. For a test example where the full input to the LLM is the concatenation of an instruction u, the task input x, and a prefix of the output y[1:m-1] (everything before the target n-gram), the probability assigned by the LLM to the output n-gram s^y is:

$P_{LLM}(s^y \mid s^x) = \prod_{t \in s^y} P_{LLM}(t \mid u \oplus x \oplus y[1:m-1])$

where s^y is the target n-gram (a sequence of tokens), t is each individual token in that n-gram, u is the instruction text (a minimal template indicating the task, e.g., "Translate English to French:"), x is the full task input, y[1:m-1] is the prefix of the output before the n-gram s^y starts at position m, and P_{LLM}(t | context) is the probability the LLM assigns to token t given the preceding context, as output by the model's softmax layer. The product runs over all tokens in s^y.

What it computes: the probability of generating the exact output n-gram s^y as part of the model's response, conditioned on the instruction, the full input, and any output tokens that precede this n-gram. If s^y is a 3-word phrase like "won the Nobel," and it appears as the 5th, 6th, and 7th tokens of the model's output, then the model's probability for this n-gram is the product of P(token_5="won" | instruction, input, tokens_1...4), P(token_6="the" | instruction, input, tokens_1...5), and P(token_7="Nobel" | instruction, input, tokens_1...6). These are the actual probabilities from the model's forward pass, not approximations.

Why this form: the chain rule decomposition is the standard way to extract sequence-level probabilities from autoregressive models. Using the actual model probabilities (rather than sampling or greedy decoding) captures the model's full uncertainty — even if the model doesn't ultimately generate this n-gram, the probability it assigns to it reflects how strongly it expects this output given the input. The inclusion of u (instruction) and the prefix y[1:m-1] as conditioning context mirrors the actual generation process: the model sees the instruction, the input, and any text it has already generated before predicting the next token.

Important subtlety — conditioning on the prefix matters. The probability P_{LLM}(s^y|s^x) is not simply the probability of generating s^y from scratch given s^x; it is the probability of generating s^y at the specific position m in the output, given everything that came before it. This means the LLM probability for a particular output n-gram depends on what other output n-grams have already been generated — there is an order dependence. This is appropriate because LLM generation is autoregressive, but it also means that two different test examples that share the same input n-gram s^x might yield different LLM probabilities for the same output n-gram s^y depending on the surrounding context.

The instruction template u. The paper states that it uses "a minimal instruction template to indicate the input and output" (Section 2), but does not provide the exact wording in the main text. The choice of template matters because the instruction affects the model's output distribution — a verbose instruction might push the model toward different output patterns than a minimal one. The paper's use of a minimal template is intentional: it avoids introducing additional distributional shift beyond what the task itself requires, keeping the analysis focused on the relationship between input content and output content rather than prompt engineering effects.

The Distributional Memorization Metric: Spearman Correlation Between Two Probability Distributions

With the task-gram LM probabilities P_{n,D}(s^y | s^x) and the LLM probabilities P_{LLM}(s^y | s^x) computed for a set of test n-gram pairs, the paper defines its central metric: distributional memorization. This is formalized in Definition 3.

Collecting test n-gram pairs. For a test set D'_T = {(x_i, y_i)}_i, the paper collects all n-gram pairs that appear in both the test examples and the task-gram table:

$\Phi = \{(s^x, s^y) \mid \forall (s^x, s^y) \in [G_n(x) \times G_n(y)] \cap H_n(T), \forall (x, y) \in D'_T\}$

where Φ is the set of all validated test n-gram pairs, G_n(x) × G_n(y) is the Cartesian product of all n-grams in the test example's input and output (same extraction as during task-gram construction), and H_n(T) is the task-gram table. The intersection ensures that only pairs that survived the embedding similarity filtering and have nonzero pretraining counts are included in the correlation analysis.

Computing the vectors. For each pair (s^x, s^y) in Φ, the paper computes two log-probabilities:

log P_{n,D}(s^y | s^x) — the log of the task-gram LM conditional probability.
log P_{LLM}(s^y | s^x) — the log of the LLM's predicted probability.

These are collected into two vectors: log P_{n,D}(Y|X) and log P_{LLM}(Y|X), each of length |Φ|. The log transform is standard practice for probability comparisons because raw probabilities are heavily skewed toward zero, making correlation metrics sensitive to outliers. Log probabilities spread out the low-probability region, making the comparison more informative.

The memorization metric. Distributional memorization is defined as the Spearman rank correlation between these two vectors:

$\text{Mem}_n(\text{LLM}, D \mid T) = \rho\left(\log P_{n,D}(Y \mid X), \log P_{LLM}(Y \mid X)\right)$

where Mem_n(LLM, D|T) is the distributional memorization of the LLM on task T using n-gram size n, ρ denotes Spearman's rank correlation coefficient, and the two arguments are the vectors of log-probabilities from the task-gram LM and the LLM respectively.

What it computes: Spearman's ρ measures the monotonic relationship between two variables — here, between the task-gram LM's estimated probability and the LLM's predicted probability across all test n-gram pairs. A Spearman correlation of +1 means that whenever the task-gram LM assigns a higher probability to one n-gram pair than another, the LLM also assigns a higher probability to the first pair (perfect monotonic agreement). A correlation of 0 means no monotonic relationship — the LLM's probability rankings are unrelated to the pretraining frequency rankings. A negative correlation would mean the LLM systematically assigns higher probability to pairs that are rarer in pretraining, which would be evidence of active suppression or generalization away from the training distribution.

Operationally, Spearman correlation works by:

Converting each probability vector to ranks (the highest probability gets rank 1, second-highest rank 2, etc.).
Computing the Pearson correlation between the ranks.
This makes the metric nonparametric — it doesn't assume linearity or any specific functional form for the relationship, only monotonicity.

Why this form — the choice of Spearman correlation over alternatives:

Spearman vs. Pearson correlation. Pearson correlation measures linear relationship. If the LLM's log-probability is a monotonically increasing but nonlinear function of the task-gram LM's log-probability (e.g., sigmoidal, step-like), Pearson correlation might be low even though there is a strong memorization signal — the LLM consistently gives higher probability to more frequent patterns, but not by a constant factor. Spearman correlation captures this by operating on ranks, detecting any monotonic relationship regardless of its functional form. This is important because there is no theoretical reason to expect the LLM's probability to be a linear function of pretraining frequency — neural network training dynamics could produce concave, convex, or threshold-like relationships.
Spearman vs. KL divergence. KL divergence measures the absolute difference between two probability distributions, not their correlation. A model could have zero KL divergence from the pretraining data distribution (perfect memorization in an absolute sense) and still have zero KL divergence — but Spearman would detect the agreement. Conversely, a model could have nonzero KL divergence (the distributions differ in absolute terms) while still being perfectly correlated with the pretraining distribution in terms of which outputs are more likely — which Spearman would capture. The paper's interest is in the latter: does the LLM assign higher probability to answers that are more common in pretraining, regardless of the absolute probability levels?
Spearman vs. counterfactual memorization. Counterfactual approaches (Feldman and Zhang, 2020) define memorization as the performance difference when a specific training example is removed, requiring model retraining. The Spearman approach is correlational — it doesn't establish that pretraining data causes the LLM's behavior, only that the two are statistically associated. This is a limitation that the paper partially addresses with the gradient influence analysis in Section 6, which provides causal evidence from training dynamics. However, the correlational nature of the main metric means that high distributional memorization could in principle arise from a confounder — both the pretraining data and the LLM might independently reflect some underlying property of language that makes certain n-gram pairs inherently more probable, without the LLM having learned it from pretraining. The paper implicitly assumes that if the LLM's probabilities correlate with pretraining frequencies, it's because the LLM learned those frequencies during training — a reasonable assumption given that neural language models are known to approximate their training distributions, but not formally proven.

Statistical significance. The paper reports statistical significance using p-values from the Spearman correlation test. A result is considered "significant" if p < 0.05, and the paper marks significant values with solid round markers in its plots and insignificant values with gray star markers. This is important because the Spearman correlation itself can be nonzero purely by chance, especially when the number of test n-gram pairs |Φ| is small. The p-value quantifies the probability of observing a correlation at least as extreme as the measured one under the null hypothesis that the true correlation is zero.

Distributional generalization as the complement. The paper defines distributional generalization as "the opposite of distributional memorization" — meaning that when memorization is low or decreasing, generalization is high or increasing. There is no separate metric for generalization; it is inferred from the absence or decline of memorization. This is a deliberate simplification: the paper focuses on quantifying one end of the spectrum (memorization) and treats the residual as generalization, rather than attempting to define and measure both independently.

The Alternative: ∞-gram Language Model for Comparison

To demonstrate that the task-gram LM captures something that standard n-gram LMs miss, the paper defines an alternative memorization metric using the ∞-gram language model (Liu et al., 2024). This serves as a baseline that represents traditional, local-context approaches to modeling pretraining data distributions.

The ∞-gram LM definition. An ∞-gram LM predicts each token based on the longest matching prefix that exists in the pretraining corpus, starting from n = ∞ and backing off when the count is zero. For a sequence of tokens t_1, t_2, ..., t_k (the concatenated instruction, input, and output prefix), the probability of an output n-gram s^y is:

$P_{\infty,D}(s^y \mid u \oplus x \oplus y) = \prod_{t_i \in s^y} P_{\infty,D}(t_i \mid t_{[1:i-1]})$

where each token-level probability is:

$P_{\infty,D}(t_i \mid t_{[1:i-1]}) = \frac{C(t_{[i-(n_i-1):i]})}{C(t_{[i-(n_i-1):i-1]})}$

Here, n_i is the size of the longest prefix of token t_i that can be found in the pretraining corpus — formally, n_i = \max\{n' \in [1, i] \mid C(t_{[i-(n'-1):i-1]}) > 0\}. If no prefix of any length exists, the count is zero, and the paper sets the probability to 1 for that token (so it doesn't zero out the entire product). The final n-gram probability is zero only if every token has zero probability — an extreme edge case.

What it computes: for each token in the output n-gram, it finds the longest sequence of preceding tokens (including the instruction, input, and already-generated output) that appears verbatim in the pretraining corpus, and uses the empirical frequency of the next token following that prefix as the probability. This is essentially a maximum-context n-gram model — it uses as much context as the training data supports, up to potentially thousands of tokens. If the sequence "The capital of France is" appears 1000 times in the pretraining data and is followed by "Paris" 900 times, then P_∞("Paris" | "The capital of France is") = 900/1000 = 0.9. If no document contains exactly this prefix, it backs off to a shorter prefix (e.g., "capital of France is") and continues backing off until it finds a match.

Why this is a meaningful comparison. The ∞-gram LM captures local sequential dependencies — it models the probability of text based on its immediate preceding context within a contiguous sequence. This is fundamentally different from the task-gram LM, which captures long-range, cross-sequence associations between input and output n-grams that may be separated by arbitrary distances within a document. The paper hypothesizes that for complex capabilities like translation and reasoning, the relevant pretraining signal is in these cross-sequence co-occurrences (e.g., seeing a question and its answer in the same document, even if separated by paragraphs), and that a local-context model like the ∞-gram LM will fail to detect this signal.

Alternative memorization metric using ∞-gram. The distributional memorization defined with the ∞-gram LM instead of the task-gram LM is:

$\text{Mem}_\infty(\text{LLM}, D \mid T) = \rho\left(\log P_{\infty,D}(Y \mid X), \log P_{LLM}(Y \mid X)\right)$

where P_{\infty,D}(Y|X) is the vector of ∞-gram LM log-probabilities for all test n-gram pairs. This is the same Spearman correlation computation, just with a different pretraining data model on the left side.

Key prediction. If the task-gram LM captures task-relevant patterns that the ∞-gram LM misses, then Mem_n should be larger than Mem_∞ for tasks where long-range co-occurrence patterns matter. The paper indeed finds that Mem_n(LLM, D|T) is consistently greater than or equal to Mem_∞(LLM, D|T) across TriviaQA and MMLU (Figure 4), confirming this prediction.

Handling zero probabilities. A practical complication: the ∞-gram LM may assign zero probability to some tokens if no prefix of any length exists in the pretraining corpus. The paper handles this by setting P_∞(t_i) = 1 for such tokens, effectively ignoring them in the product. This is a pragmatic choice — setting probability to zero would make the entire n-gram probability zero, making log-probability undefined and breaking the correlation analysis. The paper notes that P_{\infty,D}(s^y) = 0 only when all its tokens have zero probability, which is extremely rare for non-trivial n-grams.

Design Choices, Hyperparameters, and Their Justifications

Choice of n-gram sizes. The paper experiments with n = 2, 3, 4, and 5 for different tasks. The key finding is that n = 3 (trigram pairs) shows the strongest memorization signal for knowledge-intensive tasks (TriviaQA, MMLU knowledge-intensive subset, both > 0.25 Spearman correlation), while n = 5 is more informative for reasoning-intensive MMLU tasks. This is interpreted as: knowledge retrieval operates on relatively short, atomic fact units (3 words is enough to capture most entity names, relations, and properties), while reasoning involves larger conceptual units that require longer spans to be meaningful. A 3-word chunk from a math reasoning solution (e.g., "first we compute") is not informative about the reasoning step, while a 5-word chunk (e.g., "let x be the number") begins to capture the structure.

Choice of embedding models. The paper uses LASER for translation and E5 for everything else. The justification is task-specific: LASER is purpose-built for cross-lingual semantic similarity, which is exactly what translation pairs require. E5 is a general-purpose embedding model trained on diverse retrieval tasks, making it suitable for the monolingual question-answer and problem-solution relationships in TriviaQA, MMLU, and GSM8K. The paper does not ablate this choice — no experiments compare alternative embedding models — so the sensitivity of results to embedding quality is unknown.

Choice of cosine similarity thresholds. The thresholds (0.85 → 0.7 for translation, 0.75 → 0.65 for QA/MMLU, decreasing with n) are set as hyperparameters without automated tuning. The paper provides a sensitivity analysis in Figure 6 (Appendix D) showing that the distributional memorization trend on TriviaQA is "essentially unchanged" with thresholds γ ∈ {0.7, 0.75, 0.8}, suggesting the method is not highly sensitive to this parameter within a reasonable range. The decreasing threshold with n-gram size is motivated by the observation that longer n-grams are inherently more specific — the probability of randomly matching a 5-gram is much lower than randomly matching a 2-gram, so a lower similarity bar is appropriate to avoid over-filtering.

Choice of document-level co-occurrence. The paper defines co-occurrence at the document level (both n-grams appear in the same document) rather than at the sentence level, paragraph level, or within a fixed window. The justification is that pretraining data includes long documents where related information appears in different sections — a news article might introduce a concept in the opening paragraph and discuss its implications pages later. Sentence-level or window-based co-occurrence would miss these associations. The document-level choice is also practical: WIMBD indexes document-level occurrence, making it computationally feasible to count.

Choice of Spearman correlation over alternatives. Discussed in detail under the metric definition: Spearman captures monotonic relationships without assuming linearity, is nonparametric, and is robust to outliers in probability space (which is heavy-tailed due to the Zipfian nature of language). The log-transform further symmetrizes the distribution, making the correlation more informative in the low-probability regime where most n-gram pairs reside.

Choice of Pythia model family. The paper uses Pythia (Biderman et al., 2023) because it satisfies three requirements simultaneously: (1) fully open-source with available pretraining data (the Pile), (2) a wide range of model sizes from 13M to 12B parameters, enabling scaling analysis, and (3) multiple pretraining checkpoints available for the gradient influence analysis in Section 6. Most contemporary open-source LLMs (e.g., Llama, Mistral) do not release their pretraining data, making the paper's analysis impossible. The paper acknowledges that Pythia is "slightly outdated" (Section A) but argues that the availability of pretraining data and checkpoints makes it the best choice for this type of analysis. The paper also includes some results with OLMo models (1B and 7B) trained on Dolma (3T tokens) to provide a more contemporary comparison.

Task selection. The four tasks — WMT (translation), TriviaQA (factual QA), MMLU (world knowledge), and GSM8K (math reasoning) — are chosen to span a spectrum from "pure retrieval" (TriviaQA) to "pure reasoning" (GSM8K), with translation and MMLU occupying intermediate positions. This task diversity is essential for the paper's central claim that the memorization-generalization balance depends on task type. The paper also splits MMLU into knowledge-intensive and reasoning-intensive subtasks (the classification is provided in Appendix D, with 33 knowledge tasks and 24 reasoning tasks) to show that even within a single benchmark, different subtasks exhibit different memorization profiles.

Summary of the Pipeline in Action

To make the abstract definitions concrete, here is what happens when the pipeline runs on a specific task, say TriviaQA with n = 3:

Phase 1: Build the task-gram table. Take the TriviaQA training set (thousands of question-answer pairs). For each pair, extract all 3-grams from the question and all 3-grams from the answer. For a question like "What element did Marie Curie discover?" and answer "radium," the question 3-grams include "What element did," "element did Marie," "did Marie Curie," "Marie Curie discover," and "Curie discover?" — 5 trigrams. The answer "radium" might yield "radium" itself as a single token (treated as an n-gram of length 1 if it's shorter than n, though the paper's framework is designed for n-length units; for TriviaQA, the paper treats the whole answer as one n-gram since answers are typically short). Each question trigram is paired with each answer token/ngram. All pairs are embedded with E5 and filtered by cosine similarity > 0.75. The surviving pairs are then checked against the Pile: only pairs that co-occur in at least one Pile document are retained. The resulting task-gram table might contain entries like ("Marie Curie discover", "radium") with some co-occurrence count — perhaps 15 documents in the Pile contain both the trigram "Marie Curie discover" and the token "radium."

Phase 2: Build the task-gram LM. For each pair in the table, compute P(radium | "Marie Curie discover") = (number of documents containing both) / (number of documents containing "Marie Curie discover"). If "Marie Curie discover" appears in 200 documents and "radium" co-occurs in 15 of them, the probability is 15/200 = 0.075.

Phase 3: Extract LLM probabilities. For each test question where a task-gram pair is found, run the LLM with instruction "Answer the following question:" and the question text. For each position in the output where a task-gram output n-gram could appear, extract the product of token probabilities. If the LLM assigns probability 0.2 to generating "radium" at the answer position, then P_LLM("radium" | context) = 0.2.

Phase 4: Compute memorization. Collect all (log(P_taskgram), log(P_LLM)) pairs across all test examples. Compute Spearman's ρ. If the correlation is 0.35 with p < 0.001, the paper reports significant distributional memorization — the LLM's output probabilities are substantially predicted by pretraining co-occurrence frequencies.

Phase 5: Interpret. A positive, significant, and sizeable Spearman correlation means the LLM is behaving as if it has internalized the pretraining data distribution for this task — it assigns higher probability to answers that are more commonly associated with the question patterns in its training data. A near-zero or non-significant correlation means the LLM's behavior cannot be explained by pretraining frequency — it is generalizing beyond the training distribution.

4. Key Insights and Innovations

Innovation 1: A Practical, Scalable Definition of Memorization That Captures Task-Level Dependency on Pretraining Data

The paper's most fundamental intellectual contribution is the concept of distributional memorization itself — not as a method, but as a diagnostic framing that reframes the entire memorization-vs-generalization debate from a binary question ("does the model memorize?") to a continuous, measurable, and task-specific quantity ("to what degree does pretraining data frequency predict the model's output distribution on this task?").

Before this work, the field operated with a definition of memorization inherited from computer security and privacy research: a model "memorizes" a piece of data if it can reproduce it verbatim when prompted (Carlini et al., 2022; Zhang et al., 2023). This was natural for studying data extraction attacks and privacy leakage, but it was fundamentally mismatched with how memorization matters for capabilities. The paper identifies the gap precisely: when an LLM answers "radium" to "What element did Marie Curie discover?", it is clearly retrieving knowledge from pretraining, yet there may be zero verbatim overlap between any pretraining document and the test output. The verbatim recall definition would classify this as "not memorized," which is technically correct under its own terms but misses the phenomenon that actually drives task performance.

The alternative tradition — counterfactual memorization (Feldman, 2020; Feldman and Zhang, 2020) — captures a richer notion of memorization by measuring how model predictions change when specific training examples are removed. This is theoretically appealing because it's causal rather than correlational. But the paper explicitly argues this approach "is impractical for large-scale analysis of LLMs as it requires retraining an LLM from scratch to analyze one data point." The computational infeasibility of counterfactual methods has meant that memorization studies using this definition have been confined to small models and small datasets, leaving the question of memorization in production-scale LLMs essentially unaddressed.

Distributional memorization cuts through this dilemma by changing what is being measured. Instead of asking "did the model see this exact example?", it asks "does the model's behavior on this task reflect the statistical patterns of its training data?" The shift is from instance-level memorization (did the model memorize this specific training example?) to distribution-level memorization (did the model internalize the training data's conditional distribution over outputs?). This is simultaneously:

Scalable: it requires only n-gram counts from the pretraining corpus and forward passes through the trained LLM — no retraining, no gradient computation, no per-example interventions. This is what enables analysis of full-scale models (up to 12B parameters) on large test sets (thousands of examples).
Task-relevant: by conditioning on task-specific n-gram pairs (extracted from supervised task data rather than generic text), the metric captures the kind of memorization that matters for a given capability. The correlation between pretraining frequency and LLM probability on translation pairs answers a different question than the same correlation on trivia question-answer pairs, and the metric makes this distinction explicit.
Gradable: the Spearman correlation produces a continuous value rather than a binary classification, enabling the paper to compare degrees of memorization across tasks, model sizes, and n-gram types, and to track how memorization changes with scale.

This is a reframing rather than an incremental improvement. The paper doesn't argue that verbatim recall is wrong or that counterfactual memorization is invalid; it argues that we need a complementary concept to answer a different question — "where do LLM capabilities come from?" — and that distributional memorization is the right concept for that question. The distinction is analogous to the difference between studying whether a student cheated on a specific exam question (instance-level) versus studying whether their overall performance reflects the material they studied (distribution-level). Both are valid questions, but they require different measurement frameworks.

The practical significance of this reframing is demonstrated throughout the paper's results: distributional memorization reveals clear, interpretable patterns — TriviaQA shows strong and increasing memorization with model size, translation shows no significant memorization at any scale, MMLU splits along knowledge-vs-reasoning lines — that align with intuitive notions of which tasks "should" involve retrieval versus reasoning. A metric that produces intuitively plausible results across diverse tasks is not proof of correctness, but it is evidence that the concept captures something real about how LLMs work.

Innovation 2: The Task-Gram Language Model as a Method for Capturing Long-Range, Cross-Sequence Supervision Signals in Unaligned Pretraining Data

The distributional memorization metric requires a model of the pretraining data distribution — something that tells us how probable a given output is, conditioned on a given input, based on what the training corpus contains. The paper's second key innovation is recognizing that standard language models — even unbounded n-gram models — are structurally incapable of capturing this conditional distribution for task-relevant patterns, and that a fundamentally different architecture is needed, which it provides in the form of the task-gram language model.

To appreciate why this is innovative, consider what the state-of-the-art in pretraining data modeling was before this work. The ∞-gram language model (Liu et al., 2024) represented the frontier: it could use effectively unbounded context by finding the longest matching prefix in the training corpus, enabling it to capture long-range dependencies within a single contiguous sequence. This was a significant advance over fixed-width n-gram models, but it retained a fundamental limitation — it models sequential contiguity. The probability of a token depends on the tokens that immediately preceded it in the same text stream. For tasks like translation or question answering, the relevant relationship is not between adjacent tokens but between semantically related spans that appear in different parts of a document (or even in different documents), often separated by large distances and unrelated intervening text.

The task-gram LM solves this by modeling document-level co-occurrence rather than sequential adjacency. It counts how often an input n-gram and an output n-gram appear in the same document, regardless of their positions or the text between them, and converts this into a conditional probability. This is conceptually simple — it is essentially a maximum-likelihood estimate of a joint distribution using document co-occurrence as the dependency signal — but it represents a different modeling paradigm from standard language modeling. A standard LM answers "what word comes next in this sequence?" A task-gram LM answers "when this concept appears in a document, how often does that related concept also appear in the same document?"

The paper's empirical finding that the task-gram LM consistently produces higher and more significant memorization correlations than the ∞-gram LM (Figure 4, all panels where both are shown) validates that this structural difference matters in practice. The task-gram LM's Mem_n values are systematically larger than Mem_∞, and in some cases the ∞-gram version shows decreasing memorization with scale while the task-gram version shows stable or increasing memorization (e.g., TriviaQA with n=3 vs. ∞-gram in Figure 4, top middle panel). This means the ∞-gram LM is not just a noisy version of the task-gram LM — it is measuring a fundamentally different signal, one that is less aligned with what the LLM actually learns.

The intellectual move here is recognizing that task-relevant supervision in pretraining data has a particular structure — cross-sequence, semantically mediated, document-scoped — and that capturing it requires abandoning the sequential modeling paradigm that dominates language modeling. The task-gram table is explicitly inspired by phrase tables from statistical machine translation (Passban et al., 2016), but the innovation is generalizing this idea from the specialized setting of parallel corpora (where alignment is given) to the general setting of arbitrary pretraining data (where alignment must be inferred via embedding similarity and document co-occurrence). This transforms the phrase table from a task-specific engineering artifact into a general-purpose analytical tool for understanding any supervised task in terms of its pretraining data correlates.

Innovation 3: Empirical Evidence That Larger Models Improve Through Different Mechanisms for Different Tasks — A Task-Dependent Scaling Law for Capability Origins

The paper's third key contribution is an empirical finding that has the character of a scaling law for capability mechanisms: as model size increases, the source of performance improvement is fundamentally different for knowledge-intensive tasks versus reasoning-intensive tasks. This is not a metric or a method but a discovery about how scaling interacts with task type.

Prior work on memorization and scale had produced a relatively simple picture: larger models memorize more. Carlini et al. (2022) found that verbatim recall of training data is more prevalent in larger language models. Merrill et al. (2024) found that larger LLMs generate fewer novel n-grams, suggesting increased reliance on training data patterns. These findings could be interpreted as evidence that scaling primarily improves performance by increasing memorization capacity — bigger models are better at storing and retrieving their training data, and this improved retrieval drives their improved task performance.

The paper's results complicate this picture substantially. It finds:

For factual QA (TriviaQA): the simple picture holds. Distributional memorization increases with model size (Figure 4, top middle panel, solid lines), consistent with larger models better internalizing the statistical associations between questions and answers in their training data. Improved performance comes from improved memorization.
For knowledge-intensive MMLU subtasks: memorization decreases with model size (Figure 4, top right panel, solid lines). The larger models' improved performance on specialized knowledge tasks (college-level subjects, professional exams) cannot be explained by closer adherence to pretraining frequency — if anything, larger models deviate more from the pretraining distribution. The paper hypothesizes this is because MMLU involves rarer, more specialized knowledge than TriviaQA, requiring larger models to adjust or override the base frequencies from pretraining to answer correctly.
For reasoning-intensive MMLU subtasks and GSM8K: memorization is either decreasing (MMLU reasoning) or entirely absent (GSM8K, where no significant Spearman correlation is found for any model size). The policy implication is that larger models' improved reasoning comes from better generalization — the ability to combine and transform training patterns in novel ways — not from better retrieval of training patterns.
For translation (WMT): memorization is never significant at any model size (Figure 4, top left panel), yet performance improves substantially with scale (Figure 3, leftmost panel). The paper demonstrates that larger models generate more novel n-gram pairs not seen in any pretraining document (Figure 4, top left, blue line rising), providing direct evidence of increasing generalization.

This pattern — memorization increases for some tasks, decreases for others, and is absent for yet others — is the key empirical contribution. It implies that there is no universal answer to "do larger models memorize more?" The answer depends on what task you're asking about. For tasks that inherently rely on retrieving stored knowledge (factual QA), memorization is the mechanism of improvement. For tasks that require manipulating knowledge in novel ways (reasoning, translation), generalization is the mechanism, and memorization either plays no role or actually declines.

This finding has the structure of a scaling law, though the paper does not fit a parametric curve. It establishes a regularity — a predictable relationship between model size, task type, and the memorization-generalization balance — that can guide expectations about future models. If someone is building a larger model and hopes to improve it on factual recall, the evidence suggests they should expect this to come from better memorization of training data patterns. If they hope to improve it on mathematical reasoning, they should expect this to come from better generalization mechanisms, not from memorizing more problem-solution pairs.

The paper's Figure 2 captures this as a predicted trend line: memorization should be high and increasing for "knowledge-intensive" tasks, low and flat or decreasing for "generalization-intensive" tasks. The results in Figure 4 largely confirm these predictions, making this one of the rare cases in LLM research where a conceptual framework successfully predicts empirical patterns across diverse tasks before they are measured.

A notable nuance: the contrast between TriviaQA and MMLU knowledge-intensive subtasks — both are "knowledge tasks," but memorization increases with scale for the former and decreases for the latter — suggests that the knowledge-vs-reasoning axis is not the whole story. The rarity of the knowledge matters. TriviaQA covers common trivia facts that appear frequently in web text; MMLU covers specialized academic knowledge that appears sparsely. Larger models may memorize common knowledge more thoroughly while developing strategies to generalize or infer rarer knowledge, producing the divergent memorization trends. This is a hypothesis rather than a demonstrated mechanism, but it points to a more nuanced taxonomy of capability origins than the simple "knowledge = memorization, reasoning = generalization" dichotomy.

Innovation 4: The Task-Gram Table as a General-Purpose Framework for Characterizing Pretraining Data Relevance to Arbitrary Supervised Tasks

The task-gram table is not just a component of the memorization metric — it is a general analytical primitive for studying the relationship between pretraining data and downstream task capabilities. The paper demonstrates this by using the task-gram table for three distinct purposes beyond the core memorization analysis: (1) predicting overall task performance from n-gram pair counts (Section 4, Figure 3), (2) identifying which pretraining documents are most influential for test predictions via gradient analysis (Section 6, Figure 5), and (3) optimizing prompts to encourage memorization or generalization (Section 7, Table 1).

This versatility is evidence that the task-gram table captures something fundamental about the relationship between pretraining data and task behavior, not just a statistic that happens to correlate with LLM probabilities for the specific purpose of measuring distributional memorization.

The intellectual contribution is the abstraction itself: the recognition that any supervised task can be characterized by a filtered set of semantically related input-output n-gram pairs, and that this characterization enables a suite of analyses that were previously impossible at scale. Before this work, if you wanted to know which pretraining documents were relevant to a downstream task, your options were limited: you could search for exact overlaps between the test set and the pretraining corpus (decontamination analysis, as in Brown, 2020), you could train an embedding-based retriever to find semantically similar documents (a supervised approach requiring task-specific training), or you could compute training data influence via gradient-based methods (expensive and only feasible for small models and subsets of data). The task-gram table offers a middle ground: it is task-aware (unlike exact overlap), it requires no additional training (unlike supervised retrievers), and it scales to full corpus analysis (unlike gradient influence methods, which the paper applies only to a small sample of documents).

The gradient influence analysis in Section 6 provides the strongest validation of the task-gram table's utility beyond correlation. The finding that documents containing task-gram pairs exert systematically more influence on test predictions than documents containing only individual output n-grams (Figure 5, green lines consistently above blue lines across all three tasks and model sizes) is causal evidence — it shows that the n-gram pairs identified by the task-gram table correspond to data that actually shapes model behavior during training, not just data that happens to be statistically associated with test examples. This addresses the natural criticism of the correlational memorization metric: that high Spearman correlation could arise from a confounder rather than genuine learning from pretraining data. The gradient influence results make this alternative explanation less plausible by showing that the same n-gram pairs identified by the table have detectable effects on model parameters during training.

The prompt optimization application (Section 7) demonstrates the framework's practical utility: by simply rewriting a task prompt to use n-grams with higher (or lower) pretraining frequency, the paper can shift model performance on TriviaQA and GSM8K in predictable directions. For TriviaQA, maximizing n-gram overlap with pretraining data improves accuracy from 17% to 23.5% (Pythia 6.9B) or 28.7% to 23.2% for the "memorization" vs. "generalization" prompts. For GSM8K, minimizing overlap improves accuracy from 6.3% to 7.3%. These are modest effects, but their direction aligns with the paper's central thesis — knowledge tasks benefit from pretraining-like prompts, reasoning tasks benefit from pretraining-divergent prompts — and they demonstrate that the task-gram framework has actionable implications for prompt engineering.

This innovation is methodological rather than empirical. The paper doesn't claim that the task-gram table is the optimal way to characterize task-pretraining relationships; it claims that the framework is general, scalable, and useful across multiple analytical contexts. The fact that the same task-gram table construction serves three distinct purposes (memorization measurement, influence analysis, prompt optimization) with no task-specific modifications suggests that the abstraction is capturing a real structure in the data, which is the hallmark of a good analytical primitive.

5. Experimental Analysis

Evaluation Methodology

Dataset. The paper evaluates on four tasks spanning the memorization-to-generalization spectrum: WMT-09 (machine translation, 2.5K test set across six European languages: Hungarian, Czech, German, Italian, Spanish, French), TriviaQA (factual question answering, 10K test set from Joshi et al., 2017, with short-phrase answers treated as single output n-grams), MMLU (world knowledge and problem-solving, 57 subtask benchmark from Hendrycks et al., 2020, evaluated via multiple-choice accuracy with four options per question), and GSM8K (grade-school math word problems with chain-of-thought solutions, from Cobbe et al., 2021). WMT-09 was chosen over newer WMT versions because it contains more European languages prominent in the Pile pretraining corpus. For MMLU, the paper splits the 57 subtasks into 33 knowledge-intensive tasks (answerable by retrieving static knowledge, e.g., high_school_geography, astronomy, professional_psychology) and 24 reasoning-intensive tasks (requiring computation or logical inference, e.g., college_mathematics, econometrics, abstract_algebra), with the full classification provided in Appendix D.
Base model(s). The primary analysis uses the Pythia model family (Biderman et al., 2023), spanning 13M, 70M, 160M, 410M, 1B, 1.4B, 2.8B, 6.9B, and 12B parameters, all Transformer-decoder architectures trained on the Pile (Gao et al., 2020), a 207B-token diverse corpus. Pythia is chosen because it satisfies three simultaneous requirements: fully open-source weights, fully available pretraining data (the Pile), and wide model size range enabling scaling analysis. The paper also includes supplementary results with OLMo models (Groeneveld et al., 2024) at 1B and 7B parameters, trained on the Dolma corpus (Soldaini et al., 2024b) with 3T tokens, to provide more contemporary comparisons on GSM8K where Pythia's performance is very low (<5% accuracy).
Metrics. The paper uses two categories of metrics: task performance metrics and memorization metrics. For task performance: WMT is evaluated by BLEU score between greedily generated translations and reference translations; TriviaQA by exact match accuracy of the generated answer; MMLU by multiple-choice accuracy (selecting the option with the highest LM-predicted probability, with random baseline at 25%); GSM8K by BERTScore precision (Zhang* et al., 2020) between model-generated chain-of-thought and ground-truth chain-of-thought, since raw accuracy is too low (<5%) to show meaningful trends. For memorization: distributional memorization Mem_n(LLM, D|T) is the Spearman rank correlation ρ between the vector of log task-gram LM probabilities log P_{n,D}(s^y|s^x) and the vector of log LLM predicted probabilities log P_{LLM}(s^y|s^x), computed over all test n-gram pairs Φ that exist in the task-gram table H_n(T). Statistical significance is assessed via p-value from the Spearman correlation test with threshold p < 0.05. For GSM8K where memorization is not significant, the paper substitutes normalized Kendall tau ranking distance to quantify distributional generalization — the fraction of data pairs that disagree on their rankings between the task-gram LM and the LLM.
Baselines. The paper employs several comparison points for the memorization metric: (1) ∞-gram LM memorization Mem_∞ — the same Spearman correlation computation but using the ∞-gram language model (Liu et al., 2024) probabilities P_{\infty,D}(s^y|u⊕x⊕y) instead of task-gram LM probabilities, where each token's probability is based on the longest matching prefix found in the pretraining corpus with backoff — this represents the state-of-the-art in local-context pretraining data modeling and serves as a direct test of whether cross-sequence co-occurrence (task-gram) captures signal that sequential contiguity (∞-gram) misses. (2) For the causal analysis (Section 6), the baseline is documents containing only the output n-gram from task-gram pairs (blue lines in Figure 5), compared against documents containing the full n-gram pair (green lines), providing a test of whether both n-grams together matter more than the output alone. (3) For n-gram count vs. performance (Section 4, Figure 3), the implicit baseline is the trend across bins grouped by n-gram pair count.
Generation budget / compute accounting. Since this is an analysis paper rather than a method paper, there is no test-time compute budget to allocate. The computational cost lies entirely in the analysis pipeline: (1) constructing the task-gram table requires embedding all n-gram pairs from supervised data with LASER or E5 and computing pairwise cosine similarities — a one-time cost per task; (2) searching n-gram pairs in the Pile uses WIMBD API calls for exact co-occurrence counts and ∞-gram API calls for single n-gram frequency counts — the key cost, scaling with the number of task-gram pairs; (3) extracting LLM probabilities requires one forward pass per test example to collect token-level logits at the positions of interest — negligible compared to training. The paper does not report wall-clock time or FLOP counts for any part of the analysis. A significant unaccounted cost is that constructing the task-gram table for MMLU and GSM8K uses the test set directly (due to lack of large training sets), creating a potential circularity where the task is characterized by the same data used for evaluation — the paper does not discuss the implications of this cost in terms of information leakage.
Cross-validation / statistical protocol. The paper applies decontamination to ensure the Pile does not contain direct overlaps with any test set: it searches for large n-grams (n = 8 and n = 14) from test examples in the Pile and confirms no matches, following the method from the GPT-3 technical report (Brown, 2020). For the core memorization analysis, statistical significance is assessed per Spearman correlation computation — values with p < 0.05 are marked with solid round markers in Figure 4, while insignificant values (p ≧ 0.05) are marked with gray star markers. The paper does not use train/validation/test splits for the task-gram table construction itself (the mined pairs come from either a separate parallel corpus like Europarl for WMT, the training set for TriviaQA, or the test set itself for MMLU and GSM8K). For the gradient influence analysis (Section 6), only R = 50 pretraining documents are sampled per test example due to the expense of computing full parameter gradients, and the analysis uses k evenly spaced pretraining checkpoints (the exact number k is not specified in the main text). A cosine similarity sensitivity analysis (Figure 6 in Appendix D) tests γ ∈ {0.7, 0.75, 0.8} for TriviaQA trigram pairs and finds that "the trend of distributional memorization does not change with different thresholds," providing some robustness evidence though only for one task.

Main Quantitative Results

Task Performance as a Function of n-Gram Pair Counts in Pretraining Data (Section 4, Figure 3)

This analysis establishes the coarse relationship between how much task-relevant data exists in pretraining and how well models perform, before the more detailed distributional memorization analysis. The x-axis groups test examples by the number of task-gram pairs found in the pretraining corpus (as defined in Equation 4), and the y-axis shows task performance for different Pythia model sizes.

For WMT (n = 2, leftmost panel): Performance measured by BLEU score against the total n-gram pair count per language (since test set size is equal across languages, counts are summed). The six data points correspond to Hungarian, Czech, German, Italian, Spanish, and French (from left to right in order of increasing n-gram pair counts). For models smaller than 410M parameters, BLEU scores are near zero regardless of n-gram count. For the 410M model, a positive slope emerges weakly. For the 1B model and above, BLUE scores increase monotonically with n-gram pair count, with the largest model (12B) showing the steepest slope — reaching approximately 0.25–0.30 BLEU for the highest-count languages (Spanish, French) versus near zero for the lowest-count language (Hungarian). The key observation is that higher n-gram pair counts in pretraining data correlate with better translation performance, and this correlation strengthens with model size, though Section 5 will show this is not due to memorization of the specific n-gram pairs.

For TriviaQA (n = 5, second panel from left): The x-axis bins test examples by n-gram pair count (ranging from 0 to approximately 14,000). Models below 410M show near-zero accuracy across all bins. At 410M, a weak positive trend emerges from baseline. At 1B, accuracy rises from near zero in the lowest-count bin to approximately 0.15 in the highest-count bin. At 1.4B, the trend steepens, reaching approximately 0.25 accuracy in the highest bin. The largest models (6.9B, 12B) show the strongest relationship: 12B rises from approximately 0.05 accuracy at zero counts to over 0.35 accuracy in the highest-count bins. The monotonic and steep relationship is consistent with TriviaQA being the most memorization-dependent task — the counts of question-answer n-gram pairs in pretraining strongly predict whether the model can answer a given question.

For MMLU (n = 3, third panel): Performance is less smooth than TriviaQA, with notable non-monotonicity. For small models (< 410M), accuracy hovers near or below random guessing (25%). For larger models (1B–12B), accuracy generally increases with n-gram pair count, but with a conspicuous dip around 150 counts where performance drops below 25% — the paper attributes this to that interval containing "more reasoning or math problems, which appear to be harder for Pythia models." The 12B model reaches approximately 30–35% accuracy in higher-count bins versus approximately 22% in the lowest bins. The relationship is noisier and weaker than TriviaQA, consistent with MMLU being a mix of memorization-dependent and generalization-dependent subtasks.

For GSM8K (n values not specified, rightmost panel): Performance is measured by BERTScore precision (accuracy being too low to plot meaningfully). The 2.8B model shows the clearest increasing trend with n-gram pair count. Smaller models show weak or no trend. Curiously, larger models (6.9B, 12B) show a less significant trend than the 2.8B model — the paper interprets this as larger models relying less on the specific n-gram pairs and more on generalization, consistent with the distributional memorization results in Section 5.

Overall pattern: All four tasks show that performance increases when the number of task-related n-gram pairs in pretraining increases, provided the model is large enough (> 410M). The trend strengthens with model size for WMT, TriviaQA, and MMLU, but weakens for GSM8K at the largest scales. The paper notes that this correlation alone cannot distinguish memorization from generalization — "these performance curves can also be explained by the improved generalization ability of LMs when there is more relevant pretraining data" — which motivates the distributional memorization analysis in Section 5.

Distributional Memorization Results: The Core Memorization-vs-Generalization Analysis (Section 5, Figure 4)

Figure 4 is the central result figure of the paper, showing Mem_n(LLM, D|T) — the Spearman correlation between task-gram LM probabilities and LLM probabilities — as a function of Pythia model size (x-axis) for each task, computed with different n-gram sizes (n = 3 and n = 5 for most tasks; also n = 2 and n = 4 for some) and compared against Mem_∞ (the ∞-gram baseline). Solid lines show task-gram LM memorization; dashed lines show ∞-gram LM memorization. Solid round markers indicate statistical significance (p < 0.05); gray star markers indicate insignificance (p ≧ 0.05).

Translation (WMT): No significant memorization at any model size (top left panel). The paper states explicitly that no distributional memorization values for WMT are shown "because none of them are statistically significant." This is a notable null result — despite WMT showing strong positive correlation between n-gram pair counts and task performance in Figure 3, the LLM's output probability distribution is not predicted by the pretraining data frequency distribution. Instead of memorization, the panel shows the number of novel n-gram pairs generated by LLMs on WMT (a measure of generalization): larger models generate more novel n-gram pairs (the blue line rises from approximately 0.2 at 70M to approximately 0.5 at 12B), indicating that "larger LLMs generate more novel n-grams, which indicates a larger discrepancy in text distribution from the pretraining data and better distributional generalization." The paper interprets this as evidence that translation ability comes from transferable skills learned across languages, not from memorizing specific translation pairs. This finding contradicts Merrill et al. (2024), who found larger LLMs generate fewer novel n-grams; the paper attributes the discrepancy to Merrill et al. evaluating on in-distribution data (Pile validation set) that encourages memorization, while the out-of-distribution translation task encourages generalization.

TriviaQA: Strong and increasing memorization (top middle panel). This is the clearest evidence for memorization-driven capability. For n = 3 task-gram LM (dark solid line): memorization starts around ρ ≈ 0.37 at 70M and rises to approximately ρ ≈ 0.42 at 12B, consistently statistically significant (solid markers at all sizes). For n = 5 task-gram LM (lighter solid line): memorization is consistently lower, starting around ρ ≈ 0.15 at 70M and rising to approximately ρ ≈ 0.28 at 12B. The ∞-gram LM memorization (dashed lines) is lower than both task-gram variants: starting around ρ ≈ 0.1 and reaching approximately ρ ≈ 0.18 at 12B, with the larger model sizes showing less ∞-gram memorization. The key findings are: (1) Memorization increases with model size for TriviaQA — the correlation between LLM probabilities and pretraining frequencies gets stronger as models get larger, indicating that improved TriviaQA performance (Figure 3) is driven by better internalization of pretraining data patterns. (2) n = 3 task-gram LM substantially outperforms n = 5 and ∞-gram, with Mem_{n=3} (approximately 0.37–0.42) more than doubling Mem_{n=5} (approximately 0.15–0.28) at all model sizes. This means shorter, atomic fact-pair co-occurrences (trigrams) explain LLM behavior better than longer spans (5-grams) or local sequential dependencies (∞-gram) for factual QA. The paper interprets this as: "LLMs memorize small, long-range parallel data pieces more than large, local data pieces."

MMLU knowledge-intensive tasks: Memorization decreases with model size (top right panel). For n = 3 task-gram LM: memorization starts at approximately ρ ≈ 0.32 for the 70M model and declines to approximately ρ ≈ 0.28 at 12B. For n = 5 task-gram LM: memorization is lower overall (approximately ρ ≈ 0.15–0.22) and relatively flat across model sizes. The ∞-gram memorization is notably lower than the task-gram versions and decreases with model size. The stark contrast with TriviaQA — both are knowledge tasks, but memorization increases for one and decreases for the other — is the paper's key finding for MMLU. The interpretation is: "MMLU involves more specialized and less common knowledge compared to TriviaQA, making its occurrence in the pretraining corpus relatively infrequent. Consequently, for larger models to perform better on MMLU tasks, they may need to adjust the probability of recalling this knowledge, resulting in a decrease in distributional memorization." In other words, larger models deviate more from raw pretraining frequencies to answer rare knowledge questions correctly, while for common trivia, they follow pretraining frequencies more closely the larger they get.

MMLU reasoning-intensive tasks: Mixed pattern, trend toward generalization (bottom right panel). For n = 5 task-gram LM: memorization is actually the most significant here (approximately ρ ≈ 0.22 at 70M, declining to approximately ρ ≈ 0.12 at 12B), while n = 3 memorization is lower and largely insignificant. The paper interprets the n = 5 advantage as: "some concepts can be meaningfully expressed in large text segments while the small text segments are meaningless in a reasoning-intensive context" — reasoning involves conceptual units that require longer spans to be meaningful. The declining trend with model size mirrors the knowledge-intensive MMLU pattern, indicating that "memorization is not the driving force of performance improvement." The ∞-gram memorization is close to zero and mostly insignificant.

GSM8K: No significant memorization; increasing generalization with scale (bottom two panels, Pythia and OLMo). The paper found no significant Spearman correlation for GSM8K with either Pythia or OLMo models — the task-gram LM probabilities do not predict LLM probabilities. Instead, the paper uses normalized Kendall tau ranking distance as a measure of distributional generalization (how much the LLM's probability rankings deviate from the pretraining frequency rankings). For Pythia models: the Kendall tau distance with the task-gram LM is around 0.51–0.54 and generally increases with model size (larger models disagree more with pretraining rankings). The Kendall tau distance with the ∞-gram LM is higher (approximately 0.55–0.58), meaning the LLM's rankings deviate even more from local sequential frequency patterns than from task-gram co-occurrence patterns. For OLMo models (1B and 7B): the same pattern holds — task-gram LM Kendall tau distance is lower (approximately 0.52) than ∞-gram Kendall tau distance (approximately 0.55–0.56), and the 7B model shows higher distance than the 1B model. The interpretation is: "generalization is the driving force of performance improvement" for reasoning, and the task-gram LM captures the pretraining distribution better than the ∞-gram LM even when neither shows significant memorization.

Task-gram LM vs. ∞-gram LM across all tasks. A consistent pattern across every task where both are computed (TriviaQA, MMLU both splits): Mem_n (solid lines) is always greater than or equal to Mem_∞ (dashed lines) for the same n-gram size. For TriviaQA, the gap is large (approximately 0.37 vs. 0.1 for n = 3 at 70M, widening to approximately 0.42 vs. 0.13 at 12B). For both MMLU splits, the gap is present but smaller. The paper concludes that "our task-gram LM can better explain LLM predicted probabilities than ∞-gram LM" because it captures cross-sequence co-occurrence that ∞-gram's local sequential modeling misses — a key validation of the task-gram approach.

Gradient-Based Influence Analysis: Causal Evidence from Training Dynamics (Section 6, Figure 5)

The gradient influence analysis estimates how much individual pretraining documents affected test predictions through pretraining, following the method of Pruthi et al. (2020): the influence of pretraining document d on test example (x, y) is the sum over pretraining checkpoints of the dot product between the document's training loss gradient and the test example's loss gradient, both taken with respect to the target output n-gram s^y. The paper computes average influence over R = 50 randomly retrieved pretraining documents per test example, comparing two retrieval schemes: documents containing the full n-gram pair (s^x, s^y) (green lines) versus documents containing only the output n-gram s^y (blue lines).

Results across tasks (Figure 5):

WMT (left panel): The influence values are the smallest among all three tasks (y-axis scale approximately -0.02 to 0.08). N-gram pair documents (green) consistently show higher influence than output-only documents (blue), but the gap is small and both are low. The influence generally decreases as model size increases. This is consistent with the null memorization finding — translation does not depend strongly on specific pretraining documents.
TriviaQA (middle panel): The influence values are the largest among all three tasks (y-axis scale approximately 0.00 to 0.16). N-gram pair documents (green) show substantially higher influence than output-only documents (blue) across all model sizes, with the difference being most pronounced at larger scales. The influence increases slightly from approximately 0.14 (green) at 70M to approximately 0.16 at 12B. This is consistent with TriviaQA's strong memorization effect and the increasing memorization trend with model size.
MMLU (right panel): Influence values are intermediate between WMT and TriviaQA. N-gram pair documents (green) show higher influence than output-only documents (blue), with the gap widening at larger model sizes. The influence decreases with model size, consistent with MMLU's decreasing memorization trend.

Key finding: "Across all three datasets, we observe that pretraining documents containing n-gram pairs consistently contribute more to the testing examples than documents containing only the output n-gram in n-gram pairs." This provides causal evidence that the task-gram table identifies genuinely influential pretraining data — documents containing both the input and output n-grams of a pair (not just the output n-gram alone) have measurably larger impact on model behavior during training, as quantified through gradient-based influence accumulation.

Limitation acknowledged by the paper: The gradient analysis uses only R = 50 documents per test example due to computational expense ("the computation of the full gradient is relatively expensive"), meaning it covers a tiny fraction of the pretraining corpus. The paper states: "This analysis does not aim to cover the full pretraining corpus but to give complementary causal evidence to the previous findings." The small sample size means that while the direction of the effect (pairs > singles) is reliable, the absolute influence magnitudes may not be representative.

Prompt Optimization: Practical Application of the Memorization-Generalization Distinction (Section 7, Table 1)

The paper tests whether prompts can be rewritten to encourage memorization or generalization by optimizing for higher or lower n-gram overlap with pretraining data, using GPT-4o as the prompt optimizer with WIMBD n-gram count feedback as the reward signal. The meta-prompt instructs GPT-4o to rewrite a given task prompt to either maximize or minimize the average n-gram count of the rewritten prompt in the pretraining corpus.

Results from Table 1:

Model	TriviaQA (Mem)	TriviaQA (Gen)	GSM8K (Mem)	GSM8K (Gen)
Pythia 6.9B	17%	9%	2.6%	2.8%
Pythia-Instruct 6.9B	23.5%	23.2%	6.3%	7.3%
Pythia 12B	28.7%	23.2%	2.7%	2.8%
OLMo 7B	36.4%	29.8%	2.5%	3.1%
OLMo-instruct 7B	29%	10%	6.3%	7.9%

Key findings: For TriviaQA, the memorization-encouraging prompt (maximizing n-gram overlap with pretraining) consistently outperforms the generalization-encouraging prompt (minimizing overlap) across all models: 17% vs. 9% for Pythia 6.9B, 23.5% vs. 23.2% for Pythia-Instruct (smallest gap), 28.7% vs. 23.2% for Pythia 12B, 36.4% vs. 29.8% for OLMo 7B, and 29% vs. 10% for OLMo-instruct 7B. For GSM8K, the generalization-encouraging prompt consistently (though marginally) outperforms the memorization-encouraging prompt: 2.6% vs. 2.8% for Pythia 6.9B, 6.3% vs. 7.3% for Pythia-Instruct, 2.7% vs. 2.8% for Pythia 12B, 2.5% vs. 3.1% for OLMo 7B, and 6.3% vs. 7.9% for OLMo-instruct. The effect sizes are small for reasoning (0.2–1.6 percentage points for GSM8K) and moderate to large for factual QA (up to 19 percentage points for OLMo-instruct on TriviaQA). The paper notes that "the lengths of the optimized prompts are not significantly different," suggesting the effect is due to lexical choice (n-gram overlap with pretraining) rather than prompt length.

Interpretation: The asymmetric results — TriviaQA substantially benefits from pretraining-aligned prompts, while GSM8K benefits marginally from pretraining-divergent prompts — align with the paper's central thesis: knowledge-intensive tasks rely on memorization (and thus benefit from prompts that evoke pretraining-like distributions), while reasoning-intensive tasks rely on generalization (and thus benefit from prompts that avoid triggering memorized patterns). The small effect sizes on GSM8K reflect the fact that these models have very low math reasoning capability to begin with (<8% for the best model), so prompt optimization can only shift performance within a narrow range.

Ablation Studies and Robustness Checks

Cosine similarity threshold for task-gram table construction (Figure 6, Appendix D): The paper tests three threshold values γ ∈ {0.7, 0.75, 0.8} for constructing the TriviaQA trigram task-gram table and recomputes distributional memorization across Pythia model sizes. The resulting memorization curves for all three thresholds largely overlap, with the paper stating "the trend of distributional memorization does not change with different thresholds." The memorization values are approximately ρ ≈ 0.35–0.42 across the three thresholds, with the same increasing trend with model size. This suggests the method is robust to the exact threshold choice within a reasonable range, though only one task (TriviaQA) is tested.

Embedding model choice (implicit ablation across tasks): While not presented as a formal ablation, the paper uses two different embedding models — LASER for WMT and E5 for all other tasks — and both produce task-gram tables that lead to interpretable results (no memorization for WMT, strong memorization for TriviaQA, intermediate for MMLU). This provides weak evidence that the framework is not tied to a specific embedding model, though no direct comparison of LASER vs. E5 on the same task is performed.

n-gram size comparison (n = 3 vs. n = 5, throughout Figure 4): The paper systematically shows results for multiple n-gram sizes within each task panel, revealing that the optimal n for detecting memorization varies by task type: n = 3 dominates for TriviaQA and knowledge-intensive MMLU (shorter, atomic fact units), while n = 5 is more informative for reasoning-intensive MMLU (longer conceptual units). This pattern is consistent with the interpretation that different tasks involve different granularities of pretraining patterns, though the paper does not sweep n systematically across all values or provide a principled method for selecting n per task.

Decontamination check: The paper verifies that "the Pile is not contaminated by any of the datasets we used, by ensuring there are no large n-grams (n = 8 and n = 14) overlaps between the Pile and the testing data," following the GPT-3 contamination detection method (Brown, 2020). This ensures that the memorization signals detected are not artifacts of test set leakage — the n-gram pairs found in the pretraining data represent naturally occurring co-occurrences, not verbatim copies of test examples.

MMLU subtask splitting (knowledge-intensive vs. reasoning-intensive): By manually classifying the 57 MMLU subtasks into knowledge-intensive (33 subtasks) and reasoning-intensive (24 subtasks) and computing memorization separately for each group, the paper reveals that aggregating all MMLU subtasks together would obscure the divergent memorization patterns. The classification itself (provided in Appendix D) is based on whether a task "can be answered by retrieving static knowledge" versus requiring "computation or logical reasoning over the knowledge" — a reasonable but subjective division that the paper does not validate through inter-annotator agreement or alternative classification schemes.

OLMo replication on GSM8K (Figure 4, bottom panels): The null memorization result and increasing generalization trend on GSM8K are replicated with OLMo models (1B, 7B) trained on Dolma (3T tokens), a different and larger pretraining corpus than the Pile. The task-gram LM Kendall tau distances for OLMo (approximately 0.52 for task-gram, 0.55–0.56 for ∞-gram) are similar in magnitude and pattern to the Pythia results, providing evidence that the finding is not specific to the Pythia-Pile combination. However, only two OLMo model sizes are tested, limiting the scaling analysis compared to the nine Pythia sizes.

Gradient influence retrieval schemes (Figure 5, green vs. blue curves): Comparing documents containing n-gram pairs against documents containing only output n-grams serves as an ablation of the task-gram table's specificity. The consistent finding that pair-containing documents exert more influence (green > blue across all tasks and model sizes) validates that the full input-output n-gram relationship, not just the output side, is what drives training influence — a non-trivial result since one might expect the output n-gram alone to be sufficient for influencing test predictions.

Critical Assessment

This section evaluates whether the reported experiments genuinely support the paper's central claims, identifying specific strengths, weaknesses, and gaps.

Does the evidence support the core claim that distributional memorization validly measures task-level dependency on pretraining data?

What was demonstrated: The paper shows that a Spearman correlation between task-gram LM probabilities and LLM probabilities produces interpretable, task-dependent patterns: strong and increasing for TriviaQA, decreasing for MMLU knowledge tasks, absent for translation and reasoning. The task-gram LM version consistently outperforms the ∞-gram baseline, suggesting it captures something the baseline misses. The gradient influence analysis provides complementary causal evidence that documents containing task-gram pairs influence test predictions during training.

What was NOT demonstrated: The paper does not establish that the Spearman correlation captures all forms of pretraining data dependency, nor that a zero correlation implies zero dependency. There could be forms of memorization that do not manifest as monotonic relationships between pretraining frequency and LLM probability — for example, if the LLM learns to suppress certain high-frequency patterns (anti-memorization) or if it learns nonlinear transformations of pretraining frequencies that Spearman correlation cannot detect. The paper's definition of distributional memorization as a Spearman correlation is a sufficient condition for detecting one type of memorization, not a necessary condition — a negative result (no significant correlation) does not prove absence of memorization, only absence of this specific form of it. This is particularly relevant for the translation results, where the paper concludes "translation ability does not come from memorization" based on non-significant Spearman correlations. The translation model might still depend on pretraining data in ways not captured by n-gram co-occurrence frequency — for instance, through abstract grammatical rules or cross-lingual semantic mappings that are not reducible to n-gram pair statistics.

The gradient influence analysis partially addresses this concern by providing converging evidence: WMT shows the lowest and most flat gradient influence, consistent with the null memorization finding. But the gradient analysis uses only R = 50 documents per test example and covers at most a few thousand documents total, which is a minuscule fraction of the 207B-token Pile corpus. The strongest causal evidence would come from a full-scale data attribution study, which the paper acknowledges is computationally infeasible.

Is the task-gram LM a valid model of the pretraining data distribution for the purpose of measuring memorization?

Strengths: The task-gram LM's document-level co-occurrence modeling is conceptually well-motivated for capturing long-range task-relevant associations. The consistent finding that Mem_n > Mem_∞ validates that this structural difference matters empirically. The cosine similarity filtering provides a principled way to select semantically meaningful pairs from the enormous Cartesian product of all n-grams.

Weaknesses: Several aspects of the task-gram LM construction raise concerns about what exactly is being measured:

The task-gram LM is extremely sparse. The probability P_{n,D}(s^y|s^x) = C((s^x, s^y), D) / C(s^x, D) is zero for any output n-gram that does not co-occur with the input n-gram in at least one pretraining document. For a given input n-gram, there may be dozens of semantically appropriate output n-grams that never happen to co-occur in the same document, and the task-gram LM assigns them all zero probability. The LLM, by contrast, will assign some non-zero probability to many of these. This fundamental mismatch — the task-gram LM is supported on a tiny subset of the output space while the LLM distributes probability broadly — means the Spearman correlation is computed primarily over n-gram pairs that do co-occur in pretraining. The correlation therefore measures whether, among co-occurring pairs, the LLM's probability ranking aligns with the empirical frequency ranking. It does not measure whether the LLM assigns low probability to n-gram pairs that never co-occur in pretraining — those pairs are simply excluded from Φ. This means the metric is somewhat circular: it tests whether the LLM agrees with pretraining frequencies on the subset of pairs where pretraining provides a signal, but cannot detect cases where the LLM assigns high probability to outputs the task-gram LM assigns zero probability (which would be evidence of generalization).
The conditional probability normalization is by document count of the input n-gram, not by anything task-specific. C(s^x, D) is the number of documents containing the input n-gram in any context, not necessarily in a task-relevant context. If the input n-gram "Marie Curie" appears in 10,000 documents (many of which are biographies, science articles, etc.) and only 15 also contain "radium" (the output n-gram), the conditional probability is 15/10000 = 0.0015. In a task-appropriate pretraining corpus, many of those 10,000 documents might not be "question-answering-like" — they might be narrative text where the task-gram relationship doesn't apply. The task-gram LM does not distinguish between task-relevant and task-irrelevant occurrences of the input n-gram, potentially diluting the conditional probability for n-grams that appear in many non-task contexts. This could suppress memorization signals for tasks where the input n-grams are common but task-relevant co-occurrence is rare.
The task-gram table for MMLU and GSM8K is built from the test set. For MMLU, the paper mines n-gram pairs from "the test sets directly" because the training set is very small. For GSM8K, pairs are mined from "the 1K test sets directly." This means the task-gram table — which defines what counts as a task-relevant n-gram pair — is constructed using the same data that memorization is evaluated on. This creates a potential information leak: the pairs in the task-gram table are guaranteed to be present in the test set by construction. The paper acknowledges this constraint but does not discuss whether it inflates the memorization estimates (if test-set n-gram pairs happen to be more common in pretraining than a typical task-gram pair would be) or deflates them (if filtering by test-set pairs misses many task-relevant pairs that would show stronger memorization). For TriviaQA and WMT, the use of separate training data or Europarl avoids this issue, making their results more trustworthy.
Document-level co-occurrence may be too coarse for some tasks. Two n-grams appearing in the same document does not mean they appear in a task-relevant relationship — "Marie Curie" and "radium" appearing in the same Wikipedia article is counted as a co-occurrence regardless of whether the text actually states that Curie discovered radium, or simply mentions both in unrelated sections. The task-gram LM would benefit from a more sophisticated co-occurrence definition (e.g., co-occurrence within a paragraph, or co-occurrence where a dependency parse links the two spans), but this would be computationally prohibitive at the 207B-token scale. The current document-level definition is a pragmatic compromise that likely introduces noise (counting spurious co-occurrences) while still capturing genuine associations.

Does the task-dependent memorization pattern genuinely establish that "memorization drives improvement on knowledge tasks while generalization drives improvement on reasoning tasks"?

Evidence in favor: The pattern is consistent across four tasks and robust to n-gram size variation: TriviaQA shows the strongest and most consistently increasing memorization; MMLU shows intermediate memorization that decreases with scale; GSM8K shows no memorization at all; WMT shows no memorization but increasing novelty. The prompt optimization results (Table 1) provide a practical validation — knowledge tasks benefit from pretraining-aligned prompts, reasoning tasks from pretraining-divergent prompts. These converging lines of evidence make a compelling case.

Evidence that complicates the picture: The key anomaly is the MMLU knowledge-intensive vs. TriviaQA contrast. Both are knowledge tasks, yet memorization increases with scale for TriviaQA and decreases for MMLU knowledge. The paper's explanation — that MMLU involves rarer knowledge requiring generalization — is plausible but not directly tested. An experiment that would clarify this: bin MMLU questions by the pretraining frequency of their correct answers (analogous to the difficulty binning in the reference paper's compute-optimal allocation) and check whether memorization is higher and increasing for high-frequency MMLU questions (like TriviaQA) and lower/decreasing for low-frequency ones. This experiment is not performed. Without it, the decreasing memorization trend for MMLU knowledge tasks could reflect a statistical artifact (regression to the mean as the task-gram LM becomes noisier for rare knowledge) rather than a genuine shift to generalization.

The GSM8K generalization claim is based on a null result (no significant memorization) plus a weak positive trend in Kendall tau distance. Demonstrating that memorization is absent does not demonstrate that generalization is present — it demonstrates only that pretraining frequency doesn't predict LLM behavior. The paper's evidence for generalization comes from: (a) the increasing Kendall tau distance between LLM rankings and task-gram LM rankings, and (b) the slight improvement from generalization-encouraging prompts. Both are suggestive but not definitive. A stronger test would be to show that GSM8K performance correlates with some measure of the model's ability to transform or recombine pretraining patterns (e.g., performance on held-out reasoning problems whose solutions require combining multiple pretraining facts in novel ways), but no such analysis is conducted.

Are the effect sizes practically meaningful?

The 4× efficiency gain from compute-optimal scaling is a benchmark from the reference paper and is NOT present in this paper. This paper does not make efficiency claims. The practical significance of its findings lies in the interpretability and predictive value of the memorization metrics:

For TriviaQA, the Spearman correlations of ρ ≈ 0.35–0.42 for n = 3 are moderate by conventional standards (explaining roughly 12–18% of variance in LLM probability rankings). Whether this constitutes "strong" memorization is a matter of interpretation — the paper treats it as strong relative to the other tasks, but a correlation of 0.4 leaves substantial variance unexplained, even among n-gram pairs that do co-occur in pretraining.
For prompt optimization, the TriviaQA improvements (e.g., 17% → 23.5% for Pythia-Instruct 6.9B, a 6.5 percentage point gain, and 29% → 10% for OLMo-instruct 7B, a 19 percentage point gap between memorization and generalization prompts) are practically meaningful. However, these results come from a single iteration of GPT-4o-based optimization (the paper does not report how many iterations were run), and the optimized prompts are shown in Table 2 (Appendix E) but no human evaluation of prompt quality or alternative optimization methods are tested. The GSM8K improvements are very small (0.2–1.6 percentage points), often within the range of random variation for small test sets.
The gradient influence values are reported without confidence intervals, making it difficult to assess whether the differences between green and blue lines in Figure 5 are statistically significant or practically meaningful. The y-axis scales differ across the three panels, suggesting the raw influence values are not directly comparable across tasks.

What experiments would have strengthened the paper?

A test of causality beyond gradient influence. The gradient-influence analysis covers only 50 documents per test example. A stronger causal test would be to remove documents containing task-gram pairs from the pretraining corpus and retrain a (smaller) model, then measure whether performance on the corresponding task degrades. This is precisely the counterfactual approach the paper argues is impractical, but doing it at a small scale (e.g., with a 70M or 160M Pythia model) would provide a validity check on the distributional memorization framework. If removing high-count task-gram pair documents degrades TriviaQA performance more than removing random documents, it would validate that the task-gram table identifies causally important data.
A test of whether distributional memorization predicts generalization failure. If high distributional memorization means the LLM is relying on pretraining patterns, then the model should perform worse on test examples where the correct answer has low pretraining frequency (since it can't retrieve the answer from memory). The paper could bin TriviaQA test examples by the task-gram LM probability of the correct answer and show that accuracy drops for low-probability answers — this would connect the memorization metric to actual task performance in a falsifiable way.
Replication on a model family with released pretraining data and stronger capabilities. The paper acknowledges Pythia is "slightly outdated." OLMo provides a partial replication on GSM8K but uses only two model sizes. A more convincing demonstration would use a model family like OLMo across the full size range (1B, 7B, and ideally larger) and replicate the full set of findings across all four tasks. The computational cost of searching Dolma (3T tokens, 14× larger than the Pile) is cited as a limitation — "The current WIMBD system also has limitations in searching larger corpora like Dolma."
Sensitivity analysis for the MMLU knowledge-reasoning split. The classification of MMLU subtasks into knowledge-intensive and reasoning-intensive is provided as a list in Appendix D but no validation is reported (e.g., inter-annotator agreement, comparison against alternative classification schemes, or correlation with external measures of knowledge vs. reasoning requirements). The divergent memorization patterns between the two splits are a key result; if the classification is noisy or biased, the divergence may be overstated.
A calibration of the ∞-gram baseline. The paper consistently finds that Mem_n > Mem_∞, which is interpreted as evidence that the task-gram LM captures something the ∞-gram LM misses. However, the ∞-gram LM is used in a specific way — computing the probability of isolated output n-grams given the full concatenated input-output context — that is not its intended use case (it is designed for next-token prediction within a contiguous sequence). The fact that it performs poorly at the task-gram comparison may reflect this mismatch rather than a fundamental limitation of local-context modeling. A fairer baseline would be an ∞-gram LM variant that computes P(s^y | s^x) by conditioning on the longest prefix ending with the input n-gram, rather than conditioning on the entire test sequence with instruction and output prefix. The paper does not explore this variant.
Statistical rigor for the main results. The memorization values in Figure 4 are reported as point estimates of Spearman ρ with significance markers (solid vs. star), but no confidence intervals are shown. With Spearman correlations of approximately 0.3–0.4 on what may be hundreds or thousands of n-gram pairs, the 95% confidence intervals would likely be wide enough that some of the apparent trends (e.g., the increase from 0.37 to 0.42 for TriviaQA n = 3 across model sizes) might not be statistically distinguishable. The paper also does not correct for multiple comparisons despite computing memorization for multiple n-gram sizes, multiple tasks, and multiple model sizes simultaneously.

Overall assessment: The paper's experiments convincingly demonstrate that LLM output probabilities correlate with task-gram co-occurrence frequencies in pretraining data, and that this correlation varies systematically across tasks and model sizes in interpretable ways. The evidence is strongest for TriviaQA (clear, significant, increasing memorization) and for the null result on translation (consistent absence of memorization across all model sizes). The evidence for the memorization-to-generalization shift with scale is suggestive but rests more heavily on interpretation of correlational patterns and declining memorization trends that may be influenced by noise, data sparsity, or the specific construction choices of the task-gram LM. The gradient influence analysis and prompt optimization experiments provide valuable complementary evidence but are limited in scale (50 documents, single optimization iteration) and effect size (small for reasoning). The paper's central claim — that memorization drives knowledge task improvement while generalization drives reasoning task improvement — is supported but not conclusively proven; the data is at least as consistent with a more nuanced interpretation where both mechanisms operate across all tasks but in different proportions, and where the task-gram LM captures memorization more effectively for some tasks than others due to its construction biases rather than genuine differences in underlying mechanisms.

6. Limitations and Trade-offs

The Difficulty Estimation Cost Is Unaccounted for in the Memorization Metric

The assumption or constraint. The entire distributional memorization framework depends on constructing a task-gram table — a filtered set of semantically related n-gram pairs mined from supervised task data and validated against the pretraining corpus. For MMLU and GSM8K, the task-gram table is built directly from the test set, as the paper acknowledges: "For MMLU, since the training set is very small (100–500 examples for each task), we mine the n-gram pairs from the test sets directly. For GSM8K, due to the short time limitation during rebuttal, we also mine the n-gram pairs from the 1K test sets directly" (Appendix D). This means the characterization of what counts as a "task-relevant n-gram pair" — the very definition of the task for analysis purposes — is derived from the same data used to evaluate memorization.

The consequence. This creates a subtle but consequential circularity. The task-gram table for MMLU and GSM8K is guaranteed to contain n-gram pairs that are present in the test set, since those pairs were extracted from the test set to begin with. If test-set n-gram pairs happen to co-occur more frequently in the pretraining corpus than a typical task-relevant pair would (because test examples are often drawn from commonly discussed topics), the memorization estimates will be inflated — the task-gram LM will appear to predict LLM behavior better than it would on genuinely held-out task examples. Conversely, if the task-gram table built from test data misses many task-relevant n-gram pairs that would show strong memorization (because the test set is too small to capture the full task distribution), the memorization estimates will be deflated. The paper does not quantify the magnitude or direction of this bias. Because two of the four tasks (MMLU and GSM8K) — including the critical reasoning task where memorization is claimed to be absent — suffer from this circularity, the comparative conclusions about memorization-vs-generalization across tasks are potentially confounded by differences in how the task-gram table was constructed.

What evidence exists in the paper. The paper transparently states the data sources for task-gram mining in Appendix D, making the circularity visible to careful readers. For TriviaQA and WMT, separate data sources are used (TriviaQA training set and Europarl respectively), avoiding the circularity for those tasks. The paper does not discuss how the use of test data for MMLU and GSM8K might affect the memorization estimates relative to TriviaQA and WMT. No ablation compares memorization computed from test-mined task-gram tables versus memorization computed from tables mined from a held-out task dataset for MMLU or GSM8K — such an ablation is impossible given the paper's data constraints, but the absence of any discussion of the issue is notable. The GSM8K null memorization result is partially replicated with OLMo models (Figure 4, bottom panels), but the OLMo analysis also uses test-set-mined task-gram pairs, so the replication does not address the circularity concern.

Mitigation status. The paper does not attempt to mitigate this circularity, acknowledge it as a limitation, or discuss its potential impact on the results. The authors flag data constraints as a practical issue ("due to the short time limitation during rebuttal") but do not connect this to the validity of the memorization estimates. Future work could address this by constructing task-gram tables from larger, held-out supervised datasets for reasoning tasks (e.g., using GSM8K's training set of 7.5K examples rather than the 1K test set), or by cross-validating the task-gram table construction to use non-overlapping data for table building and memorization evaluation.

The Task-Gram LM Is an Extremely Sparse Model of the Pretraining Distribution, Potentially Masking Forms of Memorization

The assumption or constraint. The task-gram language model assigns zero probability to any output n-gram that does not co-occur with the input n-gram in at least one pretraining document. Formally, P_{n,D}(s^y|s^x) = 0 for any pair not in the task-gram table H_n(T), and only pairs with C((s^x, s^y), D) > 0 survive in the table after the final filtering step (Definition 1 and the post-construction filtering described in Appendix D). The distributional memorization metric Mem_n is computed only over the set Φ of test n-gram pairs that exist in both the test examples and the task-gram table (Definition 3). Pairs where the output n-gram never co-occurs with the input n-gram in pretraining — even if the LLM assigns those pairs high probability — are excluded from the correlation computation entirely.

The consequence. The Spearman correlation measures whether, among n-gram pairs that do co-occur in pretraining, the LLM's probability ranking aligns with the empirical co-occurrence frequency ranking. It is silent on whether the LLM assigns high probability to outputs that never co-occur with the input in pretraining — the very behavior that would be the strongest evidence of generalization. Imagine a scenario where for a given input n-gram, the pretraining data contains 5 co-occurring output n-grams (included in Φ) and the LLM assigns high probability to those 5 plus 20 other semantically appropriate outputs that never co-occurred. The correlation is computed over only the 5 co-occurring pairs, completely ignoring the 20 novel outputs. The actual degree to which the LLM's output distribution reflects pretraining data is therefore systematically overestimated by the metric — the LLM could be generating many more novel outputs than the correlation suggests, and these simply don't enter the computation. This limitation is most consequential for tasks where the task-gram table is small relative to the space of possible outputs, which is likely the case for reasoning (GSM8K) and translation (WMT), where the correct output can be expressed in many different ways, most of which never co-occurred with the specific input phrasing in any single pretraining document.

What evidence exists in the paper. The paper does not report what fraction of test n-gram pairs are excluded from Φ because they have zero co-occurrence counts — this number would quantify the sparsity of the task-gram LM and help readers assess how much of the LLM's output distribution is being measured. The paper does show that WMT exhibits no significant memorization and that larger models generate more novel n-gram pairs (Figure 4, top left), which is consistent with the task-gram LM being too sparse to capture the relevant pretraining signal for translation. However, this pattern could reflect either genuine generalization or a failure of the task-gram LM to capture the form of memorization that translation relies on (e.g., abstract syntactic patterns rather than specific n-gram co-occurrences). The paper does not distinguish between these possibilities.

Mitigation status. The paper does not acknowledge this sparsity limitation or discuss its implications for interpreting null memorization results. The reliance on document-level co-occurrence guarantees that many semantically related pairs will have zero counts simply because the pretraining corpus, however large, cannot contain all valid input-output pairings within individual documents. The paper does not explore alternatives that could capture broader forms of pretraining dependency, such as smoothed or interpolated n-gram models, embedding-based density estimators, or models that condition on co-occurrence within a broader context (e.g., same domain, same website) rather than strictly within the same document. The ∞-gram baseline avoids the sparsity problem (since it backs off to shorter prefixes rather than assigning zero probability), but it captures local sequential dependencies that are structurally different from the cross-sequence associations the task-gram LM targets. A hybrid approach combining the task-gram LM's cross-sequence modeling with some form of smoothing or backoff would address this limitation but is not explored.

All Results Are on a Single, Relatively Weak Model Family with a Single Pretraining Corpus; Generalizability to Stronger Models Is Unknown

The assumption or constraint. The paper's entire empirical analysis uses the Pythia model family (Biderman et al., 2023) trained on the Pile (Gao et al., 2020), with supplementary OLMo results on GSM8K only. The paper acknowledges this limitation candidly in Appendix A: "the model we use, Pythia, and the pretraining corpus, Pile, are slightly outdated and have been outperformed by many new open-source LLMs. However, most open-source LLMs lack corresponding pretraining data and have limited model sizes and pre-training checkpoints, hindering scaling effect studies." The largest Pythia model (12B) achieves less than 5% accuracy on GSM8K and approximately 35% on TriviaQA, meaning all memorization analyses on reasoning tasks are conducted on models that fundamentally cannot perform the task at a meaningful level.

The consequence. The central finding — that memorization increases with scale for knowledge tasks but generalization increases for reasoning tasks — is established on model sizes and capability levels where reasoning performance is near the floor. It is entirely possible that at larger scales (e.g., 70B, 400B parameters) or with stronger base models (e.g., Llama-3, Qwen, DeepSeek), the memorization-generalization balance shifts qualitatively. A 70B+ model with non-trivial GSM8K accuracy (say, 50%+) might achieve that accuracy through very different mechanisms than a 12B model with 2.8% accuracy. The paper's scaling trends (memorization decreasing for MMLU knowledge tasks, flat/absent for GSM8K) might plateau, reverse, or change functional form at larger scales — the range from 13M to 12B parameters represents only the early part of the scaling curve for modern LLMs. The finding that distributional memorization fails to explain reasoning performance might be an artifact of the base model being too weak to exhibit the relevant memorization behavior, rather than a fundamental property of reasoning tasks.

The reliance on Pythia also means the paper cannot disentangle model-specific effects from task-general effects. Pythia models are known to have specific weaknesses (e.g., poor multi-step reasoning, limited code generation) that may affect how they utilize pretraining data for different tasks. A model family with different architecture, training recipe, or data mixture might exhibit different memorization-generalization patterns on the same tasks. The paper provides no evidence that the findings would transfer to a different model architecture (e.g., a mixture-of-experts model) or training objective (e.g., a model trained with fill-in-the-middle, or instruction-tuned).

What evidence exists in the paper. The paper reports OLMo results on GSM8K (Figure 4, bottom panels) as a partial replication, finding similarly non-significant memorization and similarly increasing generalization (Kendall tau distance) with model scale. This is a meaningful robustness check — OLMo uses a different architecture, training recipe, and pretraining corpus (Dolma, 3T tokens) — but it is limited to one task (the hardest one for Pythia) and two model sizes (1B and 7B). The memorization values for OLMo are not directly compared to Pythia's on the same axes (GSM8K uses Kendall tau for both, while other tasks use Spearman) and only two data points per curve make it impossible to assess whether the scaling trend is consistent across the full OLMo size range. The paper does not replicate the TriviaQA, MMLU, or WMT findings on OLMo, leaving open the possibility that the key knowledge-task memorization results are specific to Pythia's training on the Pile.

Mitigation status. The paper is transparent about this limitation in Appendix A, framing it as a constraint imposed by the availability of open pretraining data rather than an oversight. The authors call for "improved searching and retrieval methods" for larger corpora, implicitly acknowledging that the current infrastructure (WIMBD) cannot yet support analysis of datasets like Dolma at the same level of detail as the Pile. The OLMo partial replication is a step toward addressing this limitation but falls short of a full multi-model, multi-task validation. Until the analysis is replicated on a contemporary, high-capability model family with released pretraining data — or until the WIMBD/∞-gram infrastructure scales to support searching multi-trillion-token corpora — the generalizability of the findings to models that practitioners actually deploy remains an open question.

The Definition of Distributional Generalization as "the Opposite of Memorization" Collapses a Spectrum into a Binary, Masking Important Nuance

The assumption or constraint. The paper defines distributional generalization entirely negatively: "The distributional generalization is then defined as the opposite of the distributional memorization: increased memorization implies decreased generalization, as the LLM predictions and the n-gram LM are more distributionally correlated, and vice versa" (Section 2, after Definition 3). There is no independent metric for generalization — no measure of how much the LLM's output distribution deviates from the pretraining distribution in structured, task-appropriate ways. When memorization is low or non-significant, the paper concludes that generalization is high; when memorization decreases with scale, the paper concludes that generalization increases. Generalization is treated as a residual category.

The consequence. This binary framing collapses several distinct phenomena into a single label. An LLM could deviate from pretraining frequencies for at least three fundamentally different reasons: (1) productive generalization — the model applies learned rules or patterns to produce correct outputs that are distributionally distinct from its training examples (the desired interpretation); (2) noise or miscalibration — the model's output distribution is effectively random with respect to the task, deviating from pretraining frequencies simply because it has not learned the task at all; or (3) overfitting to spurious patterns — the model relies on superficial cues in the prompt or formatting that correlate with correctness in some training examples but do not reflect the underlying task structure. The paper's framework cannot distinguish between these. For GSM8K, where Pythia models achieve <5% accuracy and memorization is non-significant, the "generalization" conclusion may actually reflect the second case — the model is not generalizing in any meaningful sense; it is simply failing to learn the task, and its output probabilities bear no systematic relationship to pretraining frequencies because they bear no systematic relationship to the task at all. The paper's evidence for the third interpretation (productive generalization) rests on the increasing Kendall tau distance with scale and the small improvement from generalization-encouraging prompts, neither of which distinguishes between the model genuinely learning to reason versus the model simply becoming less tethered to training data frequencies without improving task-relevant behavior.

The novel n-gram count for WMT (Figure 4, top left panel) is the paper's only direct measure of generalization — showing that larger models produce more n-gram pairs never seen in any pretraining document. However, novelty alone does not imply task-appropriate generalization. A model could produce many novel but incorrect translations, which would register as high generalization under the paper's framework but would be undesirable behavior. The paper reports BLEU scores alongside novelty (Figure 3, leftmost panel), showing that novelty and performance both increase with scale, but this is a correlation, not a demonstration that the novelty reflects useful generalization rather than benign variation.

What evidence exists in the paper. The entire interpretation of results for translation, reasoning, and the MMLU knowledge-task decrease in memorization hinges on the assumption that declining memorization implies rising generalization. For WMT, the evidence is relatively strong because BLEU scores increase meaningfully (0 to 0.25–0.30) and the novel n-gram count tracks model size, providing converging evidence that larger models are both more novel and more accurate. For GSM8K, the evidence is weak — accuracy is near floor levels and "generalization" is inferred almost entirely from the absence of memorization plus a slight upward trend in Kendall tau distance. For MMLU knowledge tasks, the decreasing memorization could reflect the task-gram LM becoming less relevant (due to the rarity of specialized knowledge in pretraining) rather than the LLM becoming better at generalization.

Mitigation status. The paper does not acknowledge this limitation or propose independent metrics for generalization. The framework would be substantially strengthened by positive measures of generalization — for example, quantifying how often the LLM produces correct answers that (a) have zero or low pretraining frequency according to the task-gram LM, (b) cannot be produced by a nearest-neighbor retrieval baseline over pretraining documents, or (c) require combining multiple distinct pretraining facts in ways not attested in any single document. Such measures would convert generalization from a residual category ("not memorization") into a directly observed phenomenon with its own scaling trends. The paper gestures toward this with the novel n-gram count for WMT but does not develop analogous measures for the other tasks.

The Analysis Cannot Distinguish Task-Relevant from Task-Irrelevant Pretraining Co-occurrence, Diluting the Memorization Signal

The assumption or constraint. The task-gram LM's conditional probability is defined as C((s^x, s^y), D) / C(s^x, D) — the fraction of documents containing the input n-gram that also contain the output n-gram. The denominator C(s^x, D) counts all documents containing the input n-gram, regardless of whether those documents are in any sense "task-relevant" or structured in a way that would support learning the input-output relationship. An input n-gram like "Marie Curie" might appear in thousands of Pile documents — Wikipedia biographies, news articles, forum discussions, academic paper references — but only a small fraction of those documents present the information in a question-answering-like format where the relationship to an output n-gram like "radium" would be learnable as a task-relevant association. The task-gram LM treats all occurrences equally, so the conditional probability is diluted by the prevalence of task-irrelevant contexts.

The consequence. The memorization metric may systematically underestimate the degree to which LLMs rely on pretraining data for tasks where the input n-grams are common but their task-relevant occurrences are rare. If "Marie Curie" appears in 10,000 documents and "radium" co-occurs in 50 of them, the task-gram LM estimates P("radium" | "Marie Curie") = 0.005. But if among the 200 documents that are Wikipedia-style encyclopedic entries (the task-relevant context for factual QA), "radium" co-occurs in 40, then the task-appropriate probability is 0.20 — a 40× difference. The LLM, trained on a mixture of document types, may learn to weight task-relevant contexts more heavily (through attention mechanisms and hierarchical representations), effectively computing something closer to the 0.20 probability for QA-style prompts. The Spearman correlation between the diluted task-gram LM probabilities and the LLM's context-aware probabilities would then underestimate the true degree of memorization. This dilution effect is likely strongest for tasks where the input n-grams are common phrases that appear in many different contexts — which describes TriviaQA question n-grams well — and weakest for tasks where input n-grams are highly specific and appear in limited contexts. The MMLU knowledge-task results, where memorization decreases with model size, could partially reflect this dilution: larger models may become better at identifying and weighting the task-relevant subset of documents for rare knowledge, making their probability assignments diverge from the diluted task-gram LM probabilities even as their reliance on the task-relevant subset increases.

What evidence exists in the paper. The paper does not measure or control for the task-relevance of pretraining documents in which co-occurrences are counted. All co-occurrences within the same document are treated equally, regardless of document type, structure, or the textual relationship between the two n-grams. The paper does not report the distribution of C(s^x, D) values (how common input n-grams are overall) or the ratio C((s^x, s^y), D) / C(s^x, D) (how concentrated the co-occurrences are relative to total occurrences), which would help readers assess the magnitude of the dilution problem. The gradient influence analysis (Section 6) partially addresses this by showing that documents containing n-gram pairs are more influential than documents containing only output n-grams, but it does not compare documents where the n-gram pair appears in a task-relevant context versus a task-irrelevant context — a much more challenging analysis that would require annotating pretraining documents for task relevance.

Mitigation status. The paper does not discuss this dilution issue or propose methods to address it. Potential mitigations could include: restricting co-occurrence counting to documents of a specific type or from specific sources likely to contain task-relevant structure (e.g., Wikipedia, textbooks, Q&A forums for factual QA; bilingual or parallel corpora for translation); using paragraph-level or passage-level co-occurrence instead of document-level to increase the likelihood that co-occurring n-grams are in a meaningful relationship; or weighting co-occurrences by some measure of textual relatedness (e.g., tf-idf weighted overlap, embedding similarity between the surrounding contexts of the two n-grams). Each of these would add computational cost and methodological complexity, but they would produce a task-gram LM that more closely approximates the distribution the LLM actually learns from — which is the paper's stated goal.

The Prompt Optimization Results Demonstrate Weak and Unreliable Effects for Reasoning Tasks, Limiting Practical Applicability

The assumption or constraint. Section 7 presents prompt optimization as a practical implication of the memorization-generalization framework: if knowledge tasks benefit from memorization, prompts should be rewritten to use high-frequency pretraining n-grams; if reasoning tasks benefit from generalization, prompts should be rewritten to use low-frequency n-grams. The optimization is performed by GPT-4o in an unspecified number of iterations, using WIMBD n-gram count as a reward signal, with the meta-prompt and optimized prompts shown in Appendix E.

The consequence. The effect sizes on GSM8K are extremely small and likely not practically actionable. Across five model configurations, the difference between memorization-encouraging and generalization-encouraging prompts is at most 1.6 percentage points (OLMo-instruct 7B: 6.3% vs. 7.9%; Pythia-instruct 6.9B: 6.3% vs. 7.3%) and as low as 0.1–0.2 percentage points for some models (Pythia 12B: 2.7% vs. 2.8%). With a test set of 1,319 GSM8K examples, a 0.2 percentage point difference corresponds to approximately 2–3 examples out of 1,319. This magnitude is well within the range of random variation from prompt wording effects, and the paper does not report statistical significance or confidence intervals for these differences. The claim that "GSM8K benefits from the prompts that are less similar" to pretraining data is technically true directionally for all five models, but the effect is so small that a practitioner rewriting prompts based on this principle would see negligible practical improvement — and might overfit to the specific test set if they iterated prompt optimization using test accuracy as implicit feedback. For TriviaQA, the effects are larger and more robust (e.g., +6.5 percentage points for Pythia-instruct 6.9B, +19 percentage points for OLMo-instruct 7B), making the practical implication meaningful for knowledge tasks. The asymmetry — large effects for knowledge tasks, negligible effects for reasoning tasks — aligns with the paper's thesis but also means the framework's practical value is limited to the subset of tasks where memorization is already known to be the dominant mechanism.

What evidence exists in the paper. Table 1 reports the raw accuracies. The paper notes that "more sophisticated prompt optimization algorithms with more detailed distributional memorization feedback can be designed based on a similar idea," implying that the current results are a proof of concept rather than a mature method. However, the paper does not report how many optimization iterations were run, whether the process converged, the variance across different optimization runs, or the computational cost of the optimization process (GPT-4o API calls plus WIMBD queries). The optimized prompts themselves (Table 2, Appendix E) show lexical differences — the memorization TriviaQA prompt ("Deliver an exact single word or concise phrase in response to the factual question") uses higher-frequency words while the generalization prompt ("Formulate a distinctive and concise term or phrase to clearly answer the factual question") uses lower-frequency words like "distinctive" — but the paper does not analyze which specific n-gram choices drive the performance difference. Without such analysis, a practitioner cannot manually craft effective prompts without access to the GPT-4o + WIMBD optimization pipeline.

Mitigation status. The paper presents the prompt optimization results as a demonstration of practical implications rather than as a fully validated method, and explicitly calls for future work on "more sophisticated prompt optimization algorithms." This is a reasonable scope for a paper whose primary contribution is analytical rather than methodological. However, the paper could have strengthened this section by reporting the statistical reliability of the differences, testing alternative optimization approaches (e.g., simpler frequency-based word substitution without GPT-4o), or evaluating whether the optimized prompts transfer to held-out test sets or different model families. As presented, the prompt optimization results are suggestive but not sufficiently robust to guide practical prompt engineering decisions, particularly for reasoning tasks where the framework's predictions are least certain.

7. Implications and Future Directions

How This Work Changes the Landscape

This paper introduces a new diagnostic paradigm for understanding LLM capabilities — not a new model or training method, but a measurement framework that makes visible something previously invisible at scale: the degree to which a model's output distribution on a specific task reflects the statistical patterns of its pretraining data. The magnitude of the contribution is best characterized as a reframing with practical tooling, not a paradigm shift. The paper does not overthrow the memorization-versus-generalization debate or settle it definitively; rather, it provides the first scalable methodology for gathering quantitative evidence on the question, transforming it from a philosophical argument into an empirical one.

The methodological shift is from instance-level to distribution-level analysis. Prior work on memorization in LLMs operated at the granularity of individual training examples — "did the model see this exact string?" (Carlini et al., 2022; Zhang et al., 2023) or "would removing this example change the model's prediction?" (Feldman and Zhang, 2020). These approaches produced binary or counterfactual answers that were theoretically clean but computationally prohibitive at scale, effectively restricting memorization studies to small models, small datasets, or approximate methods. This paper's distributional memorization metric — the Spearman correlation between pretraining n-gram co-occurrence frequencies and LLM output probabilities — operates at the granularity of the entire task distribution. It does not ask whether the model memorized a specific question-answer pair; it asks whether the model's overall pattern of assigning probability to different answers on this task is predicted by how often those answers co-occurred with question patterns during training.

This shift has three cascading consequences for the field:

1. It makes memorization analysis tractable at deployment scale. The computational cost of the approach is dominated by n-gram search over the pretraining corpus (one-time per task, parallelizable via WIMBD) and forward passes through the trained LLM (negligible per example). There is no retraining, no gradient computation for the main metric, and no per-example intervention. This means the framework can, in principle, be applied to any model for which the pretraining data is available — a constraint, certainly, but one satisfied by the growing ecosystem of fully open-source LLMs (Pythia, OLMo, Amber, DCLM). The paper demonstrates this scalability by analyzing nine model sizes (13M to 12B parameters) across four tasks, computing memorization over hundreds to thousands of test n-gram pairs per task — a scale that would be entirely infeasible with counterfactual methods.

2. It creates a common language for comparing memorization across tasks and models. Because distributional memorization is a single scalar (Spearman ρ) computed through a standardized pipeline, it enables direct quantitative comparisons: "TriviaQA memorization at 12B (ρ ≈ 0.42 for n=3) is roughly 3× MMLU reasoning memorization at the same scale (ρ ≈ 0.12 for n=5)" is a statement that was not possible before this work. This common metric allows researchers to build taxonomies of tasks by their memorization profiles, track how memorization evolves through training (by applying the metric to intermediate checkpoints), and compare different model architectures or training recipes in terms of their reliance on pretraining data patterns. The paper's Figure 4 is a first step toward such a taxonomy, and the clear separation between knowledge tasks (TriviaQA) and reasoning tasks (GSM8K) on this metric suggests the measure captures something structurally meaningful.

3. It partly reconciles conflicting findings in the literature by reframing the question. Prior work produced apparently contradictory results: Carlini et al. (2022) found that larger models exhibit more verbatim memorization, while Merrill et al. (2024) found that larger models generate fewer novel n-grams. This paper's framework suggests both can be true simultaneously, because they measure different things at different levels of analysis. A model can increasingly memorize specific training strings (verbatim recall) while also developing the capacity to deviate from training data distributions when the task demands it (distributional generalization). The paper's finding that TriviaQA memorization increases with scale while WMT generalization (novel n-gram generation) also increases with scale — within the same model family — is direct evidence that memorization and generalization are not zero-sum at the model level, only at the task level. This reframing turns an apparent contradiction into a coherent picture: LLMs become both better at retrieving training patterns and better at recombining them as they scale, and which capability dominates depends on what the task demands.

Where the landscape does NOT shift. The paper does not provide a causal theory of how memorization and generalization interact during training, nor does it prescribe interventions to shift a model's reliance from one to the other. The gradient influence analysis (Section 6) is a step toward causality but is limited to 50 documents per test example — a tiny fraction of the training data. The paper also does not address the normative question at the heart of the memorization debate: is memorization desirable or problematic for a given application? High memorization on factual QA might be exactly what we want from a knowledge retrieval system; high memorization on a creative writing task might indicate undesirable regurgitation. The framework measures "how much" but does not evaluate "how appropriate," leaving that judgment to downstream users.

The research directions that become more attractive after this work include: systematic audits of pretraining data influence on specific model behaviors (the task-gram approach provides a template), comparative studies across model families to understand how architecture and training recipe affect the memorization-generalization balance, and tool development for interactive exploration of pretraining data relevance (a "WIMBD for task analysis" that researchers can use without writing custom search pipelines). The directions that arguably become less attractive — or at least, more clearly bounded — include: studies that define memorization purely as verbatim recall and generalize from that to claims about overall model behavior (the paper shows verbatim recall misses the most consequential form of memorization for task performance), and counterfactual memorization studies that retrain models from scratch to analyze individual examples (the paper provides a path to similar insights at vastly lower cost).

Follow-Up Research This Work Enables

Testing causality through controlled pretraining data ablation at small scale. The paper's distributional memorization metric is correlational — it shows that LLM output probabilities align with pretraining frequencies, but does not prove that pretraining data causes those outputs. The gradient influence analysis (Section 6) provides partial causal evidence but covers only 50 documents per example. A direct follow-up would be: take a small Pythia model (70M or 160M, where retraining is feasible), identify all Pile documents containing task-gram pairs for TriviaQA above some frequency threshold (e.g., pairs with top-10% co-occurrence counts), remove those documents from the training corpus, retrain from scratch, and measure whether TriviaQA accuracy drops more than a control condition where an equal number of random documents are removed. The distributional memorization framework predicts that removing high-count task-gram documents should specifically impair performance on test examples containing those n-gram pairs, while leaving performance on test examples with low-frequency pairs relatively unchanged. A null result — equal degradation regardless of n-gram pair frequency — would suggest the correlational memorization signal reflects a confounder rather than genuine learning from those documents. This experiment is feasible at the 70M–410M scale (training Pythia takes days on a small cluster) and would provide the strongest causal validation of the task-gram table's relevance.

Building a difficulty-aware memorization map: which specific examples within a task show the strongest memorization? The paper aggregates memorization across entire test sets, but the metric is computed per n-gram pair. A natural extension would be to report memorization at the level of individual test examples (e.g., average Spearman correlation across all task-gram pairs in that example) and analyze which example characteristics predict high versus low memorization. For TriviaQA, do questions about popular entities (high pretraining frequency) show stronger memorization than questions about obscure entities? Does question phrasing that closely matches Wikipedia text show stronger memorization than rephrased questions? For MMLU, does memorization vary by subtask in ways that align with expert judgments of knowledge-vs-reasoning requirements, or does it cut across subtask boundaries in unexpected ways? The paper already takes a step in this direction by splitting MMLU into knowledge and reasoning subtasks, but the within-subtask variance is unexplored. A heatmap of MMLU subtasks with memorization on one axis and task accuracy on the other would reveal whether high-memorization subtasks are also the ones where models perform best — if so, it would strengthen the claim that memorization drives performance for knowledge tasks; if not, it would complicate the picture and suggest additional factors (reasoning difficulty, answer format, distractor quality) matter independently of memorization.

Using the task-gram framework to compare pretraining data mixtures. The paper uses a single pretraining corpus (the Pile) and a single model family (Pythia). The OLMo-on-Dolma partial replication for GSM8K hints at what a multi-corpus comparison could reveal. A systematic follow-up would apply the full memorization pipeline to models trained on different data mixtures: Pythia on the Pile vs. OLMo on Dolma vs. any available model trained on a known corpus, measuring distributional memorization on a fixed benchmark (TriviaQA, MMLU, WMT, GSM8K). The key question: does a corpus with higher representation of task-relevant documents (e.g., more Wikipedia, textbooks, Q&A forums) produce models with higher distributional memorization on knowledge tasks, or does the model compensate by developing stronger generalization? The task-gram LM provides a natural way to quantify "task-relevance" of a corpus — the average P_{n,D}(s^y|s^x) across task-gram pairs — and the distributional memorization metric then measures whether this higher relevance translates into stronger dependency. This experiment could inform corpus curation for specific downstream tasks: if higher task-relevance in pretraining data produces models that memorize task patterns rather than generalizing from them, that has implications for whether to include or exclude task-like data from pretraining.

Stress-testing the framework on a task where memorization is demonstrably harmful. The paper's task selection spans memorization-favoring (TriviaQA) to generalization-favoring (GSM8K), but all tasks have objectively correct answers. What happens on a task where pretraining data patterns are systematically misleading? For example, a dataset of common misconceptions where the pretraining corpus overwhelmingly contains the misconception rather than the correction (e.g., "What animal has the largest eyes?" — pretraining data frequently discusses the giant squid's large eyes in hyperbolic terms, but the correct answer for largest eye-to-body ratio is a different animal). The distributional memorization framework would predict that models with high memorization on this task would tend to reproduce the misconception, while models with lower memorization (or higher generalization) would better learn the correction from supervised fine-tuning or prompting. A positive result — memorization correlates with misconception reproduction — would validate the framework's ability to detect undesirable memorization; a negative result — memorization does not predict misconception rate — would suggest the task-gram LM captures factual association frequencies but not their truth value, and that LLMs may learn to distinguish these even when both are frequent in pretraining.

Developing positive metrics for distributional generalization. The paper's largest conceptual gap is defining generalization only as the absence or decline of memorization. A well-motivated follow-up would develop independent metrics for generalization that can be computed from the same pipeline infrastructure. One candidate: measure how often the LLM assigns high probability to correct output n-grams that have zero co-occurrence count with the input n-gram in pretraining (i.e., pairs outside the task-gram table Φ). A model that frequently gets correct answers using n-gram pairs never seen together in any pretraining document is demonstrably generalizing, and the rate of such "zero-count correct answers" could be tracked across model sizes and tasks. On TriviaQA, this rate should be low (correct answers mostly come from attested co-occurrence pairs); on GSM8K, it should be high (correct solutions rarely co-occur with specific problem phrasings in pretraining). This metric would convert generalization from a residual into a positive measurement and would be directly computable from the existing WIMBD-LLM comparison pipeline with minimal modification: for each test example, check whether the correct answer contains any n-gram pairs with zero pretraining co-occurrence, and measure what probability the LLM assigns to those zero-count pairs. This is a short-term, low-cost extension that would substantially strengthen the paper's framework.

Applying the framework to intermediate training checkpoints to study the temporal dynamics of memorization versus generalization. The paper analyzes only final trained models across different sizes. The Pythia model suite releases intermediate checkpoints (every 1,000–10,000 training steps), making it possible to compute distributional memorization Mem_n(LLM_checkpoint, D|T) at each checkpoint throughout training. This would reveal whether memorization emerges early and plateaus while generalization continues to improve (suggesting a phase transition in learning dynamics), or whether both develop concurrently with memorization leading generalization by a fixed offset, or whether the relationship is task-dependent (e.g., TriviaQA memorization grows monotonically while GSM8K memorization rises then falls as the model begins to generalize). This connects directly to the literature on grokking and phase transitions in neural network training, and the task-gram LM provides a novel lens — tracking not just performance but the mechanism underlying performance — that existing accuracy-only analyses miss. A finding that memorization plateaus while generalization continues to improve on reasoning tasks would be evidence that test-time compute or inference-time strategies can extract additional capability from a fixed model without requiring more pretraining data memorization.

Practical Applications and Downstream Use Cases

Pretraining data auditing for regulated or high-stakes deployments. Organizations deploying LLMs in domains with legal or ethical requirements around data usage (healthcare, legal, finance, education) often need to understand whether model outputs on specific tasks are derived from memorized training data. The distributional memorization framework provides a concrete auditing tool: for a task of concern (e.g., answering medical licensing exam questions), compute Mem_n(LLM, D|T) using the task-gram pipeline. A high memorization score (e.g., ρ > 0.4, as seen for TriviaQA) indicates the LLM's output distribution closely tracks pretraining data frequencies, suggesting answers are predominantly retrieved rather than reasoned. This can inform decisions about whether the model is suitable for deployment (memorization of copyrighted medical textbooks might create liability), whether additional fine-tuning or alignment is needed to shift the model toward generalization, or whether certain test examples — those where the correct answer has extremely high P_{n,D}(s^y|s^x) — should be flagged as potential training data extractions. The paper's 12B Pythia results on TriviaQA, where accuracy reaches ~35% and memorization is ρ ≈ 0.42, provide a concrete benchmark: a model of this scale with this task accuracy exhibits substantial memorization, and a similar analysis on a proprietary 70B+ model with 80%+ TriviaQA accuracy would reveal whether that higher performance comes with proportionally higher memorization or with a shift toward generalization.

Prompt engineering guided by pretraining data frequency, not just intuition. The prompt optimization results in Section 7, while preliminary, demonstrate a practically actionable principle: for knowledge-intensive tasks, rewrite prompts to use n-grams with higher pretraining frequency to nudge the model toward retrieval behavior; for reasoning tasks, use lower-frequency n-grams to encourage novel recombination. The effect sizes are modest for reasoning (0.2–1.6 percentage points for GSM8K) but meaningful for factual QA (6.5 percentage points for Pythia-instruct 6.9B, up to 19 percentage points for OLMo-instruct). A practitioner building a trivia bot could implement this today: use the WIMBD or ∞-gram API to score candidate prompts by the average pretraining frequency of their n-grams, select the highest-scoring prompt, and expect a non-trivial accuracy improvement for knowledge retrieval tasks. The paper's specific optimized prompts (Table 2, Appendix E) — e.g., "Deliver an exact single word or concise phrase in response to the factual question" for memorization-encouraging TriviaQA — could be used directly or adapted. The key practical insight is that the optimization criterion (n-gram frequency in pretraining) is task-agnostic and automatable — it does not require task-specific human judgment about prompt quality, only access to a pretraining data search index. This makes it suitable for automated prompt optimization pipelines where human evaluation is a bottleneck.

Corpus curation for training models with specific memorization-generalization profiles. The finding that task performance on TriviaQA increases monotonically with n-gram pair count in pretraining data (Figure 3, second panel — 12B accuracy rising from ~5% at zero counts to >35% at high counts) provides direct guidance for corpus design: if you are pretraining a model intended primarily for factual QA, invest in including more documents with high task-gram pair density — encyclopedic sources, Q&A archives, fact-oriented textbooks — rather than narrative fiction or social media. Conversely, if you are building a model for mathematical reasoning, the paper's GSM8K results suggest that simply including more problem-solution pairs in pretraining may not improve reasoning performance proportionally, and may even be counterproductive if it shifts the model toward memorization of specific problem patterns rather than acquisition of general reasoning strategies. This is a testable hypothesis: pretrain two models on corpora with equal total size but different densities of math problem-solution n-gram pairs, and evaluate whether the high-density model shows higher GSM8K accuracy or merely higher distributional memorization without accuracy gains. The task-gram framework provides both the measure (n-gram pair density) and the outcome metric (distributional memorization) to run this experiment. For organizations with the resources to curate their own pretraining data, this represents a principled alternative to the current practice of including "as much diverse data as possible" without task-specific guidance.

Supporting decisions about whether to fine-tune, prompt, or rely on base model capabilities. A recurring practical question in LLM deployment is whether a given task is best addressed by few-shot prompting, supervised fine-tuning, or relying on the base model's zero-shot capabilities. The distributional memorization framework offers a diagnostic: compute Mem_n for the base model on the task. If memorization is high and increasing with model size (like TriviaQA), the base model already internalizes the relevant pretraining patterns, and few-shot prompting with pretraining-aligned prompts (Section 7) may be sufficient, with fine-tuning providing diminishing returns. If memorization is low or absent (like GSM8K), the base model has not internalized the task patterns from pretraining, and supervised fine-tuning on task-specific data — or investing in better reasoning-time strategies — is likely necessary for meaningful performance. The paper's MMLU results add nuance: knowledge-intensive subtasks might show moderate memorization and benefit from retrieval-augmented prompting, while reasoning-intensive subtasks within the same benchmark might require different interventions entirely. This diagnostic use of the framework treats distributional memorization as a task capability audit that precedes and informs deployment decisions, rather than as a post-hoc analysis. A team evaluating a new model on their internal benchmark could run the task-gram pipeline to estimate whether the model "already knows" their task from pretraining or needs explicit teaching — a question that currently relies on intuition and trial-and-error.

Facilitating data attribution for copyright and licensing compliance. As copyright lawsuits against LLM developers progress through courts, the ability to trace model outputs to specific training documents becomes increasingly important from both legal and ethical standpoints. The paper's task-gram table provides a lightweight attribution primitive: for a given output, identify which input-output n-gram pairs it contains, look up which pretraining documents contain those pairs, and rank documents by the number of matching pairs or by their co-occurrence frequency. The gradient influence analysis (Section 6) shows that documents containing task-gram pairs exert measurably more influence on test predictions than documents containing only individual output n-grams (Figure 5, green lines systematically above blue lines), validating that pair-containing documents are genuinely more causally relevant. While this is far from a complete attribution system — it cannot distinguish between a document that is the "source" of a fact versus one that merely mentions it incidentally — it provides a scalable first-pass filter for identifying candidate documents that may have contributed to a given model output. In a legal discovery context, this could narrow the search space from trillions of tokens to thousands of candidate documents that contain the specific n-gram associations manifested in the model's output, making subsequent human review or more expensive attribution methods feasible. The paper's demonstration that this works on a 207B-token corpus using commodity infrastructure (WIMBD API calls) makes it practically deployable without hyperscaler resources.

When to Prefer This Method

The paper does not position its framework as a "method" to be preferred over alternatives in the sense of choosing a model architecture, training recipe, or inference strategy. It is a diagnostic framework — a way to measure and understand memorization — not a model to deploy. However, within the narrow domain of choosing how to measure an LLM's reliance on pretraining data for a specific task, the paper's approach does occupy a distinct position in the design space relative to the two dominant alternatives (verbatim recall detection and counterfactual memorization), and the paper provides enough evidence to characterize when each is appropriate.

Prefer the distributional memorization framework (task-gram LM + Spearman correlation) when:

You need to measure pretraining data dependency at the task level, not for individual examples, and you care about the overall pattern rather than specific instances.
You are working with models at scales where retraining is infeasible (hundreds of millions to hundreds of billions of parameters) — the approach requires only forward passes and n-gram search, no gradient computation or weight modification.
The capability you are studying manifests as short, structured outputs (answers, translations, completions of specific n-gram spans) rather than long, open-ended generation, since the task-gram LM conditions on specific n-gram spans and requires identifying discrete output n-grams to match against pretraining co-occurrence counts.
You have access to the pretraining corpus or a searchable index thereof (WIMBD, ∞-gram, or equivalent) — without this, the empirical frequencies cannot be computed.
You are interested in distinguishing between types of tasks (knowledge retrieval vs. reasoning) by their memorization profiles, rather than making binary memorized/not-memorized classifications.

Prefer verbatim recall analysis when:

You are specifically concerned with privacy leakage, data extraction attacks, or copyright infringement where exact reproduction of training text is the legally or ethically relevant standard.
The capability of interest involves generation of long, creative, or open-ended text where output n-gram matching against a task-gram table would be too sparse — verbatim recall detects exact overlaps without requiring semantic pairing.
You need a binary, auditable criterion ("this exact 50-token sequence appears in the training data") rather than a continuous correlation.

Prefer counterfactual memorization analysis when:

You are working with models small enough that retraining from scratch is feasible (typically under 1B parameters with academic compute budgets), and you need causal rather than correlational evidence.
You are studying the effect of individual training examples — "does this specific Wikipedia article about Marie Curie cause the model to answer questions about radium correctly?" — rather than aggregate statistical patterns.
The cost of retraining can be amortized across many analyses (e.g., studying memorization of a class of examples rather than a single instance).

The paper's key insight is that these three approaches answer different questions with different cost-reliability tradeoffs, and that the field's over-reliance on verbatim recall (for historical reasons related to security research) has left a gap in our ability to answer the question most relevant to understanding capabilities: "does the model's behavior on this task reflect its training data, and if so, how much?" Distributional memorization fills that gap at a cost that scales gracefully to current LLM sizes, making it the appropriate tool for that specific question — but not a replacement for verbatim recall in privacy contexts or counterfactual analysis in high-certainty causal studies.