The Pitfalls of Memorization: When Memorization Hurts Generalization

Thank you for your detailed review and constructive feedback. Below we address your questions. We are happy to continue the discussion in case of further follow-up questions. We have split our response into two parts for better readability.

Memorization Types in Section 5:

In Section 5, towards the end of the paper, we wanted to provide a thought-provoking discussion on the other side of memorization. As explicitly posed in line 377, we ask the question, “Is memorization always bad?” We apologize if our response to this question was unclear. The short answer is no, memorization is a double-edged sword. While it can sometimes harm generalization, there are situations where it plays a positive role, and ERM might even outperform methods like MAT in such cases.

To address your specific questions:

Does it mean the memorization at the presence of spurious features is always good/bad/ugly?

Memorization can indeed be categorized as good, bad, or ugly depending on the context, the inductive biases of the learning algorithm, and the very definition of spurious features (e.g., the type of leaves an insect is sitting on is spuriously correlated with the type of bug, yet entomologists use this signal to identify the insect). In this section, we only scratched the surface of this complex topic, hoping to inspire future research to disentangle the interplay between memorization and generalization more deeply.

Is the proposed training method going to help mitigate any type of memorization discussed here?

As stated in line 439, MAT specifically addresses “bad memorization”. In scenarios where memorization is not harmful, or is even beneficial, MAT may hurt the performance. To see this, take the Waterbirds task: suppose that the task were to classify the background rather than the bird type. ERM would perform better by memorizing the minority (which are now considered mislabeled), while MAT would (incorrectly) learn the (now spurious) foreground feature.

Can those types of memorization be easily identified in real-world scenarios?

As discussed in lines 60–63, the primary signal for identifying harmful memorization is the model's test performance, particularly on worst-group. This true generalization signal can guide us in determining which method to use. For standard tasks where harmful memorization is unlikely, ERM is generally sufficient. However, in scenarios where spurious correlations are anticipated, or when test performance suggests that, methods like MAT or other invariant learning techniques should be considered.

We will revise the manuscript to make these points more explicit.

Results Discussion about Table 1:

Thank you for raising this point. We have now added a new section in Appendix E (highlighted in green for better visibility) titled "Further Discussion on Experimental Results", to provide more details. In the final version, we will move parts of it to the main text.

The newly added section clarifies the different levels of assumptions about access to group annotations and highlights their implications for methods like MAT and GroupDRO. Specifically, we discuss how:

• GroupDRO relies heavily on group annotations, either predefined or inferred, for its reweighting strategy and performs well when full or partial access to group annotations is available.
• MAT, by contrast, is group-agnostic, requiring no explicit group annotations during training. This makes it more robust in scenarios where such annotations are unavailable or ambiguous, although it does rely on inferred validation group annotations for model selection.
• We also emphasize that MAT’s ability to perform competitively without relying on training group annotations is a key strength, especially in real-world applications where such information is often unavailable.

Ablation on Logits Shift Estimation:

Thank you for the suggestion. In the newly added section in Appnedix F (highlighted in green for better visibility) titled "Ablation Study on Estimations in $p_a(y \mid x_i)$", we have now conducted an ablation study to evaluate the accuracy and impact of the components used in the logit shift estimation $p(a \mid **x**_i)$ and $p(y \mid a)$. Specifically, we compare configurations where these components are either replaced by their ground-truth values or estimated using Equations (4) and (5). The results, detailed in the newly added table in the paper, show that while using ground-truth values yields the best performance, the fully estimated version of MAT still significantly outperforms ERM. This highlights the robustness of the estimation process and demonstrates MAT’s effectiveness even without access to ground-truth annotations. We hope this addresses your question.

Further Self-influence Analysis:

Thank you for the suggestion. In the revised manuscript, we have added a self-influence analysis for GroupDRO in Appendix G (highlighted in green for better visibility) titled "Additional Influence-Score Experiments," to provide a direct comparison. The results confirm that MAT is more effective at reducing self-influence scores across all subpopulations (groups). GroupDRO also shows some level reduction in the self-influence score of groups, specifically for the minority groups. This further indicates that memorization is limiting generalization, and mitigating it contributes to improved generalization in subpopulation shift settings.

Validation Set and the Amount Shift:

To clarify, "MAT does not use the validation sets to shift the logits". MAT uses the held-out predections from the training set. These held-out predictions are obtained by randomly splitting the training data into two halves and training two models simultaneously. Consequently, each training sample is seen and trained on by only one model, while the other model provides the held-out predictions for it. This enables us to compute held-out predictions, which serve as a strong indicator of where models may fail, enabling logit shifting for each individual training sample. We will clarify this in the paper.

