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

Thanks for your comments. We appreciate the reviewer found our investigation thorough with a wide range of tasks. Below, we provide a point-by-point response to your review.

Weaknesses:

W1. Contribution significance.

We want to point out that the mentioned workshop is non-archival, and it is generally permissible to resubmit a paper accepted to a non-archival workshop. Regarding the contribution of our paper, we define a new type of memorization, beyond any existing type of memorization. Note that our conclusion of LLMs performing knowledge-intensive tasks primarily by memorization and reasoning-intensive primarily by generalization is new and has not been found by any previous work to the best of our knowledge. Our analysis presents important new evidences regarding when LLMs memorize and when they generalize. Our proposed novel analysis technique also has interesting practical implications as shown in point 4 of the general response.

W2. Novelty compared to knowledge localization work.

Thanks for bringing up the knowledge localization papers. The storage and retrieval of factual knowledge have been studied by numerous papers. However, our paper studies the interplay between generalization and memorization, which is a different topic that was not studied before, as far as we know. "Factual question answering shows an increase in memorization with increasing LLM size” is just one of our conclusions. We also draw conclusions on LLMs performing harder, reasoning-intensive tasks primarily based on generalization, which is not studied by the knowledge localization papers. Our paper also presents a novel definition of memorization, which better captures LLMs’ generalization and memorization behaviors at scale, compared to other possible alternatives like infini-gram.

W3. Practical implication.

To show the practical implication of our study, we propose a prompt optimization technique based on our observations in point 4 of our general response. Our findings suggest that LLMs’ task performance can be improved by encouraging distributional memorization or generalization depending on the task type. So we propose to use GPT4o to rewrite the task prompts and optimize memorization/generalization, which is shown to be effective on TriviaQA and GSM8K with the Pythia 12B model.

W4. Only used Pile and Pythia.

We include some new results on Dolma with GSM8K in the revision. We would love to extend our results to more pretraining corpora but most of the open-sourced LLMs trained on other corpora do not have as many model sizes as the Pythia model family, making it difficult to see the relationship between memorization and model size/task performance.

W5. More suitable reasoning tasks.

Thank you for suggesting GSM8K as a new reasoning dataset. We have included new experiments with GSM8K in the revision. For a detailed description of the new results, please see point 1 of our general response. We found both the performance v.s. n-gram count results and the distributional generalization results on GSM8K align with our previous findings.

W6. Why use WMT09, not newer.

We have performed a data contamination search with respect to all testing data, as stated in the first paragraph of Section 4. More specifically, we also perform n-gram searches with n=8 and n=14. Neither of them found any match. Checking possible testing data contamination using large n is adopted from the GPT3 technical report: they remove contaminated training texts by searching for n-gram with n=13. We chose WMT09 as our testing data instead of newer versions of WMT because WMT09 contains more European languages, which is more prominent in the Pile. Thanks for bringing this up, we have added this explanation to the revised paper.

W7. Justification of using Spearman correlation to measure memorization.

Memory and generation themselves are not well-defined: they can have a lot of different meanings. There is not a single way to quantify all possible aspects of memorization. In our paper, as stated in Definition 3, we only want to study one type of memorization: distributional memorization, which is defined by Spearman correlation. As the general memorization is not well-defined, it is not possible to provide a mathematical argument to justify such a definition.

To empirically support our correlation-based analysis, in Section 6, we have shown a gradient-based experiment, which estimates the causal influence of the task-gram data throughout pretraining. In the last paragraph of this section, we discuss how these results support our correlation-based study: task-grams consistently have a higher influence than infini-gram. Translation and MMLU show decreased influence while TriviaQA shows increased influence, which has the same trend as the correlation-based distributional memorization.

Thanks for your positive comments! We appreciate that the reviewer found the question we studied important and our paper “structured, clear, and easy to read”. The reviewer also found our experimental methodology sound, and liked that we confirmed our findings with a gradient-based method and checked possible test set contamination. Below, we provide a point-by-point response to your review.

Weaknesses:

W1. Hyperparameter choice and sensitivity.

