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

Thank you very much for your constructive review! We first notify you that our equations 21, 22, 23 in Appendix B have typos: $\exp{(-\frac{1}{2} (s - \mu_k)^\top \Sigma^{-1}_k (s - \mu_k))}$ must be $\exp{(-\frac{1}{2} s^\top \Sigma^{-1}_k s)}$. We are sorry for the inconvenience.

1. The paper does not include error bars or statistical tests for empirical results. This undermines confidence in the generality of findings, especially on small or noisy datasets.

Though the availability of a shortcut label was consistently higher than the availability of the ground-truth label, the values of availability varied due to the randomness in picking samples for measuring NTK. Thus, we did not include error bars for the clarity of the graphs. But we will try to run important experiments such as Waterbirds and CelebA on CE loss to show you that our results are consistent. Though the values can vary due to randomness of picking training samples for measurement of NTK, the availability of shortcut labels was consistently higher than the one of ground-truth labels. Also, in the experiment measuring the strength of a shortcut (Fig. 7), a stronger shortcut always exhibited higher availability. The experiments below were run with CE loss. The availability was measured at the 1000-th iteration to observe the convergence behaviour in the early phase of training (which was also our focus in Fig. 6, 9, and 10)

Waterbirds

CelebA

Colored-MNIST

Patched-MNIST

On the other hand, for the experiments measuring the strength of a shortcut in Patched-MNIST (Fig. 7), the variance of availability was small (across 5 runs):

If given a chance to revise our paper, we will reflect your comment and include an error bar for this experiment. The experiments below were run with MSE loss.

Waterbirds

CelebA

Colored-MNIST

Patched-MNIST

Waterbirds

CelebA

Colored-MNIST

Patched-MNIST

We also answer your questions below:

1. The theoretical analysis relies heavily on the NTK framework (infinite-width, lazy training). In practice, finite-width networks often exhibit feature learning. Is there any empirical observations or theoretical intuitions about how the findings generalize to feature-learning regimes?

All the plots in Fig. 6, 9, and 10 are measurements from empirical NTK of finite-width networks. They were measured at >= 100 epoch (130 epoch for CelebA, 200 epoch for the rest). This empirically implies that our findings could be generalizable to feature-learning regimes. Also, one might extend our results with perturbation theory and dynamical mean-field theory, which might need more effort.

2. The shortcut labels in real-world datasets are manually constructed based on model predictions. How consistent are these across runs or architectures? Could a data-driven or automated method for shortcut label discovery be developed using availability scores?

The shortcut labels were consistent across runs. As in an example in Waterbirds, when the label was manually found, the match rate of the test predictions to the found label was 0.7819 ± 0.0018 (measured at > 100 epochs across 5 runs). On the other hand, when the label was composed solely from the shortcut feature, the match rate of the test predictions to the composed label was 0.7488 ± 0.0029 (across 5 runs) which was lower than the aforementioned one. In an example in CelebA, when the label was manually found, the match rate of the test predictions to the found label was 0.9527 ± 0.0003 (measured at > 100 epochs across 5 runs). On the other hand, when the label was the ground-truth label, the match rate of the test predictions to the ground-truth label was 0.9510 ± 0.0004 (across 5 runs) which was lower than the aforementioned one. In CelebA, we compared the match ratio of the shortcut label to the match ratio of the ground-truth label since the match rate to the label solely from the shortcut feature was much lower than the one to the ground-truth label. This shows that our shortcut label is reliable. But the shortcut labels were not consistent across parameters or architectures. For example, as shown in A.4, the shortcut labels might be different if the network is pretrained. And yes, it might be possible to construct a shortcut label from availability scores or from top eigenvectors of NTK (by constructing top eigenvectors of NTK as shortcut labels). But there is a limitation which was also noted in Appendix A.6: if there are multiple shortcuts in the datasets, it might not be possible to detect all the shortcut labels from the datasets.

Thank you very much for your constructive review! We first notify you that our equations 21, 22, 23 in Appendix B have typos: $\exp{(-\frac{1}{2} (s - \mu_k)^\top \Sigma^{-1}_k (s - \mu_k))}$ must be $\exp{(-\frac{1}{2} s^\top \Sigma^{-1}_k s)}$. We are sorry for the inconvenience.

1. Unclear logical structure. The manuscript would benefit from a more layered structured, which makes clearer claims.