Learning Order of Features:

In Figure 1 (left and middle panels), you can observe that both the training and test accuracy for the majority groups (red and black solid lines) increase rapidly to nearly 100% accuracy after a few epochs. These majority groups are where the spurious features $x_a$ dominate the learning dynamics, meaning the model first learns $x_a$. Meanwhile, we observe a delayed learning process (at epoch ~200) for the minority groups (red and black dashed line).

In the left panel, where noise features are off ($\sigma_{\epsilon} = 0$) and memorization through these features is not possible, the model eventually learns the generalizable pattern (by learning the core feature $x_y$ from the minority groups, achieving 100% accuracy in both training and test sets for the minority groups.

In contrast, in the middle panel, where noise features are present ($\sigma_{\epsilon} = 1$), memorization is possible. Here, the model tends to rely on these features instead of the core feature $x_y$ to reduce loss, resulting in 100% training accuracy, but significantly lower test accuracy.

We’d be happy to clarify any remaining questions or engage further in discussion. Thanks again for your review and consideration!

Dear Reviewer,

Thank you once again for taking the time to provide your valuable feedback on our paper. With the discussion phase nearing its end, and to ensure that our responses are fully considered, we request that you kindly confirm whether you’ve had a chance to review them.

We hope our responses have addressed your concerns and clarified any ambiguities. If there are any further questions or suggestions, we would be glad to address them promptly.

Additionally, if our responses have adequately resolved your concerns, we kindly ask you to consider reflecting this in your review score.

Thank you!

Regarding the rightmost plot in Figure 3:

Sorry for the confusion! Figure 3 and the experiments in Section 5 are based on standard ERM in a linear regression task, and MAT is not being used here. To clarify, the main goal of Figure 3 and Section 5 is to discuss the broader implications of memorization and to highlight the need to study memorization under various setups. We aimed to provide a thought-provoking and balanced discussion on the other side of memorization, arguing that while it can sometimes harm generalization, there are situations where it plays a positive role, and ERM might even outperform methods like MAT in such cases. We will make this more clear in the next revision.

Regarding Influence scores distributions:

This is, in fact, related to the characteristics of the Waterbirds dataset. Specifically, "LB on Land" and "WB on Water" are the majority groups within their respective classes, but the dataset is also class-imbalanced. For example, the "LB on Land" group has 3,498 samples, while "WB on Water" has only 1,057 samples.

As expected, memorization is more severe for smaller groups because there are fewer samples to support feature learning—i.e., the model focuses more on learning features that represent larger groups. This makes memorization less severe for "LB on Land" compared to "WB on Water." This is why, in Figure 2, the self-influence score is higher for "WB on Water" than for "LB on Land." It also supports the idea that memorization is inversely related to group size: smaller groups are more vulnerable to memorization effects.

We have also extended our self-influence analysis for GroupDRO in Appendix G (highlighted in green for better visibility) titled "Additional Influence-Score Experiments" to provide a direct comparison. The results confirm that MAT is more effective at reducing self-influence scores across all subpopulations (groups). GroupDRO also demonstrates some reduction in the self-influence scores, particularly for minority groups. This further supports the idea that memorization limits generalization and that mitigating it contributes to improved generalization in subpopulation shift settings.

Thanks again for your review and feedback. We hope our responses help address your concerns and show the value of our work. Let us know if you have any other questions, we’d be happy to discuss further!

Thank you for your thoughtful feedback and kind words about the introduction and experiments! Below we address your remarks and provide further contexts. We have split our response into three parts for better readability.

Why and when should MAT be used?

This is a valid question. As you pointed out, one major challenge in subpopulation shift problems is group annotations, which face two main issues: (1) Group annotation could be expensive, and (2) The second issue, which is somewhat hidden in this setup, is defining the groups themselves. In many real-world problems, the types and number of groups may not be easily defined. While datasets like Waterbirds have clear group structures (e.g., [Waterbird, Landbird] x [Water, Land]), many real-world scenarios lack such clarity.

One MAT’s advantage is that it is group-agnostic, which eliminates the need for predefined or inferred group labels. Unlike GroupDRO, which depends on explicit group annotations (human-labeled or XRM-inferred), MAT uses soft held-out predictions for individual examples, enabling it to train without any notion of group. While MAT’s performance on current benchmarks is comparable to GroupDRO with XRM-inferred annotations, its group-agnostic nature holds promise for more complex, less structured setups, a direction we leave for future work.

Another advantage of MAT is its reduced sensitivity to early stopping or model selection. Unlike reweighting methods such as GroupDRO, which converge to the same max-margin solution as ERM unless carefully regularized or stopped early, MAT fundamentally changes both the optimization trajectory and the final solution. This topic is discussed further in the newly added section in Appendix D (highlighted in yellow for better visibility) titled "Empirical Insights into the Effects of Reweighting and Logit-Shifting".

That being said, as discussed in the paper, we only expect MAT to excels in setups where memorization is "bad". When memorization is good, even ERM may outperform MAT and any invariant learning method. For example, in the Waterbirds dataset, if the task were background classification, ERM’s reliance on background features would be advantageous, whereas MAT might learn a non-generalizable pattern (the foreground in this case).

On the Role of Noise Dimensions:

Thanks for pointing this out! We would like to clarify that the noise features we add to the inputs play a different role from their traditional usage, where input noise is employed as a form of regularization (e.g., Bishop, 1995 [Training with Noise is Equivalent to Tikhonov Regularization]). In our setup, these features are used to model any feature that, in expectation (given infinite data), is independent (and thus decorrelated) from the target labels. However, due to the finite sample size, they exhibit apparent empirical correlations and can be leveraged for memorization. We believe our setup closely resembles real-world scenarios, and we will illustrate this similarity by examining one of the datasets in our study: the Waterbird dataset.

In our setup (Section 2.1), we define three types of features:

1. Spurious feature: Partially correlated with the label; learning it doesn't generalize to perfect accuracy in testing.
2. Core feature: Fully correlated with the label; learning it ensures perfect accuracy in testing.
3. Noise features: Uncorrelated with the label in expectation but may show empirical correlations in finite data, e.g., irrelevant pixels in an image.

Now consider the Waterbird dataset. In the generation process of this dataset, we have two features: the foreground (the type of bird, such as "waterbird" or "landbird") and the background ("water" or "land"). The type of bird is considered the core feature because it is 100% correlated with the target label, while the background is considered the spurious feature, as it is not fully correlated. For instance, a small fraction of the "waterbirds" appear with a "land" background.

However, in high-dimensional settings such as images, this distinction becomes less clear. Essentially, every pixel in an image dataset can be considered a feature, allowing a model to exploit any arbitrary combination of pixels for shortcut learning. For instance, a random pixel at position (x, y) might be distinguishable enough for the model to memorize a specific sample. Therefore, similar to our setup, such features can be considered noise features (also referred to as example-specific features): they are not correlated with the label in expectation but may be used to reduce the training loss for specific samples through memorization.

Additionally, although we concatenate the noise feature to the inputs for simplicity, once the data passes through the network, all features (including the noise) interact and become mixed together in the representations. The network does not treat the noise as separate. For this reason, if we rotate the features (e.g., by multiplying them by an orthogonal matrix) and then add the noise, the results and conclusions remain unchanged.

Lastly, we plan to rename "input noise" to "example-specific features" to avoid any confusion between its usage in our setup and the traditional notion of input noise.

Dear Reviewer,

Thank you once again for taking the time to provide your valuable feedback on our paper. With the discussion phase nearing its end, and to ensure that our responses are fully considered, we request that you kindly confirm whether you’ve had a chance to review them.

We hope our responses have addressed your concerns and clarified any ambiguities. If there are any further questions or suggestions, we would be glad to address them promptly.

Additionally, if our responses have adequately resolved your concerns, we kindly ask you to consider reflecting this in your review score.

Thank you!

Thank you for your thoughtful and constructive review of our submission. Below, we address each of your comments and clarify aspects of the paper as requested.

On the theoretical analysis:

Our theoretical analysis is indeed limited to the case of a linear model, which is a common and well-established practice in the field. The primary goal of this approach is to simplify the analysis and gain interpretable insights that can then be tested in more complex scenarios. While extending the theory to non-linear models is challenging due to their complexity, many of the findings from our linear analysis carry over to real-world, non-linear setups, as demonstrated by our experimental results.

On how logit shift enables memorization-aware training:

As you stated, “the shift down-weights the gradient descent updates coming from the training samples with spurious correlations.” In subpopulation shift setups, these samples are typically from the majority group (e.g., waterbird on a water background (majority) vs. landbird on a land background (minority)), and this learning dynamic often leads to the memorization of minority groups. However, under the assumption that the core features are present in both groups, the goal is to promote learning these features while avoiding harmful memorization.

The logit shift makes training "memorization-aware" by adjusting the model's focus based on its performance on held-out data. When a sample is heavily memorized (e.g., when the model fits noise or spurious correlations), its held-out prediction deviates significantly, and, as a result, the logit shift places greater emphasis on these samples. On the other hand, for samples where the model relies on core features, the held-out prediction aligns closely with the target, and the emphasis is lower. Thus, the extent of the logit shift implicitly depends on the degree of memorization for each sample (e.g., emphasizing minority samples), steering the model’s training toward learning invariant (i.e., core) features.

We will add an explanation of this mechanism to Section 3.1 in the revised version of the paper to improve clarity. Thank you for pointing this out.

Duplication of a reference:

Thank you for the precise observation; we have fixed this in the current revision.

Code:

We will release our codebase as soon as possible in our next revision of the paper.

On the general spurious correlation - (Multiple spurious features setup)

Thank you for your interesting suggestion to explore multiple spurious features. To investigate this, we have designed an experiment and added a new section in Appendix H (highlighted in purple for better visibility), titled "Multiple spurious features setup". This new design is adapted from our original setup described in Section 2.1 of the paper. We introduce a second spurious feature, resulting in one main feature ($x_y$) and two spurious features ($x_{a_1}$ and $x_{a_2}$), making the task more challenging. For more details, please refer to Appendix H.

We illustrate the training dynamics of ERM (with and without memorization capacity) and MAT with memorization capacity. As shown in Figure 7, MAT proves highly effective in this scenario, achieving near-perfect test accuracy across all groups. These results demonstrate the robustness of MAT in multiple spurious features setup, and we hope this addresses your questions.

Thank you once again for your review and feedback. We hope our responses have addressed your concerns and highlighted the value of our work. Please feel free to reach out if you have any additional questions—we’d be glad to discuss them further!

On the Role of Noise Dimensions:

We would like to first clarify that the "input noise" used in our setup differs from its traditional usage, where input noise is used as a form of regularization (e.g., Bishop, 1995 [Training with Noise is Equivalent to Tikhonov Regularization]). Instead, we use input noise to model any feature that, in expectation (infinite data), is independent (thus decorrelated) of the target labels but, due to finite samples, exhibits apparent empirical correlations. For instance, in an image classification task, irrelevant pixels might appear correlated and thus be used for memorization.

In Section 5.1, we do not claim that adding noise to inputs prevents overfitting in the traditional sense. Rather, we argue that memorization through input noise dimensions can prevent catastrophic overfitting by redirecting spurious patterns into auxiliary features, leaving core features free to learn generalizable patterns.

Now to answer your question, in the noiseless setting ($\sigma = 0$), the model is constrained to rely solely on a single feature (along the x-axis) to minimize the training loss to zero. Regardless of the model's over-parameterization, it must memorize every detail in the training data using this single feature. This leads to catastrophic overfitting because there are no auxiliary noise dimensions available to "absorb" small variations or non-generalizable patterns..

In the small-noise setting ($\sigma = 1e-4$), however, there are noise dimensions that the model can use to fit any remaining variation that the single core feature couldn't explain. Here, memorization is not harmful; instead, it helps the model fit the last bit of unexplained variation, serving as a subtle regularizer.

Planned Updates

To address your remarks, we will revise the draft to:

• Clarify differences between our work and Sagawa et al. (2020), emphasizing our focus on memorization and differing assumptions about core and spurious feature noise.
• Already added appendix section D (highlighted in yellow for better visibility), "Empirical Insights into the Effects of Reweighting and Logit-Shifting", to show reweighting converges to the max-margin solution while logit shifting changes the fixed point.
• Highlight MAT’s group-agnostic advantage in the main text, emphasizing its flexibility to handle arbitrary groups without requiring predefined groups.
• Rename "input noise" to "example-specific features" to avoid misunderstandings. Expand Section 5.1 to explain how these example-specific features prevent catastrophic overfitting.

We remain at your disposal, should you have further questions. Thank you for your time and attention.

References:

[1] Sagawa et al. Section 5.3: "... show how the minimum-norm inductive bias translates into a bias against memorization."

[2] Sagawa et al. Page 2: "... the spurious features ... entail less memorization, and therefore suffer high worst-group test error."

[3] Our work, Line 67: "... the combination of spurious correlations with memorization that leads to poor generalization."

[4] Our work, Figure 2: "Models trained with ERM exhibit higher self-influence scores (higher memorization) for minority subpopulations."

[5] Sagawa et al., Page 9: "... the model prefers the spurious feature because it is less noisy than the core feature."

[6] Our work, Lines 124-126: $ρ^{tr} = 0.9$ makes the spurious feature correlated with y, whereas the core feature is fully correlated.

[7] Walczak & Cerpa (1999), Heuristic Principles for the Design of Artificial Neural Networks, Page 22: "... to memorize the training data set, by providing a one-to-one mapping between input values and the corresponding output values, which decreased generalization performance."

[8] Sagawa et al. (2020), Distributionally Robust Neural Networks for Group Shifts

Thank you for taking the time to review our work. Below, we provide clarifications and hope they address your questions. We have split our response into three parts for better readability.

On the Novelty:

As stated in lines 99-100 of our paper, our setup adapts that of Sagawa et al. (2020). However, our analysis leads to fundamentally different findings that cannot be derived from their results. Below, we detail the core differences:

1. Different Focus and Findings:

   • Sagawa et al. focus on the effect of over-parameterization. They show that over-parameterization leads to a bias against memorization ([1]) and then conclude that: less memorization leads to higher worst-group test error ([2]).
   • Our work centers on memorization itself. We conclude the opposite: more memorization leads to higher worst-group test error ([3, 4]).

   This apparent contradiction arises from different setups between Sagawa’s work and ours:

2. Different Setups:

   • Sagawa et al. assume that the spurious feature is less noisy than the core feature ([5]). Thus, if a model captures the core feature, it must memorize more due to the core feature’s higher noise.
   • Our work assumes the core feature is fully correlated with the target class, while the spurious feature is less correlated ([6]) with it. In our setup, if learning manages to focus on the core feature, it requires no memorization of data points. This adjustment reflects more realistic settings, where spurious features, despite being easier to learn, often have weaker correlations with labels. For example, in the Waterbirds dataset, the background ("water" or "land") is less correlated with the label than the foreground ("waterbird" or "landbird").

   Both conclusions are valid within each setup. The apparent contradiction stems from Sagawa et al.'s setup in which the core feature has more noise. As noted in Sagawa et al. (Section 9), "This assumption need not always apply...". Their focus is on over-parameterization, so this setup is suitable for their analysis but limits its generalizability.

3. Memorization as a Starting Point: The intuition that “neural networks memorize the training data, which decreases generalization” is not new and has been discussed in earlier literature, including Shah et al. (2020), Tsipras et al. (2018), and even back in 1999 in Walczak & Cerpa (1999) ([7]). Our work formalizes this and then builds upon that, introducing MAT to leverage this insight to learn an invariant predictor.

On Shifting the Logits (through MAT) vs. Reweighting (through GroupDRO):

While XRM+GroupDRO (per-group reweighting) and MAT (per-example logit shifting) achieve comparable performances in table 1, they are fundamentally different:

1. Reweighting converges to max-margin while shifting changes the loss landscape:

   Any reweighting scheme, including GroupDRO, eventually falls back to ERM if example weights are non-zero. To avoid this, as noted in the original GroupDRO paper ([8]), "stronger-than-typical l2 regularization or early stopping" with surgical precision is required. Without such measures, as shown in Figure 2 of Sagawa et al., GroupDRO’s performance reverts to ERM (see Figure 2, ERM vs DRO with Standard Regularization).

   This behavior arises because, greadient decent dynamics will converge to the same max-margin solution, even with different weights per example or per group. This result follows from the first part of Theorem 3 in Soudry et al. ("The Implicit Bias of Gradient Descent on Separable Data", JMLR, 2018). To see why, note that applying group weights is equivalent to repeating data points from each group a certain number of times based on the group label, which does not affect the separability assumption of the theorem. As such, gradient descent with per-group weighting will eventually converges to the max-margin solution.

   In contrast, the logit-shifting changes the fixed point of gradient descent, resulting in a different predictor. To show that, we have now added a new section in the appendix D (highlighted in yellow for better visibility) titled "Empirical Insights into the Effects of Reweighting and Logit-Shifting". There, we show that while reweighting changes the optimization path, it ultimately converges to the max-margin solution. Logit shifts, however, modify both the trajectory and the fixed point, making shifting more robust and alleviating the need for precise early stopping.

2. MAT is group-agnostic:

   An advantage of MAT's logit shifting over GroupDRO is its ability to operate without relying on any notion of groups. GroupDRO requires predefined (human-labeled or XRM-inferred) groups, with the number of groups dictated by the group-inference algorithm (e.g., XRM always identifies two groups). In contrast, MAT uses soft held-out predictions for individual examples, without requring any group-specific information in the logit adjustment phase.

   This flexibility allows MAT to handle an arbitrary number of groups in subpopulation shift scenarios. While standard benchmarks on which we have reported our results often involve two groups, MAT’s group-agnostic approach could offer greater adaptability in real-world problems, where explicit group definitions are often unavailable.