Think or Remember? Detecting and Directing LLMs Towards Memorization or Generalization

We appreciate the reviewer’s thoughtful feedback, which prompted us to conduct additional experiments, yielding encouraging initial findings that further support our hypothesis. While we acknowledge that applying this approach to broader real-world contexts (e.g., larger models, diverse tasks) requires further exploration and will be addressed in future work, we believe the current study’s contributions are significant. Sharing these findings now can provide valuable insights to the community and inspire further research to advance our understanding and control of memorization and generalization behaviors in LLMs. Below, we address the concerns and questions raised.

Weakness

1. Definition of Generalization and “Correct Reasoning”

We agree that the definition of generalization requires further clarification. In our study, we define generalization as the ability of the model to generate outputs based on reasoning derived from patterns learned during training, without directly recalling specific instances. The datasets we used were designed so that the outputs can be clearly categorized as either memorization or generalization, minimizing ambiguity in the distinction. We will ensure this is made clearer in the revised paper.

2. Memorization and Generalization within Designed Dataset

We recognize that reliance on specialized datasets may limit immediate practical applicability. However, these datasets are designed to provide a controlled environment for isolating and studying behaviors, a necessary step before extending to real-world, multi-task settings. As noted in the paper, we are actively exploring how to generalize these findings to larger models and multiple tasks, for example, we further found that common neurons for generalization/memorization are shared among multiple tasks, which we plan to report in future work with comprehensive analysis.

3. Low Rate of Behavior Change

The modest rate of behavior change is expected, as we do not anticipate large shifts in LLM behavior with only minor input re-phrasing. Additionally, while generating multiple test instances may introduce some randomness in the timing of behavior transitions (from memorization to generalization or vice versa), once a behavior is determined, the model’s internal representation at that point remains stable and informative for subsequent experiments. We believe this approach still provides valuable insights for the analysis.

4. Neuron-Wise Mean Difference and Detailed Analysis

In the original paper, we select neurons with a Neuron-Wise Mean Difference (NMD) that are strongly correlated with the distinction between memorization and generalization behaviors, as this method yields more effective results than random selection of neurons. In the revised paper, we will provide additional insights into how specific neurons behave in relation to these two mechanisms. Regarding concerns about the limited datasets, we have recently obtained positive findings showing that certain neurons exhibit consistent behavior across multiple tasks. We will include these results in the future work with more comprehensive experiments.

5. Generalization the Results

While our current study demonstrates behavior differentiation primarily in two tasks, we agree that further validation across a wider range of tasks is necessary to establish the robustness of spatial differentiation and predictability of behaviors, and we acknowledge this is only a starting point. In future iterations, we plan to explore whether these patterns generalize across diverse datasets and task domains, ensuring a more comprehensive understanding. However, we think sharing current findings can still provide valuable insights to the community and inspire further research to advance our understanding and control of memorization and generalization behaviors in LLMs.

Questions

1. Model Performance After Training

In our study, we do not define "performance" in the traditional sense, as we do not specify memorization or generalization as the correct answer. Instead, our goal is to trigger both behaviors and investigate the detection and control mechanisms that differentiate them. Thus, we focus on understanding the dynamics between memorization and generalization, rather than optimizing for one behavior over the other.

2. Cluster and Interactions Between Neurons

The influence of clusters or interactions between neurons is a valid point. While our initial analysis focuses on neuron-wise mean differences, we plan to investigate cluster-level interactions in future work. This will allow us to better understand how groups of neurons contribute to the overall behavior change.

3. Meaning of “Other” in Tables

It means the result is neither memorization nor generalization, e.g. irrelevant output. We will clarify it in the revised version.

Dear Reviewer wpuf,

Thank you for your valuable feedback. We have carefully addressed your comments in our response and would greatly appreciate it if you could kindly review it and share any further thoughts.

Thank you for your time and consideration.

Thank you for answering my questions and addressing my concerns. I think the additional experiments contribute to a deeper analysis. However, my concerns regarding the use of restricted datasets and models remain, as these limitations affect the generalizability and broader applicability of the findings. While I recognize the value of sharing these initial insights, I believe further work is needed to address these issues. Therefore, I will maintain my original score.

We appreciate the reviewer’s thoughtful feedback, which prompted us to conduct additional experiments, yielding encouraging initial findings that further support our hypothesis. While we acknowledge that applying this approach to broader real-world contexts (e.g., larger models, diverse tasks) requires further exploration and will be addressed in future work, we believe the current study’s contributions are significant. Sharing these findings now can provide valuable insights to the community and inspire further research to advance our understanding and control of memorization and generalization behaviors in LLMs. Below, we address the concerns and questions raised.