We are sorry for our unclear statements. Your statements you wrote about the logical argument are correct. For more information, we clarify our results as below: Theorem 3.1 and Theorem 3.2 are based on data following the Gaussian mixture model, while Corollary 3.3 does not necessarily need data to be on a Gaussian mixture model. Those three results need an assumption that the trained model is a linear neural network. Meanwhile, the toy experiments from Fig. 4, Fig. 5, and Fig. 7 are trained on a two-layer ReLU FC network. Saliency maps (Fig. 3) were generated from a two-layer ReLU CNN network and predictability and availability (Fig 6) were measured from a pretrained ResNet-18.

2. The predictability and availability metrics seem reasonable, to the point that I would guess they already have a name or previous usage.

Yes, especially the availability metric has been widely used to measure the alignment of NTK to the ground-truth label [1], but it has not been used to investigate the phenomenon of shortcut learning.

3. The NTK analysis is relevant only in the lazy regime, which is not what happens when neural networks get trained on real-world samples. The derived results provide only an intuitive understanding that needs to be tested.

Theorem 3.1 is relevant in the lazy regime and our analysis on the convergence speeds of features is dependent on the lazy regime assumption. We acknowledge that it is our limitation. However, it is known that the neural network can sometimes exhibit the lazy regime behaviour in the early phase of training, which is our main focus of our analysis on the convergence speeds of features [2]. Also, Theorem 3.2 and Corollary 3.3 do not need an assumption of lazy regime.

4. While it is note that "the data variance of each cluster is not really small", the contribution of this phenomenon is not well investigated. The point of Figure 4 is unclear.

The fact that the data variances of the clusters are not small is used in Theorem 3.2, which implies that shortcut features dominate the decision boundary when the variances are not zero. The point of Fig. 4 is to empirically prove Theorem 3.2 and show the effect of data variances on the decision boundary.

We also answer your questions below:

1. Can you show that the Gaussian mixture model exhibits shortcut features without using the NTK analysis?

In Theorem 3.2, we show that the Gaussian mixture model exhibits shortcut features after convergence without using any NTK analysis. (Unfortunately, we currently do not provide an analysis on convergence speeds of features during convergence without NTK analysis.)

2. Can you quantitatively show the parameters of the GMM under which shortcut feature do or do not show, based on the NTK analysis?

If the data distribution follows the Gaussian mixture model of $p(x) = \sum^K_{k=1} \pi_k \mathcal{N}(\mu_k, \sigma^2_K I)$ and $\sigma_k > 0$, then the imbalance in the weights $\pi_k$ can trigger shortcut learning after convergence. Spectral bias can occur even when $\sigma_k = 0$.

[1] Baratin, Aristide, et al. "Implicit regularization via neural feature alignment." International Conference on Artificial Intelligence and Statistics. PMLR, 2021.

[2] Lyu, Kaifeng, et al. "Dichotomy of Early and Late Phase Implicit Biases Can Provably Induce Grokking." The Twelfth International Conference on Learning Representations.

Thank you very much for your constructive review! We first notify you that our equations 21, 22, 23 in Appendix B have typos: $\exp{(-\frac{1}{2} (s - \mu_k)^\top \Sigma^{-1}_k (s - \mu_k))}$ must be $\exp{(-\frac{1}{2} s^\top \Sigma^{-1}_k s)}$. We are sorry for the inconvenience.

1. I recommend clarifying this aspect in the analysis, as it potentially restricts the theoretical setup even further to datasets where the mean vectors have approximately equal norms, further drifting away from practice.

Firstly, both toy experiments in Fig. 4 and 5 were conducted with orthogonal clusters with means of norm 1. The norms of the means of clusters were identical and the effect of the norm was decoupled. While we were also aware of the effect of the norm of the mean, we did not elaborate on the subject since our focus was on the weight $\pi_k$. If given a chance to revise our paper, we will reflect your comment and emphasize more on the effect of norms of the means to the decision boundary of a neural network.

2. Theory relies on static NTK assumption that only holds in the infinite-width limit.

Theorem 3.1 assumes the static NTK assumption, which is the limitation of our work. However, it is known that the neural network can sometimes exhibit the lazy regime behaviour in the early phase of training, which is a focus of our analysis on the convergence speeds of features [3]. On the other hand, Theorem 3.2 and Corollary 3.3, which show the decision boundary of a network after convergence of the network, are not based on the constant NTK assumption.

We also answer your questions below:

