The Pitfalls of Memorization: When Memorization Hurts Generalization
Authors: Reza Bayat
Mohammad Pezeshki
Elvis Dohmatob
David Lopez-Paz
Pascal Vincent
Venue: ICLR 2025 Poster (2025)
Keywords: Memorization
Generalization
Spurious Correlations
Primary area: unsupervised, self-supervised, semi-supervised, and supervised representation learning

Abstract: Neural networks often learn simple explanations that fit the majority of the data while memorizing exceptions that deviate from these explanations. This behavior leads to poor generalization when the learned explanations rely on spurious correlations. In this work, we formalize $\textit{the interplay between memorization and generalization}$, showing that spurious correlations would particularly lead to poor generalization when are combined with memorization. Memorization can reduce training loss to zero, leaving no incentive to learn robust, generalizable patterns. To address this, we propose $\textit{memorization-aware training}$ (MAT), which uses held-out predictions as a signal of memorization to shift a model's logits. MAT encourages learning robust patterns invariant across distributions, improving generalization under distribution shifts.

Published as a conference paper at ICLR 2025
THE PITFALLS OF MEMORIZATION: WHEN MEMORIZATION HURTS GENERALIZATION
Reza Bayat 1,2, Mohammad Pezeshki*3
Elvis Dohmatob 1,3,1 David Lopez-Pazz Pascal Vincent1.2.3
Mila Université de Montréal
Meta FAIR
'Concordia University
ABSTRACT
Neural networks often learn simple explanations that fit the majority of the data while memorizing exceptions that deviate from these explanations. This behav- ior leads poor generalization when the learned explanations rely on spurious correlations. In this work, we formalize the interplay between memorization and generalization, showing that spurious correlations would particularly lead to poor generalization when are combined with memorization. Memorization can reduce training loss zero, leaving no incentive to learn robust, generalizable patterns. To fraining (MAT), which uses held-out
address this, we propose memorization-awwen predictions as a signal of memorization to shift model's logits. MAT encourages learning robust patterns invariant across distributions, improving generalization under distribution shifts.
INTRODUCTION
Watching the stars at night gives the illusion that they orbit around the Earth, leading to the longstand- ing belief that our planet is the center of the universe-a a view supported by great philosophers such as Aristotle and Plato. This Earth-centric model could explain the trajectories of nearly all celestial bodies. However, there were exceptions; notably, five celestial bodies, including Mars and Venus, occasionally appeared to move backward in their trajectories (Kuhn, 1992). To account for these anomalies, Egyptian astronomer Claudius Ptolemy introduced "'picycles" orbits within orbits-a complex yet effective system for predicting these movements. Over 1400 years later, Nicolaus Copernicus proposed an alternative model that placed the Sun at the center. This Sun-centric view not only simplified the model but also naturally explained the previously perplexing backward motions.
Drawing a parallel to modern machine learning, we observe a similar phenomenon in neural networks. Just as the Earth-centric model provided incomplete explanation that required epicycles to account for exceptions, neural networks can learn simple explanations that work for the majority of their training data (Geirhos et al., 2020; Shah et al., 2020; Dherin et al., 2022). These models might then treat minority examples- those that do not conform the learned explanation- as exceptions (Zhang et al., 2021) 1). This becomes particularly problematic if the learned explanation is spurious, meaning does not hold in general or is not representative of the true data distribution (Idrissi et al., 2022; Sagawa et al., 2020; Pezeshki et al., 2021; Puli et al., 2023).
Empirical Risk Minimization (ERM), the standard learning algorithm for neural networks, can exacerbate this issue. ERM enables neural networks to quickly capture spurious correlations and, with sufficient capacity, memorize the remaining examples rather than learning the true patterns that explain the entire dataset. This has real-world implications; for example, neural networks designed to detect COVID-19 from x-ray images have been found to rely on spurious correlations, such as whether a patient is standing or lying down (Roberts al., 2021). This could be dangerously misleading, as a model that appears to excel in most cases may have actually captured a spurious correlation. Memorization of the remaining minority examples can fully mask neural network's failure to grasp the true patterns in the data, giving a false sense of reliability. Again, the Earth-centri model
"Equal contribution.
Code: http:///iikkknmecccceeeeees///st

---

Published as a conference paper at ICLR 2025
of the universe able explain all celestial trajectories with complex epicycles but ultimately failed to reveal the true nature of our solar system.
Identifying whether model with nearly perfect accuracy on the training data has learned generalizable patterns or merely relics on a mix of spurious correlations and memorization is critical. The answer lies in the model's 's performance held-out average accuracy or more fine-gr geee-grrindd group accuracies can help us identify a better model. question that arises is: How can one use held-out performance signals to proactively guide a model toward learning generalizable patterns?
held-out data, partieularly minority examples. Metries such as
Traditionally, held-out performance signals are mainly used for hyperparameter tuning and model selection. However, in this work, we propose a novel approach that leverages these signals strategically to guide the learning process. Towards this goal, our paper makes the following contributions:
Formalizing the interplay between memorization and generalization: We study how mem- orization affects generalization in an interpretable setup. We show that while spurious correlations inherently problematic, they by themselves do not always lead to poor generalization in neural networks. Instead, it is the combination of spurious correlations with memorization that leads this problem. Our analysis shows that models trained with empirical risk minimization (ERM) tend to rely on spurious features for the majority of the data while memorizing exceptions, achieving zero training loss but failing to generalize on minority examples.
Introducing memorization-awee aaaere training (MAT): MAT is a novel learning algorithm that leverages the flip side of memorization by using held-out predictions to shift model's logits during training. This shift guides the model toward learning invariant features that generalize better under distribution shifts. Unlike ERM, which relies on the i.i.d. assumption, MAT is built upon an alternative assumption that takes into account the instability of spurious correlations across different data distributions.
THE INTERPLAY BETWEEN MEMORIZATION AND SPURIOUS CORRELATIONS IN ERM
Problem Setup and Preliminaries. We consider standard supervised learning setup for K-class classification problem. The data consists of input-label pairs (( y;) J where æ; is the input vector and y/ E (1, (1,..., K} is the class label. Let a; denote any attribute combination of attributes within x; that may or may not be relevant for predicting the target yı. The objective is to train model p(y x; w) parameterized by w. Given an input H,, let f(x;; w) e RK represent the output logits of the model, then:
ply xy; u) offtmax(((ax;)
Under the i.i.d. assumption that p(y, æ) is invariant between training and test sets, empirical risk minimization (ERM) seeks to minimize the following loss over the training dataset:
CERM 1 Σ 1(bly x;; w), v.).) +
1
w wwl
(2)
where 1(.,-) is the cross-entropy loss, and > | 2 the weight-decay regularization.
2.1 MEMORIZATION CAN EXACERBATE SPURIOUS CORRELATIONS
Spurious correlations violate the i.i.d. assumption, and when combined with memorization, can hurt generalization. We now study such scenario. Adapting the frameworks introduced in Sagawa et al. (2020) and Puli et al. (2023), we look into the interplay between memorization and spurious correlations in an interpretable setup.
Setup 2.1 (Spurious correlations and memorization). Consider a binary classification problem with labels y E [-1,+1} and an unknown spurious attribute a E (-1,+11 +1). Each input x E Rd42 is given by x (Ing TFa, é), where Ty E R is core feature dependent only on y. To a sparious feature dependent only a, and E TRd are caample-speciffe noise features uncorrelated with both y
2

---

Published as a conference paper at ICLR 2025
and The scalar r R. modulates the rate which the model learns rely the spurious feature Eas effectively acting as a scaling factor that increases the feature's learning rate relative to the core feature zy. The attribute is considered spurious; it is assumed to be correlated with the labels y training but has no correlation with y at test time, potentially leading to poor generalization if the model relies Ta- Specifically, the data generation process is defined as:
Ty oyy
yıa ((a, E~~ E~N((0,021)
ER where,
w.p. w.p.
and p
p
(train) 0.5 (test)
To better understand this setup, one can think of a classification task between cows and camels, In this example, z represents the pixel data, E {cow, camel} are the class labels, and E {grass, sand] the background labels. Here, Ty represents the pixels associated with the animal itself (either cow camel), represents the pixels associated with the background (grass or sand), and represents irrelevant pixels that varies from one example to another. The key assumption is that the joint distribution of class labels and attribute labels differs between training and test datasets, ie.. p" (a. y) # p'((a, y). For example, in the training set, most cows (camels) might appear on grass (sand), while in the test cows (camels) appear equally each background.
Illustrative Scenarios. We first empirically study configuration of the above setup where p" = 0.9 makes the spuriously correlated with y. We set y 5 making the spurious feature easier for the model to learn. In contrast, the core feature Ty is fully correlated with but due to smaller norm, it learned more slowly. Here we consider two cases:
memorization)
ERM
n
1.0
1.0
Accurwcy
Accuraey
Majority Trom
Accuracy
Minsety Taali
0.6
Majority
Majority Train
Minoelty Tex:
0.6
Maaerity Train
Minocily Traid
0.2
Majority
2.2
Majority
Minurity Teet
Minority Text.
0.0
Epochs
Egochs
Epochs
7.5
7.5
7.5
3.4
5.0
24 3.-
0.0
25
235
25*
-7.
-7.5
-7.5
10.0
12.
3ig
Figure Illustration of scenarios in the interpretable classification setup involving spurious comelatiins and memorization. The left panel represents scenario without example-specific features (ae - 0), where memorization not possible. this case, the model trained with ERM initially learns the spurious feature I. serving the majority, but eventually adjusts the decision boundary to the core feature xy, resulting in good generalization on minority test examples. The middle and right panels depict scenario with example-specific features (a 0), where memorization is possible. In the middle plot, the model trained with ERM fails to generalize as memorizes the minorities using the example example-specifes features leaving no more incentive for the model to learn the core feature. In contrast, the model trained with MAT successfully learns the invariant features, and generalizes well in the presence of example-specific features.

---

Published as a conference paper at ICLR 2025
1. No example-specific features = No Memorization = ERM generalizes well. Figure I-(left) presents case where there are no input example-speciit features (Øe + 0). As training progresses, the model first learns Za due its larger norm, resulting in perfect accuracy on the majority examples. Once the model achieve nearly perfect accuracy on the majority examples, it starts to learn the minority examples. At this point, the model must adjust its decision boundary to place more emphasis on the core feature Zy, ultimately achieving perfect generalization both majority and minority examples.
2. With example-specific features + Spurious Features + Memorization = ERM fails to generalize. Figure 1-(middle) presents similar setup to the former, but this time with input example-specific features (o, > 0). Again, initially, the model learns to rely on the spurious feature Va- However, unlike Case 1, the example-specific features provides the model an opportunity to memorize minority examples directly. As a result, the model achieves zero training loss by memorizing minority examples using the example-specific dimensions instead of learning to rely on the core feature Ty. Consequently, the model fails to adjust its decision boundary align with Ty. and does not generalize held-out minority examples.
These results illustrate that the combination of spurious correlations and memorization creates a 'loophole' for the model. When memorization happens, it leaves no more incentives for the model to learn the true, underlying patterns necessary for robust generalization.
Theoretical Analysis. We now provide formal analysis to formalize our empirical observations. Complete proofs are provided in Appendix F.
Theorem 2.2. Consider binary classification problem under the setup described in Setup 2.1. where linear model fa; w) x w is trained using ERM dataset D Let ÖMRM = (8.y. Way we) Rd42 be the learned parameters, and f(x)) argmax, fyl (x;; WERM) the learned classifier.
Under the conditions where, X - 0+, 00, X/n - co, p² 0.501, have,
0. Perfect training accuracy: The classifier achieves perfect accuracy on all training examples:
(Ø(() =y) 1, Vo D
No example-specific features: of + 0+, the classifier converges to one that relies solely the core feature Ty. For test point x::
( v)
Vo D
2. With example-specific features: If > is bounded away from zero and >> log n and 7 S dev/ d m. where m == p0nn is the number of majority samples the training set. Then, for a test point x, the classifier relies pathologically on the spurious feature Fas iee,
p (ū(x) a) + 1,, Va & D
The condition d >> log ensures that mately orthogonal, and x > dey /e/d/mm guarantees that the spurious feature Ia learned faster by gradient descent than other features.
features from different samples are approxi-
Theorem 2.2 suggests that in the presence of spurious correlations and example-specific features, p(ū(x)) = a) + for held-out examples. This indicates that the model confidently predicts as the label for any held-out input Dd leading to poor generalization. The probability that model assigns to the held-out examples is then defined as:
f(z; WERM)
x)) softmax
Vx Dir
(4)
where the superseript 'ho' stands for held-out and is temperature parameter that controls the sharpness of the distribution. This probability serves as the basis for the shifts introduced in the next section.
4