Weakness

1. Generality of conclusions

We acknowledge that our current experiments are limited to synthetic datasets and relatively small model scales. However, our findings represent an initial step toward understanding the localization of memorization and generalization behaviors in LLMs. While these results are derived in a controlled setting, we believe they provide a solid foundation for future exploration in more complex, real-world scenarios. In fact, recently we have found positive findings on the pretrained LLM such as Llama3-instruct, and also shared neurons for memorization/generalization across multi-tasks, and we plan to report it in the future work with more comprehensive analysis. However, we think sharing current findings now can provide valuable insights to the community and inspire more research directions.

2. Technical contributions

We agree that some techniques, such as inference-time intervention (ITI) and causal analysis, build upon prior works. However, our primary contribution lies in their novel application to investigate a previously unexplored research question: understanding the interplay between neuron-level behaviors and memorization/generalization in LLMs. We believe this unique perspective adds value to the field and provides a starting point for deeper investigation.

Questions

1. training mechanism

We tried both fine-tuned and from-scratch mechanism. For in-context inference, we fine-tuned the pretrained GPT-2-Medium with our designed datasets, and for arithmetic addition, we trained the GPT2 model architecture from scratch. The experimental results reveal memorization/generalization behavior differentiation in both scenarios.

2. re-phrase the input prompts still introduce additional variables

The reviewer's concern is valid and appreciated. However, in order to examine the differences in the model’s internal representations between memorization and generalization behaviors, we need to ensure that the model’s internal states differ between the two behaviors (with the same input, the model internal representations are identical). To achieve this, we set the temperature to 0 to generate deterministic outputs. And the variability in behavior and internal representations is then introduced through the re-phrasing mechanism. While we acknowledge that re-phrasing could introduce additional variables, we believe these effects are minimal, allowing us to isolate the specific behaviors of memorization and generalization.

Dear Reviewer 6PR6,

Thank you for your valuable feedback. We have carefully addressed your comments in our response and would greatly appreciate it if you could kindly review it and share any further thoughts.

Thank you for your time and consideration.

We appreciate the reviewer’s thoughtful feedback, which prompted us to conduct additional experiments, yielding encouraging initial findings that further support our hypothesis. While we acknowledge that applying this approach to broader real-world contexts (e.g., larger models, diverse tasks) requires further exploration and will be addressed in future work, we believe the current study’s contributions are significant. Sharing these findings now can provide valuable insights to the community and inspire further research to advance our understanding and control of memorization and generalization behaviors in LLMs. Below, we address the concerns and questions raised.

Weaknesses

1. Focus on Small Models and Synthetic Datasets:

We fully acknowledge the limitations of our study's scope, focusing on smaller models (GPT-2 and GPT-2-medium) and synthetic datasets. These choices were intentional to maintain experimental control and clarity in identifying neuron-level behavior patterns. However, we agree that expanding the analysis to larger models and more complex, real-world tasks would provide stronger empirical support and improve the generalizability of our findings. In fact, we are actively working on experiments involving larger models (e.g., ＬLlama3) and broader tasks, and currently with some positive findings, which we hope to report in follow-up work with comprehensive analysis. However, we still think it beneficial to share the current work with the community and potentially inspire other researchers.

2. Generalization of Observed Patterns Across Tasks:

While our current study demonstrates behavior differentiation primarily in two tasks, we agree that further validation across a wider range of tasks is necessary to establish the robustness of spatial differentiation and predictability of behaviors. Our current findings suggest that behavior differentiation is more pronounced in the later layers, but we acknowledge this is only a starting point. In future iterations, we plan to explore whether these patterns generalize across diverse datasets and task domains, ensuring a more comprehensive understanding. In fact, we have found that common neurons for generalization/memorization are shared among multiple tasks, which we plan to report in future work with more comprehensive analysis.

3. Novelty and Practical Implications at Scale:

We appreciate the reviewer’s suggestion to clarify the broader implications of our findings. While the method for detecting hidden layer activations is relatively straightforward, the novelty lies in the combination of analysis and intervention techniques to steer memorization/generalization behaviors. The practical implications are significant, particularly for tasks requiring control over memorization and generalization, such as privacy-preserving language generation. In the revised version, we will explicitly connect our findings to potential real-world applications to better highlight their impact.