1. Given the spectral bias analysis, I am curious whether you have thought about how the new spectral bias insights (and the metrics for availability and predictability) might inform designing further debiasing strategies?

Our new spectral bias insights might aid finding possible shortcut labels inherent in the dataset. Since shortcut features are learned faster than core features, one can construct shortcut labels by looking at the predictions of the network in the early phase of training, which is already done in a famous debiasing method called LfF [1].

2. Since the theory assumes a static NTK (which holds for infinite-width limit), have you observed a particular width threshold beyond which the spectral bias observations tend to emerge in the experiments?

We could not particularly find a case where our observations were not valid. We show our results on Patched-MNIST and a two-layer ReLU FC network when the width of the network is small as 32 or 64: availability of the shortcut label was consistently higher than the availability of the ground-truth label.

3. As mentioned in the Weaknesses section, for the toy experiments highlighting theorems 3.1 and 3.2, has the norm of the cluster means been controlled in some way in order to decouple its effect on the eigenvalue scales?

Both toy experiments in Fig. 4 and 5 were conducted with orthogonal clusters with means of norm 1. The norms of the means of clusters were identical and the effect of the norm was decoupled.

4. Can some aspects from the theoretical argument be directly applied from linear NTK to other types of kernels?

You cannot directly apply the results to NTK from a ReLU network, but in some cases where eigenfunctions of a ReLU NTK include linear functions [2], our results might be valid for linear eigenfunctions in those simple cases.

[1] Nam, Junhyun, et al. "Learning from failure: De-biasing classifier from biased classifier." Advances in Neural Information Processing Systems 33 (2020): 20673-20684. [2] Hermann, Katherine L., et al. "On the foundations of shortcut learning." arXiv preprint arXiv:2310.16228 (2023). [3] Lyu, Kaifeng, et al. "Dichotomy of Early and Late Phase Implicit Biases Can Provably Induce Grokking." The Twelfth International Conference on Learning Representations.

I thank the authors for their responses! In the light of the discussions, I have increased my score.

Regarding control for the mean-vectors, although the experiment accounts for equal mean-vector norms, my original curiosity (in hindsight, unclearly formulated, apologies for that) was rather about systematically controlling the drift between mean-vector norms such that its effect on the network outputs is observed.

While explicitly evaluating the effect of the mean-vector norms could more strongly highlight the theoretical-empirical link, I recognize that the focus of the experiment is on the mixing weights. I welcome and support the decision of the authors to emphasize more clearly on the effect of the mean-vector norms on the eigenvalues in an updated version.

Thank you very much for your constructive review! We first notify you that our equations 21, 22, 23 in Appendix B have typos: $\exp{(-\frac{1}{2} (s - \mu_k)^\top \Sigma^{-1}_k (s - \mu_k))}$ must be $\exp{(-\frac{1}{2} s^\top \Sigma^{-1}_k s)}$. We are sorry for the inconvenience.

1. All curves come from single runs; no confidence intervals reported.

Though the availability of a shortcut label was consistently higher than the availability of the ground-truth label, the values of availability varied due to the randomness in picking samples for measuring NTK. Thus, we did not include error bars for the clarity of the graphs. Due to the limited time, we cannot run all experiments with multiple runs during our rebuttal period: but we will try to run important experiments such as Waterbirds and CelebA on CE loss to show you that our results are consistent. Though the values can vary due to the randomness of picking training samples for measurement of NTK, the availability of a shortcut label was consistently higher than that of ground-truth label. Also, in the experiment measuring the strength of a shortcut (Fig. 7), a stronger shortcut always exhibited higher availability. The experiments below were run with CE loss. The availability was measured at the 1000-th iteration to observe the convergence behaviour in the early phase of training (which was also our focus in Fig. 6, 9, and 10)

Waterbirds

CelebA

Colored-MNIST

Patched-MNIST

On the other hand, for the experiments measuring the strength of a shortcut in Patched-MNIST (Fig. 7), the variance of availability was small (across 5 runs):

If given a chance to revise our paper, we will reflect your comment and include an error bar for this experiment. The experiments below were run with MSE loss.

Waterbirds

CelebA

Colored-MNIST

Patched-MNIST

2. Manual construction of shortcut labels may inject subjectivity.

