Emergence of Alignment and Local Elasticity in Two-Layer Neural Networks

Thank you for your review. We respond to the weaknesses you pointed out as follows.

W1 -> We support our motivation through experiments measuring the clustering performance (recall@1), which implies the prediction on unseen data is available through $\beta^\top \mu$.

We observe that the alignment increases as $\beta^\top \mu$ increases in two-layer networks, and multi-class networks also demonstrate a positive correlation between alignment and recall@1 with low p-value.

Explanations on intuition including margin-based method is also added in Introduction section.

W2 -> The assumption over data distribution makes our proof unique to previous works.

Existing studies analyzing two-layer networks [Ba et al. (2022), Moniri et al. (2024)] assume standard gaussian distribution over their data distributions, with their analysis scope on regression problems.

On the other hand, we generalize such assumption into non centered Sub-Gaussian distribution, which is applicable to classification problems thanks to its finite supports representable as separable class distributions.

At proof technique level, we apply approximation methods by Ba et al. (2022) and Moniri et al. (2024) but with new assumptions applicable to a classification problem. Especially, we expand Gaussian distribution assumption on data into Sub-Gaussian and substitute Lipschitz teacher function into binary class labels. Such differences on conditions derive the general statements to previous works, but the definitions and the proof between statements require novel delineations. For instance, the definition of $\beta$ exploits the opposite signs of $y_1$ and $y_2$ to compose different gradient signals in the two-label classification, and we utilize SVD to acquire the norm bound of arbitrary labels without Lipschitz assumption.

W3 -> For Theorems 5.3 and 5.4, we modify their symbols and clarify the assumptions, which are now readable without appendices.

Additionally, we cleared the randomness and made Theorem 5.3 and 5.4 deterministic statements, so as informal statements Thm 1.3, 1.4.

Through Sections 5.3 and 5.4, we observe that as the similarity $\beta^\top \mu$ between the train data and unseen data increases, the metrics of alignment and LE also increase. To clarify this further, we have added a discussion section.

Additional experiments also verify the influence of $\beta^\top \mu$ to alignment and LE in both two-layer and multi-layer networks.

We have improved the paper based on your comments. We would appreciate it if you could refer to the revised version. Additionally, below are the contributions and meaning of our informal theorem we are asserting.

-> The main contribution is that we can predict the feature structure after training, expressed through metrics such as alignment and LE, based on information from the data before training. We will clarify this in the introduction.

The key idea of the informal theorem is that the alignment and LE of the feature vectors learned using MSE can be expressed in terms of the statistics of the training and test data, namely $\mu$ and $\beta$. We have rewritten this more clearly. Apologies for any confusion caused.

Thank you for the helpful review that contributed to improving the paper. We are now addressing your questions.

We have clarified the theorem, particularly providing a clearer that “we provide analysis of the features after one step of training”. Additionally, we explicitly stated that this study focuses solely on classification problems. Apologies for any confusion caused.

We have claimed that after one step of GD, alignment and LE are increased by the same variable, $\beta^\top \mu$. To express this more precisely, we will replace the term "connection" with the statement "both phenomena occur after one step." Furthermore, in the final "further works" section, we plan to mention the need to explore the connection of these two phenomena as events that occur after one step.

We have made it clear in the problem setup that only the first layer is trained, while the initialized second layer affects the gradient. Reflecting this, we revised the overall description of the problem setup and did not depict the second layer before and after training in Fig. 1, thus enhancing clarity.

Incorporating your comment, we experimented with the CARS196 and CUB200 datasets using a multi-layer network after sufficient gradient steps(20 epoch), and confirmed that $\beta^\top \mu$, alignment, and LE are strongly rank correlated with Kendall’s W statistic and move together.

Q. Can the alignment part of the result be related to the alignment phenomena discussed in https://arxiv.org/abs/2410.03006？ -> This is a paper that we have not yet reviewed. It has provided us with a great deal of inspiration, and we are grateful for it. The paper above argues that alignment between latent representations, weights, and gradients during training enables the neural network to naturally learn compact representations, supporting this claim through covariance and correlation analysis. In contrast, we focus on the alignment between representations within the same class (intra-class representations) and have discovered conditions under which alignment occurs more strongly in a two-layer neural network.

We believe there is a strong connection between these two papers when considering recent trends in Neural Collapse (NC) research. NC suggests that in a classifier, variability collapse—alignment between representations—occurs, as well as convergence to self-duality, meaning alignment between representations and weights. Our paper addresses the alignment related to variability collapse in a two-layer structure that can learn features, whereas the above paper extends the analysis by including gradients in addition to the alignment between representations and weights, as seen in NC. This distinction is significant, and we will incorporate this correlation in our paper.

