Shortcut Features as Top Eigenfunctions of NTK: A Linear Neural Network Case and More

AI summary

Overview

Research area: Deep learning theory — specifically shortcut learning analyzed through the Neural Tangent Kernel (NTK) framework.

Technical level: Advanced. The paper relies on kernel operator theory, eigendecompositions, gradient-flow ODEs, and Gaussian mixture model analysis.

Scope: A theoretical and empirical study showing that shortcut features correspond to eigenfunctions of the NTK with large eigenvalues, using a linear neural network on a Gaussian mixture as the analytical case and extending the results to a two-layer ReLU network and a ResNet-18.

What This Paper Is About

Neural networks trained on biased data often latch onto shortcut features — attributes that correlate strongly with labels inside the training set but do not generalize outside it. This paper asks why that happens and whether the usual explanation (the max-margin bias) is sufficient, by defining a "feature" as an eigenfunction of the NTK and tracking how much each eigenfunction contributes to the network's output during and after training.

Key Contributions

1. Eigenvalue analysis for shortcut clusters. Proposition 3.1 derives the eigenfunctions and eigenvalues of the inner-product kernel k(x, y) = ⟨x, y⟩ under a Gaussian mixture model p(x) = Σ πₖ N(μₖ, σₖ²I), showing that eigenvalues depend on cluster weights πₖ through the eigendecomposition of Σ πₖ μₖ μₖᵀ.

2. Post-convergence influence analysis. Proposition 3.2 gives a closed-form expression for the weight wₖ that each eigenfunction contributes to the converged network output under MSE loss, showing wₖ increases with πₖ from 0 to 1 and that this dependence arises from within-cluster variances σᵢ.

3. Ablation of the max-margin explanation. Corollary 3.3 shows that under strong SD / Marg-Ctrl regularization (λ → ∞), the ratios of feature weights converge to those of an MSE-trained network, so controlling the margin does not remove shortcut dominance.

4. New empirical metrics. Definitions 4.1 and 4.2 introduce predictability (alignment of a feature with ground-truth labels, yᵀg / |X|) and availability (alignment of a label with the empirical NTK, computed from its eigendecomposition), then apply them to real-world datasets.

Main Findings

• Large clusters produce large eigenvalues. Eigenvectors close to a cluster with larger weight πₖ have larger eigenvalues aᵢ, so inner products with vectors from larger clusters converge faster under the spectral bias. Since the assumption is Σᵢ π_{B_{y,i}} ≫ Σᵢ π_{C_{y,i}}, shortcut clusters converge faster than others.

• Large eigenvalues also mean large post-convergence influence. Because wₖ increases with πₖ, eigenfunctions corresponding to larger clusters have more influence on the converged neural network output. Under the orthogonality assumption μᵢ ⊥ μⱼ and vₖ = μₖ/‖μₖ‖, wₖ = ± πₖ‖μₖ‖₂ / (Σᵢ πᵢ σᵢ² + πₖ‖μₖ‖₂²).

• Data variance is the mechanism. If within-cluster variances were very small, wₖ would not change when the distribution (πₖ) changed. The toy 2D experiment with a two-layer fully-connected ReLU network classifying four clusters matches this: at σ = 1/3 and σ = 1/5 the decision boundary tilted toward the larger-weight clusters, while at σ = 1/20 the boundary depended only on cluster positions, not weights.

• Max-margin bias is not the only cause. In the 2D toy example, applying SD to the MSE loss left the decision boundary unchanged, while applying SD to the cross-entropy loss changed it and converged toward the MSE boundary at λ = 1.0 (λ values tested: 0.1 and 1.0). The authors state this means shortcut bias arises from learning the label itself.

• Saliency maps confirm the theory visually. In a two-layer ReLU CNN, features with larger eigenvalues (ranked by eigenvalue magnitude) focus on biased attributes: the patch in Patched-MNIST, the background in Waterbirds, and face edges or background rather than hair in CelebA.

• Shortcut labels have lower predictability but higher availability. On Waterbirds and CelebA with a pretrained ResNet-18, the availability of manually identified shortcut labels exceeded that of ground-truth labels, and test accuracy was also higher for the shortcut labels. Availability in Figure 6 was measured from 500 randomly sampled training data points.