We chose the cosine similarity threshold based on the human inspection of the task-gram quality (i.e. how well they represent the current task). We use smaller thresholds for larger n-grams to tolerate more noises contained in them. We have added a cosine similarity threshold sensitivity experiment in Figure 5 in the Appendix as described in point 3 of the general response. The trend of distributional memorization does not change across different cosine similarity thresholds (0.7, 0.75, and 0.8), and our method is in general robust to the choice of cosine similarity threshold. The 0.05 difference between WMT and other datasets is only a result of human preference, which does not affect the final conclusion.

The choice of dataset comes from the consideration of balancing different types of tasks, and we believe the combination of translation + factual QA + reasoning can well represent the spectrum of possible tasks. To enhance the selection of reasoning tasks, we also add GSM8K as a new task as described in point 1 of the general response. We found both the performance v.s. n-gram count results and the distributional generalization results on GSM8K align with our previous findings.

The choice of embedding model comes from accommodating the need for different tasks. We use the LASER embedding model for the Translation task because it is specialized for translation. We use E5 for other tasks because it is a general-purpose sentence embedding model that performs cosine-similarity-based retrieval.

Thanks for bringing this up. We have added the above discussion to Appendix D.

W2. The possible influence of generation fluency.

The LLM predicted probability is not evaluated on LLM-generated outputs, but on the ground truth outputs, thus does not have the fluency issue. The LLM-generated outputs are used for evaluating the task performance and counting the number of novel/unseen n-gram pairs.

Questions:

Q1. Stop word filtering and NER-based post-filtering.

For the main paper results, we only perform stop word filtering with Translation tasks, as the translation of stop words in European languages is very similar and can result in identical pairs. For other tasks, we do not perform stop word filtering. In Appendix F, we showcase the matched task-gram pairs in Pythia model generations, while we mine these pairs from the ground truth data. For displaying per-instance results we used both techniques interchangeably, just to showcase their occurence in the pretraining data.

For additional filtering of the mined n-gram pairs, we have not performed any such (e.g. NER based) post-filtering. Due to the short time constraint of the rebuttal period, we are not able to extensively modify and rerun our pipeline at this time, but we would like to investigate this factor after the rebuttal.

Q2. MMLU performance drops at 150 frequency.

Yes. As stated in the second last paragraph of Section 4: “A closer inspection of the test examples in this interval reveals that they contain more reasoning or math problems, which appear to be harder for Pythia models.” More specifically, this interval contains more n-grams with numbers and math operations. The questions containing these n-grams likely contain math operations.

Q3. Choice of large n-gram for decontamination.

For contamination check, we do not remove stop words. In addition to n=14, we also perform a check with n=8. Neither of them found any match. Checking possible testing data contamination using large n is adopted from the GPT3 technical report: they remove contaminated training texts by searching for n-gram with n=13.

Thank you for your positive comments! We truly appreciate that the reviewer found our paper to be “exceptionally well written” and “very enjoyable to read it end to end”. The reviewer also found our “methodological contribution is strong” and our “method can likely be extended to multiple different tasks easily”. Comparing with existing literature, the reviewer found our paper “makes an important contribution by computing the statistics at scale”. The reviewer also found we have a “good choice of downstream tasks”. Below, we have provided a point-by-point response to your review.

Weaknesses:

W1&2. More downstream tasks.

Thank you for suggesting GSM8K as a new reasoning dataset. We have included new experiments with GSM8K in the revision. For a detailed description of the new results, please see point 1 of our general response.

W3. Only use Pile.

We include some new results on Dolma with GSM8K in the revision. However, since OLMo only has two model sizes, we can only show results with these two sizes. We would love to extend our results to more pretraining corpora but most of the open-sourced LLMs trained on other corpora do not have as many model sizes as the Pythia model family, making it difficult to see the relationship between memorization and model size/task performance.

Questions:

Q1. Percentage of zero denominator.

The percentage of n-gram pairs that we need to drop due to the zero denominator is relatively small, around 5-10%, because single n-grams are usually easy to find in Pile.

Q2. Synonym n-grams.

We do observe the existence of synonyms in mined n-gram pairs, especially in MMLU and GSM8K, likely because many of the input and output texts are related logically instead of semantically. However, some synonyms are potentially meaningful, for example, the same phrase in question form and not in question form. We currently do not have a filter for synonyms, but we would like to investigate this factor after the rebuttal as there is not enough time to modify and rerun the searches before the rebuttal ends.