---

Published as a conference paper at ICLR 2025
3 MEMORIZATION-AWARE TRAINING (MAT)
As exemplified in the previous section, spurious correlations between labels and certain attributes lead to poor generalization, particularly when models memorize specific training examples rather than learning robust patterns. To mitigate this issue, we propose memorization-avalt training (MAT), which leverages calibrated held-out probabilities to shift the logits, prioritizing learning of examples with worse generalization performance.
3.1 SHIFTING THE LOGITS WITH CALIBRATED HELD-OUT PROBABILITIES
MAT modifies the empirical risk minimization (ERM) objective by introducing a per-example logit shift based on calibrated probabilities. The training loss is redefined as:
1 Σ (sootmaxx(f( w) log (h
MAT
(5)
where f(x;; w) are the logits for input 24, and pho æ;) are the calibrated held-out probabilities, defined as:
o x)
Σ plu y pp("" ar),
Ep(y
yoo
where ply"o x) is derived from Eq. equation 4, and ply yho) is the calibration matrix, given by: p(u,yho) Ey p(y', yhoj'
ply 1h)
where p(u.v"o) 1 {(
ply xx).
(7)
The calibration matrix ply )² represents the confusion matrix of true labels versus held-out probabilities y0 The calibrated probabilities p x) guide the learning such that correct held-out predictions lower the loss for an example, while incorrect predictions increase it, prompting the model focus challenging examples. This ensures MAT prioritizes minority examples with poor generalization, improving generalization under distribution shifts. See Appendix B for discussion on the effect of loss reweighting and logit shifting.
Estimating poy" x.). The final challenge in implementing MAT is that we require an auxiliary model provide held-out probabilities for every single example. This auxiliary model can be derived using Cross-Risk Minimization (XRM) (Pezeshki et al., 2023).
XRM, originally proposed as an environment discovery method, trains two networks on random halves of the training data, encouraging each to learn biased classifier. It uses cross- mistakess (errors made by one model on the other's data) annotate training and validation examples. Here, MAT does not require environment annotations but uses held-out logits, (==), from pretrained XRM model. For model selection using the validation set, either ground-truth or inferred environment annotations from XRM can be used.
The pseudo-code in Algorithm summarizes the MAT algorithm. Note that MAT introduces a single hyper-parameter the softmax temperature in Eq. equation 4.
4 EXPERIMENTS
We validate the effectiveness of MAT in improving generalization under subpopulation shift. Then, we provide a detailed analysis of the memorization behaviors of models trained with ERM and MAT.
4.1 EXPERIMENTS ON SUBPOPULATION SHIFT
We evaluate approach on four datasets under subpopulation shift, as detailed in Appendix A. In all experiments, we assume that training environment annotations are not available. For the validation set and for the purpose of model selection, consider settings: (1) group annotations available in the validation set for model selection, and (2) no annotations are available even the validation set.
5

---

Published as a conference paper at ICLR 2025
Algorithm I Memorization-Awarc Training (MAT)
Input: Training set {( wi) Validation set (( v1 m
=== Pre-trained XRM model (xRM(a)) Optional: Validation environment annotations {0Ji= if not available, infer āj argmax (x;) + Y/
Model Selection: Early stopping based on validation gorst-group accuracy (using df if provided, or de if inferred)
Initialize classifier f(z).
Compute p(y/o wx) softmax (x.)7T)
Compute ply, / = concat ( (v xi)ly = vil) ys ((...K....k)
Compute p(y ) plu."")
Compute ( x;)= = (p(y ae.). p(y )
Repeat until early stopping:
Update the loss: Σ (f + log ( x.),
Track worst-group accuracy and update the best model
Stop if no improvement is observed after P iterations
Table 1: Average and worst-group accuracies (avg/wga) comparing methods. We specify access to group annotations training (tr) and validation (va) data. Symbol denotes original numbers.
Waterbirds
CelebA
MultiNII
CivilComments
fr UQ
Avg
WGA
Avg
WGA
Avg
WGA
Avg
WGA
/
GroupDRO
90.2:
86.5
93.1 03
88.3
80.6
73.4 4:48
84.2 a 0.
73.8 0.6
ERM
97.3
72.6
95.6
47.2
82.4
67.9
83,1
69.5
LFF:
91.2
78.0
85.1
77.2
80.8
70.2
68.2
$0.3
X
JTT
93.3
86.7
88.0
81.1
78.6
72.6
83.3
64.3
LCI
90.5 1.1
88.1 0.8
70,3 1.2
AFR
94.2
90.4
82.0
81.4
73.4
89.8
68.7
MAT
90.4
88.1
92.4 044
90.5
79.4
74.6
84.3
74.0
ERM
83.5
66.4
95.4
54.3
82.1
67.9
67.2
uLAi
91.5 0.7
86.1
93.9
86.5
x
XRMI
89.3
88.1 1.09
91.4 0.
89.1 1.1
75.8
72,1 1.
84.4
72.2
MAT
90.4 1: 0.7
88.1 1.
92.3 0:
89.9
79.6 0.2
73.0 a
85.7 a: 0.1
68.1 0.7
For evaluation, we report two key metrics the test set: (1) average test accuracy and (2) worst-group test accuracy, the latter being computed using ground-truth annotations.
Table compares the performance of MAT with several bascline methods, including ERM, GroupDRO (Sagawa et al., 2019), and other invariant methods like LiF (Nam et al., 2020), JTT (Liu et al., 2021), LC (Liu et al., 2022), uLA (Tsirigotis et al., 2024), AFR (Qiu et al., 2023), XRM++GrouDROO (Pezeshki et al., 2023). These methods vary in their assumptions about access to annotations, both in training and validation for model selection. For instance, ERM does assume any training group annotations, while GroupDRO has full access to group annotations for training and validation data. Further details the experimental setup and methods are in Appendix A.
4.2 ANALYSIS OF MEMORIZATION SCORES
To understand the extent memorization in models trained with ERM, we analyze the distribution of memorization scores across subpopulations. We focus the Waterbirds dataset, which includes two main classes - and Landbird eaach divided into majority and minority subpopulations based their background (e.g., Waterbird on water vs. Waterbird on land).
The memorization score is derived from the influence function, which measures the effect of each training sample model's prediction. Formally, the influence of training sample on target

---

Published as a conference paper at ICLR 2025
sample j under training algorithm A is defined as:
iA, D, i.)) ((AA) P (uj xt)) PD-(e...))
(ys xzj)
(8)
where D is the training dataset, D-(m,.u::) denotes the dataset with the sample (I,, y4) removed. The memorization score a specific case of this function where the target sample (æzz Mj) is the same as the training sample. It measures the difference between model's performance on a training sample when that sample is included in the training set (held-in) versus when is excluded (held-out).
Calculating self-influence scores using naive leave-one-out approach is computationally expensive. However, recent methods, such as TRAK (Park et al., 2023), provide an efficient alternative. TRAK approximates the data attribution matrix, and the diagonal of this matrix directly gives the self- influence scores (see Appendix A.2 for more details).
Figure 2 depicts the distribution of self-influence scores across subpopulations in the Waterbird dataset. We note that minority subpopulations (e.g., Waterbirds on land) show higher self-influence scores compared their majority counterparts (e.g., Waterbirds water) in model trained with ERM. However, a model trained with MAT shows a similar distribution of self-influence scores for both the majority and minority examples, with overall lower scores compared to ERM. These results show that MAT effectively reduced memorization, while leading to improved generalization.
ERM Waterbird
ERM Landbird
ERM -MMo-Proortton Samples
Minority
Minority
1.6
Land
Density
Mejority
Density
Majority
Proportion
WB on Water
LB on Water
0.8
WB
0.4
0.00
0,25
0.50
0.75
1:00
0.00
0,25
0.50
0.75
1.00
0.0
0.1
0.3
0.5
0.7
0.9
Seifinfluence Score
Selfinfluence Score
Self Influennee Interval
MAT Waterbird
MAT Landbird
MAT Proportion Samples
Minority Majority
Minority Majority
3.2
land
Density
Density
Progortion
Water
WB on Water
1.6
WB on Land
0.8
0.00
0.25
0.50
0.75
1.00
0.00
0.25
0.50
0.75
1.00
0.0
0.1
0.3
0.5
0.7
0.9
Self-Influence Score
Self-ifluence Score
Self-lnfluence Interval
Figure 2: Self-Influence estimation of the Waterbird groups by ERM and MAT. The distribution of self-influence scores is shown for both the majority and minority subpopulations (e.g.. Waterbirds on water vs. Waterbirds land). Models trained with ERM exhibit higher self-influence scores for minority subpopulations, suggesting increased memorization in these groups In contrast, models trained with MAT show more uniform self-influence distributions across both majority and minority subpopulations. The rightmost plots display the proportion of samples in different self-influence intervals, with MAT producing a more balanced distribution compared ERM. Further details can be found in Appendix D.
MEMORIZATION: THE GOOD, THE BAD, AND THE UGLY
In this work, showed that the combination of memorization and spurious correlations, could be key reason for poor generalization. Neural networks can exploit spurious features and memorize exceptions to achieve zero training loss, thereby avoiding learning more generalizable patterns. However, an interesting and somewhat controversial question arises: Is memorization always bad?
To explore this, we look into simple regression task understand different types of memorization and their effects on generalization. We argue that the impact of memorization on generalization can vary depending on the nature the data and the model's learning dynamics, and we categorize these types of memorization into three distinct forms. The task is defined as follows.
Setup 5.1. Let By E R be a scalar feature that determines the true target, 3" = f(z,). Let n 1 be dataset consisting of input-target pairs (x. y). Define the imput vector as D (a ve) )
7

---

Published as a conference paper at ICLR 2025
Good memorization d, 10-1
Bad memorization -, =1110
Ugly memorization a, 0.0
(Noig) Triining mamples
and prodicctoo furettons
l'redated fanction sta) 141
tree and pretiction Furattios
tree and prodtetton furettons
1.0
-3.3
-3.7
6.2
(osce (bature)
(0oce Seature)
(osce Seaturo)
Figure 3: Three types of memorization in regression models trained with different levels of example- specific features (oe). The plots show the ERM-trained model g(x) g( E) (solid blue line) versus the true underlying function ff (zyy (dashed gray line) and the noisy training examples. In all the three, the models are trained until the training loss goes helow 10-6 Good memorization (Lefft, 0e d== ===10 10-1): Model learns the function f (3x, well but slightly memorizes residual noise in the training data using the input example-specific features e. This type of memorization is benign, as it does not compromise generalization. Bad memorization (Middle, ae 10 10-3): The model relies more scample-specifie features than learning the true function f(yy). leading to partial learning of (zyy and fitting of diisinddmdttted input features. This type of memorization impedes learning of generalizable patterns and is considered malign. Ugly memorization (Right, a, 0.0): Without example-specine features, the model overfits the training data, including label noise, resulting in a highly non- non-linear and complex model that fails to generalize to new data. This type is referred to as catastrophic overfitting.
concat(x,, <) E Rd+1, where c ~ NO, o21) E Rd represents input example-specific features concatenated with the true feature Ty. The larget is defined as y = y + where & (1 aq) represents additive target noise.
In this context, Ty can be interpreted the core feature (c.g., the object in an object classification task), e as irrelevant random example-specific features, and & as labeling noise or error. Now, consider training linear regression models ý = g(x) on this dataset. Fixing dg. we train three models under three different input example-speciff features levels: or E 10,10-1,10 10-3J. The results, summarized Figure 3, showcases three types of memorization:
when memorization benefits generalization. At an intermediate level of example- specific features, Oe 10 the model effectively captures the true underlying function, f(xy). However, due to the label noise, the model cannot achieve a zero training loss solely by learning (*yy) As a result, it begins to memorize the residual noise the training data by using the example- specific features This is evidenced by sharp spikes at each training point, where the model, g(x), precisely predicts the noisy label if given the exact same input as during training. Nevertheless, for a neighboring test example with no example-specific features, the model's predictions align well with (z,y). demonstrating good generalization.
The Good:
This phenomenon is often referred to as "benign overfitting" where model can perfectly fit (overfit in fact) the training data while relying on example-specific fcatures, yet still generalize well to unseen data (Belkin et al., 2019a; Muthukumar et al., 2020; Bartlett et al., 2020). The key insight is that the overfitting in this case is "benign" because the model's memorization by relying on example-specific features does not compromise the underlying structure of the true signal. Instead, the model retains a close approximation to the true function on test data, even though it memorizes specific noise in the training data. This has been shown occur particularly in over-parameterized neural networks (Belkin et al., 2019b; Nakkiran et al., 2021).
The Bad: when memorization prevents generalization. At higher level of example-specife features, oe = 10 3 the model increasingly rely on thesefeatures e rather than fully learning the true underlying function (zyy). In this case, memorization is more tempting for the model because the
8

---

Published as a conference paper at ICLR 2025
example-specifer features dominates the input, making it difficult to recover the true signal. As result, the model g(x) might achieve zero training loss by only partially learning f f(zy) and instead relying heavily on the example-speciss features in the inputs to fit the remaining variance in the training data.
This is instance of bad memorization as it hinders the learning of generalizable patterns, the case we studied in this work. This phenomenon is referred to as "malign overfitting" in Wald et al. (2022), where model fits the training data but in way that compromises its ability to generalize, especially in situations where robustness, faimess,, or invariance are critical.
It is important to note that both good and bad memorization stem from the same learning dynamics. ERM, and the SGD that drives it, do not differentiate between the types of correlations or features they are learning. Whether features contributes to generalization or memorization is only revealed when the model is evaluated on held-out data. If the features learned are generalizable, the model will perform well on new data; if they are not, the model will struggle, showing its reliance on memorized, non-generallizable patterns.
The Ugly: Catastrophic overfitting Finally, consider the case where there is no example-specific features, de = 0.0. In this case, the model may initially capture the true function [(rgg), but due to the presence of label noise, camnot achieve zero training loss by learning only f(zy). Unlike the previous cases, the absence of example-specifie features means the model has additional features to leverage in explaining the residual error. As a result, the model is forced to learn a highly non-linear and complex function of the input x Zy fit the noisy labels.
In this situation, memorization is ugly: The model may achieve perfect predictions on the training data, but this comes at the cost of catastrophic overfitingg where the model overfits severely that it not only memorizes every detail of the training data, including noise, but also loses its ability to generalize new data (Mallinar al., 2022). Early stopping generally can prevent this type of severe overfitting.
6 RELATED WORK
Detecting Spurious Correlations. Early methods for detecting spurious correlations rely on human annotations (Kim et al., 2019; Sagawa al., 2019; Li & Vasconcelos, 2019), which are costly and susceptible to bias. Without explicit annotations, detecting spurious correlations requires assumptions. common assumption is that spurious correlations are learned more quickly or are simpler to learn than core features (Geirhos et al., 2020; Arjovsky et al., 2019; Sagawa et al., 2020). Based on this, methods like Just Train Twice (JTT) (Liu et al., 2021), Environment Inference for Invariant Learning (EIIL) (Creager et al., 2021), Too-Good-Tor
Prior (Dagaev et al., 2023), and Correct-n-Contrall (CnC) (Zhang et al., 2022) train models with limited capacity to identify "hard" (minority) examples. Other methods such as Learning from Failure (LfF) (Nam et al., 2020) and Logit Correction (LC) (Liu et al., 2022) use generalized cross-entropy to bias classifiers toward spurious features. Closely related to this work is Cross-Risk Minimization (XRM) Pezeshki et al. (2023), where uses the held-out mistakes as a signal for the spurious correlations.
Mitigating Spurious Correlations. Reweighting, resampling, and retraining techniques are widely used to enhance minority group performance by adjusting weights or sampling rates (Idrissi et al., 2022; Nagarajan
al., 2020; Ren et al., 2018). Methods like Deep Feature Reweighting (DFR) (Kirichenko et al., 2022) and Selective Last-Layer Finetuning (SELF) (LaBonte et al., 2024) retrain the last layer on balanced or selectively sampled data. Automatic Feature Reweighting (AFR) (Qiu et al., 2023) extends these methods by automatically upweighting poorly predicted examples without needing explicit group labels. GroupDRO (Sagawa al., 2019) minimizes worst-case group loss, while approaches like LiF and JTT increase loss weights for likely minority examples. Data balancing can also be achieved through data synthesis, feature augmentation, or domain mixing (Hemmat et al., 2023; Yao et al., 2022; Han et al., 2022).
Logit adjustment methods are another line of work for robust classification under imbalanced data. Menon et al. (2020) propose method that corrects model predictions based on class frequencies, building on prior work in post-hoc adjustments (Collell et al., 2016; Kim & Kim, 2020; Kang et al., 2019). Other methods, such as Distell-iriibbtoo--AAe Aware Margin (LDAM) loss (Cao al., 2019),

---

Published as a conference paper at ICLR 2025
Balanced Softmax (Ren et al., 2020), Logit Correction (LC) (Liu al., 2022), and Unsupervised Logit Adjustment (uLA) (Tsirigotis et al., 2024), adjust classifier margins to handle class or group imbalance effectively.
Memorization Memorization in neural networks has received significant attention since the work of Zhang et al. (2021), which showed that these models can perfectly fit the training data, even with completely random labels. Several studies have studied the nuances of memorization across various scenarios (Zhang et al., 2019; Feldman, 2020; Feldman & Zhang, 2020; Brown al., 2021; 2022; Anagnostidis et al., 2022; Garg et al., 2023; Attias et al., 2024). Feldman (2020); Feldman & Zhang (2020) examined how memorization contributes improved performance on the tail of the distribution, especially visually similar training and test examples, and highlighted the role of memorizing outliers and mislabeled examples in preserving generalization, akin to our concept of "good memorization" terminology (Section 5).
Memorization and Spurious Correlations. Research has shown that memorization in neural networks can significantly affect model robustness and generalization. Arpit al. (2017); Maini et al. (2022); Stephenson et al. ;; Maini et al. (2023); Krueger et al. (2017) explore memorization's ' impact on neural networks, examining aspects like loss sensitivity, curvature, and the layer where memorization occurs. Yang et al. (2022) investigate "rare spurious correlations," which are akin to example-specific noise features that models memorize. Bombari & Mondelli (2024) provide a theoretical framework quantifying the memorization of spurious features, differentiating between model stability with respect to individual samples and alignment with spurious patterns. Finally, Yang et al. (2024) propose k-nearest neighbor- baaebbor-based regression to fit residuals, enhancing test performance across benchmarks.
(ResMem), which combines neural networks with
7 CONCLUSION
In this work, we show that while spurious correlations are inherently problematic, they become particularly harmful when paired with memorization. This combination drives the training loss to zero too early, stopping learning before the model can capture more meaningful patterns. To address this, we propose Memorization-Aware Training (MAT), which leverages the negative effects of memorization to mitigate the influence of spurious correlations. Notably, we highlight that memorization not always detrimental; its impact varies with the nature of the data. While MAT mitigates the negative effects of memorization in the presence of spurious correlations, there are cases where memorization can benefit generalization or even be essential (Feldman & Zhang, 2020). Future work could focus on distinguishing these scenarios and exploring the nuanced role of memorization in large language models (LLMs). Recent work (Carlini al., 2022; Schwarzschild et al., 2024; Antoniades et al., 2024) have highlighted the importance of defining and understanding memorization in LLMs, as it can inform how these models balance between storing training data and learning generalizable patterns.
ACKNOWLEDGEMENTS
We thank Kartik Ahuja, Andrei Nicolicioiu, and Reyhane Askari Hemmat for their invaluable feedback and discussions. RB acknowledges support from the Canada CIFAR AI Chair Program and the Canada Excellence Research Chairs (CERC) Program, well as additional support from the UNIQUE scholarship. Part of the early experiments were conducted using computational resources provided by Mila Quebec AI Institute. Finally, we thank the FAIR leadership for their support in publishing this work.
10

---

Published as a conference paper at ICLR 2025
REFERENCES
Sotiris Anagnostidis, Gregor Bachmann, Lorenzo Noci, and Thomas Hofmann. The curious case of benign memorization. arXiv preprint arXiv:2210. 14019, 2022.
Antonis Antoniades, Xinyi Wang, Yanai Elazar, Alfonso Amayuelas, Alon Albalak, Kexun Zhang, and William Yang Wang. Generalization vs memorization: Tracing language models' capabilities back to pretraining data, arXiv preprint arXiv:2407.14985, 2024.
Martin Arjovsky, Léon Bottou, Ishaan Gulrajani, and David Lopez-Paz. Invariant risk minimization. arXiv preprint arXiv:1907.02893; 2019.
Devansh Arpit, Stanistaw Jastrzebski, Nicolas Ballas, David Krueger, Emmanuel Bengio, Maxinder Kanwal, Tegan Maharaj. Asja Fischer, Aaron Courville, Yoshua Bengio, et closer look at memorization in deep networks. In International conference on machine learning, pp. 233-242.
PMLR, 2017.
Idan Attias, Gintare Karolina Dziugaite, Mahdi Haghifam, Roi Livni, and Daniel M. Roy. Information complexity of stochastic convex optimization: Applications to generalization and memorization. arXiv preprint arXiv: 2402.09327, 2024.
Peter I. Bartlett, Philip M Long. Gábor Lugosi, and Alexander Tsigler. Benign overfitting in linear regression. Proceedings of the National Academy of Sciences, 117(48)::0063-30077) 2020.
Mikhail Belkin, Daniel Hsu, Siyuan Ma, and Soumik Mandal. Reconciling modern machine-learning practice and the classical varianae trade-off. Proceedings of the National Academy of Sciences, 116632):5544-11554 2019a.
Mikhail Belkin, Alexander Rakhlin, and Alexandre B Tsybakov. Does data interpolation contradict statistical optimality? In The 22nd International Conference on Artificial Intelligence and Statistics, pp. 1611-1619. PMLR, 2019h.
Simone Bombari and Marco Mondelli. How spurious features memorized: Precise analysis for random and ntk features. In Forty-first International Conference on Machine Learning, 2024.
Daniel Borkan, Lucas Dixon, Jeffrey Sorensen, Nithum Thain, and Lucy Vasserman. Nuanced metrics for measuring unintended bias with real data for text classification. WWW, 2019. URL https: araii.ii....//... org/ otg/abss////ooooooo/ooooo//.sa.r 1903. 04561.
Gavin Brown, Mark Bun, Vitaly Feldman, Adam Smith, and Kunal Talwar. When memorization of irrelevant training data necessary for high-accuracy learning? In Proceedings of the 53rd annual ACM SIGACT symposium on theory of computing, pp. 123-132, 2021.
Gavin Brown, Mark Bun, and Adam Smith. Strong memory lower bounds for learning natural models. In Conference on Learning Theory, pp. 4989-5029. PMLR, 2022.
Kaidi Cao, Colin Wei, Adrien Gaidon, Nikos Arechiga, and Tengyu Ma. Learning imbalanced datasets with tabel-distribution-aaaa margin loss. Advances in neural information processing systems, 32, 2019.
Nicholas Carlini, Daphne Ippolito, Matthew Jagielski, Katherine Lee, Florian Tramer, and Chiyuan Zhang. Quantifying memorization across neural language models. arXiv preprint arXiv:2202.07644 076666 2022.
Guillem Collell, Drazen Prelec, and Kaustubh Patil. Reviving threshold-moving: a simple plug-in bagging ensemble for binary and multiclass imbalanced data. arXiv preprint arXiv:1606.0860 08698, 2016.
Elliot Creager, Jömn-Henrik Jacobsen, and Richard Zemel. Environment inference for invariant learning. In International Conference Machine Learning, pp. 2189-2200, PMLR, 2021.
Nikolay Dagaev, Brett D Roads, Xiaoliang Luo. Daniel N Barry, Kaustubh R Patil, and Bradley C Love. A co--oodd-to-be-tee prior to reduce shortcut reliance. Pattern recognition letters, 166: 164-171, 2023.
11

---

Published as a conference paper at ICLR 2025
Jacob Devlin, Ming-Wei Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprins arXiv: 04800, 2018.
B. Dherin, M. Munn, M. Rosca, and D. Barrett. Why neural networks find simple solutions: the many regularizers of geometric complexity. Neural Information Processing Systems, 2022. doi: 10.485500/r2xv.2220... 130833
Vitaly Feldman. Does learning require memorization? a short tale about a long tail. In Proceedings of the 52nd Annual ACM SIGACT Symposium on Theory of Computing, pp. 954-959, 2020.
Vitaly Feldman and Chiyuan Zhang. What neural networks memorize and why: Discovering the long tail via influence estimation. Advances in Neural Information Processing Systems, 33:2881-2891,
2020.
Isha Garg, Deepak Ravikumar, and Kaushik Roy. Memorization through the lens of curvature of loss function around samples. arXiv preprint arXiv:2337.05881, 2023.
Robert Geirhos, Jörn-Henrik Jacobsen, Claudio Michaelis, Richard Zemel, Wieland Brendel, Matthias Bethge, and Felix A Wichmann. Shortcut learning in deep neural networks. Nature Machine Intelligence, 2(11):655-673 2020.
Zongbo Han, Zhipeng Liang, Fan Yang, Liu Liu, Lanqing Li, Yatao Bian, Peilin Zhao, Bingzhe Wu, Changqing Zhang, and Jianhua Yao. Umix: Improving importance weighting for subpopulation shift via incertainty-awant mixup. Advances in Neural Information Processing Systems, 35: 37704-37718, 2022.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 770-778, 2016.
Reyhane Askari Hemmat, Mohammad Pezeshki, Florian Bordes, Michal Drozdzal, and Adriana Conero-Soriano, Feedback- guiided data synthesis for imbalanced classification. arXiv preprint arXiv:2310.00158, 2023.
Badr Youbi Idrissi, Martin Arjovsky, Mohammad Pezeshki, and David Lopez-Paz. Simple data balaneing achieves competitive worst-group- araat-rruuuaaarccc...... CLeaR, 2022. URL hittts://aaxiis org/abs/211 2210 14503.
Bingyi Kang, Saining Xie, Marcus Rohrbach, Zhicheng Yan, Albert Gordo, Jiashi Feng, and Yannis Kalantidis. Decoupling representation and classifier for tai-tailed recognition. arXiv preprint arXiv: 1910.9291 2019.
Byungju Kim and Junmo Kim. Adjusting decision boundary for class imbalanced learning. IEEE Access, 81685, 2020.
Byungju Kim, Hyunwoo Kim, Kyungsu Kim, Sungjin Kim, and Junmo Kim. Learning not to learn: Training deep neural networks with biased data. In Proceedings of the IEEE/CVE conference on computer vision and pattern recognition, pp. 9012-9020, 2019.
Polina Kirichenko, Pavel Izmailov, and Andrew Gordon Wilson. Last layer re-training is sufficient for robustness to spurious correlations. arXiv preprint arXiv:2204. 02937, 2022.
David Krueger, Nicolas Ballas, Stanislaw Jastrzebski, Devansh Arpit, Maxinder S Kanwal, Tegan Maharaj, Emmanuel Bengio, Asja Fischer, and Aaron Courville. Deep nets don't learn via memorization. 2017.
Thomas S Kuhn. The Copernican revolution: Planetary astronomy in the development of Western thought. Harvard University Press, 1992.
Tyler LaBonte, Vidya Muthukumar, and Abhishek Kumar. Towards last-layer retraining for group robustness with fewer annotations. Advances in Neural Information Processing Systems, 36, 2024.
Yi Li and Nuno Vasconcelos. Repair: Removing representation bias by dataset resampling. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pp. 9572- 2019.
12

---

Published as a conference paper at ICLR 2025
Evan Z Liu, Behzad Haghgoo, Annie S Chen, Aditi Raghunathan, Pang Wei Koh, Shiori Sagawa, Perey Liang, and Chelsea Finn. Just train twice: Improving group robustness without training group information. In International Conference on Machine Learning, pp. 6781-6792. PMLR, 2021.
Sheng Liu, Xu Zhang, Nitesh Sekhar, Yue Wu, Prateek Singhal, and Carlos Fernandez-Granda. Avoiding spurious correlations via logit correction. arXiv preprint arXiv:2212.01433, 2022.
Ziwei Liu, Ping Luo, Xiaogang Wang, and Xiaoou Tang. Deep learning face attributes the wild. In ICCV, 2015. URL https: http:///aavvvvoiie abs 11114 7766.
Iyya Loshchilov and Frank Hutter. Decoupled weight decay regularization. arXiv preprint arXiv:1711.05101, 2017.
Pratyush Maini, Saurabh Garg, Zachary Lipton, and J Zico Kolter. Characterizing datapoints via second-split forgetting. Advances in Neural Information Processing Systems, 35:30044-30057, 2022.
Pratyush Maini, Michael C Mozer, Hanie Sedghi, Zachary C Lipton, J Zico Kolter, and Chiyuan Zhang. Can neural network memorization be localized? arXiv preprint WrXiv:2307.09542 2023.
Neil Mallinar, James Simon, Amirhesam Abedsoltan, Parthe Pandit, Misha Belkin, and Preetum Nakkiran. Benign, tempered, or catastrophic: Toward refined taxonomy of overfitting. Advances in Neural Information Processing Systems, 35:1 35:1182-1195 1195, 2022.
Aditya Krishna Menon, Sadeep Jayasumana, Ankit Singh Rawat, Himanshu Jain, Andreas Veit, and Sanjiv Kumar. Long-tail learning via logit adjustment. arXiv preprint arXiv:2007.07314 2020.
Vidya Muthukumar, Kailas Vodrahalli, Vignesh Subramanian, and Anant Sahai. Harmless interpo- lation of noisy data in regression. IEEE Journal on Selected Areas in Information Theory, 1(1): 67-83, 2020.
Vaishnavh Nagarajan, Anders Andreassen, and Behnam Neyshabur. Understanding the failure modes of out-of-distribution generalization. arXiv preprint arXiv:2010.15773 2020.
Preetum Nakkiran, Gal Kaplun, Yamini Bansal, Tristan Yang. Boaz Barak, and llya Sutskever. Deep double descent: Where bigger models and more data hurt. Journal of Statistical Mechanics: Theory and Experiment, (11))))224003 2021.
Junhyun Nam, Hyuntak Cha, Sungsoo Ahn, Jaeho Lee, and Jinwoo Shin. Learning from failure: De-biasing classifier from biased classifier. Advances in Neural Information Processing Systems, 2020.
Sung Min Park, Kristian Georgiev, Andrew Ilyas, Guillaume Leclere, and Madry. Trak: Attributing model behavior at scale. International Conference on Machine Learning, 2023, doi: 10.48550/
arXiv.2303.1 14186.
Mohammad Pezeshki, Oumar Kaba, Yoshua Bengio, Aaron C Courville, Doina Precup, and Guil- laume Lajoie. Gradient starvation: A learning proclivity in neural networks. Advances in Neural Information Processing Systems, 34: 1es111252222 1272, 2021.
Mohammad Pezeshki, Diane Bouchacourt, Mark Ibrahim, Nicolas Ballas, Pascal Vincent, and David Lopez-Paz. Discovering environments with xrm, arXiv preprint arXiv:2309. 16748, 2023.
Aahlad Manas Puli, Lily Zhang, Yoav Wald, and Rajesh Ranganath. Don't blame dataset shift! shortcut learning due to gradients and cross entropy. Advances in Neural Information Processing Systems, 36.71874-71910, 2023.
Shikai Qiu, Andres Potapezynski, Pavel Izmailov, and Andrew Gordon Wilson. Simple and fast group robustness by automatic feature reweighting. In International Conference Machine Learning, pp. 28448-28461 PMLR, 2023.
Jiawei Ren, Cunjun Yu, Xiao Ma, Haiyu Zhao, Shuai Yi, et al. Balanced meta-softmax for long-tailes visual recognition. Advances in neural information processing systems, 33:4175-4186 2020.
13

---

Published as a conference paper at ICLR 2025
Mengye Ren, Wenyuan Zeng, Bin Yang, and Raquel Urtasun. Learning to reweight examples for robust deep learning. In International conference on machine learning, pp. 4334 4343. PMLR, 2018.
Michael Roberts, Derek Driggs, Matthew Thorpe, Julian Gilbey, Michael Yeung, Stephan Ursprung. Angelica I Aviles-Rivero, Christian Etmann, Cathal McCague, Lucian Beer, et al. Common pitfalls and recommendations for using machine learning to detect and prognosticate for covid-19 using chest radiographs and ct scans. Nature Machine Intelligence, (3):199-217, 2021.
Shiori Sagawa, Pang Wei Koh, Tatsunori B Hashimoto, and Percy Liang. Distributionally robust
neural networks for group shifts: On the importance of regularization for worst-case generalization. arXiv preprini arXiv: 1119911.000 1.08731, 2019.
Shiori Sagawa, Aditi Raghunathan, Pang Wei Koh, and Percy Liang. An investigation of why overparameterizatioo exacerbates spurious correlations. In International Conference on Machine Learning, pp. 6346-8356. PMLR, 2020.
Avi Schwarzschild, Zhili Feng, Pratyush Maini, Zachary C Lipton, and J Zico Kolter. Rethinking IIm memorization through the lens of adversarial compression. arXiv preprint arXiv:2404. 15146, 2024.
Harshay Shah, Kaustav Tamuly, Aditi Raghunathan, Prateek Jain, and Praneeth Netrapalli, The pitfalls simplicity bias in neural networks. Advances in Neural Information Processing Systems, 33:9573-9585, 2020.
Daniel Soudry, Elad Hoffer, Mor Shpigel Nacson, Suriya Gunasekar, and Nathan Srebro. The implicit bias of gradient descent on separable data. Journal of Machine Learning Research, 19(70):1-57, 2018.
Cory Stephenson, Suchismita Padhy, Abhinav Ganesh, Yue Hui, Hanlin Tang, and Yuon Chung. On the geometry of generalization and memorization in deep neural networks. arXiv preprint arXiv:210S. 14602, 2021.
Christos Tsirigotis, Joao Monteiro, Pau Rodriguez, David Vazquez, and Aaron C Courville. Group robust classification without any group information. Advances in Neural Information Processing Systems, 36, 2024.
Catherine Wah, Steve Branson, Peter Welinder, Pietro Perona, and Serge Belongie. The Caltech- UCSD birds-200-2011 dataset. Califormia Institute of Technology, 2011. URL https: : /www. vision. caltech. wet/ datasets//ssaaaaaa cub 200 200_2011/.
Yoav Wald, Gal Yona, Uri Shalit, and Yair Carmon. Malign overfitting: Interpolation and invariance are fundamentally at odds. In NewrIPS 2022 Workshop on Distribution Shifts: Connecting Methods and Applications, 2022.
Adina Williams, Nikita Nangia, and Samuel R Bowman. A broad-coverage challenge corpus for sentence understanding through inference. ACL, 2017. htl.tttss///ccaattott org/ N1 N18-1101/.
Yao- Yuo- Yang. Chi-Ning Chou, and Kamalika Chaudhuri. Understanding rare spurious correlations in neural networks. arXiv preprint arXiv:2202.05/89, 2022.
Zitong Yang. Michal Lukasik, Vaishnavh Nagarajan, Zonglin Li, Ankit Rawat, Manzil Zaheer, Aditya K Menon, and Sanjiv Kumar. Resmem: Learn what you and memorize the rest. Advances in Neural Information Processing Systems, 36, 2024.
Huaxiu Yao, Yu Wang, Sai Li, Linjun Zhang, Weixin Liang, James Zou, and Chelsea Finn. Improving out-of-distribution robustness via selective augmentation. International Conference on Machine Learning, pp. 25407-25437. PMLR, 2022.
Chiyuan Zhang, Samy Bengio, Moritz Hardt, Michael C Mozer, and Yoram Singer. Identity crisis: Memorization and generalization under extreme overparameterizztton. arXiv preprint arXiv:1902 04698, 2019.
14

---

Published as a conference paper at ICLR 2025
Chiyuan Zhang, Samy Bengio, Moritz Hardt, Benjamin Recht, and Oriol Vinyals. Understanding deep learning (still) requires rethinking generalization. Communications of the ACM, 64(3):107-115, 2021.
Michael Zhang, Nimit S Sohoni, Hongyang Zhang, Chelsea Finn, and Christopher Ré. Correct-n-- contrast: A contrastive approach for improving robustness to spurious correlations. arXiv preprint arXiv:2203.01517, 2022.
Bolei Zhou, Agata Lapedriza, Aditya Khosla, Aude Oliva, and Antonio Torralba. Places: A 10 million image database for scene recognition. IEEE transactions pattern analysis and machine intelligence, 40(6))1452-1444 2017.
15

---

Published as a conference paper at ICLR 2025
A EXPERIMENTAL DETAILS
A.1 EXPERIMENTS ON SUBPOPULATION SHIFT
For the results presented in Table 1, we follow the experimental settings in Pezeshki et al. (2023). The hyperparameter search involves testing 16 random hyperparameter combinations sampled from the search space described in Table 3, using a single random seed. We select the hyperparameter combination and the carly-stopping iteration that achieve the highest validation worst-group accuracy, either with ground truth group annotations pseudo annotations, depending worst-class accuracy if groups are not available. Then, to calculate the results with mean and standard deviation, repeat the best-chosen hyperparameter experiment 10 times with different random seeds. Finally, to ensure a fair comparison between different methods, we always report the test worst-group accuracy based on the ground-truth (human annotated) group annotations provided by each dataset.
the method, or the
We tested our method on 4 standard datasets: two image datasets, Waterbirds (Sagawa et al., 2019) and CelebA (Liu et al., 2015), and two natural language datasets, MultiNLI (Williams et al., 2017) and CivilComments (Borkan al., 2019). The configuration of each dataset provided below. For CelebA, predictors map pixel intensities into a binary "blonde/not-blonee characteristics, landmarks, keypoints, facial mapping, metadata, or any other information was used to train our CelebA predictors. We use pre-trained ResNet-50 (He et al., 2016) for image datasets. For text datasets, we use a pre-trained BERT (Devlin et al., 2018). We initialized the weights of the linear layer added on top of the pre- pre-trained model with zero. All image datasets have the same pre- pree-prceesing scheme, which involves resizing and center-cropping to 224 x 224 pixels without any data augmentation. We use SGD with momentum of 0.9 for the Waterbirds dataset, and we employ AdamW (Loshchilov & Hutter, 2017) with default values of 31 = 0.9 and 3a = 0.999 for the other datasets.
label. No individual face
Dataset. The detailed statistics for all datasets are provided in Table 2, together with the descriptions of each task below.
Waterbirds (Sagawa et al., 2019): The Waterbirds dataset a combination of the Caltech- UCSD Birds 200 dataset (Wah et al., 201 and the Places dataset (Zhou et al., 2017). It consists of images where two types of birds (Waterbirds and Landbirds) are placed on either water or land backgrounds. The objective is to classify the type of bird as either "Waterbird" or "Landbird", with the background (water or land) introducing a spurious correlation.
CelebA (Liu et al., 2015): The CelebA dataset is a binary classification task where the objective is to classify hair either "Blond" or "Not-Blond", with gender considered spurious correlation.
MultiNLI (Williams et al., 2017): The MultiNLI dataset is natural language inference task in which the objective is to determine whether the second sentence in a given pair is "entailed "neutral with", or "contradicts" the first sentence. The spurious correlation is the presence of negation words.
CivilComments (Borkan et al., 2019): The CivilComments dataset is natural language inference task in which the objective is classify whether a sentence is "Toxic" "Non- Toxic".
Table 2: Summary of datasets used, including their data types, the number of classes, the number of groups, and the total dataset size for each.
Dataset
Data type Num. of classes
Num. of groups
Train size
Waterbirds
Image
2
2
4795
CelebA
Image
2
2
162770
MultiNLI
Text
3
6
206175
CivilComments
Text
2
8
269038
16

---

Published as a conference paper at ICLR 2025
The baseline methods compared our method with include ERM, GroupDRO (Sagawa et al., 2019), LfF (Nam et al., 2020). JTT (Liu et al., 2021), ) LC (Liu et al., 2022). uLA (Tsirigotis et al., 2024), AFR (Qiu et al., 2023), XRM+GroupDRO (Pezeshki et al., 2023). The results for ERM. GroupDRO are based on our own implementation, while the results for rest of the methods are adapted from the respective papers.
Table Hyperparameter search space, ERM and MAT share the same hyperparameter search space, except that MAT has one additional hyperparameter. т, which is used in the softmax function as the temperature parameter control the sharpness/smoothness of the output distribution.
algorithm
prper-parameter
ResNet
BERT
learning rate weight decay batch size dropout
10 Uniformm( ( 55)33)
10 Uniforom( -6
ERM & MAT
10 Uniformm-66 (-6.-3)
Uniiommm -6,
2 2
2 2nitiomm (4
Random (0, 0.5)
MAT specific
Random([0.001, 0.01, 0.1])
Random Random((0.001,00 (0.001,, 0.01, 0.1))
GroupDRO
Un
10 Uniform -3,-1)
Methods.
We compared our method with variety of baseline approaches presented in the follow-
ing.
Empirical Risk Minimization (ERM) trains model by minimizing the average loss over the entire training dataset.
ERM:
GroupDRO (Sagawa et al., 2019): Group Distributionally Robust Optimization (GroupDRO)
minimizes the worst-group loss across different predetermined groups in the training data using ground-truth -tuth group annotation.
LfF (Nam et al., 2020): Learning from Failure (LfF) trains two models, one biased and one debiased. The biased model is trained to amplify reliance on spurious correlations using generalized cross-entropy, while the debiased model focuses on samples where the biased model fails.
JTT (Liu et al., 2021): Just Train Twice (JTT) is two-stage method that first identifies misclassified examples with ERM, then upweights them in the second stage to improve performance on hard-to-learn groups.
LC (Liu et al., 2022): Logit Correction (LC) trains two models, one biased and one debiased. The biased model, trained with generalized cross-entropy, produces a correction term for the logits of the debiased model. Additionally, LC uses Group MixUp strategy between minority and majority groups further enrich the representation of the minority groups.
uLA (Tsirigotis et al., 2024): Unsupervised Logit Adjustment (uLA) uses pretrained self- supervised model to generate biased predictions, which are then used to adjust the logits of another model, improving robustness without needing group information.
AFR (Qiu et al., 2023): Automatic Feature Reweighting (AFR) retrains the last layer of model by up-weighting examples that the base model poorly predicted to reduce reliance on spurious features.
XRM+GroupDRO (Pezeshki al., 2023): Cross-Risk Minimization (XRM) automatically discovers environments a dataset by training twin networks on disjoint parts of the training data that learn from each other's errors. XRM employs label-flippig strategy amplify model biases and better identify spurious correlations. In the next stage, an invariant learning method like GroupDRO uses the discovered environments (i.e., groups) to improve the worst group performance.
17

---

Published as a conference paper at ICLR 2025
A.2 ANALYSIS OF MEMORIZATION SCORES
target
For the results presented in Figure 2, we used the TRAK frame- work (Park et al., 2023) to calculate self-influence scores for individual training data points. The TRAK framework com- putes the influence of each training point on a target dataset, generating an influence matrix, as illustrated in Figure 4. By setting the target dataset as the training dataset itself, we com- pute pairwise influence scores for all training data points, where the diagonal of this matrix represents the self-influence each point on its own prediction. We visualized the distribution of these scores for both the majority and minority groups of the Waterbird dataset.
Sij
Influence Matrix
Figure 4: Influence Matrix: In- fluence scores of all the train and target data point pairs. Setting target train and reveals the self-influence.
To utilize the framework, we provided the necessary model checkpoints saved during training. These checkpoints enabled TRAK to evaluate the contributions of specific training exam- ples the model's predictions.
B EMPIRICAL INSIGHTS INTO THE EFFECTS OF REWEIGHTING AND LOGIT-SHIFTING
In this section, we empirically illustrate how gradient descent dynamics in logistic regression are influenced by different weighting and shifting schemes. ERM with uniform weights converges the max-margin solution, aligning with the theoretical expectation (Soudry et al., 2018). Reweighting examples changes the optimization path but ultimately converges to the max-margin fixed point. In contrast, introducing a shift that scales with the logits changes the optimization landscape, leading to a distinct fixed point and different classifier. These results highlight how both reweighting and logit-scaling shifts impact the intermediate learning dynamics, but only shifts fundamentally alter the convergence behavior. Below we provide details on the dataset, model, and experimental configurations, with results shown in Figure 5.
The dataset consists of two-class synthetic data generated using Gaussian mixtures with overlapping clusters. We generate n 100 samples TR² with each class drawn from distinet Gaussian distributions. The features are standardized zero mean and unit variance using standard scaler, and the labels are assigned as y E 1}.
Dataset.
Model and Optimization.
The logistic regression model is parameterized as:
f(; w. b) Xww
where X R2x2 is the input data, E 2 is the weight vector, and b IR is the bias term. The gradient descent update is computed for the cross-entropy loss:
- T1
C
w. log
i (1=
where w; are per-example weights. We compare three configurations:
ERM: Uniform weights are applied to all examples, i.e., 1 for all
Reweight: We reduce the weights of examples for half of the examples of one class (e.g.. ! = 1) by setting w; for these examples and w; = for the rest. This simulates scenarios where certain groups are more important during training, e.g. are minority.
Shift: A margin-dependent shift is introduced by modifying the logits as:
fo; w, b) X + b + 8, -|w|||
where = 2 for the same half of class y = 1 and 6, = for the other. The scaling with the norm of the weight vector effectively adjusts the margin each example.
18

---

Published as a conference paper at ICLR 2025
To analyze the optimization dynamics, plot the trajectories of the normalized weight vectors during training. Specifically, we plot: w/ w/|lw|| x log(t), where is the iteration step, to ensure the trajectories visually distinguishable since all points would otherwise lie the umit circle.
As shown in Figure 5, ERM converges to the max-margin solution. Reweighting changes the optimization trajectory initially but still converges the max-margin fixed point. In contrast, the logit-shifting changes the fixed point of the optimization, making logit-shifting more robust method and alleviating the need for precise early stopping.
Trajectory Normalized Weight Vector
2-
Normatzzed seggit
trajectory
trajectory
trajectory
Nax margin solution
Normaizzed my logic)
Figure 5: Weight vector trajectories for ERM (blue), reweighting (orange), and shifting (green) compared with the max-margin solution (black dot). Weight vector trajectories normalized as w/||w|| x log( where is the iteration step. The scaling by log(t) is done for better visualization only.
C FURTHER DISCUSSION ON EXPERIMENTAL RESULTS
Subpopulation shift methods operate under three setups regarding access to group annotations, cach progressively increasing the task's 's difficulty.
The least restrictive scenario assumes full access to group annotations for both training and validation sets. This allows methods such GroupDRO to leverage loss re-weighting with group-aware model selection and early stopping. While effective, this setup requires strong assumptions about the availability of group annotations. As noted in the original GroupDRO paper (Sagawa al., 2020), achieving consistent performance in this setup often necessitates "stronger-than-typiee L2 regularization early stopping". Without such measures, as shown Sagawa al. (2020) (Figure 2), GroupDRO's behavior can revert to ERM-like performance due to the implicit bias of gradient descent toward the max-me -maggin solution.
A slightly more restrictive scenario assumes access to group annotations only for the validation set, which are used for model selection and early stopping. However, most algorithms, including GroupDRO, still rely on explicit group annotations during training, which must often be inferred using methods such as XRM. contrast, MAT does not depend on training group annotations, making it more robust in cases where such information is unavailable.
The most restrictive and realistic scenario assumes no access to group annotations at any stage, either for training validation. In this context, methods like MAT, which per-example logits from held-out predictions, become particularly interesting. MAT operates without explicit group-specific information, making it flexible and robust in scenarios with ambiguous undefined group structures. However, for the purpose model selection, MAT does rely on explicit validation group annotations being inferred.
19

---

Published as a conference paper at ICLR 2025
MAT's group-agnostic nature makes it particularly suitable for real-world scenarios where group definitions may be complex or unavailable. While GroupDRO relies on predefined or inferred group labels for its reweighting strategy, MAT uses per-example and soft adjustments without explicit notion of group
D ADDITIONAL INFLUENCE-SCORE EXPERIMENTS
In this section, we extend our experiments on influence-score analysis (refer to 4.2) across three methods: ERM, GroupDRO, and MAT. Figure 6 demonstrates that MAT effectively reduces self- influence scores relative to the other methods. Additionally, GroupDRO also shows some level reduction in the self-influence score of groups, specifically for the minority groups. This further indicates that memorization limiting generalization, and mitigating it contributes to improved generalization in subpopulation shift settings.
ERM Waterbird
ERM Landbird
ERM Proportion Samples
Minority Majarity
Minority Majarity
LB Land
Progortion
LB Water
Density
Density
50.8
WB Land
0.4
0.00
0.50
0.75
1.00
0.00
0.50
0.75
1.00
0.7 Self-Influence Interval
0.5
0.9
Seifinfluence Score
Seifinfluence Score
GroupDRO Waterbird
GroupDRO Landbird
GroupDRO Proportion of Samples
Minority Majority
Minority Majority
Prepertion
2.4
LB Water
Density
Density
1.6
WB on Water
WB
0.00
0,25
0.50
0.75
1.00
0.00
0.25
0.50
0.75
1.00
0.0
0.1
0.3
0.5
0.7
0.9
Self Influence Score
Seif-Influence Score
Self Influence Interval
Waterbird
Landbird
Proportion Samples
Minority Majerity
Minority Majority
3.2
LB on Land Water
Density
Density
Prosortion
2.44
WB Water
1.6
WB Land
Figure 6: Self-Influence estimation of the Waterbird groups by ERM, GroupDRO and MAT, Each row corresponds to one of the training methods (ERM, GroupDRO, and MAT). The left and middle columns show the density distributions of self-influence scores for minority and majority subpop- ulations within each method, separately for Waterbirds and Landbirds. The right column depicts the proportion of samples across different self-influence score intervals for each subpopulation (LB on Land, LB Water, WB on Water, WB on Land). The results illustrate that MAT achieves the most uniform and reduced self-influence distribution, indicating effective mitigation of memorization across subpopulations, especially in minority groups. GroupDRO also shows some reduction in self-influence scores for the critical group in this dataset, WB on Land.
MULTIPLE SPURIOUS FEATURES SETUP
To test the capability of MAT in handling multiple spurious features, have adapted our original setup in Section 2.1 and added a second spurious feature. Specifically, the data generation process is defined
20

---

Published as a conference paper at ICLR 2025
Ty ~ Ny,, og)
a
Jo. NNaa,
E GR++2 where, a
y
y w.p. - and p
w.p. P
p 0.5 (test)
(train)
Faz N(a, o2)
N(0,0
Illustrative Scenarios. Similarly, we chose p² = 0.9 (i.e., the correlation of spurious features with labels) and y = 5 (i.e., the strength of the spurious features) for both spurious features. In this new setup, with one main feature and two spurious features, there are four groups within each class, resulting in total of eight groups. For simplicity in our visualization, we identified the group with the highest number of samples as the majority group and the group with the fewest samples as the minority group (i.e., the perss-pprforinn group). Note that the number of samples in this minority group is even smaller than before, as the addition of the second spurious feature splits the previous minority group, introduced by the first spurious feature, into two even smaller groups.
Consistent with the single spurious feature setup (Section 2.1). Figure shows that ERM, with memorization capacity, struggles to learn generalizable features and achieves only random guesss performance for the minority samples. In contrast, MAT effectively handles this scenario and nearly achieves perfect test accuracy for groups.
6, memorization)
1.0
1.0
Accarscy
Aecurary
64
Majority
Majority Truis
Majority Trais
Minseity muin
MinortlyImmn
Minectly ruin
Majority Test
Majority Text
Majarity Test
Minuetty Text
60
Minuetiy Text
60
Minuetty Text
Epochs
Epochs
Figure 7: Similar to Figure 1, this figure illustrates the performance of ERM and MAT across different memorization setups in a scenario with multiple spurious features. The left panel shows ERM without memorization, where the model generalizes well both majority and minority groups. The middle panel shows ERM with memorization capacity, where the model exhibits poor generalization for minority test samples. The right panel illustrates MAT, where it effectively learns invariant features and generalizes well to all groups.
F PROOFS
Σ Et eda be a softmax over the logits (;(x), and define P2(y = x) such that p2(y = x) x w(j, pily = x) for some weighting function
Lemma F.1. Let (1 = x)) = mf(m)
w(ja x). Then:
(l
P2(U = x))
Σ (o0g
Proof. Starting with the definition:
Pz(U = j xx) x w(i pi (v x))
Substituting the expression for pi (y = x):
(m))
pa = x)) w(j
Σ 1==1 (0)(1)
21

---

Published as a conference paper at ICLR 2025
Since the denominator Σ - ( is constant with respect can be absorbed into the propor-
tionality constant:
p2 = j x)) x w(j. x) eds(a),
Using the property wojj m) ebslm) (00 have:
Pz(y= = x)) x ( w(j.))
To obtain a valid probability distribution, we normalize by computing the normalization constant Z:
Z =
(ug w(i.))
Therefore, the normalized p2(y = x:) is:
p2 = j m²) (m)( ((j.)))))))))))))))))))))
(m) log m(
Z
Li=1 ( w(im))
This completes the proof.
G PROOF OF THEOREM 2.2
Setting the derivative of the objective function [ERM (Equation equation 2) with respect w to zero gives the normal equation according to Appendix G.4
CCEMM
1
(5(2, w) W1)a, Aw O,
dw
n
where s(x]Tw) w) = (y xe; w), and solving for w then gives
w
Cg
with oq Ri- with [y.00), s(u),, vi == x2 W. 17 nd. 7
Note that the vy's correspond to logits, while a (an...., an) ) E RD should be thought of as dual representation of the weights vector w. Indeed, by construction, one has
T = TTaa a
where X Ra the design matrix.
Notation. Henceforth, with abuse of notation but WLOG, we will write Eepu = a for the spurious feature, and therefore write x; = (ys, ag ag. G)., where e; are the example-specifie noise features for the example and E (+1) } is its label.
Our mission is then to derive necessary and sufficient conditions for > O, where
= Wpuprr Weore =
Σ (22a; - y: yi
(10)
i=1
22

---

Published as a conference paper at ICLR 2025
G.1 FIXED-POINT EQUATIONS
Define subsets 1,, L C [n] and integers m, k E [n] by
4 := E [n] yi = =1},
(11)
S {i [n] a; 79:7,
(12)
L = {i E [n] a = -196)
(13)
E |S|
pn, pn,
== B (1-
(14)
Thus, S (resp. L) corresponds to the sample indices in the majority (resp. the minority) class.
One computes the logits as follows
0 x. w0 Σ
a oi Σ
ag + y²
Σ f=1
ajaaa +
Σ 3=1
age
T
6i
j=1
3=1
a + Σ = 9g6 T
ifiEsnL.
(15)
b j= Og9 Gr ifiESnL,
c+ j=1 aye, 6/ ifie
e 1リー1 ag9 4, ifiELnI,
where a, b, R are defined by
= Yuspu + Woore
Σ
+
==1
b = spupp
Weore
Σ
(7aag + Ub Wjj
j=1
(16)
= spspu Weore
Σ (v²ag yj)ag,
C = Weore ypppp
Σ
(-72a, + y;)aa.
5=1
Observe that
a0
c= -c.
(17)
The following lemma will be crucial our proof. Lemma G.1. a 2 0 and e 2 O, then part 2 of Theorem 2.2 holds. On the other hand, if a 2 0 and S. O, then part I of Theorem 2.2 holds.
Proof. Indeed, for random test (i.e., held-out) point (x, v) we have
p(CFRA((z) Capu p(æ x ④ 2 0)
p(yas Werre + @ + spp x We 2 0)
= p( Dp x The S 22 + sp Weore)
Now, independent of y, the random variable sp T ④ has distribution (O,, o² (.1l' 1))). Now, because w = X a by construction, the variance can be written as 12 o itself chi-squared random variable which concentrates around its mean ||] Furthermore, thanks to equation 21, ||a]]2 < 1/(n>?). which vanishes in the limit
_ ail² which is
X + ++ + 00, Avn + D0.
(18)
We deduce that
p(CERM(Φ)) Copu (x)) - (Spp + yo you 1 core 0)
-p + Seore 201 + p) -Bore =pl(a20) + ()=2))(220).
23

---

Published as a conference paper at ICLR 2025
Thus, if a 2 0 and 2 O, must have ((CERM)( Copu p 1 - p 1, that is, part of Theorem 2.2 holds. In other words,
ho VERM (z) =a) 1.
p
On the other hand, one has
p(CERM((z)) - Coore (x)) (core x x ④ 2 0) =p( (pp((uaan Weore + y/sspu Wap 0)
(Beare ty0 20) + (1 - (10rre
gl (a20) + 4e500
where q : p(a y)-
We deduce that if a 2 0 and e < 0, then p(CERM = Ccore + - -q 1, ie part 2 of Theorem 2.2 holds. In other words,
(+) = v) + 1.
G.2 STRUCTURE THE DUAL WEIGHTS
The following result shows that the dual weights 01,-..;0y cluster into 4 lumps corresponding the following 4 sets of indices $ n [., S S T_, Ln T,, and L I_.
Lemma G.2. There exist positive constants A, C,C,E. E 0 such that the following holds with large probability uniformly over all indices i [n]
A fiieSLL,
- B, fiienL.
(19)
C, fii
-F,, Ini.
Furthermore, the empirical probabilities predicted by ERM are given by
- nA fitesL nB, fnE 1-nC, fnE.
T. [v.=1) 7/00
(20)
EE inL..LL
Proof. First observe that
||||| < AVn
(21)
Indeed, one computes
|al12
Σ 1=1
(T: - )2 <
5
n1
1
172
22m
1=1
Next, observe that Σ, ays, e = ailley11 + Ejer ayeJe es ~ aaadd.
This is because as/|e,]|2
concentrates around it mean which equals afadd, while w.h.p.
logn n
log
log n
sup ie(n]
Σ Jaysee e i71
$ |all
= oe aellallVa
< σ.
d = o(1)
2dd
d
d
24

---

Published as a conference paper at ICLR 2025
The above is because À 0 and (log n)/d - by assumption. Henceforth simply ignore the contributions the terms Σxzt ajez c,. We get the following equations in the limit 18
aa a,
oa + b,
fieSnL.
1/3
o afa + e,
Ln
o a te,
ffiE
(22)
ss(oaa + a),
i
s(o +b),
i Snl_,
11, = y1 - s(U))
slofa,d +c).
i LNI+
s(02ad + e),
if LNL.
Now, because monotonicity of a, we can find A, B, C, E such that
A, E
B, ifieSnL
04
C ffieLLI, I+,
E,
LnL,
as claimed.
We will make use of the following lemma. Lemma G.3. In the unregularized limit X + 0+, it holds that nA,BB, n,, E [0, 1/2].
Proof. Indeed, in that unregularized limit, ERM attains zero classification error on the training dataset (part O of Theorem 2.2). This means mean that म. 1/2 iff y; 1, and the result follows.
G.3 FINAL ToUCH (PROOF THEOREM 2.2)
We resume the proof of Theorem 2.2. The scalars A, B, C, E must verify
7 s0 Ad + a) s( s(-0 Ad a),
BB = o
- Bd + b) = s(-0, a) = - Bd
(23)
E = s( s(-o² Ed + e),
nC = 1 - sd c) = 1- slo.Cd - e) = s(( s(-o,Cd+e))
We deduce that
A B, C=
(24)
71A = s(-o Ad - a), qE = 8(-0² y===(-o3 Ed + e).
(25)
Proof of Part Theorem 2.2, In particular, for the case with no example-specific features where
G. + 0+, we have nA s(-a) and nE ~ s(e). We know from Lemma G.3 that nA n < 1/2. This implies a 2 0 and e < 0, and thanks to Lemma G.I, we deduce part 1 of Theorem 2.2
Proof of Part of Theorem 2.2. In remains to show that 2 0 and e > in the noisy regime G. 0, and then conclude via Lemma G.1.
Define N₁ == |5 I- |, N₂ IS J- N₂ 41 N₂ = I-. Note that from the definition of b, c, in equation 16, one has
(-2 + 1) + N2)A - -6-2 1 N4)E,
6 = (n² 1) + No)A - -622 +1)(s + +N)E,
c, n s( o(-o" Ad a),
b -a,
(26)
EE s s( Ed e).
B = A, C' E.
25

---

Published as a conference paper at ICLR 2025
We now show that and 0 under the conditions d >> logn and y >> 0ev d/m. Indeed, under the second condition, the following holds w.h.p
add d + (n2 + 1)(N1 + N2) = ((>² + 1) + N2) + +o2d) = ((72 + 1)m + o2.)) (7² + 1)m= (7² 1) N2),
where we have used the fact that N1 + N2 concentrates around its mean m = pn. We deduce that
Ad + a= (o.d (7 + 1)( + No)) (0- 1))(N3 + Na)
~ (772 1)(N + NN) - + NN)
from which we get.
1/2 nA 1 s( o(- a) = ((-(1+ = s(-a) + o1),
Le s-aa) 2 1/2 (1)). But this only happen if 0.
Finally, the conditions d log and y > O. v d/m imply 7 > Koevdo d/k and 9 2 K log(31) for any constant K > 0. Theorem of Puli et al. (2023) then gives = ypp W weore > 0, and we are done.
G.4 DERIVATIVE THE OBJECTIVE FUNCTION
Given the objective function (Equation equation 2):
Uy,, (y (g + X |11011
CERM
where U( (v (g:")) is the cross-entropy loss for binary classification, we can write:
1y, p' (y xy; w)) - y/ log(p"((s (y: ay w)) (1 y:) log(1 13t (ya x,;; w)),
where p"(v w; s(x] w) is the predicted probability of class y = 1, and s(.) is the sigmoid function:
s(x, w)
1 x
Now, the derivative of the LRR w.r.t. w consists of two parts:
Derivative of the cross-entropy loss: we compute the derivative of the cross-entropy term w.r.t. wr
Iyys, s(x; w)) Ow [- log(s( w)) - (1 y/ log(1 - s(x w))]
Let P; = =s w). Using the chain rule, we first calculate the derivative of the cross-entropy loss
w.r.t. P::
alur,
Oy: P: + y/
Ps
Next, we compute the derivative of pt = s(x0. w) w.r.t. w:
8p.
sax w1 sax. w))a:; ps(1 pt) p:)a..
Sw
Using the chain rule:
Øl(y,, s(x] (xxTw)) w))
(- V/1
s(x] w)(1 s(x) w))
Ow
s(x,, w)
s(x=, [2, w)
26

---

Published as a conference paper at ICLR 2025
Simplifying the expression:
al,,, s0, w))
(s(2)) w) 2:)
dw
Thus, the derivative of the cross-entropy loss term is:
a
1
Uy,, p0 (y w:)w))
1
(s w) y.)
ow
2. Derivative of the regularization term: the second part of the objective function is the b-- regularization term:
w
2
The derivative of this term w.r.t w is straightforward:
Aw.
Sw
Final derivative: combining both terms, the derivative of the objective function LRRM w.r.t. w is:
FCR
n
Σ
(s(aTw) yx + Aw.
dw
This is the required derivative of the regularized cross-entropy loss function with respect w.
27


Reviews:
[{"review_id": "8lT0zycAmk", "reviewer": "Reviewer_1LhZ", "summary": "This paper studies memorization in the presence of spurious features and input noise in the training data. To improve generalization the paper proposes memorization aware training (MAT). MAT computes held-out predictions, obtained using XRM (prior work), and then uses the predictions to shift model logits. The shifted logits force the model to learn generalizable features, empirically improving worst-group accuracy across common distribution shift benchmarks with spurious correlations like CelebA and Waterbirds.", "questions": "- In the \"ugly\" memorization setting in Section 5, it is unclear to me why in the noiseless setting, the solution learnt is different from the setting with $\\sigma=10^{-4}$. The class of over-parameterized networks contain both the function learnt in the $\\sigma=10^{-4}$ setting and $\\sigma=0$ setting. So, why is the learning biased to the bad solution?", "limitations": "", "rating": 5, "confidence": 5, "soundness": 1, "presentation": 2, "contribution": 1, "strengths": "- The analysis in Section 5 is interesting, particularly the distinction between the noiseless and small noise setting, where the latter actually generalizes better. Essentially, the results indicate that a small amount of independent noise in the input is essential to fit the label noise. It would have been really nice to see an expanded analysis on this claim, with some real world experiments and formal analysis.", "weaknesses": "- My main concern with the paper is lack of novelty in the main technical claim, i.e., memorization (afforded by overparameterization)+spurious correlations = poor generalization. There are multiple works like Shah et al. 2020, and Sagawa et al. 2020 that make the same claim with supporting evidence. \n- The connection between the objective (shifting the logits) and the final goal of preventing memorization is not very clear. In my understanding, shifting the logits should have the same affect as adding more weight in the loss on minority groups. \n- The theoretical analysis in Section 2.1 is very similar to Sagawa et al. 2020 (An investigation of why overparameterization exacerbates spurious correlations). In the prior work, the analysis in the same toy setup shows that memorization is exacerbated by spurious correlations, and thus the results in this work are almost directly implied by the results in Sagawa et al. 2020.\n- The empirical results are not very strong, since XRM+GroupDRO is comparable to MAT.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "oWug8LDA5M", "reviewer": "Reviewer_23ej", "summary": "This paper investigates the effect of the combination of spurious correlation and memorization and points out that the latter exacerbates the poor generalization caused by the former. The claim has basis on an experimental and theoretical analysis with a linear model. Furthermore, a method to mitigate the problem is proposed and its effectiveness is checked in experimental setups of subpopulation shift.", "questions": "- Please consider providing clarification regarding the weaknesses.\n- Is it possible to extend MAT for more general spurious correlation? In general situations, it can happen that multiple spurious features may correlate with each other, and the indirect path mentioned in Section 3 may change to the one involving intermediate correrations.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "- A theoretical analysis gives us an example model where the training with noiseless data leads to a test performance asymptotic to perfect accuracy and the training with noise data leads a test performance asymptotic to perfect fitting to spurious correlation(Theorem 2.2).\n\n- A new training loss for mitigating the harmful effect of spurious correlation without requiring annotations on spurious features is proposed based on the theoretical analysis mentioned above (Section 3.1).\n\n- The effect of the proposed method is checked experimentally with the subpopulation shift problem. Although superiority against existing methods for this problem is not observed clearly, the improvement from naive empirical risk minimization is obvious(Table 1). Analysis with influence function is also conducted, and it is observed that memorization is suppressed by the proposed method(Figure 2).", "weaknesses": "- The theoretical analysis is limited to the case of a linear model.\n\n- It is not clear to me how the logit shift proposed in 3.1 makes the training more \"memorization-aware\". In my understanding, the shift down-weights the gradient descent updates coming from the training samples with spurious correlation, and it should work, but it is not clear how the shift values change depending on the extent of memorization. It is preferable if a comment about this is added in 3.1. \n\n- (L691, L694) Duplication of a reference.\n\n- The code is not available with the first submission.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "AM0IjSdZCv", "reviewer": "Reviewer_USfi", "summary": "This paper examines the relationship between spurious correlations and memorization. A model trained with Empirical Risk Minimization (ERM), which has learned spurious correlations and has memorized irrelevant patterns, can have poor generalization. Spurious correlations result from patterns within the input data that are closely associated with the output, but not necessarily predictive (e.g., grass in a scene with cows). Memorization occurs when a model memorizes specific patterns rather than learning robust features that generalize to new examples.\n\nTo address this, the authors propose a novel approach, Memorization-Aware Training (MAT), a method that modifies the logits of the cross-entropy loss to discourage the model from relying on spurious features (\"the indirect path\"). By minimizing the log probability of the \"indirect path\", where the output y depends on a spurious feature a, MAT encourages the model to learn more generalizable patterns. \n\nThe accuracy of MAT is compared to various baseline methods under different sub-population label settings. Additionally, the paper shows that MAT effectively shifts the self-influence distribution, reducing the reliance on spurious correlations.\n\nThe authors additionally provide a thorough theoretical analysis and a detailed description MAT algorithm.", "questions": "1. (Continuation of weakness #2) Why did you chose to concatenate the noise to the input, rather than add it to the input? How do you think your results would differ if you added the noise to the input, simulating noisy real-world data? Your results, which rely on this simplification, may not hold true to real-world tasks.\n2. (Confusion) Looking at the results depicted in Figure 2, I noted that when training with ERM, the Waterbird on Water has a right-tailed self-influence score (e.g., in the right most figure, WB on Water has a large proportion of samples with a self-influence score of 0.3, yet LB on Land has a very tiny proportion of samples with a self-influence score of 0.3). Do you have any ideas on why this is the case? Is this due to a characteristic of the dataset? Of the sub-population images?\n3. You can also address any weaknesses.", "limitations": "", "rating": 6, "confidence": 2, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "1. I really appreciated the first paragraph of the introduction and how it related to your work. This was very creative and smart, and helps the reader in understanding.\n2. I like some the experiments in the paper and I have questions about others. First, for Figure 1, setting gamma to 5 creates a type of \"worst-case\" scenario with respect to spurious features. MAT is able to overcome this learning a generalizable function in the presence of noise. The experiments in Figure 2 show an improved self-influence score distribution with the use of MAT, which is a convincing way to present the effectiveness of your method.", "weaknesses": "1. One thing that the paper is missing is arguments as to why MAT is a better approach than other invariant methods? Or, alternatively, why should one use MAT over these other methods? What are the advantages and disadvantages, aside from sub-population labelling? The average accuracies among these methods seem to be similar. That is, there is no consistent 'best' approach. This doesn't invalidate the results, but I would like to know why and when I should choose MAT over another approach.\n2. I think that this paper would benefit from a more realistic scenario where noise modifies the input features. The concatenation of noise to the input is a simplification, which serves a clear purpose; however, it is not representative of real-world noise, in which the input features are directly modified.\n3. The rightmost image in Figure 3, with noisy labels, confused me in the sense that I don't understand how it demonstrates the benefit of using MAT.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "qlkqCmg0BR", "reviewer": "Reviewer_ttKz", "summary": "This paper studies the interplay between learning with spurious features and memorizing the noises in the training data. It shows that the combination of the two can be harmful for model generalization because the model lacks further incentive to learn the actual underlying generalizable solution. Based on this, the paper proposed a training paradigm called memorization-aware training (MAT) by utilizing a model trained with heldout data to shift the current model's logits.", "questions": "1. In the experiments, does the validation set used to shift the model logits contain spurious feature?\n2. L160: the paper says \"the model *first* learns $x_a$ ... Once the model achieve nearly perfect accuracy on the majority examples, it *starts* to learn the minority examples\". Could you explain where do we observe such order of learning from Figure 1?\n\n\n==============\nPost rebuttal: Thanks the authors for the response, especially the additional results and analysis. I've raised my rating.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "1. This paper presents a study of how memorization impact generalization, with a specific focus on the presence of spurious features and distribution shift between training and testing. \n2. The intuition is presented with both a synthetic data example and a theoretical construction.\n3. A new method is proposed based on the intuition that perform competitively compared to the previous baselines.", "weaknesses": "1. Spurious features and distribution shift is at the core of the most of the paper. I believe it is better to reflect this in the paper title, which currently is framed quite broadly.\n2. The presentation of the paper could be improved. For example\n\n    1. The synthetic example and discussions in Section 5, while interesting, feels disconnected and unrelated to the rest of the paper. Does it mean the memorization at the presence of spurious feature are always good / bad / ugly type of memorization? Or is the proposed training method going to help mitigate any type of memorization discussed here? Can those types of memorization be easily identified in real world scenarios? \n\n    2. Section 4.1 ends abruptly with a reference to Table 1 and a list of the baseline methods. Having some discussions of the results would be good since the proposed MAT algorithm is an important contribution of the paper. Right now looking at Table 1 it seems the proposed method is not outperforming previous methods when the group annotations for the validation data is available.\n\n3. The logits shift is estimated with Eq (4) when annotation is not available. It would be great to have an ablation study of the accuracy of such estimation in one of the experiments.\n\n4. The explanation of learning spurious feature + memorizing noise is interesting. However, it is still unclear from the results of the paper if this is what is happening in more general setting (i.e. without explicit spurious features and subpopulation shift). So it would be great to have some results on other standard ML benchmark as well.\n\n5. The analysis in Section 4.2 show that the proposed method reduces self influence in the minority subpopulations. It would be great to include the same analysis on some baseline methods compared in Section 4.1 as well.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)