We deeply appreciate your warm encouragement and thoughtful feedback. To address your concerns regarding the number of layers, steps, and tasks, we have conducted extensive validations of our theory using ResNet architectures with various depths and a wide range of steps. As detailed in the revised manuscript, we have confirmed that our theoretical predictions continue to hold for neural networks trained on real-world data.

That said, we find your concern about limitations regarding the number of layers, steps, and tasks still exist. To provide additional clarity and to better convey our intentions, we provide additional background information and a concise theoretical analysis.

We hope these additions address your concerns and help you view the manuscript more favorably. If there are any other aspects you believe require further clarification or additional analysis, please do not hesitate to let us know.

Additionally, we have been informed by ICLR that the discussion period has been extended, and we look forward to engaging in meaningful discussions to address any further points you may have. Thank you once again for your time and constructive evaluation.

Response to Concerns Regarding Two-Layer Networks: Neural Network Theories Based on the Two-Layer Assumption

Several prior studies have effectively utilized the feature extractor assumption same to our, to interpret phenomena observed in practical neural networks. For example, Damian et al., 2022 analyzed the efficient generalization and transfer performance of neural networks, while [1] used this framework as a tool to study robustness to input distribution shifts. Similarly, [2] employed it to analyze out-of-distribution inputs, and [3] utilized it to investigate adversarial robustness. These studies focused on understanding phenomena of neural network representations, particularly the hidden representations allowing them to model and explain behaviors observed in practical deep learning scenarios. Based on this body of work, we argue that the assumption of a two-layer network capable of learning hidden representations is a reasonable and effective framework for analyzing neural networks without significant loss of generality.

[1] Nilesh Tripuraneni, Ben Adlam, and Jeffrey Pennington. Overparameterization improves robustness to covariate shift in high dimensions. In A. Beygelzimer, Y. Dauphin, P. Liang, and J. Wortman Vaughan (eds.), Advances in Neural Information Processing Systems, 2021. URL https:// openreview.net/forum?id=PxMfDdPnTfV.

[2] Donghwan Lee, Behrad Moniri, Xinmeng Huang, Edgar Dobriban, and Hamed Hassani. Demysti- fying disagreement-on-the-line in high dimensions, 2023. URL https://arxiv.org/abs/ 2301.13371.

[3] Simone Bombari, Shayan Kiyani, and Marco Mondelli. Beyond the universal law of robustness: Sharper laws for random features and neural tangent kernels, 2023. URL https://arxiv. org/abs/2302.01629.

Concerns Regarding One-Step Learning: Theoretical Justification of the One-Step Learning Assumption and Discussion of the Extensibility of Our Analysis

As you pointed out, the assumption that a neural network learns in a single step may seem restrictive. However, as mentioned by Ba et al., 2022, numerous studies on the early phase of training in deep learning ([4,5,6,7]) suggest that neural networks are significantly influenced by their initial learning stages. Moreover, it has been theoretically established that a sufficiently large single gradient step can sharply reduce the training loss ([8]), enable the learning of task-relevant features ([9,10]), and move the network away from its initialization ([11]). Based on these findings, we argue that analyzing feature learning under the one-step assumption provides meaningful insights into neural network behavior. At the same time, Dandi et al., 2023 analyzed the staircase property of neural networks in the context of multi-step two-layer neural networks with a large (non-constant) learning rate, similar to our setting. However, such an analysis requires delving into the dynamics of the gradual expansion of the feature subspace, which is beyond the scope of this paper due to volume constraints. We acknowledge the importance of this direction and plan to address it in future work.

[4] Aditya Golatkar, Alessandro Achille, and Stefano Soatto. Time matters in regularizing deep net- works: Weight decay and data augmentation affect early learning dynamics, matter little near convergence, 2019. URL https://arxiv.org/abs/1905.13277.

[5] Guillaume Leclerc and Aleksander Madry. The two regimes of deep network training, 2020. URL https://arxiv.org/abs/2002.10376.

[6] Scott Pesme, Loucas Pillaud-Vivien, and Nicolas Flammarion. Implicit bias of sgd for diagonal linear networks: a provable benefit of stochasticity, 2021. URL https://arxiv.org/abs/ 2106.09524.

[7] Stanislav Fort, Gintare Karolina Dziugaite, Mansheej Paul, Sepideh Kharaghani, Daniel M. Roy, and Surya Ganguli. Deep learning versus kernel learning: an empirical study of loss landscape geometry and the time evolution of the neural tangent kernel, 2020. URL https://arxiv. org/abs/2010.15110.

