[Paper Review] Gradient Starvation: A Learning Proclivity in Neural Networks

Mohammad Pezeshki, Sékou-Oumar Kaba+4 more|arXiv (Cornell University)|Nov 18, 2020

Adversarial Robustness in Machine Learning102 references42 citations

TL;DR

The paper formalizes Gradient Starvation under cross-entropy in over-parameterized networks, showing how strong features hinder learning of weaker ones, and introduces Spectral Decoupling to decouple feature learning and improve robustness and OOD generalization.

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.

Motivation & Objective

• Motivate and formalize why gradient descent emphasizes a subset of predictive features under cross-entropy loss.
• Develop a theoretical framework for learning dynamics in the NTK regime to explain feature imbalances.
• Propose Spectral Decoupling as a simple regularizer to decouple feature learning and mitigate Gradient Starvation.
• Provide theoretical guarantees and empirical evidence across classification and OOD tasks.
• Discuss implications for robustness and generalization under data with spurious correlations.

Proposed method

• Model the neural network in the Neural Tangent Kernel regime to linearize training dynamics.
• Decompose learning dynamics along orthogonal feature directions via SVD of Y Phi0 to define features and responses.
• Formulate Gradient Starvation as a coupling between feature directions that slows learning of some features when others are stronger.
• Derive a dual formulation of cross-entropy via a variational bound to obtain tractable dynamics and fixed points.
• Introduce Spectral Decoupling by replacing weight decay with a logit-based L2 penalty to decouple dual feature dynamics.
• Provide theoretical results (e.g., fixed-point analysis and perturbation results) and validate with simple analytic cases and experiments.

Experimental results

Research questions

• RQ1 Under what conditions does Gradient Starvation emerge in cross-entropy trained networks?
• RQ2 How do feature strength disparities and coupling between feature directions affect learning dynamics?
• RQ3 Can a simple regularizer decouple feature learning and mitigate Gradient Starvation without hurting clean accuracy?
• RQ4 Does Spectral Decoupling improve robustness and out-of-distribution performance across tasks?

Key findings

• Gradient Starvation occurs when stronger features suppress learning of weaker but predictive features under cross-entropy.
• Coupled learning dynamics due to non-diagonal interactions in the feature space drive GS, especially with unequal feature strengths.
• A Spectral Decoupling regularizer makes dual dynamics independent, mitigating GS and enabling learning of multiple features.
• SD improves adversarial robustness and OOD performance on CIFAR-2/10/100 without sacrificing clean accuracy in reported experiments.
• SD yields larger classification margins in CIFAR-2 and improves worst-group accuracy in CelebA hair-color classification with spurious gender-color correlations.
• Experiments on colored MNIST demonstrate that SD helps learn robust features beyond color, improving test performance in non-training environments.