Q3. Testing sample size as a confounder.

Since the distributional memorization is calculated on an n-gram pairs level, and usually an n-gram pair only appears once in the test set, the effective example number should be the number of non-zero count n-gram pairs, which is actually similar for MMLU and TriviaQA. We will add an experiment with the same number of n-gram pairs for all datasets.

Q4. Why use both infini-gram and WIMBD.

We use infini-gram for counting the denominator mainly because it is much faster than WIMBD, which enables faster experiment iteration. The occurrence of single n-grams is more frequent and thus can tolerate the counting approximation used in infini-gram. The cooccurrence of n-gram pairs is rare and infini-gram’s approximation ends up overlooking them, thus we have to use the full search of WIMBD.

Q5. Different memorization trends between infini-gram and task-gram.

The infini-gram only produces probability for single n-grams, while the task-gram produces conditional probability for task-related n-gram pairs. The infini-gram reducing with scale for TriviaQA reflects that larger LMs do not plainly memorize the local probability distribution, but a task-relevant long-range probability distribution.

Q6. Progressively adding task-related data.

The shape of the distributional memorization curve would depend on the effect of the task-related data: whether its effect of helping the LM to generalize to other unseen cases is stronger or its effect of grounding the LM to memorize its distribution is stronger. The accurate answer can only be confirmed with experiments, but from the existing observations, we expect that knowledge-intensive tasks should display increasing memorization, while reasoning-intensive tasks should display increasing generalization. From our current results, the task-related data for reasoning-intensive tasks seem to help the task performance by encouraging LMs to produce more novel generations, while for knowledge-intensive tasks, the task-relate data encourages memorization.

Suggestions:

S1: Thanks for pointing this out. We have combined the MMLU figures together under the same scale in the revised paper.

S2: Thanks for the suggestion. We have added such a plot in the revision.

Questions:

Q1. n-gram merging.

The overlapping n-grams are counted individually. Note that the distributional memorization is computed on a n-gram level instead of on an example/sentence level, so the probabilities of individual n-grams do not need to be merged. Thus, the text length does not impact the memorization score computed for the current task.

Q2. Short task input description.

We assume that the input description referred to by the reviewers means the task instruction. In our experiments, we only use task descriptions within one or two words, like “Question: “. Also, the task instruction is not involved in mining the task-gram pairs, only the raw task input-outputs are used. We do not see how short input descriptions would impact our method. On the other hand, very short task inputs or outputs that are smaller than the used n-gram would have an impact. In this case, we use the whole input or output as the n-gram if their length is smaller than n. For the sparse n-gram distribution issue, we only use the task-gram pairs with non-zero counts to compute the distribution memorization score. As shown in point 2 of the general response, there are in general enough task-grams for each task.

Q3. Memorization v.s. Task-related patterns.

Please see our response to W2. While the task-gram pairs counts are indeed important for task performance as shown in Figure 2, our proposed distributional memorization does not correlate with task performance.

Q4. Decreased ability in capturing rare knowledge v.s. Better generalization.

We conclude that the decreased distributional memorization in MMLU implies better generalization because the performance is also increased for larger models. It is not reasonable if both memorization and generalization decreased but the performance increased.

Q5. Alternative method in Appendix.

The scope and setting of the project were very different from our current paper in our early exploration. This alternative method is not used for identifying parallel n-gram structures, but for identifying any general parallel pieces. In this alternative method, we prompted GPT3.5 or fine-tuned Llama with a pretraining document and its entity tags produced by another NER model, and instructed the model to classify whether this document contains any “parallel structure”. Our initial hypothesis was that the general “parallel structure” presented in the pretraining data, like QA or translation-like text pieces, can be viewed as a form of implicit supervision of the LLMs, which is the reason why the LLMs can perform diverse tasks in zero-shot. However, since the “parallel structure” itself is not well-defined, our attempt to identify it failed.

Q6. Choice cosine similarity threshold.