[8] Niladri S. Chatterji, Philip M. Long, and Peter L. Bartlett. When does gradient descent with logistic loss find interpolating two-layer networks?, 2021. URL https://arxiv.org/abs/2012. 02409.

[9] Amit Daniely and Eran Malach. Learning parities with neural networks, 2020. URL https: //arxiv.org/abs/2002.07400.

[10] Spencer Frei, Niladri S. Chatterji, and Peter L. Bartlett. Random feature amplification: Feature learning and generalization in neural networks, 2023. URL https://arxiv.org/abs/ 2202.07626.

[11] Karl Hajjar, L´ena¨ıc Chizat, and Christophe Giraud. Training integrable parameterizations of deep neural networks in the infinite-width limit, 2021. URL https://arxiv.org/abs/2110. 15596.

Concerns Regarding Classification Tasks: Extend Classification Settings to Regression Settings.

We chose a classification setup to analyze network learning in settings where classes are distinct. However, our analysis is not limited to classification tasks alone. To address your concerns, we have prepared additional clarification. Inspired by works like Ba et al., 2022, we will incorporate a setting that reflects an regression form to demonstrate that our proof techniques can straightforwardly extend to scenarios involving regression setting. This straightforward adaptability is possible because our analysis applies to any loss or model that satisfies the condition of Proposition 3.1 and Lemma 3.2 in the main text. We argue that this is a key aspect showcasing the extensibility of our study, and we will include this discussion in the Appendix.

Under all the assumptions stated in our paper, we define a new random variable $a \sim \mathcal{N}(0, \frac{1}{N} I)$, and modify the assumptions from Ba et al. 2022 for regression by replacing the centered Gaussian assumption with a non-centered sub-Gaussian assumption, leading to the following problem setup:

In the above setup, we define the loss as follows:

$\theta = \{ W, a\}$, $L(X, y; \theta) = \frac{1}{2n} ||y - \sigma(XW)a||^2$

