Source-linked AI summary
Gradient Starvation: A Learning Proclivity in Neural Networks
Mohammad Pezeshki, Sékou-Oumar Kaba, Yoshua Bengio, Aaron Courville, Doina Precup, Guillaume Lajoie
TL;DR
The paper addresses why cross-entropy training can learn only a subset of predictive features, leaving other useful features undiscovered and harming robustness or OOD generalization. It formalizes this Gradient Starvation through dual-space learning dynamics, derives Spectral Decoupling, and reports improved margins, robustness, and performance on biased-data settings. The paper also notes that Gradient Starvation can be beneficial when the dominant feature generalizes.
Problem
Cross-entropy-trained neural networks may favor superficial, fast-to-learn correlations while failing to discover other predictive features, with consequences for OOD generalization and robustness.
Method
The paper uses dynamical-systems analysis of linearized neural networks in a dual feature space to formalize Gradient Starvation and derive Spectral Decoupling.
Results
Spectral Decoupling improves classification margins and robustness without apparently compromising noise-free test performance, and improves performance on the CelebA gender-bias task.
Takeaways & Limitations
Controlling feature-learning dynamics can mitigate Gradient Starvation in settings where predictive features are imbalanced or spuriously correlated.
Takeaways & Limitations
The method assumes feature imbalance, although it does not require identifying the imbalanced features exactly.
Abstract
from arXiv · showhide
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.
1 Introduction
The paper studies Gradient Starvation, where cross-entropy training favors strongly correlated, fast-to-learn features while other predictive features remain undiscovered. It formalizes this imbalance, proposes Spectral Decoupling, and supports the analysis with experiments.
- Motivation: Gradient descent can prioritize strongly correlated, fast-to-learn features, leaving other predictive features undiscovered.This imbalance is associated with poor OOD generalization, reliance on spurious correlations, and reduced robustness.
- Motivation: Prior analyses largely focus on squared-error loss, whose learning dynamics do not readily generalize to cross-entropy classification.The paper argues that cross-entropy produces distinct dynamics requiring separate analysis.
- Contributions: The paper provides a theoretical framework for cross-entropy learning dynamics in linearized neural networks and formalizes Gradient Starvation through coupled orthogonal feature directions.It uses a dual-space analysis and perturbation analysis to characterize the imbalance.
- Contributions: Spectral Decoupling is introduced as a regularizer intended to decouple feature-learning dynamics and mitigate Gradient Starvation.The paper supports the proposal with theoretical guarantees and empirical results across classification and adversarial attack tasks.
2 Gradient Starvation: A simple example
A small shift can make the training data linearly separable, allowing cross-entropy optimization to achieve nearly zero loss by learning one feature while ignoring another. This creates narrow, vulnerable decision boundaries, whereas Spectral Decoupling learns a larger-margin boundary.
- Simple example: A small shift makes the 2-D training data linearly separable, enabling negligible loss by discriminating along the horizontal axis while ignoring the vertical axis.The example uses a two-layer ReLU network with 500 hidden units trained with cross-entropy.
- Simple example: When one feature is learned faster, correctly processed examples contribute less gradient information, starving the remaining features of learning signal.This mechanism explains how cross-entropy can preserve feature imbalance during training.
- Consequences: Near-zero training loss can coexist with a decision boundary close to the data points, creating adversarial vulnerability and weak OOD robustness.High confidence on the training data does not ensure a wide-margin boundary.
- Consequences: Gradient Starvation can produce excessive invariance, leaving low-probability inputs classified with high confidence despite task-relevant changes.The example shows that ignored features need not affect predictions strongly.
- Consequences: Gradient Starvation may instead act as an implicit regularizer when training and test data share the same dominant feature.Under that distributional condition, suppressing weaker features can prevent overfitting.
3 Theoretical Results
The paper analyzes cross-entropy learning in a dual feature space, identifies coupled dynamics and unequal feature strengths as conditions for Gradient Starvation, and derives Spectral Decoupling to promote balanced learning.
- Problem Setup: The theoretical analysis assumes the Neural Tangent Kernel regime, binary classification, and interactions between two feature directions.These assumptions provide analytical tractability for the learning-dynamics study.
- Problem Setup: In the NTK regime, the network output is approximated as a linear function of parameters, with dominant feature directions obtained from the NTRF singular-value decomposition.The feature responses are defined through the SVD of YΦ0.
- Gradient Starvation Definition: Feature responses aggregate example logits weighted by their feature coordinates, enabling Gradient Starvation to be defined through cross-feature effects on optimal responses.An increase in one feature’s strength can have a detrimental effect on another feature’s learning.
- Training Dynamics: A variational Legendre-transform approach converts difficult coupled gradient-descent dynamics into tractable dynamics over dual variables and feature directions.The resulting autonomous system incorporates different learning rates across feature directions.
- Training Dynamics: Theorem 1 states that fixed points of the approximated system are attractive for α_i ∈ (0, 1).The theorem characterizes stability of the fixed points used in the analysis.
- Gradient Starvation Regime: Gradient Starvation arises when feature strengths differ and coupled learning dynamics connect their directions, with stronger discrepancies producing detrimental effects.The theorem’s two-feature regime extends pairwise to higher-dimensional and multi-class settings.
- Spectral Decoupling: Spectral Decoupling cancels the coupling matrix in the dual dynamics, making each α_i independent of the other α_j and promoting balanced feature learning.The regularizer uses a penalty on output weights and is intended to mitigate Gradient Starvation without requiring exact identification of imbalanced features.
4 Experiments
Experiments evaluate Spectral Decoupling across synthetic and real-world settings, including adversarial robustness, spurious correlations, and OOD performance. Across these tasks, SD improves robustness or worst-group accuracy while retaining strong standard performance, although the Colored MNIST probe relies on test-environment hyperparameter access.
- SD produces a curved decision boundary with a larger input-space margin in the 2-D classification example.
- CIFAR classification and adversarial robustness: SD consistently improves CIFAR adversarial OOD performance, and its margin improvement does not compromise noise-free test performance.The CIFAR evaluation compares ERM and SD on CIFAR-2, CIFAR-10, and CIFAR-100 under perturbed test inputs.
- Colored MNIST with color bias: 68.4% accuracy is achieved by SD on Colored MNIST while training on a single aggregated environment.Unlike IRM and REx, SD does not require access to multiple training environments.
- Colored MNIST with color bias: On Colored MNIST, SD learns color but also digit-shape features, whose predictions prevail when color conflicts with the label.In the colored-digit environment, SD reaches 67.2% accuracy versus 23.9% for ERM; IRM ignores color but requires multiple training environments.
- Colored MNIST with color bias: The Colored MNIST results should be interpreted as a probe because all reported methods use test-environment access for hyperparameter tuning.The authors state that this assumption is not generally valid.
- CelebA with gender bias: On CelebA, SD improves performance more than twofold for the rare blond-male group, where ERM fails despite generalizing well on average.Blond males comprise 0.85% of the training data, and the task contains a spurious Gender–HairColor correlation.
- CelebA with gender bias: The CelebA SD variant provably decouples learning dynamics and appears to outperform the original SD formulation, while Group DRO and LfF require additional information or simultaneous network training.
5 Related Work and Discussion
The paper situates Gradient Starvation within work on implicit bias, spurious correlations, and OOD generalization, emphasizing that feature selection can help or harm depending on distributional assumptions. It also discusses societal risks from superficial dataset biases and the need to control reliance on non-generalizable features.
- On reliance upon spurious correlations and robustness: OOD robustness remains difficult because networks may rely on superficial or spurious correlations rather than abstract, robustly informative features.The paper connects Gradient Starvation to this reliance and to possible adversarial vulnerability.
- On implicit bias: Related work explains neural-network generalization through simplicity, implicit regularization, spectral bias, NTK alignment, or convergence to maximum-margin solutions.These accounts include learning simpler functions first and aligning with only a few task-relevant directions.
- Gradient Starvation: friend or foe?: Gradient Starvation may act as an implicit regularizer when the learned features suffice for test generalization, but otherwise can prevent predictive features from being learned.The paper presents GS as potentially beneficial or adverse depending on the relationship between training and test distributions.
- GS social impact: Dataset biases can produce superficially high performance in critical applications, motivating methods that uncover and mitigate reliance on biased features.The CelebA experiment is cited as an example of gender-specific bias that can be difficult to detect.
6 Conclusion
The paper formalizes Gradient Starvation for cross-entropy training, analyzes it through dual-space learning dynamics, and derives Spectral Decoupling as a remedy. The analysis shows that GS can slow learning of features present in the training set.
- 6 Conclusion: Gradient Starvation emerges during cross-entropy training when coupled dual-space dynamics slow the learning of certain features despite their presence in the training set.The paper derives Spectral Decoupling as a possible remedy.
A Further discussions
This discussion explains how Spectral Decoupling is motivated, derived, and interpreted, while clarifying the paper’s analytical assumptions and future scope. SD penalizes output confidence to encourage learning additional features and decouple feature dynamics.
- The intuition behind Spectral Decoupling (SD): Spectral Decoupling favors learning a new feature over becoming more confident on features already learned by other examples.Its L2 output penalty penalizes over-confident predictions and encourages broader feature learning.
- Derivation of Spectral Decoupling: The SD penalty is derived by removing a coupling term from the dual dynamics, yielding decoupled feature responses across the dual parameters.The construction uses ||SV^Tθ||2, which equals ||ŷ|| under the stated factorization.
- Loss choice: Squared-error loss may avoid gradient starvation and improve robustness, but selecting suitable classification targets remains a major obstacle.The paper presents the robustness explanation as a conjecture.
- Analytical assumptions: The analysis assumes the Neural Tangent Kernel regime, binary classification, and interaction between two features for analytical tractability.NTK analysis is restricted primarily by its accuracy in the large-width limit.
- Future Directions: Future work should identify reliance on spurious correlations in sensitive applications before applying methods such as Spectral Decoupling to mitigate it.The discussion cites failures to generalize across hospitals as an example of hospital-specific feature reliance.
B.1 A Simple Experiment Summarizing the Theory
The toy experiment compares primal and dual learning dynamics with and without Spectral Decoupling. Increasing one feature’s singular value intensifies starvation of the other, whereas SD decouples their dynamics.
- Experiment setup: The experiment compares primal dynamics, dual dynamics, and Spectral Decoupling on a two-datapoint toy classification task.The primal and dual formulations arise from separate optimization problems, while SD uses a modified optimization.
- Primal versus dual dynamics: Although primal and dual dynamics differ, they share the same fixed points.This comparison appears in the left panel of Figure 5.
- With Spectral Decoupling: Spectral Decoupling makes the z1 and z2 learning dynamics independent of changes to the other feature’s corresponding singular value.This behavior is shown in the right panel of Figure 5.
B.2 Two-Moon Classification: Comparison with other regularization methods
In the two-moon task, training data lets cross-entropy achieve negligible loss by learning only the horizontal feature, while common regularization methods do not recover a larger-margin curved boundary. Spectral Decoupling is evaluated across regularization strengths, with excessive regularization harming IID generalization.
- Two-Moon Classification: Both axes are relevant, but the small class margin lets the network reach negligible loss using only the horizontal feature.The second dimension is overlooked because the training data permits explaining labels with one feature.
- Comparison with regularization methods: Weight Decay, Dropout, and Batch Normalization do not help achieve a larger-margin classifier on the two-moon task.The comparison uses full-batch gradient descent with learning rate 1e −2, momentum 0.9, and 10k iterations unless otherwise stated.
- Spectral Decoupling: Sweeping λ for Spectral Decoupling produces a smooth transition from green to orange, but larger λ values hurt IID test generalization.The selected λ is cross-validated using the average of IID and OOD generalization performance.
B.4 Colored MNIST with color bias
The Colored MNIST analysis examines feature learning under color bias, comparing SD with ERM and IRM. SD retains color sensitivity but reduces confidence when expected digit features are absent and performs better than ERM when color-label relationships flip, although the evaluation uses test-environment hyperparameter tuning.
- Results: SD achieves significantly better performance than ERM when both color and digit features are present in an environment with flipped colors.The result is consistent with SD mitigating Gradient Starvation from color features onto digit-shape features.
- Results: In the environment containing no digit features, SD matches ERM’s low accuracy but has much lower confidence.This lower confidence indicates that SD expects additional features beyond color for classification.
- Comparison with IRM: IRM appears not to capture the color feature, predicting 50% accuracy with low confidence when only color is presented.The low-confidence behavior is described as consistent with IRM being invariant to color.
- Scope and assumption: The task assumes access to the test environment for hyperparameter tuning, so results should be interpreted as a probe of SD’s feature-control capability.The paper states that this assumption is not generally valid.
- Hyperparameter selection: Hyperparameter search applies SD at the 450th step, coinciding with the observed onset of traditional in-distribution overfitting.This suggests hyperparameters might be tuned without monitoring the test set.
C Proofs of the Theories and Lemmas
The proofs recast cross-entropy through a Legendre transform and analyze gradient dynamics in a dual space. They establish attraction of fixed points, characterize feature responses in special cases, formalize Gradient Starvation under feature coupling, and show that Spectral Decoupling makes coordinate dynamics independent.
- Dual dynamics: The transformed formulation supports analysis of learning dynamics as an autonomous gradient system in the dual variable.The continuous gradient-ascent learning-rate parameter affects the dynamics but not the fixed point.
- Legendre transformation: The Legendre transformation expresses cross-entropy as a linear lower bound in ω, maximized over a variational parameter α.For binary cross-entropy, the intercept function equals the binary entropy H(α), with equality at α∗= −∇ωL.
- Fixed-point stability: Any fixed point of Eq. 10 is attractive for α_i ∈(0, 1).The proof uses linearization and shows the Jacobian is negative-definite.
- Gradient Starvation regime: When feature coupling is a perturbation and s1^2 > s2^2, the theory identifies a Gradient Starvation regime.The analysis treats off-diagonal coupling terms as a small perturbation of a decoupled system and studies its effect on steady-state feature responses.
- Spectral Decoupling: Spectral Decoupling replaces L2 weight decay with an L2 penalty on logits and yields independent dynamics for each dual coordinate α_i.The loss formulation makes the dynamics of α_i independent of α_j for j ≠ i.