Gradient Starvation: A Learning Proclivity in Neural Networks
Authors: Mohammad Pezeshki
Sékou-Oumar Kaba
Yoshua Bengio
Aaron Courville
Doina Precup
Guillaume Lajoie
Venue: NeurIPS 2021 Poster (2021)
Keywords: generalization
neural networks
dynamics
OOD

Abstract: We identify and formalize a fundamental gradient descent phenomenon resulting in a learning proclivity in over-parameterized neural networks. Gradient Starvation arises when cross-entropy loss is minimized by capturing only a subset of features relevant for the task, despite the presence of other predictive features that fail to be discovered. This work provides a theoretical explanation for the emergence of such feature imbalance in neural networks. Using tools from Dynamical Systems theory, we identify simple properties of learning dynamics during gradient descent that lead to this imbalance, and prove that such a situation can be expected given certain statistical structure in training data. Based on our proposed formalism, we develop guarantees for a novel regularization method aimed at decoupling feature learning dynamics, improving accuracy and robustness in cases hindered by gradient starvation. We illustrate our findings with simple and real-world out-of-distribution (OOD) generalization experiments.

Gradient Starvation: A Learning Proclivity in Neural Networks
Mohammad Pezeshki
Sékou-Oumar Kaba
Yoshua B.
Aaron Courville
Doina Precup
Guillaume Lajoie
Mila
Université de Montréal
McGill University
*Google DeepMind
corresponding authors: {pezeshki,
guillaume. lajoie}
@ajoie]@mia.. @mila. emila.quebec
Abstract
We identify and formalize a fundamental gradient descent phenomenon leading a learning proclivity in pere-arameteri neural networks. Gradient Starvation arises when cross-entropy loss is minimized by capturing only a subset of features relevant for the task, despite the presence of other predictive features that fail to discovered. This work provides a theoretical explanation for the emergence such feature imbalances in neural networks. Using tools from Dynamical Systems theory, we identify simple properties of learning dynamics during gradient descent that lead to this imbalance, and prove that such situation can be expected given certain statistical structure in training data. Based on our proposed formalism, develop guarantees for novel but simple regularization method aimed at decoupling feature learning dynamics, improving accuracy and robustness in cases hindered by gradient starvation. We illustrate findings with simple and real- world distriibiitiiioo (OOD) generalization experiments.
Introduction
In 1904, horse named Hans attracted worldwide attention due the belief that it was capable of doing arithmetic calculations [81]. Its trainer would ask Hans question, and Hans would reply by tapping the ground with its hoof. However, it was later revealed that the horse was only noticing subtle but distinctive signals its trainer's unconscious behavior, unbeknown to him, and actually performing arithmetic. An analogous phenomenon has been noticed when training neural networks [e.g. 85, 109, 54, 39, 17, 14, 37, 51, 107, 76, 48, 19, 61, 77]. In many cases, state-of-the-an neural networks appear to focus low-level superficial correlations, rather than more abstract and robustly informative features of interest [16, 88, 40, 68, 30).
The rationale behind this phenomenon is well known by practitioners: given strongly-correlate and fast-to-learn features training data, gradient descent biased towards learning them first. However, the precise conditions leading to such learning dynamics, and how one might intervene to control this feature imbalance are not entirely understood. Recent work aims at identifying the reasons behind this phenomenon [97, 70, 22, 71, 76, 100, 92, 83, 105, 42, 79, 4], while complementary work quantifies resulting shortcomings, including poor generalization to test data, reliance upon spurious correlations, and lack of robustness [30, 68, 77, 41, 63, 64, 9]. However most established work focuses on squared-error loss and its particularities, where results do not readily generalize to other objective forms. This is especially prohlematic since for several classification applications, cross-entropy is the loss function choice, yielding very distinct learning dynamics. In this paper, we argue that Gradient Starvation. first coined in [26], is leading cause for this feature imbalance in neural networks trained with cross-entropy, and propose a simple approach mitigate it.
(OOD)
35th Conference on Neural Information Processing Systems (NcurIPS 2021).

---