$G = - \frac{\partial L}{\partial W} = -\frac{1}{ n} \Bigg[X^\top \bigg[ \big(\frac{1}{\sqrt{N}} (\frac{1}{\sqrt{N}} \sigma(XW) a - y) a^\top \big) \odot \sigma'(XW) \bigg] \Bigg]$

In this case, if we define $\alpha = a$ as opposed to the main text, $\alpha$ becomes a Gaussian with zero mean and variance halved, so it follows the same bound structure.

Specifically, based on the sub-Gaussian bound results in Section D, $\mathbb{A}$ can be bounded using the second equation of Lemma 14(i) from Ba et al. 2022 and the fact that $\mathbb{A}$ is rank-1, so $\lVert A \rVert = \lVert A\rVert_F$. For $\mathbb{B}$, we can remove $||a_2||_\infty$ from equation 20 in our Lemma E.1 proof. For $\mathbb{C}$, by removing $a_2 a_2^T$ from equation 29, we obtain the same bounds as the previous results.

In conclusion, these three bounds satisfy our Proposition 3.1 under the same conditions, and the same conclusion holds even outside of classification tasks. In this case, $\beta$ is defined as in the previous setting.

To explain in the context of previous research, the original paper on Local Elasticity (He & Su, 2019) also analyzes the difference between linear classifiers and neural networks in the setting of a "binary classifier." We were greatly motivated by He & Su, 2019 in formulating our analysis setting, and we would like to clarify that our classification setup assumption is simply the same as the original setting and not an arbitrary scope limitation.

Furthermore, we would like to emphasize again that one of the key contributions of our work is demonstrating that, based on our assumptions, the class-input similarity in He & Su, 2019 clearly operates as $|\beta^\top \mu|$. This is also a significant insight, as it highlights a limitation in the work of Dan et al., 2023, who extended LE to regression tasks. While their paper shows that, under certain assumptions, the LE-based similarity metric proposed by He & Su, 2019 can be defined for regression problems, they did not theoretically analyze what factors contribute to increasing elasticity, as seen in Hypothesis 1.1(a).

Thank you for your comments, we reply your questions and concerns

1. Experiments are limited to simulation data only. It would be very helpful to verify if the conclusion still holds when going beyond the simple two-layer neural network settings. -> We only substitute two-layer networks to pretrained ResNet18 from our theoretic configurations and verify the equivalent results in additional experiments.

2. Why do Theorem 5.3 and 5.4 both appear in high probability formulation? The equation for the expectation is deterministic on both sides, what is the exact randomness here?

->Apologies, we are already using the approximated $F_l$ in this theorem, so it is not a probabilistic technique. We have removed it. Additionally, we have improved the theoretical techniques, so I would appreciate it if you could review it again.

3. Section 4 seems to be purely technical lemmas and has no direct implication on the main conclusion of the paper, therefore I suggest the authors defer it to the appendix. -> Except for the parts necessary to understand the final theorem, unnecessary part has been moved to the appendix.

4. Is it essential to use two-layer neural networks to conclude the relationship between $\beta^\top \mu$ and the alignment and local elasticity? The conclusion is rather simple so I suspect onc may reach the same conclusion even in the linear model.

-> We studied on the occurrence of alignment and LE in the context of classification tasks using a two-layer network. This two-layer network is particularly relevant to the practice of using feature extractors in practical settings, especially when features are obtained from a network with a removed task head. While two-layer analysis is not essential, the intention behind this approach exists as modeling feature extraction, and since many previous studies have used the same Conjugate Kernel setting, we can also explain that the CK was used in terms of comparison and extension.

12. In the discussion of the Recall @ 1 performance, the authors say that "the second phase provides a dataset which effectively supports training for retrieval tasks and is explainable by our theory": could the author clarify this point? From the plots, it looks like the Recall @ 1 performance is not obviously correlated with the Alignment results.

To avoid confusion, we added a new dataset, which does not change the L2 distance between the two training classes, which is different to original dataset. This allowed us to obtain experimental results that eliminate the need for phase analysis in the recall@1 discussion within the two-layer setup. We found that as $|\beta^\top \mu|$ increases, the recall@1 performance after training improves, while the recall@1 performance based on the initialized features remains unchanged regardless of $|\beta^\top \mu|$. Based on this, we revised the explanation in the main text. Additionally, we conducted practical metric learning experiments using ResNet50 as the backbone and confirmed the positive correlation between recall@1 and alignment.

13. what is the role of \psi_1 and \psi_2 ? -> This is a constant needed to define the proportional regime. The proportional regime is necessary to bound the gradient decomposition. The value itself is not important for the purposes of this paper.

Thank you for the extensive comments, which have greatly improved the paper. Below are the revisions made based on your feedback. We would appreciate it if you could evaluate it once again.

1. Is the orthogonal decomposition required for at line 167 the one explicitly stated soon after in terms of Hermite polynomials?

-> Yes, that is correct. To minimize misunderstandings, we have consistently expressed all the expressions using Hermitian decomposition. Orthogonal decomposition can be considered as applying Hermitian decomposition only to the linear components, while treating the rest as a single function.

2. In line 184, the authors refer to an extension that might open "new avenues for analysis of deep representation of classification". Are they referring to the sub-gaussian assumption? Is it not a very common assumption for asymptotic analysis in particular when a Gaussian equivalence principle is applied or proved?

-> The sub-Gaussian assumption on the data, to the best of our knowledge, is introduced here for the first time in the context of classification problems. Existing works analyzing two-layer models, such as Ba et al. (2022) and Moniri et al. (2024), are based on the standard Gaussian data assumption. We extend this by highlighting the importance of the sub-Gaussian class, which, owing to its finite support, can represent multiple separable data distributions.

3.Fig 2's relation with the product similarity to is not clear and its presence here.

-> I have taken your advice and moved the figure to Section 6. It has helped in separating the specific examples from the theoretical assumptions. Thank you for your valuable suggestion.

4. At line 202, it seems that the label vector y of each datapoint is a row or column vector depending on the class (?!).... -> It should be understood that $\beta$ is defined as X^Ty_1. I have made this correction in the main text as well. Additionally, in order to set up the neural network for classification with the MSE loss, we made $y$ as symmetric. Therefore, any form of $y$ is acceptable. We have added a discussion on this point after the main theorem.

5. The reference to the training algorithm is not clear at line 204 as the loss used has not yet been specified. -> I have clarified the loss function part. Apologies for the confusion. I have removed the term f_NN(x)and revised the loss to accept weight and train data.

6. Isn't Lemma 5.1 a special case of Theorem 5.2? You are correct in your observation. We have rephrased Lemma 5.1 as a corollary and moved both Lemma 5.1 and Theorem 5.2, which are only used in the intermediate steps of the theoretical proof, to the appendix in order to add experiments and discussions.

7. In this sense, I would have expected the theory or experiments to be simplified by some concentration result, as the one obtained in the proportional regime for classification tasks, e.g., by Mignacco et al. (2020) or Loureiro et al. (2021) for the case of convex losses. Could the authors comment on this?

-> We do not analyze the concentration in this work, but will discuss it in further works. Besides, we averaged Alignment and LE with 30 different seeds in two-layer networks and every experiment show consistent result.

8. What do the authors mean by "Considering randomness, we assume that the sole term C_jk represents the Alignment"?

-> The revised theorem, which now also take expectation over the parameter. Even in this revision, we demonstrates that $\beta^\top\mu$ is the only term dependent to $\beta$ and $\mu$ as original statement, deriving the equivalent result to our previous claim.

9. Questions about truncated Gaussian (The appendix is written very poorly and, in some points, looks more like a collection of sketchy notes, rich in typos and poor in comments. In some points, the text is just confusing. For example, in line 890 the authors write ... )

-> We intended to emphasize that the truncated Gaussian dataset as a example satisfies our theoretic assumption on Sub-Gaussian condition, which only used for experiments through Lemma D.1. This dataset follows the truncated Gaussian in element-wise. We will details it in the revised main paper to clarify the confusion. Besides, we will greatly appreciate to the reviewer if they let us know the “usual definition” of truncated-Gaussians.

10. What is the role played by the separability condition on the training dataset in the theoretical analysis? -> We have confirmed that the assumption of separable sets is not necessary. We will remove it.

11. It is not clear to me if Fig 3 and Fig 4 represent the results of the experiments only or if any quantitative comparison with the theoretical prediction is given. Could the authors clarify this point?

->We will rewrite it more clearly. We have only presented the experimental results.

Thank you for your valuable comments. we reply your questions and concerns.

The paper heavily relies on techniques already proven in other works, this fact induces some concern on the clear elements of novelty of this submission. -> Though our work depends on two-layer networks analysis suggested by Ba et al. (2022), it still has a novelty on expanding the problem domain into classification and investigating the emergence condition of alignment and LE theoretically.

At a proof technique level, we employ the approximation methods by Ba et al. (2022) and Moniri et al. (2024) but with new assumptions applicable to a classification problem. Especially, we expand Gaussian distribution assumption on data into Sub-Gaussian and substitute Lipschitz teacher function into binary class labels. Such differences on conditions derive the general statements to previous works, but the definitions and the proof between statements require novel delineations. For instance, the definition of $\beta$ exploits the opposite signs of $y_1$ and $y_2$ to compose different gradient signals in the two-label classification, and we utilize SVD to acquire the norm bound of arbitrary labels without Lipschitz assumption.

For Figure 1

We reflect your comments by designating the origin on Figure 1 and add explanation about the train datasets and evaluation data “a, b, c, d.” The second layer weights are omitted from the Figure 1, which do not engage with the gradient analysis.

For definition of CK We replace “CK” in Line 76 to “neural network feature” and specify “Conjugate Kernel” in Line 82 instead of its abbreviation.

Organization of paper As you suggested, we supplement additional experiments to verify the theoretic relation between $|\beta^\top \mu|$, Alignment, and LE in two-layer networks. We also conduct the analysis on practical neural networks such as ResNet with the mitigated assumption. The lemmas in Page 6 and 8 except essential ones to understand the main paper are moved to appendices.

discussion of Question: How crucial is the learning rate assumption scaling?

Our analysis follows the framework by Moniri et al. (2024) which determines the interval of learning rate by parameter $l$, and it is related to the degree of Hermite expansion in our work. Both Alignment and LE are expanded by polynomials of $|\beta^\top \mu|$, and the max degree of the polynomial is determined by the aforementioned $l$. Thus the learning rate grows as the $l$ grows by the framework, and the measurements of LE and Alignment is representable in the polynomials with the bigger degree. It matches the intuition that the loss is minimized more as the feature is learned with the larger learning rate, resulting features are updated and clustered strongly.

The second layer weight $a_c$ Unlike previous studies requiring a single output layer for regression, two classifier heads compose the two-class classifier networks in our work.

Thank you

Thank you so much for taking the time to carefully reconsider your evaluation and for your thoughtful feedback. We truly appreciate your engagement with the paper. In response to the concerns raised, we have prepared thorough answers tailored to each reviewer’s comments, and we are currently awaiting responses from the others. As noted by ICLR, we are fortunate to have an extended discussion period, which provides us with an excellent opportunity to further address any remaining questions or concerns you may have. Over the coming period, We are fully committed to clarifying any points of uncertainty and ensuring that our paper meets your expectations. We sincerely hope that this process will not only lead to a more favorable evaluation from you, but also allow us to further enhance the manuscript. Given the extended timeframe, We would greatly appreciate it if you could share any specific issues or concerns you still have regarding the paper—separate from the feedback of other reviewers. This will allow us to provide you with the clearest possible explanations and make further improvements, with the goal of increasing your satisfaction with our work. Once again, We want to express our deep gratitude for your valuable contribution to the paper. We look forward to continuing this discussion and developing this paper from your insight. Warm regards,