Though we manually constructed shortcut labels, the assignment of labels does not depend on our subjective decision: for real-world datasets, we grouped the samples into four groups: [Bias-aligned samples labelled as 1], [Bias-aligned samples labelled as -1], [Bias-conflicting samples labelled as 1], and [Bias-conflicting samples labelled as -1]. We observed how each group was predicted in the test sets and assigned a shortcut label to each group as the predicted label of the majority of samples in the group (if the majority of the group was predicted as 1, then the shortcut label becomes 1 for this group). In Fig. 6, 9, and 10, the test predictions of the network were closer to our constructed shortcut labels rather than the original ground-truth labels, which shows that our constructed shortcut labels are valid. More explanation is in Appendix A.1. This is purely based on the statistics of the predictions not our subjective decisions. As in an example in Waterbirds, when the label was manually found, the match rate of the test predictions to the found label was 0.7819 ± 0.0018 (measured at > 100 epochs across 5 runs). On the other hand, when the label was composed solely from the shortcut feature, the match rate of the test predictions to the composed label was 0.7488 ± 0.0029 (across 5 runs) which was lower than the aforementioned one. In an example in CelebA, when the label was manually found, the match rate of the test predictions to the found label was 0.9527 ± 0.0003 (measured at > 100 epochs across 5 runs). On the other hand, when the label was the ground-truth label, the match rate of the test predictions to the ground-truth label was 0.9510 ± 0.0004 (across 5 runs) which was lower than the aforementioned one. In CelebA, we compared the match ratio of the shortcut label to the match ratio of the ground-truth label since the match rate to the label solely from the shortcut feature was much lower than the one to the ground-truth label. This shows that our shortcut label is reliable.

3. Results remain in infinite-width/lazy-training regime; impact on feature-learning settings uncertain.

Theorem 3.1 assumes the infinite-width/lazy-training regime, which is the limitation of our work. However, it is known that the neural network can sometimes exhibit the lazy regime behaviour in the early phase of training, which is the reason our analysis focuses on the convergence speeds of features [1]. On the other hand, Theorem 3.2 and Corollary 3.3, which show the behaviour of shortcut features after convergence of the network, are not based on the constant NTK assumption.

We also answer your questions below:

1. Finite-width drift – Could you run a width-sweep (e.g. 256-2 k channels) on Patched-MNIST to show when NTK approximation breaks?

We measured availability for two-layer ReLU FC networks with various widths: 32, 64, 256, 512, 1024, and 2048. Availability of the shortcut label was consistently higher than the availability of the ground-truth label.

2. Statistical confidence – Please repeat each benchmark over ≥5 seeds and report mean ± s.d.; how stable is the availability ranking?

We showed the partial results above: though the value of the availability itself could be unstable, the availability of a shortcut label was consistently larger than the availability of the ground-truth label.

3. Shortcut-label sensitivity – How do availability curves change if shortcut labels are derived automatically (e.g. k-means on top eigen-projections) instead of manual inspection?

We are not sure if we understood your question correctly, but if you mean automatic derivation by composing the label from k-means clustering on data projected to principal components, then the availability varied across runs. We experimented on Waterbirds and the label was constructed from 2-means clustering on data projected to principal components. Below shows the result:

Waterbirds

4. Variance term in Thm 3.2 – Provide a concrete bound showing when a small but high-variance cluster can outweigh a large low-variance one.

If the means of the clusters are orthogonal to each other and the norms of the means are identical, a small but high-variance cluster cannot outweigh a large low-variance one. The variance term $\sum^K_{i=1} \pi_i \sigma^2_i$ in $w_k$ includes the variance of all clusters and $w_k$ is not dependent only on the variance of a single cluster. If at least one cluster has a positive variance, then a large cluster always outweighs small clusters.

5. Late-epoch NTK – Have you measured availability using the empirical NTK after feature learning (e.g. epoch 50 ResNet-18) to see if shortcut alignment grows?

All the experiments in Fig. 6, 9, and 10 show the alignment of empirical NTK after epoch 100 (except CelebA that was run for 130 epochs, all the experiments were run for 200 epochs). The availability grew in experiments on synthetic datasets, but it did not in experiments on real-world datasets.

[1] Lyu, Kaifeng, et al. "Dichotomy of Early and Late Phase Implicit Biases Can Provably Induce Grokking." The Twelfth International Conference on Learning Representations.

Thanks for self-reporting the error. Seems like an honest mistake, and, importantly, does not seem to affect the main claims in the paper.

It goes without saying but if your paper were accepted, the final version must be updated to reflect the corrected and clarified results presented during the review. I would be explicitly looking out for it. Note, of course, that the discussions are still ongoing, and no decisions have been made.