Venue: NeurIPS 2021 Conference

Reviews:

final_score: rating: 7: Good paper, accept
review_content: 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.
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. 

Clarity/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.


Major Areas for Improvement:
- 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.
- 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.
- 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.
- Explaining how future work could either build on your analysis or implications of SD regularization would be helpful.


Specific Comments:
- 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).
- 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?
- It would be helpful to describe the distinction between "feature space" and "parameter space" and the learning dynamics in each?
- In equation (4) I am assuming you are missing an index for j?  Should is say $z_j = (U)_j^\intercal Y \hat{y}$?
- In equation (8) is the $c^\intercal$ supposed to be an $\alpha^\intercal$?
- 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?
- Why don't you merge assumption (1) as a condition to your lemma (1)?  It is only ever used in this lemma correct? 
- Would you consider describing	intuition 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.
- 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?
- 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.
limitations_and_societal_impact: Yes.
review_title: null

final_score: rating: 6: Marginally above the acceptance threshold
review_content: 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.
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. 

For 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.

I 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. 

Some other questions I have:
1. 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
2. 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?)
3. 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
4. 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?


Other lesser problems include:
1. theta_0 introduced before definition in line 105
2. what "system 3" refers to needs clarification. my guess is that it refers to "the system prescribed by equation 3"
3. "training optimum theta^*" should be defined more formally in lines 117-118
4. the "circle-dot" symbol between 131-132 undefined

limitations_and_societal_impact: Sufficiently addressed.
review_title: null

final_score: rating: 6: Marginally above the acceptance threshold
review_content: 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:
- Formally defining GS in the NTK regime, and theoretically deriving the presence of GS when training with the cross-entropy loss. 
- Proposing a novel Spectral Decoupling (SD) regularizer that promotes decoupled learning dynamics and mitigates GS.
- Empirically showing that the proposed SD regularization can effectively mitigate GS and improve generalization.

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.

Quality: This paper is technically sound in general. However, here are some critiques / questions that I would like the authors to address:
- 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?
- 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.
- Why is gamma needed in line 274?
- 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?
- How important is the NTK assumption in practice? Does the SD regularization still work well when the network is not overparameterized?

Clarity: In general, the paper is very well-written. The following are a few typos and suggestions:
- Line 23: “among the authors’ personal experiences” sounds unprofessional. Either remove or provide citations.
- 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).
- Equation (8): c transpose -> alpha transpose.
- Line 259: missing period mark at the end.

Significance:
The theoretical formalization of GS and the proposed SD regularizer are significant contributions to the community.

limitations_and_societal_impact: I think the authors have adequately addressed the limitations and potential negative societal impact of the work.
review_title: null

final_score: rating: 7: Good paper, accept
review_content: 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.

Although 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.
main_review: 
## Contributions:

1. 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.

2. The authors explicitly show how gradient starvation happens by analyzing the training dynamics.

3. The authors propose using spectral decoupling (a simple regularizer) to deal with the gradient starvation problems.

4. Experiments show that such a regularizer works when there are spurious features. Besides, it performs well on colored MNIST.

##Concerns:

1. 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?

2. 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? 

## Other limitations:

1. 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.
2. 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).
3. 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.
limitations_and_societal_impact: No potential negative societal impact.
review_title: null

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)