We chose the cosine similarity threshold based on the human inspection of the task-gram quality (i.e. how well they represent the current task). We use smaller thresholds for larger n-grams to tolerate more noises contained in them. We have added a cosine similarity threshold sensitivity experiment in Figure 5 in the Appendix as described in point 3 of the general response. The trend of distributional memorization does not change across different cosine similarity thresholds, and our method is in general robust to the choice of cosine similarity threshold.

Q7. Cross-task memorization.

We primarily focus on the dynamic between memorization and generalization within the same task in this paper. While the cross-task memorization effect is an interesting direction that our proposed method can be used to study, we leave it to future research due to space and time limitations.

Q8. Practical implication.

Yes. We propose a prompt optimization technique based on our study in point 4 of our general response. we propose to use GPT4o to rewrite the task prompts and optimize memorization/generalization, which is shown to be effective on TriviaQA and GSM8K with the Pythia 12B model.

Thanks for your comments. We appreciate that the reviewer found our finding “an insight useful for understanding scaling effects”. Below, we provide a point-by-point response to your review.

Weaknesses:

W1. Task-gram sufficiency in capturing complex tasks.

While task-gram is not as complex as LLMs, we found it strikes a good balance between expressiveness of the task and faithfulness to the pretraining data. Our extensive empirical analysis also shows that task-gram can explain a good portion of LLMs’ task performance by distributional memorization or generalization. As suggested by the reviewer, we have added a new experiment on GSM8K, which is detailed in the general response and the revised script.

W2. Memorization v.s. task features.

The task-grams mined by our method are indeed essential to the task performance, as shown in Figure 2 of the paper, that all task performance is positively correlated with more task-gram counts. However, distributional memorization as defined in Definition 3 with Spearman correlation, does not necessarily positively correlate with the task performance. In Figure 3, Translation and GSM8K do not show a significant distributional memorization effect. MMLU shows a negative correlation between task performance (implied by the model size) and distributional memorization. In summary, our results show that distributional memorization does not correlate with task performance.

W3. Related work discussing similar topics.

Thanks for pointing out the related work, which is a nice position/survey paper that proposes a taxonomy that summarizes all existing LLM memorization works together. We have cited it in our related work section in lines 516-517. Note that our paper defines a new type of memorization, beyond any type of memorization discussed in the reference paper. Note that our conclusion of LLMs performing knowledge-intensive tasks primarily by memorization and reasoning-intensive primarily by generalization is new and has not been concluded by any previous work, to the best of our knowledge. Whether LLMs primarily memorize or generalize when performing different tasks still does not have a good consensus in the community. Our analysis presents important new evidence regarding when LLMs memorize and when they generalize. Our proposed novel analysis technique also has interesting practical implications as shown in point 4 of the general response.

4. Practical implication:

An important observation of our study is that knowledge-intensive tasks benefit from LLMs’ distributional memorization, while reasoning-intensive tasks benefit from LLMs’ distributional generalization. Then it is possible to design or rewrite the prompt according to this principle to improve an LLM’s task performance, based on the hypothesis that the LLM generation distribution is strongly affected by the prompt distribution. More specifically, to encourage memorization, we can rewrite the task instruction to be more similar to pretraining data in terms of n-gram counts. To encourage generalization, we can rewrite the prompt to be less similar to training data.

We implement a simple prompt optimizer based on GPT4o and the WIMBD n-gram count feedback. More specifically, we instruct GPT4o to rewrite a given task prompt at each iteration, and give the average n-gram count in the pretraining corpus of the rewritten prompt to GPT4o in the next iteration as the reward. We instruct GPT4o to maximize this reward if we want to encourage memorization, and instruct it to minimize this reward if we want to encourage generalization. Here, we show a maximization and a minimization result for TriviaQA and GSM8K respectively. We report zero-shot testing accuracy with the Pythia models and OLMo models below. The meta prompt we used to perform such optimization is included in Appendix E.

We also show the optimized prompts below:

Note that the lengths of the optimized prompts are not significantly different, while TriviaQA significantly benefits from the prompts that are more similar to the pretraining data, and GSM8K benefits from the prompts that are less similar. More sophisticated prompt optimization algorithms with more detailed distributional memorization feedback can be designed based on a similar idea. We leave the investigation of other possibilities for future work.