Questions

1. Training Objective and Examples:

Yes, we fine-tuned the models using a causal language modeling objective. For example, in one synthetic dataset, the input might be: "Input: 2+3+5\nTarget: 10\n" This format allows the model to infer either the memorized mapping (e.g., for seen inputs) or generalize (e.g., for unseen inputs) depending on the training setup. We will include additional concrete examples in the revision to clarify this process.

2. Monitoring Memorization and Generalization:

We chose to monitor behaviors during training and preserve the model when it exhibited both memorization and generalization on a held-out test set because these behaviors form the focus of our study. Training to convergence or optimizing validation set performance may result in models biased toward one behavior (either memorization or generalization), which would undermine the balanced analysis we aimed to achieve. While this approach deviates from standard practices, it was specifically designed to align with our study’s objectives. We will provide a more detailed rationale for this choice in the revised manuscript.

Dear Reviewer PREZ,

Thank you for your valuable feedback. We have carefully addressed your comments in our response and would greatly appreciate it if you could kindly review it and share any further thoughts.

Thank you for your time and consideration.

We appreciate the reviewer’s thoughtful feedback, which prompted us to conduct additional experiments, yielding encouraging initial findings that further support our hypothesis. While we acknowledge that applying this approach to broader real-world contexts (e.g., larger models, diverse tasks) requires further exploration and will be addressed in future work, we believe the current study’s contributions are significant. Sharing these findings now can provide valuable insights to the community and inspire further research to advance our understanding and control of memorization and generalization behaviors in LLMs. Below, we address the concerns and questions raised.

Weaknesses

1. Correlation to Neuroscience

Our inspiration stems from foundational work like Brodmann's cortical localization (Garey, 1999) and early computational models such as the perceptron (Rosenblatt, 1958), which established the link between neural and artificial systems. Recent studies further explore this connection, investigating alignment between LLMs and human brain activity (e.g., Tang et al., 2023; Aw et al., 2023; Ren et al., 2024) and drawing on neuroscience-inspired methods like neural cluster knockout (Bhile & Maes, 2024). Given this context, we believe it is intuitive and valuable to explore whether localization principles in the brain extend to LLMs, as our study contributes to this growing interdisciplinary dialogue. We will add more context regarding to the correlation to neuroscience.

2. Dataset specialization and single-task limitation

We recognize that reliance on specialized datasets may limit immediate practical applicability. However, these datasets are designed to provide a controlled environment for isolating and studying behaviors, a necessary step before extending to real-world, multi-task settings. As noted in the paper, we are actively exploring how to generalize these findings to larger models and multiple tasks, for example, we further found that common neurons for generalization/memorization are shared among multiple tasks, which we plan to report in future work with comprehensive analysis.

3. Binary behavior shift

Based on our experimental results, we have observed that the difficulty of transitioning from memorization to generalization differs from the reverse, suggesting that these are not strictly binary scenarios. We will clarify this observation in the paper and highlight it as an area for future exploration, where we plan to conduct more comprehensive experiments to further investigate this phenomenon.

4. Granularity of neuron-level interventions

We agree that higher-level interventions, such as prompt engineering or input modifications, may complement neuron-level adjustments. However, our work focuses on neuron-level interventions as a first step to empirically identify the mechanistic underpinnings of behavior in LLMs. This granularity allows precise control and a better understanding of the internal workings, which can inform future higher-level approaches.

Questions

1. Choosing values for topN and alpha

Selecting optimal values for topN and alpha is indeed crucial. In our experiments, we empirically determine these values through grid search on the validation set. We acknowledge that optimal values may vary across architectures, datasets, and tasks. We will emphasize this in our revision.

2. Scope of Neuron Adjustments

For the current study, adjustments were applied to neurons across all layers. However, the impacted neurons are targeted according to the correlation with behavior shifts, and our analysis shows that neurons in the later layers exhibit the most pronounced correlation with behavior shifts. We will include visualizations in the revised paper to highlight the distribution of impacted neurons and the specific layers they belong to.

3. Definition of "Spatial" Differentiation

By “spatial differentiation,” we refer to the observation that neurons associated with memorization and generalization tend to activate in distinct regions of the latent space. This term highlights the structural separation in the model’s internal representations. We will clarify this terminology in the paper for better understanding.

Dear Reviewer ct23,

Thank you for your valuable feedback. We have carefully addressed your comments in our response and would greatly appreciate it if you could kindly review it and share any further thoughts.

Thank you for your time and consideration.