festuure
(0) A= 00.1
A= 0.1
(c) A = 0.1 with SD
iterations
Figure 1: Diagram illustrating the effect of gradient starvation in simple 2-D classification task. (a) Data is not linearly separable and the learned decision boundary curved. (b) Data is lincarly separable by a small margin (A = 0.1). This small margin allows the network to discriminate confidently only along the horizontal axis and ignore the vertical axis. (c) Data is linearly separable as in (b). However, with the proposed Spectral decoupling (SD), a curved decision boundary with a large margin is learned. (d) Diagram shows the evolution of two of the features (Eqq. of the dynamics in three cases shown dotted, dashed and solid lines. Analysis: (dotted) vs (dashed): Linear separability of the results in increase in 21 and a decrease (starvation) 22. (dashed) vs (solid): SD suppresses 21 and hence allows 22 to grow. Decision boundaries are averaged over ten runs. More experiments with common regularization methods are provided in App. B.
Here we summarize our contributions:
We provide a theoretical framework to study the learning dynamics of linearized neural networks trained with cross-entropy loss in a dual space.
Using perturbation analysis, we formalize Gradient Starvation (GS) in view of the coupling between the dynamics of orthogonal directions the feature space (Thm. 2).
We leverage our theory to introduce Spectral Decoupling (SD) (Eq. 17) and prove this simple regularizer helps to decouple learning dynamics, mitigating GS.
We support our findings with extensive empirical results a variety of classification and adversarial attack tasks. All code and experiment details available at Gi GitHub repository.
In the rest of the paper, we first present simple example to outline the consequences of GS. We then present our theoretical results before outlining number of numerical experiments. We close with review of related work followed by discussion.
2 Gradient Starvation: A simple example
Consider 2-D classification task with training set consisting of two classes, as shown in Figure A two-layer ReLU network with 500 hidden units is trained with cross-entropy loss for two different arrangements of the training points. The difference between the arrangements is that, in one setting, the data is not linearly separable, but a slight shift makes linearly separable in the other setting. This small shift allows the network to achieve negligible loss by only learning to discriminate along the horizontal axis, ignoring the other. This contrasts with the other case, where both features contribute the learned classification boundary, which arguably matches the data structure better. We observe that training longer or using different regularizers, including weight decay [58], dropout [95], batch normalization [49], as well as changing the optimization algorithm to Adam [56] or changing the network architecture or the coordinate system, do not encourage the network to learn a curved decision boundary. (See App. for more details.)
We argue that this occurs because cross-entropy loss leads to gradients "starved" of information from vertical features. Simply put, when one feature is learned faster than the others, the gradient contribution of examples containing that feature diminished (i.e., they are correctly processed based on that feature alone). This results lack of sufficient gradient signal, and hence prevents any remaining features from being learned. This simple mechanism has potential consequences, which we outline below.
2.1 Consequences of Gradient Starvation
Lack of robustness. the example above, even in the right plot, the training loss nearly zero, and the network is very confident in its predictions. However, the decision boundary is located very

---

close to the data points. This could lead to adversarial vulnerability as well as lack of robustness when generalizing to -uttr--diittibutiion data.
Excessive invariance. GS could also result in neural networks that are invariant to task-relevant changes in the input. In the example above, it is possible to obtaín a data point with low probability under the data distribution, but that would still be classified with high confidence.
Implicit regularization. One might argue that according to Occam's razor, a simpler decision boundary should generalize better. In fact, if both training and test sets share the same dominant feature (in this example, the feature along the horizontal axis), GS naturally prevents the learning of less dominant features that could otherwise result in overfitting Therefore, depending our assumptions on the training and test distributions, GS could also act as an implicit regularizer. We provide further discussion on the implicit regularization aspect of GS in Section 5.
3 Theoretical Results
In this section, study the learning dynamics of neural networks trained with cross-entropy loss. Particularly, we seek to decompose the learning dynamics along orthogonal directions in the feature space of neural networks, to provide formal definition of GS, and to derive simple regularization method to mitigate it. For analytical tractability, we make three key assumptions: (1) we study deep networks in the Neural Tangent Kernel (NTK) regime, (2) we treat binary classification task, (3) we decompose the interaction between two features. In Section 4, we demonstrate our results hold beyond these simplifying assumptions, for wide range of practical settings. All derivation details can be found in SM C.
3.1 Problem Setup and Gradient Starvation Definition
Let D = {X.y} denote a training set containing n datapoints with d dimensions, where, X [x..... .... x,] e "xd and their corresponding class labely E {-1,+11 Also let yXX) == f( (X):
Raxd _ represent the logits of an L-layer fully-connected neural network where each hidden
layer (0) (x) E Rd is defined as follows,
( (x.)) h00 = -
LE (0. L}.
in which (()) E d is weight matrix drawn from (O, I) and y is scaling factor to ensure that norm of each (111) preserved at initialization (See [28] for formal treatment). The function S(.) is also an element-wise non-linear activation function.
Let 0 = concat( 1/21 vec( E R be the concatenation of all vectorized weight matrices with m as the total number of parameters. In the NTK regime [52], in the limit of infinite width, the output of the neural network can be approximated as linear function of its parameters governed by the neural tangent random feature (NTRF) matrix [23],
σy (X
4 Xx, 0)
E IR
(2)
30
In the wide-width regime, the NTRF changes very little during training [62], and the output of the neural network can be approximated by first order Taylor expansion around the initialization parameters 0g. Setting & = & (X, 0o) and then, without loss of generality, centering parameters and the output coordinates to their value at the initialization (0o and ýo), we get
y (X, 0) Pp0.
Dominant directions in the feature space as well as the parameter space are given by principal components of the NTRF matrix o. which are the same those of the NTK Gram matrix [106].
We therefore introduce the following definition.
Definition (Features and Responses). Consider the singular value decomposition (SVD) of the matrix Y USVT, where Y diag (y). The jth feature is given by (VT);. The strength of jth feature is represented by s, (S) Also, (U),, contains the weights of this feature in all examples. A neural network's 's response to feature is given by zj where,
z : --= SV
(4)
3

---

In Eq. 4, the response to feature j is the sum of the responses to every example in (Y multiplied by the weight of the feature in that example (U"). For example, all elements of (U),, are positive. indicates a perfect correlation between this feature and class labels. We are now equipped to formally define GS.
Definition 2 (Gradient Starvation). Recall the the model prescribed by Eq. 3. Let =] denote the model's response to feature j at training optimum 0+ Feature starves the gradient for feature if dzž/ [s()) < 0.
This definition of GS implies that an increase in the strength of feature i has a detrimental effect on the learning of feature j. We now derive conditions for which learning dynamics of system suffer from GS.
3.2 Training Dynamics
We consider the widely used ridge-regularized cross-entropy loss function,
L (0) = .log [1 + exp (- YY) + 2 101
(5)
where 1 is a vector of size with all its elements equal to 1. This vector form simply represents a summation over all the elements of the vector multiplied to. X [0. oo) denotes the weight decay coefficient.
Direct minimization of this loss function using the gradient descent obeys coupled dynamics and is difficult to treat directly [26]. To overcome this problem, we call on a variational approach that leverages the Legendre transformation of the loss function. This allows tractable dynamics that can directly incorporate rates of learning in different feature directions. Following [50], we note the following inequality,
log + (- 2 H(a)
a Yý,
(6)
where H(x) = -a log a + (1 ax) log (1 -x] is Shannon's binary entropy function, a E (0, 9 is variational parameter defined for each training example, and @ denotes the element-wise vector product. Crucially, the equality holds when the maximum of r.h.s. w.r.t a is achieved
ac a" TYyyr, r which leads to the following optimization problem,
min L (0) = min max 0
1 (1-1(a)) aaý + =11
(7)
0 &x
where the order of min and max can be swapped (see Lemma 3 of [50]). Since the neural network's output is approximated by a linear function of 6, the minimization can be performed analytically with an critical value 0 T = Y Po. given by a weighted sum of the training examples. This results in the following maximization problem on the dual variable, i.c., ming L (0) is equivalent to,
min L (0) = max 0
1 1-H(a)
1 H(a) - 21
a "
(8)
at
By applying continuous- -immo gradient ascent on this optimization problem, we derive an autonomous differential equation for the evolution of a. which can be written in terms of features (see Definition 1),
à = (- log a log (1 a)
1 AUUSUU
where 77 is the learning rate (see SM C.1 for more details). For this dynamical system, we see that the logarithm term acts as barriers that keep a; E (0, 1). The other term depends the matrix US" U", which is positive definite, and thus pushes the system towards the origin and therefore drives learning.
When X & 2 where k is an index over the singular values, the linear term dominates Eq. 9, and the fixed point is drawn closer towards the origin. Approximating dynamics with first order Taylor expansion around the origin of the second term in Eq. 9, we get
1 Lau (S + AI) +I) U²).
④ lolo log
(10)
with stability given by the following theorem with proof in SM C.
Training optimum refers to the solution VgC(O) 0.

---

Theorem 1. Any fixed points of the system in Eq. 10 is attractive in the domain a; E (0, 1).
At the fixed point a", corresponding to the optimum of Eq. 8, the feature response of the neural network is given by,
z' S
(11)
See App. A for further discussions on the distinction between "feature space" and "parameter space". Below, we study how the strength of one feature could impact the response of the network to another feature which leads to GS.
3.3 Gradient Starvation Regime
In general, we do not expect to find an analytical solution for the dynamics of the coupled non-linear dynamical system of Eq. 10. However, there are at least two cases where a decoupled form for the dynamics allows to find an exact solution. We first introduce these cases and then study their perturbation to outline general lessons.
If the matrix of singular values S² is proportional the identity: This is the case where all the features have the same strength s². The fixed points are then given by,
a AW(A 10-
+
2j
32 +
Wij:
(12)
where W is the Lambert W function.
the matrix is permutation matrix: This the case in which each feature is associated with single example only. The fixed points are then given by, W 1)
af XW(A A 1))
(13)
+
To study minimal case of starvation, consider variation of case 2 with the following assumption which implies that each feature not associated with single example anymore.
Lemma 1. Assume U is a perturbed identity matrix (a special case of permutation matrix) in which the off-diagonal elements are proportional to a small parameter ò > 0. Then, the fixed point of the dynamical system in Eq. 10 can be approximated by.
a = log (ao)) [A diag (a) "
(14)
where A = U + AI and on is the fixed point of the uncoupled system with 5 0.
For sake of ease of derivations, we consider the two dimensional case where,
U= ( 82
-8
(15)
V1-82
which is equivalent to a U matrix with two blocks of features with no intra-block coupling and amount of inter-block coupling.
Theorem (Gradient Starvation Regime). Consider a neural network in the linear regime, trained under cross-entropy loss for a binary classification task. With definition 1, assuming coupling between features and 2 as in Eq. 15 and st > si, we have,
dæ
< 0,
(16)
ds}
which implies GS.
While Thm. 2 outlines conditions for GS in two dimensional feature space, we note that the same rationale naturally extends to higher dimensions, where GS defined pairwise over feature directions. For classification task, Thm. 2 indicates that gradient starvation occurs when the data admits different feature strengths, and coupled learning dynamics. GS is thus naturally expected with cross-entropy loss. Its detrimental effects however (as outlined in Sect. 2) arise in settings with large discrepancies between feature strengths, along with network connectivity that couples these features' directions. This phenomenon readily extends to multi-class settings, and we validate this case with experiments in Sect. 4. Next, we introduce a simple regularizer that encourages feature decoupling, thus mitigating GS by insulating strong features from weaker ones.
5

---

3.4 Spectral Decoupling
By tracing back the equations of the previous section, one may realize that the term S in Eq. not diagonal in the general case, and consequently introduces coupling between a, 's and hence, between the features 23 ss We would like to discourage solutions that couple features this way. To that end, introduce simple regularizer: Spectral Decoupling (SD). SD replaces the general I.2 weight decay term in Eq. with an 1.2 penalty exclusively on the network's s logits, yielding
L (0) = 1 log + exp (- Yy) 11
(17)
Repeating the same analysis steps taken above, but with SD instead of general 1.2 penalty, the critical value for 0" becomes 0 &aY 201/--22 -y- This new expression for 0* results in the following modification of Eq. 9,
à = log
aUSSSs U log a
x2)
(18)
where as earlier, log and division are taken element-wise on the coordinates of a.
Note that in contrast to Eq. the matrix multiplication involving U and S in Eq. 18 cancels out, leaving a; independent of other ajziss ss We point out this is true for any initial coupling, without simplifying assumptions. Thus, simple penalty output weights promotes decoupled dynamics across the dual parameter a; which track learning dynamics of feature responses (see Eq. 7). Together with Thm. 2, Eq. 18 suggests SD should mitigate GS and promote balanced learning dynamics across features. We now verify this numerical experiments. For further intuition, we provide simple experiment, summarized in Fig. 5, where directly visualizes the primal the dual dynamics as well as the effect of the proposed spectral decoupling method.
4 Experiments
The experiments presented here designed to outline the presence of GS and its consequences, as well as the efficacy of our proposed regularization method to alleviate them. Consequently, we highlight that achieving state-of-the-art results is not the objective. For more details including the scheme for syper-parameter tuning, see App. B.
4.1 Two-Moon classification and the margin
Recall the simple 2-D classification task between red and blue data points in Fig. 1. Fig. I (c) demonstrates the learned decision boundary when SD is used. SD leads to learning a curved decision boundary with larger margin in the input space. See App. B for additional details and experiments.
4.2 CIFAR classification and adversarial robustness
To study the classification margin in deeper networks, we conduct classification experiment on CIFAR-10, CIFAR- 100, and CIFAR-2 (cats dogs CIFAR 10) [57] using a convolutional network with ReLU non-linearity. Unlike linear models, the margin to non-lineat decision boundary cannot be computed analytically. Therefore, following the approach [72], we use "the norm of input- disturbance required to cross the decision boundary" proxy for the margin. The disturbance on the input computed by projected gradient descent (PGD) [84], a well-known adversarial attack.
Miclossfieol (ieation peroutage
Dataset Method
Train
Test IID
Test OOD+
80
Cifar-2 SD
100.0s
95.2% 25.2%scoo 2011
42.3%
w/ 1Diamen100:00 95.38 was
69.7%
PGD attack
99.9 99.9% nmm 92.8% t.15
30.18
Cifar-10 w/ SDuiman 9% oni 92,9ts 16
15
Vanilla BCE
20
wio SD
99.7% 7% oun 69.2
14.3% 2o
Cifar-100 wi SD no 99.79 70.5%
19
0
0.2
0.4
0.6
0.8
+ Accuracy (#t sid) nans.
( norm of the disturbance applied
Table I: Table compares adversarial robustness of ERM (vanilla cross-entropy) vs SD with CNN trained CIFAR-2, and 100 (setup of [72].. SD consistently achieves better OOD performance.
Figure 2: The plot shows the cumulative distribution function (CDF) of the margin for the CIFAR-2 binary classification. appears to improve the margin con- siderably.

---

Table I includes the results for IID (original test set) and OOD (perturbed test set by <PGD = 0.05). Fig. 2 shows the percentange of missclassification as the norm of disturbance is increased for the Cifar-2 dataset. This plot can be interpreted as the cumulative distribution function (CDF) of the margin and hence a lower curve reads as a more robust network with a larger margin. This experiment suggests that when trained with vanilla cross-entropy, even slight disturbances in the input deteriorates the network's classification accuracy. That is while spectral decoupling (SD) improves the margin considerably. Importantly, this improvement in robustness does not seem to compromise the noise-free test performance. It should also be highlighted that SD does not explicitly aim at maximizing the margin and the observed improvement is in fact a by-product of decoupled learning of latent features. See Section 5 for a discussion on why cross- cross-entropy results in poor margin while being considered a max- mar-mn classifier in the literature [94].
4.3 Colored MNIST with color bias
We conduct experiments on the Colored MNIST Dataset, proposed in [9]. The task is to predict binary labels y = -1 for digits 0 to 4 and y +1 for digits 5 to 9. A color channel (red, green) is artificially added to each example to deliberately impose a spurious correlation between the color and the label. The task has three environments:
Training env. 1: Color is correlated with the labels with 0.9 probability.
Training env. 2: Color is correlated with the labels with 0.8 probability.
Testing env.: Color is correlated with the labels with 0.1 probability (0.9 reversely correlated).
Because of the opposite correlation between the color and the label in the test set, only learning to classify based on color would be disastrous at testing. For this reason, Empirical Risk Minimization (ERM) performs very poorly on the test set (23.7 % accuracy) shown in Tab. 2.
Method
Train Accuracy
Test Accuracy
ERM (Vanilla Cross- Enttssyytttoy)
91.1 % (1:04)
23.7 % (10.8)
REx [59]
71.5 % (+1.0)
68.7 % (+0.9)
IRM [9]
70.5 % (+E0.)
67.1% (#T.4)
SD (this work)
70.00 % (+009)
68.4 % (81.%)
Oracle (grayscale images)
73.5 % (+-02)
73.0 % (+0.4)
Random Guess
50 %
50 %
Table 2:
Test accuracy on test examples of the Colored MNIST after training for IK epochs. The standard deviation over 10 runs is reported in parenthesis. ERM stands for the empirical risk minimization. Oracle is an ERM trained on grayscale images. Note that due to 25 % label noise, hypothetical optimum achieves 75 % accuracy (the upper bound).
Invariant Risk Minimization (IRM) [9] on the other hand, performs well on the test set with (67.1 % accuracy). However, IRM requires access multiple (two in this case) separate training environments with varying amount of spurious correlations. IRM uses the variance between environments as a signal for learning to be "invariant" to spurious correlations. Risk Extrapolation (REx) [59] is related training method that encourages learning invariant representations. Similar to IRM, it requires access to multiple training environments in order quantify the concept of "iivariance".
SD achieves an accuracy of 68.4 %. Its performance is remarkable because unlike IRM and REx, SD does not require access to multiple environments and yet performs well when trained on single environment (in this case the aggregation of both of the training environments).
A natural question that arises is "How does SD learn to ignore the color feature without having access to multiple environments?" The short answer that it does not! In fact, we argue that SD learns the color feature but it also learns other predictive features, i.e., the digit shape features. At test time, the predictions resulting from the shape features prevail over the color feature. To validate this hypothesis, we study a trained model with each of these methods (ERM, IRM, SD) on four variants of the test environment: 1) grayscale-digits: No color channel is provided and the network should rely shape features only. 2) colored-digit Both color and digit are provided however the color is negatively correlated (opposite of the training set) with the label. 3) grayscale- blank: All images are grayscale and blank and hence do not provide any information. 4) colored- blank: Digit features are removed and only the color feature is kept, also with reverse label compared to training. Fig. 3 summarizes the results. For more discussions see SM B.
As a final remark, we should highlight that, by design, this task assumes access to the test environment for hyperparameter tuning for all the reported methods. This is not a valid assumption in general, and
7

---

ERM
SD
IRM
graysooll
colored
jariescaal
colored
colored
100
digt
dggit
digi
90
graysaale
oslened
3
6
3
Honk
49.9
blank
15
Mook
400
5
grayicale
colored
gropscale
colored
grayscale
colored
Lce
blnn
0.6
0.54
0.44
0.66
0.63
0.60
0.59
Adoaurt
-04
Noak 0.65
0.00
blonk
0.66
0.41
Sbak 0.65
0.64
0.0 contideet
Figure 3: Diagram comparing ERM, SD, and IRM on four different test environments on which we evaluate a pre-trained model. Top and bottom rows show the accuracy and the entropy (inverse of confidence), respectively. Analysis: Compare three values of 9.4 % 9.4 % and 49.6 % : Both ERM
have learned the color feature but since it is inversely correlated with the label, when only the color feature is provided, as expected hoth FRM and SD performs poorly. Now compare 0.00 and 0.41 Although both ERM and SD have learned the color feature, ERM is much more confident on predictions (zero entropy). digit features are provided along with the color feature (colored-digt) eavironment), ERM still performs poorly (23.9 % but SD) achieves significantly better results (6..2 % IRM ignores the color feature altogether but requires access to multiple training environments.
consequence, when
CelebA ERM
CelebA SD
100
Method
Average Acc.
Worst Group Acc.
Accucayy
60
60
ERM
94.61 (+0.67)
40.35 (+ (1:6))
SD (this work)
91.64 % (20.61)
83.24 % (#2.01)
LIF
N/A
81.24 % (#138)
Group DRO
91.76 % (+0.28)
87.78 (+0.96)
vald, wam
9
iterafions
iterations
Figure 4: CelebA: blond vs dark hair classification. The HairColor and the Gender are spuriously cor- related which leads poor OOD performance with ERM, however SD significantly improves performance. ERM's worst group accuracy significantly lower than SD.
Table 3: CelebA: blond vs dark hair classification with spurious correlation. We report test performance runs. SD significantly improves upon ERM. "Group DRO [89] requires explicit information about the spurious correlation. LfF [71] requires simultane- ous training of two networks.
hence the results should be only interpreted as a probe that shows that SD could provide an important level of control over what features are learned.
4.4 CelebA with gender bias
The CelebA dataset [65] contains 162k celebrity faces with binary attributes associated with each image. Following the setup [89], the task is to classify images with respect to their hair color into two classes of blond or dark hair. However, the Gender E [Male, Female] is spuriously correlated with the HairColor E (Blond, Dark) } in the training data. The rarest group which blond males represents only 0.85 % of the training data (1387 out of 162k examples). We train ResNet-50 model [38] on this task. Tab. 3 summarizes the results and compares the performance of several methods. A model with vanilla cross-entropy (ERM) appears generalize well average but fails generalize to the rarest group (blond males) which can be considered as "weakly" diito---dittibutionn (OOD). Our proposed SD improves the performance more than twofold. should be highlighted that for this task, we The hyper- parameeessssrs and y are tuned separately for each class (a total of four hyper parametrttttssss This variant of SD does provably decouple the dynamics too but appears perform better than the original SD in Eq. 17 in this task.
variant of SD in which, A ll 7112 is added to the original cross- cross-entropy loss.
Other proposed methods presented in Tab. 3 also show significant improvements the performance

---

of the worst group accuracy. The recently proposed "Learning from failure" (LfF) [71] achieves comparable results to SD. but it requires simultaneous training of two networks. Group DRO [89] is another successful method for this task. However, unlike SD, Group DRO requires explicit information about the spuriously correlated attributes. In most practical tasks, information about the spurious correlations is not provided and, dependence on the spurious correlation goes unrecognized.
5 Related Work and Discussion
Here, we discuss the related work. Due to space constraints, further discussions are in App. A.
On learning dynamics and Loss Choice. Several works including [90, 91, 1, 60] investigate the dynamics of deep linear networks trained with squared-error loss. Different decompositions of the learning process for neural networks have been used: [83, 104, 87, 105] study the learning in the Fourier domain and show that low- ow-frequency functions are learned earlier than high-frequency ones. [90, 2, 32] provide closed-form equations for the dynamics of linear networks terms of the principal components of the input covariance matrix. More recently, with the introduction of neural tangent kernel (NTK) [52, 62], a new line of research is to study the convergence properties of gradient descent [e.g. 8, 69, 25, 29, 7, 44, 33, 110, 11, 99]. Among them, [12, 106, 18, 22] decompose the learning process along the principal components of the NTK. The message in these works is that the training process can be decomposed into independent learning dynamics along the orthogonal directions.
Most of the studies mentioned above focus on the particular squared-error loss. For a linearized network, the squared-error loss results in linear learning dynamies, which often admit an analytical solution. However, the de-facto loss function for many of the practical applications of neural networks is the cross-entropy. Using the cross-entropy as the loss function leads to significantly more complicated and non-linear dynamics, even for linear neural network. In this work, our focus was the cross-entropy loss.
On reliance upon spurious correlations and robustness. In the context of robustness in neural networks, state-of-the-art neural networks appear to naturally focus on low-level superficial correla- tions rather than more abstract and robustly informative features of interest (e.g. [30]). As we argue in this work, Gradient Starvation is likely an important factor contributing to this phenomenon and can result in adversarial vulnerability. There is rich research literature on adversarial attacks and neural networks' vulnerability [96, 34, 48, 67, 5, 47]. Interestingly, [73], [72] and [51] draw a similar conclusion and argue that "an insufficiency of the cross-entropy loss" causes excessive invariances to predictive features. Perhaps [92] is the closest to our work in which authors study the simplicity bias (SB) in stochastic gradient descent. They demonstrate that neural networks exhibit extreme bias that could lead to adversarial vulnerability.
On implicit bias. Despite being oerlirreeeerrrraeeeriieddd modern neural networks seem to gener- alize very well [108]. Modern neural networks generalize surprisingly well in numerous machine tasks. This is despite the fact that neural networks typically contain orders of magnitude more parameters than the number of examples in a training set and have sufficient capacity to fit totally randomized dataset perfectly [108]. The widespread explanation is that the gradient descent has form of implicit bias towards learning simpler functions that generalize better according to Occam's 's razor. Our exposition of GS reinforces this explanation. In essence, when training and test data points are drawn from the same distribution, the top salient features are predictive in both sets. We conjecture that in such a scenario, by not learning the less salient features, GS naturally protects the network from overfitting.
The same phenomenon is referred to implicit bias, implicit regularization, simplicity bias and spectral bias in several works [83, 75, 36, 74, 70, 53, 94, 10, 13, 35, 82, 66).
As active line of research, numerous studies have provided different explanations for this phe- nomenon. For example, [70] justifies the implicit bias of neural networks by showing that stochastic gradient descent learns simpler functions first. [15, 78] suggests that form of implicit regularization is induced by an alignment between NTK's principal components and only a few relkvatl
Recall that it took 3 years for the psychologist, Oskar Pfungst, to realize that Clever Hans was not capable of doing any arithmetic.

---

directions. Several other works such as [20, 35, 94, 25] recognize the convergence of gradient descent to maximum-maree solution as the essential factor for the generalizability of neural networks. should be stressed that these work refer to the margin in the hidden space and not in the input space classifier the
as pointed out in [55]. Indeed, as observed in our experiments, the hidden space can be achieved at the expense of small margin in the input space.
On Gradient Starvation and no free lunch theorem. The no free lunch theorem [93, 102] states that "learning is impossible without making assumptions about training and test distributions". Perhaps, the most commonly used assumption of machine learning is the i.i.d. assumption [98], which assumes that training and test data are identically distributed. However, in general, this assumption might not hold, and in many practical applications, there predictive features in the training set that do not generalize to the test set. A natural question that arises is how to favor generalizable features over spurious features? ? The most common approaches include data augmentation, controlling the inductive biases, using regularizations, and more recently training using multiple environments.
Here, would like to elaborate network is provided with a chess book containing examples of chess games with the best movements indicated by red arrow. The network can take two approaches: 1) learn how to play chess, or 2) learn just the red arrows. Either of these solutions results in zero training loss on the games in the book while only the former is generalizable to new games. With external knowledge, the network typically learns the simpler solution.
interesting thought experiment of [79]: Suppose neural
Recent work aims to leverage the invariance principle across several environments to improve robust learning. This is akin to present several chess books to a network, each with markings indicating the best moves for different sets of games. In several studies [9, 59, 79, 3], methods are developed to aggregate information from multiple training environments in a way that favors the generalizable / Noomain-agnosia / invariant solution. We argue that even with having access to only one training environment, there is useful information in the training set that fails to be discovered due to Gradient Starvation. The information books. Still, as soon as the network learns the red arrows, the network has no incentive for further learning. Therefore, learning the red arrows is not an issue per se, but not learning play chess is.
how to actually play chess is already available in any of the chess
Gradient Starvation: friend or foe? Here, we would like remind the reader that GS can have both adverse and beneficial consequences. If the learned features are sufficient to generalize to the test data, gradient starvation can be viewed as implicit regularizer. Otherwise, Gradient Starvation could have an unfavorable effect, which we observe empirically when some predictive features fail to be learned. A better understanding and control of Gradient Starvation and its impact on generalization offers promising avenues to address this issue with minimal assumptions. Indeed, our Spectral Decoupling method requires an assumption about feature imbalance but not pinpoint them exactly, relying on modulated learning dynamics to achieve balance.
GS social impact Modern neural networks are being deployed extensively in numerous machine learning tasks. Our models are used in critical applications such autonomous driving, medical prediction, and even justice system where human lives are at stake. However, neural networks appear base their predictions on superficial biases in the dataset. Unfortunately, biases in datasets could be neglected and pose negative impacts on our society. In fact, our Celeb-A experiment is an example of the existence of such bias in the data. As shown in the paper, the gender specific bias could lead to a superficial high performance and is indeed very hard to detect. Our analysis, although mostly on the theory side, could pave the path for researchers to build machine learning systems that are robust to biases and helps towards fairness in our predictions.
Conclusion
In this paper, we formalized Gradient Starvation (GS) as phenomenon that emerges when training with cross-entropy loss in neural networks. By analyzing the dynamical system corresponding to the learning process in a dual space, we showed that GS could slow down the learning of certain features, even if they are present in the training set. We derived spectral decoupling (SD) regularization as possible remedy to GS.
10

---

Acknowledgments and Disclosure of Funding
The authors are grateful to Samsung Electronics Co., Ldt., CIFAR, and IVADO for their funding and Calcul Québec and Compute Canada for providing us with the computing resources. We would further like to acknowledge the significance of discussions and supports from Reyhane Askari Hemmat and Faruk Ahmed. MP would like to thank Aristide Baratin, Kostiantyn Lapchevskyi, Seyed Mohammad Mehdi Ahmadpanah, Milad Aghajohari, Kartik Ahuja, Shagun Sodhani, and Emmanuel Bengio for their invaluable help.
References
[1] Madhu S Advani and Andrew M Saxe. High-dimensional dynamics of generalization error in neural networks. arXiv preprint arXiv:1710.03667, 2017.
[2] Madhu S Advani, Andrew Saxe, and Haim Sompolinsky. High-dimensional dynamics of generalization error in neural networks. Neural Networks, 2020.
[3] Kartik Ahuja, Karthikeyan Shanmugam, and Amit Dhurandhar. Linear regression games: Convergence guarantees to approximate out-of-distribution solutions. arXiv preprint arXiv:2010.15234 2020.
[4] Kartik Ahuja, Karthikeyan Shanmugam, Kush Varshney, and Amit Dhurandhar. Invariant risk minimization games. arXiv preprint arXiv:2002.04669 2020.
[5] Naveed Akhtar and Ajmal Mian. Threat of adversarial attacks on deep learning in computer vision: A survey. IEEE Access, 6:14410-14430, 2018.
[6] Adelin Albert and John A Anderson. On the existence of maximum likelihood estimates in logistic regression models. Biometrika, 71 71(1):1-10, 1984.
Zeyuan Allen-Zhu, Yuanzhi Li, and Yingyu Liang. Learning and generalization in overpa- rameterized neural networks, going beyond two layers. In Advances in neural information processing systems, pages 6158-6169, 2019.
[8] Zeyuan Allen-Zhu, Yuanzhi Li, and Zhao Song. A convergence theory for deep learning via over-parameterinn wverppramettrizzatonn In International Conference on Machine Learning, pages 242-252. PMLR, 2019.
[9] Martin Arjovsky, Léon Bottou, Ishaan Guirajani, and David Lopez-Paz. Invariant risk mini- mization. arXiv preprint 190.00997..02893 2019.
[10] Sanjeev Arora, Nadav Cohen, Wei Hu, and Yuping Luo. Implicit regularization in deep matrix factorization. In Advances Neural Information Processing Systems, pages 7413-7424, 2019.
[11] Sanjeev Arora, Simon S Du, Wei Hu, Zhiyuan Li, Russ Salakhutdinov, and Ruosong Wang. On exact computation with infinitely wide neural net. In Advances in Neural Information Processing Systems, pages 8141-8150, 2019.
[12] Sanjeev Arora, Simon S Du, Wei Hu, Zhiyuan Li, and Ruosong Wang. Fine-grained analysis of optimization and generalization for overparameterized two-layer neural networks. arXiv preprint arXiv: 2019.
[13] Devansh Arpit, Stanistaw Jastrzebski, Nicolas Ballas, David Krueger, Emmanuel Bengio, Maxinder S Kanwal, Tegan Maharaj. Asja Fischer, Aaron Courville, Yoshua Bengio, et al. A closer look at memorization in deep networks. arXiv preprint
2017.
[14] Nicholas Baker, Hongjing Lu, Gennady Erlikhman, and Philip J Kellman. Deep convolu- tional networks do classify based on global object shape. PLoS computational biology. 14 14(12)xe1006613 2018.
[15] Aristide Baratin, Thomas George, César Laurent, R Devon Hjelm, Guillaume Lajoie, Pascal Vincent, and Simon Lacoste-Julien. Implicit regularization in deep learning: view from function space. arXiv preprint arXiv:2008.0099 00938, 2020.

---

[16] Sara Beery, Grant Van Horn, and Pietro Perona. Recognition in terra incognita. In Proceedings of the European Conference on Computer Vision (ECCV), pages 456 2018.
[17] Yonatan Belinkov and Yonatan Bisk. Synthetic and natural noise both break neural machine translation. arXiv preprint arXiv:1711.02173, 2017.
[18] Alberto Bietti and Julien Mairal. On the inductive bias of neural tangent kernels. In Advances in Neural Information Processing Systems, pages 12893-12904, 2019.
[19] Wieland Brendel and Matthias Bethge. Approximating cnns with low-ooo-ooo---atuuue models works surprisingly well on imagenet. arXiv preprint urXiv:/904.007000 2019.
[20] Alon Brutzkus, Amir Globerson, Eran Malach, and Shai Shalev-Shwartz. Sgd learns over- parameterized networks that provably generalize linearly separable data. arXiv preprint a7Xiv: 10174 2017.
[21] Christopher JC Burges and David J Crisp. Uniqueness of the svm solution. Advances in neural information processing systems, 12:223-229, 2000.
[22] Yuan Cao, Zhiying Fang. Yue Wu, Ding- Xuan Zhou, and Quanquan Gu. Towards understand- ing the spectral bias of deep learning. arXiv preprint arXiv:1912.01198, 2019.
[23] Yuan Cao and Quanquan Gu. Generalization bounds of stochastic gradient descent for wide and deep neural networks. In Advances in Neural Information Processing Systems, pages 10836-10846, 2019.
[24] Zixiang Chen, Yuan Cao, Quanquan Gu, and Tong Zhang. generalized neural tangent kernel analysis for two-layer neural networks. arXiv preprint arXiv:2002.044020 2020.
[25] Lenaic Chizat and Francis Bach. A note lazy training in supervised differentiable program- ming. arXiv preprint arXiv: 1, 2018.
[26] Remi Tachet des Combes, Mohammad Pezeshki, Samira Shabanian, Aaron Courville, and Yoshua Bengio. On the learning dynamics of deep neural networks. arXiv preprint arXiv:1809 (00.::88,,,0888,, 2018.
[27] Thomas M Cover.
Geometrical and statistical properties of systems of linear inequalities with applications in pattern recognition. IEEE transactions on electronic computers, 023):366-334, 1965.
[28) Simon S Du, Wei Hu, and Jason D Lee. Algorithmic regularization in learning deep homo- geneous models: Layers are automatically balanced. In Advances in Neural Information Processing Systems, pages 384-395, 2018
[29] Simon S Du, Xiyu Zhai, Barnabas Poczos, and Aarti Singh. Gradient descent provably optimizes over-parameterized neural networks. arXiv preprini arXiv: arXiv: 1810.02054 2018.
[30] Robert Geirhos, Hoon-Kerrik Jacobsen, Claudio Michaelis, Richard Zemel, Wieland Brendel, Matthias Bethge, and Felix A Wichmann. Shortcut learning in deep neural networks. arXiv preprint arXiv:2004.0778 2020.
[31] Thomas George. Nngeometry: Easy and fast fisher information matrices and neural tangent kernels in pytorch. 0, 2020.
[32] Gauthier Gidel, Francis Bach, and Simon Lacoste-Julien. Implicit regularization of discrete gradient dynamics in linear neural networks. In Advances in Neural Information Processing Systems, pages 3202-3211, 2019.
[33] Sebastian Goldt, Madhu Advani, Andrew M Saxe, Florent Krzakala, and Lenka Zdeborová. Dynamics of stochastic gradient descent for two-layer neural networks in the teacher-student setup. In Advances in Neural Information Processing Systems, pages 6981-6991, 2019.
[34] Ian Goodfellow, Jonathon Shlens, and Christian Szegedy. Explaining and harnessing adver- sarial examples. arXiv preprint arXiv:1412.65777 2014. 2014.
12

---

[35] Suriya Gunasekar, Jason D Lee, Daniel Soudry, and Nati Srebro. Implicit bias of gradient descent on linear convolutional networks. In Advances in Neural Information Processing Systems, pages 9461-9471, 2018.
[36] Suriya Gunasekar, Blake E Woodworth, Srinadh Bhojanapalli, Behnam Neyshabur, and Nati Srebro. Implicit regularization in matrix factorization. In Advances in Neural Information Processing Systems, pages 6151-6159, 2017.
[37] Suchin Gururangan, Swabha Swayamdipta, Omer Levy, Roy Schwartz, Samuel Bowman, and Noah A Smith. Annotation artifacts in natural language inference data. arXiv preprint XxXi::8033023324 2018.
[38] Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Proceedings of the IEEE conference computer vision and pattern recognition, pages 770-778, 2016.
[39] Christina Heinze-Deml and Nicolai Meinshausen. Conditional variance penalties and domain shift robustness. arXiv preprint arXiv:1710.1 1469, 2017.
[40] Dan Hendrycks and Thomas Dietterich. Benchmarking neural network robustness to common corruptions and perturbations. arXiv preprint arXiv: /903.9926,, 2019.
[41] Dan Hendrycks and Kevin Gimpel. A baseline for detecting misclassified and out-of- distribution examples in neural networks. arXiv preprint 16.0002 136, 2016.
[42] Katherine L Hermann and Andrew K Lampinen. What shapes feature representations? exploring datasets, architectures, and training. arXiv preprint arXiv:2006. 12433, 2020.
[43] Cho-Jui Hsieh, Kai- Chang, Chih-Jen Lin, S Sathiya Keerthi, and Sellamanickam Sun- dararajan. A dual coordinate descent method for large-scale linear svm. In Proceedings of the 25th international conference on Muchine learning, pages 408-415, 2008
[44] Jiaoyang Huang and Horng-Tze Yau. Dynamics of deep neural networks and neural tangent hierarchy, arXiv preprint arXiv:1909) 1909.08156, 2019.
[45] Kaixuan Huang, Yuqing Wang, Molei Tao, and Tuo Zhao. Why do deep residual networks generalize better than deep feedforward networks? neural tangent kernel perspective. Advances in Neural Information Processing Systems, 33, 2020.
[46] Like Hui and Mikhail Belkin. Evaluation of neural architectures trained with square loss vs cross-entropy in classification tasks. arXiv preprint arXiv:2006.07322, 2020.
[47] Andrew Ilyas, Logan Engstrom, Anish Athalye, and Jessy Lin. Black-box adversarial attacks with limited queries and information. arXiv preprint arXiv:1804.08593 2018.
[48] Andrew Ilyas. Shibani Santurkar, Dimitris Tsipras, Logan Engstrom, Brandon Tran, and Aleksander Madry. Adversarial examples are not bugs, they are features. In Advances in Neural Information Processing Systems, pages 125-136. 2019.
[49] Sergey Ioffe and Christian Szegedy. Batch normalization: Accelerating deep network training by reducing internal covariate shift. arXiv preprini arXiv:1502.03366 2015.
[50] Tommi S Jaakkola and David Haussler. Probabilistic kernel regression models. In AISTATS, 1999.
[51] Jörn-Henrik Jacobsen, Jens Behrmann, Richard Zemel, and Matthias Bethge. Excessive invariance causes adversarial vulnerability. arXiv preprint arXiv: 2018.
[52] Arthur Jacot, Franck Gabriel, and Clément Hongler. Neural tangent kernel: Convergence and generalization in neural networks. In Advances in neural information processing systems, pages 857 1-8580, 2018
[53] Ziwei and Matus Telgarsky. The implicit bias of gradient descent on nonseparable data, In Conference on Learning Theory, pages 1772-1798, 2019.
13

---

[54] Jason Jo and Yoshua Bengio. Measuring the tendency of cnns to learn surface statistical regularities. arXiv preprint arXiv:1711.115611 2017.
[55] Alexia machines, wasserstein distance and
and Ioannis Mitliagkas. Connections between support vector
gans. arXiv preprint arXiv: 101.092910.06922,
2019.
[56] Diederik Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint 1xXiv:4412.6980 2014.
[57] Alex Krizhevsky, Geoffrey Hinton, et al. Learning multiple layers of features from tiny images.
0, 2009.
[58] Anders Krogh and John A Hertz. A simple weight decay can improve generalization. In Advances in neural information processing systems, pages 950-957, 1992.
[59] David Krueger, Ethan Caballero, Beern-Herrik Jacobsen, Amy Zhang. Jonathan Binas, Remi Le Priol, and Aaron Courville. Dut-of-distributin generalization via risk extrapolation (rex). arXiv preprint arXiv:2003.00663 2020.
[60] Andrew K Lampinen and Surya Ganguli. An analytic theory generalization dynamies and transfer learning in deep linear networks. arXiv preprint arXiv:1809.10374 2018.
[61] Sebastian Lapuschkin, Stephan Wäldchen, Alexander Binder, Grégoire Montavon, Wojciech Samek, and Klaus-Robert Müller. Unmasking clever hans predictors and assessing what machines really learn. Nature communications, 10(1):1-8, 2019.
[62] Jaehoon Lee, Lechao Xiao, Samuel Schoenholz, Yasaman Bahri, Roman Novak, Jascha Sohl-Diekstein, and Jeffrey Pennington. Wide neural networks of any depth evolve as linear models under gradient descent. In Advances in newral information processing systems, pages 8570-8581, 2019.
[63] Kimin Lee, Kibok Lee, Honglak Lee, and Jinwoo Shin. A simple unified framework for de- tecting out-of-distribution samples and adversarial attacks. In Advances in Neural Information Processing Systems, pages 7167-7177, 2018
[64] Shiyu Liang, Yixuan Li, and Rayadurgam Srikant. Enhancing the reliability of out-of- distribution image detection in neural networks. arXiv preprint (rXXiv:706.006690 2017.
[65] Ziwei Liu, Ping Luo, Xiaogang Wang, and Xiaoou Tang. Deep learning face attributes in the wild. In Proceedings of International Conference on Computer Vision (ICCV), December 2015.
[66] Cong Ma, Kaizheng Wang, Yuejie Chi, and Yuxin Chen. Implicit regularization in nonconvex statistical estimation: Gradient descent converges linearly for phase retrieval and matrix completion. In International Conference on Machine Learning, pages 3345-3354. PMLR, 2018.
[67] Aleksander Madry, Aleksandar Makelov, Ludwig Schmidt, Dimitris Tsipras, and Adrian Vladu. Towards deep learning models resistant to adversarial attacks. arXiv preprint arviv: 700:::0066008333 2017.
[68] R Thomas McCoy, Ellie Pavlick, and Tal Linzen. Right for the wrong reasons: Diagnosing syntactic heuristics in natural language inference. arXiv preprint arXiv: 2019.
[69] Song Mei and Andrea Montanari. The generalization error of random features regression: Precise asymptotics and double descent curve. arXiv preprint arXiv:1908.055: 2019.
[70] Preetum Nakkiran, Gal Kaplun, Dimitris Kalimeris, Tristan Yang, Benjamin L Edelman, Fred Zhang, and Boaz Barak. Sgd on neural networks learns functions of increasing complexity. arXiv preprint 19Xiv:
2019.
[71] Junhyun Nam, Hyuntak Cha, Sungsoo Ahn, Jaeho Lee, and Jinwoo Shin. Learning from failure: Training debiased classifier from biased classifier. arXiv preprint vxXvv2007.00566/ 2020.
14

---

[72] Kamil Nar. Orhan Ocal, S Shankar Sastry, and Kannan Ramchandran. Cross-entropy loss and low-rank features have responsibility for adversarial examples. urXiv:/901. 900000000,,0...0.10000000000000, 2019.
arXiv preprint
[73] Kamil Nar and S Shankar Sastry. Persistency of excitation for robustness of neural networks arXiv preprint arXiv:1911.0104. 2019.
[74] Behnam Neyshabur, Ryota Tomioka, Ruslan Salakhutdinov, and Nathan Srebro. Geometry of optimization and implicit regularization in deep learning, arXiv preprint arXiv:1705.03071, 2017.
[75] Behnam Neyshabur, Ryota Tomioka, and Nathan Srebro. In search of the real inductive bias: On the role of implicit regularization in deep learning. arXiv preprint arXiv: 1412.6614, 2014.
[76] Timothy Niven and Hung-Yo Yu Kao. Probing neural network comprehension of natural language arguments. arXiv preprint arXivv:99..07355, 2019.
[77] Luke Oakden- Oakden-Rayner, Jared Dunnmon, Gustavo Carneiro, and Christopher Ré. Hidden stratification causes clinically meaningful failures in machine learning for medical imaging. In Proceedings of the ACM Conference Health, Inference, and Learning, pages 151-159,
2020.
[78] Samet Oymak, Zalan Fabian, Mingchen Li, and Mahdi Soltanolkotabi. Generalization guaran- tees for neural networks via harnessing the low-rank structure of the jacobian. arXiv preprint arXiv:1906.05320 2019.
[79] Giambattista Parascandolo, Alexander Neitz, Antonio Orvieto, Luigi Gresele, and Bernhard Schölkopf. Learning explanations that are hard to vary. arXiv preprint arXiv:2009.00322, 2020.
[80] Adam Paszke, Sam Gross, Soumith Chintala, Gregory Chanan, Edward Yang, Zachary DeVito, Zeming Lin, Alban Desmaison, Luca Antiga, and Adam Lerer. Automatic differentiation in pytorch. openreview id=BIJsrm/CZ, 2017.
[81] Oskar Pfungst. Clever Hans:(the horse of Mr. Von Osten.) a contribution to experimental animal and human psychology. Holt, Rinehart and Winston, 1911.
[82] Tomaso Poggio, Kenji Kawaguchi, Qianli Liao, Brando Miranda, Lorenzo Rosasco, Xavier Boix, Jack Hidary, and Hrushikesh Mhaskar. Theory of deep learning ii:: explaining the non-overfitting puzzle. arXiv preprint
2017.
[83] Nasim Rahaman, Aristide Baratin, Devansh Arpit, Felix Draxler, Min Lin, Fred Hamprecht, Yoshua Bengio, and Aaron Courville. On the spectral bias of neural networks. In International Conference on Machine Learning, pages 5301-5310. PMLR, 2019.
[84] Jonas Rauber, Wieland Brendel, and Matthias Bethge. Foolbox: A python toolbox to bench- mark the robustness of machine learning models. arXiv preprint 1oXiv:1707.044 04131, 2017.
[85] Marco Tulio Ribeiro, Sameer Singh, and Carlos Guestrin." " "why should trust you?" explaining the predictions of any classifier. In Proceedings of the 22nd ACM SIGKDD international conference on knowledge discovery and data mining, pages 1135-1144, 2016.
[86] Michael Roberts. Machine learning for covid-19 diagnosis: Promising, but still too flawed, Mar 2021.
[87] Basri Ronen, David Jacobs, Yoni Kasten, and Shira Kritchman. The convergence rate neural networks for learned functions of different frequencies. In Advances in Neural Information Processing Systems, pages 4761-4771 2019.
[88] Amir Rosenfeld, Richard Zemel, and John K Tsotsos. The elephant in the room. arXiv preprint arXiv:1808.088 00t0:/8000003005, 2018.
[89] 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 preprint 19I1v:1111.077//1 2019.
15

---

[90] Andrew M Saxe, James L McClelland, and Surya Ganguli. Exact solutions to the nonlinear dynamics of learning in deep linear neural networks. arXiv preprint arXiv:1312.6120, 2013.
[91] Andrew M Saxe, James L McClelland, and Surya Ganguli. A mathematical theory of semantic development in deep neural networks. Proceedings of the National Academy of Sciences, 1 166(23): 2019.
[92] Harshay Shah, Kaustav Tamuly, Aditi Raghunathan, Prateek Jain, and Praneeth Netrapalli. The pitfalls of simplicity bias in neural networks. arXiv preprint vrXiv:2006.07710, 2020.
[93] Shai Shalev-Shwartz and Shai Ben-Dar David. Understanding machine learning: From theory to algorithms. Cambridge university press, 2014.
[94] Daniel Soudry, Elad Hoffer, Mor Shpigel Nacson, Suriya Gunasekar, and Nathan Srebro. The implicit bias of gradient descent separable data. The Journal of Machine Learning Research, 19(1)28822-288 2018.
[95] Nitish Srivastava, Geoffrey Hinton, Alex Krizhevsky, Ilya Sutskever, and Ruslan Salakhutdinov. Dropout: simple way to prevent neural networks from overfitting. The journal of machine learning research, 15())1199--1988 2014.
[96] Christian Szegedy, Wojciech Zaremba, Ilya Sutskever, Joan Bruna, Dumitru Erhan, Ian Goodfellow, and Rob Fergus. Intriguing properties of neural networks. arXiv preprint arXiv:1312.6199, 2013.
[97] Guillermo Valle-Pérez, Chico Q Camargo, and Ard A Louis. Deep learning generalizes because the parameter-function map is biased towards simple functions. arXiv preprint XXXi::8805008222 2018.
[98] Vladimir Vapnik and Vlamimir Vapnik. Statistical learning theory wiley, New York, 1:624, 1998.
[99] Santosh Vempala and John Wilmes. Gradient descent for one-hidden-layer neural networks: Polynomial convergence and sq lower bounds. In Conference on Learning Theory, pages 3115-3117, 2019.
[100] Haohan Wang, Zexue He, Zachary Lipton, and Eric P Xing. Learning robust representations by projecting superficial statistics out. arXiv preprint (oXXiv:993300220 2019.
[101] Sifan Wang, Xinling Yu, and Paris Perdikaris. When and why pinns fail to train: A neural tangent kernel perspective. Journal of Computational Physies, page 110768, 2021.
[102] David H Wolpert. The lack of a priori distinctions between learning algorithms. Neural computation, 8(7):1341-139. 1996.
[103] Xin Xu and Eibe Frank. Logistic regression and boosting for labeled bags of instances. In Pacific-Asia conference on knowledge discovery and data mining, pages 272-281. Springer, 2004.
[104] Zhi-Qin John Xu, Yaoyu Zhang, Tao Luo, Yanyang Xiao, and Zheng Ma. Frequency principle: Fourier analysis sheds light on deep neural networks. arXiv preprint arXiv:1901.06:: 2019.
[105] Zhi- hi-- John Xu, Yaoyu Zhang, and Yanyang Xiao. Training behavior of deep neural network in frequency domain. In International Conference on Neural Information Processing, pages 264 124. Springer, 2019.
[106] Greg Yang and Hadi Salman. A fine lin-ggrined spectral perspective on neural networks. arXiv preprint arXiv: 1907. 10599, 2019.
[107] John R Zech, Marcus A Badgeley, Manway Liu, Anthony B Costa, Joseph J Titano, and Eric K Oermann. Confounding variables can degrade generalization performance of radiological deep learning models. arXiv preprint arXiv:1807.0044 2018.
16

---

[108] Chiyuan Zhang, Samy Bengio, Moritz Hardt, Benjamin Recht, and Oriol Vinyals. Under- standing deep learning requires rethinking generalization. arXiv preprint arXiv: 2016.
[109] Jieyu Zhao, Tianlu Wang, Mark Yatskar, Vicente Ordonez, and Kai-Wei Chang. Men also like shopping: Reducing gender bias amplification using corpus-level constraints. arXiv preprint arXiv:1707.094577 2017.
[110] Difan Zou, Yuan Cao, Dongruo Zhou, and Quanquan Gu. Gradient descent optimizes over- parameterized deep relu networks. Machine Learning, 109(3):467-492, 2020.
17


Reviews:
[{"review_id": "uKhY1N3G5yx", "reviewer": "Reviewer_3LiL", "summary": "The paper provides an NTK based theoretical framework to study the learning dynamics of neural networks. The topic of this paper is gradient starvation, the phenomenon that learning of one feature may prevent the learning of another.", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "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": "", "main_review": "While the paper is clear and I did enjoy reading this paper, I feel that one crucial experiment is lacking. Since the paper is about learning one feature may prevent the learning of another, and this has been formalized in section 3. The authors should perform one direct experiment to show this. \n\nFor example, one good choice would be to directly verify the case studied in theorem 2: show that z_2 does decrease as s_1^2 increases in the same prescribed setting. Without this type of experiment, I feel that there is a huge gap between the theory section and the experiment section. I deem this experiment crucial.\n\nI will consider raising the score if the above experiment is provided (and if the result is as expected) and the following questions are satisfactorily addressed. \n\nSome other questions I have:\n1. does inequality (6) holds for every possible alpha? if so, maybe it is good to explicitly say this, the authors might also want to include proof of this inequality in the appendix\n2. why does equality (7) hold? It seems that it should only hold if L(theta) = max_alpha E, where E is the expression following the max in equation (7), but the authors did not show this is the case in the paper (or did I miss anything?)\n3. one motivation that the authors have is that the previous work mainly studies quadratic loss but not cross-entropy loss. However, the authors also did not tackle the cross-entropy loss directly in equation (5). I feel that this is unjustified. If the authors do not directly deal with the cross-entropy loss, I feel that the authors should not motivate the paper this way\n4. the theory is done in the limit of a very small lambda, which corresponds to the case when the weight decay is small. Then why do we need the weight decay term in the first place in equation 5?\n\n\nOther lesser problems include:\n1. theta_0 introduced before definition in line 105\n2. what \"system 3\" refers to needs clarification. my guess is that it refers to \"the system prescribed by equation 3\"\n3. \"training optimum theta^*\" should be defined more formally in lines 117-118\n4. the \"circle-dot\" symbol between 131-132 undefined\n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "Sufficiently addressed.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "1.5 hour"}, "extra_text": {}}, {"review_id": "gtJgoijIXPA", "reviewer": "Reviewer_U7r7", "summary": "In summary, this paper proposes to use Spectral Decoupling to deal with Gradient Starvation problems under NTK regimes. Concretely, by analyzing the whole training dynamics, the authors first show that under NTK regimes, gradient starvation problems arise due to the correlation between different features during the training process. Therefore, the authors propose spectral decoupling to decouple the features. The authors conducted several experiments to show that this decoupling can help solve gradient starvation.\n\nAlthough I have some small concerns about the theoretical part, I still tend to accept the paper due to its solid and exciting experiments. I would like to raise my score if the authors could clarify more intuition behind the theory.", "questions": "", "limitations": "", "rating": 7, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "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": "", "main_review": "\n## Contributions:\n\n1. Due to the difficulties when analyzing cross-entropy loss, the authors formulate the training dynamics and its dual problem via a variational form approach, which provides some interesting insights.\n\n2. The authors explicitly show how gradient starvation happens by analyzing the training dynamics.\n\n3. The authors propose using spectral decoupling (a simple regularizer) to deal with the gradient starvation problems.\n\n4. Experiments show that such a regularizer works when there are spurious features. Besides, it performs well on colored MNIST.\n\n##Concerns:\n\n1. In definition 2, can the authors explain why they define gradient starvation by the gradient at the optimal point theta* (in my opinion, it is optimal via training)? Should we evaluate gradient starvation via the whole training process? Moreover, what if we cannot train the parameter to the optimal?\n\n2. Since the authors claim the regularizer solves the gradient starvation. I am wondering, with the new regularizer, can the gradient starvation vanish under the assumptions of Theorem 2? \n\n## Other limitations:\n\n1. The paper only focuses on the NTK regimes. However, since the theory mainly aims to provide some intuition to the algorithm, the NTK regime seems already enough.\n2. In Line 164-165, the authors claim that Theorem 2 is a variation of case 2. However, case 2 is a permutation matrix, but assumption 1 is a perturbed identity matrix (I believe there are no fundamental differences, just that the statement is not precise).\n3. The authors provide little intuition on how to derive Equation (17) [which is the crucial point in my opinion]. Although I believe it is true, I am still puzzled how the authors directly propose the regularizer, which decouples the features.", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "No potential negative societal impact.", "review_text": "", "extra_scores": {"ethical_concerns": "No ethical issues.", "needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "7"}, "extra_text": {}}, {"review_id": "WeYImKiOto", "reviewer": "Reviewer_3tUD", "summary": "This work studied the learning dynamics of neural networks for classification tasks trained with a cross-entropy loss. In particular the authors highlight a common failure mode \"feature imbalance\" where superficial features highly correlated with the labels, such as a water mark, will be learned quickly, but other features more important for generalization or the task will not be learned. They hypothesize that feature imbalance is caused by \"gradient starvation\", where the learning of the spurious features actually slows the learning of the other features by dampening the gradient signal for these features. They demonstrate theoretically the rational of this hypothesis in an binary classification setting and in the NTK regime. From this analysis they derive a simple regularization technique, called \"spectral decoupling\", which aims to decouple features preventing gradient starvation. Finally they empirically test their regularization technique on a variety of experiments to demonstrate that it is preventing the failure mode of feature imbalance.", "questions": "", "limitations": "", "rating": 7, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "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": "", "main_review": "Originality/Significance - I think the focus on learning dynamics with cross-entropy loss is important, as the authors mentioned, most learning dynamics analysis has focused on MSE loss.  My take-away message from this work was that the learning dynamics of cross-entropy loss do not decouple, as in the dynamics of MSE loss, which can lead to learning interference (\"gradient starvation\").  I was not aware of this idea prior to reading this work, nor had I seen the dual analysis of the CE loss.  I thought it was interesting to think about dynamics in the dual space, but it was unclear how these dynamics (equation 9) would relate to the dynamics in parameter space due to SGD. \n\nClarity/Quality - Generally the writing is clear, but at times I found myself asking what certain terms meant (such as \"feature imbalance\") or wishing a definition came earlier (such as the definition of \"feature\"). The overall structure of the paper is: (1) describe a failure mode of learning, (2) determine cause, (3) find solution, (4) test.  The failure mode motivation could be explained more clearly. Some of the derivations are hard to follow or lacking completely and it wasn't always clear how hyperparameters were chosen in experiments. For example it was strange that in one of the experimental sections 4.4 the authors use \"a variant\" of their algorithm without explanation.\n\n\nMajor Areas for Improvement:\n- I think it would be insightful to see the analysis repeated with MSE loss.  Shouldn't it be possible to take the same analysis steps with MSE loss and come to the conclusion that there is no gradient starvation? This analysis should be relatively simple to follow and might both help the reader follow your analysis and drive home the point that gradient starvation is a phenomenon unique to CE loss.\n- Shouldn't the analysis and implications apply to logistic regression without the need for certain assumptions?  I think this perspective would be helpful and provide a simple setting to test theoretical ideas.\n- Currently your experimental section neither directly tests the theoretical statements nor demonstrates that SD regularization should be used as an alternative to other regularization strategies.  Are there more direct ways of testing the hypothesis of gradient starvation (potentially with synthetic data)?  What happens when you apply SD regularization to non-binary tasks?  Clearly explaining the weaknesses of SD (such as hyperparameter tuning?) could motivate future algorithmic work building off SD even if \"achieving state-of-the-art results is not the objective\" of this work.\n- Explaining how future work could either build on your analysis or implications of SD regularization would be helpful.\n\n\nSpecific Comments:\n- In Fig. 1 are the dotted axis lines in (a),(b),(c) visualizing the features $z_1$ and $z_2$?  Initially I confused these lines with the dotted lines in the title used as a label for plot (d).  In general, I think the description of $z_1$ and $z_2$ is a little opaque in this caption and it would be great if you could describe or visualize \"features\" intuitively before defining in equation (4).\n- line 39 - \"We provide a theoretical framework to study learning dynamics of wide, over-parametrized neural networks trained with cross-entropy loss\" seems to strong of a statement.  If indeed section 3 describes a framework for studying learning dynamic how could it be applied to study for example dynamics in parameter space?  Maybe make this contribution sentence more specific to your work?\n- It would be helpful to describe the distinction between \"feature space\" and \"parameter space\" and the learning dynamics in each?\n- In equation (4) I am assuming you are missing an index for j?  Should is say $z_j = (U)_j^\\intercal Y \\hat{y}$?\n- In equation (8) is the $c^\\intercal$ supposed to be an $\\alpha^\\intercal$?\n- I could not find a derivation of equation (9) and I found it strange to see the learning rate appear.  If this is gradient flow on the objective in equation (8) then why would the learning rate appear?  Is this the learning rate used in the SGD dynamics in parameter space?\n- Why don't you merge assumption (1) as a condition to your lemma (1)?  It is only ever used in this lemma correct? \n- Would you consider describing\tintuition for spectral decoupling in the 2D classification problem for figure 1?  The analysis nicely built up the derivation of SD, but intuition for how it would mitigate GS in the simple setting where you motivated feature imbalance as a problem would be nice before describing experiments.  For that matter I wouldn't consider 4.1 a part of \"experiments\". I think it should be seen as a simple problem to build intuition.\n- Throughout your experiments I was left wondering how you choose the regularization rate $\\lambda$ for SD.  This seems like it would be crucial.  I would like to see experiments where this rate is swept over a range of values. For example shouldn't you be able to smoothly sweep from orange to green curve in fig 2 by sweeping lambda?\n- In table 2 it appears REx method is strictly outperforming SD, but you don't discuss this method?  It would be good to have a least a discussion of this as you do for IRM.", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "Yes.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "5"}, "extra_text": {}}, {"review_id": "J3L1fnRjHPQ", "reviewer": "Reviewer_MbCs", "summary": "The paper identifies and formalizes the Gradient Starvation (GS) phenomenon, which often occurs during training of overparameterized neural networks with the cross-entropy loss. GS causes the network to quickly learn features with higher strength, but fail to learn some features with lower strength due to the coupled learning dynamics. The main contributions are:\n- Formally defining GS in the NTK regime, and theoretically deriving the presence of GS when training with the cross-entropy loss. \n- Proposing a novel Spectral Decoupling (SD) regularizer that promotes decoupled learning dynamics and mitigates GS.\n- Empirically showing that the proposed SD regularization can effectively mitigate GS and improve generalization.\n", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "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": "", "main_review": "Originality: The paper has several novel contributions, including the formalization of GS, derivation of GS with cross-entropy loss, and proposal of SD that mitigates GS.\n\nQuality: This paper is technically sound in general. However, here are some critiques / questions that I would like the authors to address:\n- The analysis in Section 3.2 is based on the dual variable alpha, with the assumption that theta is optimized as a function of alpha. Does this assumption affect the learning dynamics? In other words, does the learning dynamics for alpha (equation 9) hold when we do gradient optimization on theta (instead of alpha) in practice?\n- The code for the CIFAR experiments doesn’t seem to be in the github repo. Also, I couldn’t find details about the CIFAR experiments in the main text or supplementary material. Please include these for better reproducibility.\n- Why is gamma needed in line 274?\n- It’s highlighted that hyperparameters are tuned on the test sets, which is not a valid assumption in practice. How sensitive is SD to hyperparameters, and what if they are not tuned on the test set?\n- How important is the NTK assumption in practice? Does the SD regularization still work well when the network is not overparameterized?\n\nClarity: In general, the paper is very well-written. The following are a few typos and suggestions:\n- Line 23: “among the authors’ personal experiences” sounds unprofessional. Either remove or provide citations.\n- In Definition 1 (line 109): the definition is for the NN’s response for feature j, but there is no mention of j in equation (4).\n- Equation (8): c transpose -> alpha transpose.\n- Line 259: missing period mark at the end.\n\nSignificance:\nThe theoretical formalization of GS and the proposed SD regularizer are significant contributions to the community.\n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "I think the authors have adequately addressed the limitations and potential negative societal impact of the work.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "6"}, "extra_text": {}}]

Author rebuttal:

Meta-review: After authors response, most reviewers updated their scores to a 6 or 7. Most reviewers did not feel strongly about this work and felt the paper would benefit from more discussion about its implications and more clarity about experiments. They also noted some limitations of the theoretical results presented. Due to the lack of enthusiasm, I do not feel strongly about the merit for acceptance. However, reviewers have uniformly agreed upon some of the merits of the paper - number of experiments showing gradient starvation is a phenomenon and the theoretical results the paper presents. Therefore, this paper warrants a borderline acceptance. 


Decision: Accept (Poster)