• Availability tracks shortcut strength. On Patched-MNIST with a two-layer ReLU fully-connected network, availability was larger as patch size grew across 1×1, 3×3, 5×5, and 7×7 pixels.

Methodology in Plain English

The authors start with a deliberately simple model: a linear neural network with no activation function, trained on data drawn from a Gaussian mixture where the majority of samples cluster around biased attributes. Because NTK theory says the network output can be written as a sum of eigenfunctions weighted by their eigenvalues, they can derive in closed form how fast each feature is learned and how much it contributes to the final output. They then test whether the small-scale conclusions survive on richer models. For a two-layer ReLU CNN they generate saliency maps by taking the magnitude of the loss gradient with respect to the input for individual features. For real datasets they train a pretrained ResNet-18 with last-layer weights initialized to zero and the bias parameter removed, using SGD with learning rate 0.001 and weight decay 0.0001, and measure availability from the empirical NTK's eigendecomposition. Datasets used are Patched-MNIST, Colored-MNIST, Waterbirds, CelebA, and Dogs and Cats. In Patched-MNIST and Colored-MNIST, 95% of digits labelled -1 and only 5% of digits labelled 1 carry the bias in training, with the test ratio balanced at 50%; Waterbirds has 95% of waterbirds on watery backgrounds and 95% of landbirds on land backgrounds; Dogs and Cats has 95% of cats and 5% of dogs dark-coloured.

Why This Matters

This work reframes shortcut learning as a property of the NTK spectrum rather than purely as an artifact of margin maximization. That reframing matters because common debiasing approaches built around margin control may leave the underlying preference for shortcut features intact — the paper's Corollary 3.3 shows the SD-regularized boundary converging back to the MSE boundary.

Potential real-world applications:

• Medical imaging, where a scanner-specific marker or hospital artifact can correlate with a diagnosis label and dominate training data.
• Wildlife monitoring and ecological image classification, where background context (as in Waterbirds) substitutes for the animal's appearance.
• Facial attribute recognition and demographic analysis, where correlated attributes such as gender stand in for the target attribute (as in CelebA).
• General product or content classification, where illumination or colour (as in Dogs and Cats) acts as a non-generalizable cue.

Industry relevance: The predictability and availability metrics give practitioners a way to measure, before deployment, whether a candidate label is aligned with the top eigenspaces of a model's empirical NTK. If it is, the model is likely to learn it quickly and rely on it heavily, which is a signal to investigate the data distribution rather than the optimization schedule.

Future Directions

• Extending the theoretical analysis beyond the simple case of a linear neural network trained on a Gaussian mixture model.
• Relaxing the NTK framework's infinite-width assumption, which Proposition 3.1 depends on.
• Resolving the tension between Proposition 3.1, where eigenvalues scale with the magnitude of mean vectors, and Proposition 3.2, where the contribution to the decision boundary scales inversely with it — a mismatch the authors note may limit the use of the NTK spectrum for assessing shortcut influence after convergence.
• Addressing the failure case noted in Appendix A.7, where availability might not detect certain shortcut labels if the ground-truth label is predominated by a stronger shortcut.

Target Audience

Researchers working on shortcut learning, bias mitigation, and debiasing methods; theorists interested in NTK and kernel-based analyses of generalization; and machine learning practitioners who train models on datasets with imbalanced or correlated attributes and need diagnostic tools for detecting reliance on non-generalizable features.

Authors’ abstract

One of the chronic problems of deep-learning models is shortcut learning. In a case where the majority of training data are dominated by a certain feature, neural networks prefer to learn such a feature even if the feature is not generalizable outside the training set. Based on the framework of Neural Tangent Kernel (NTK), we analyzed the case of linear neural networks to derive some important properties of shortcut learning. We defined a feature of a neural network as an eigenfunction of NTK. Then, we found that shortcut features correspond to features with larger eigenvalues when the shortcuts stem from the imbalanced number of samples in the clustered distribution. We also showed that the features with larger eigenvalues still have a large influence on the neural network output even after training, due to data variances in the clusters. Such a preference for certain features remains even when a margin of a neural network output is controlled, which shows that the max-margin bias is not the only major reason for shortcut learning. These properties of linear neural networks are empirically extended for more complex neural networks as a two-layer fully-connected ReLU network and a ResNet-18.