Early learning of the optimal constant solution in neural networks and humans

… Section 6 makes the (extremely strong) claim that such reversion can equivalently arise from input correlations.

While it might appear like a strong claim we do believe that it is true that fairly typical correlations in the input data gives rise to early learning of the optimal constant solution. For nonlinear networks, we conduct experiments without bias terms in the architecture (Fig. 5 of the revised draft), so that the effect of data is clearly isolated in this case.
However, the initial submission did not contain a figure which examines a case in deep linear networks without bias term, but with input correlations which do exhibit the $1\_N$ eigenvector. We have added a solvable case where these conditions hold true and have added it to the paper in Fig. 17 in Appendix A.11. In the bottom row of the figure, we can see that the first SVD mode is indeed exactly equivalent to the OCS mode, i.e. $\\bar{\\mathbf{y}} \\bar{\\mathbf{x}}^T$. The right panel highlights how the network is driven towards the OCS up until the time-point when the second effective singular value $a\_2(t)$ (which is quite close in time) is learned. We hope that the additional results in linear networks makes our claim about the role of input correlations in OCS learning more precise.

If it is indeed true that very general covariate structure can give rise to this behavior, then why have you spent most of the paper reasoning about the contrived and simple setting of linear networks with bias?

A bias towards the OCS is observed over the course of learning in human and artificial learners. Prior work by Saxe et al. (2013) has established linear networks as a powerful but theoretically transparent model to study such learning dynamics. Its power stems from the fact that it is agnostic about the specifics of the data: instead, it makes predictions about the learning dynamics solely in terms of the broad statistical features, in our case the leading eigenmode of the data correlation structure.
The inclusion of bias terms presents the minimal modification to observe OCS learning. Only after developing the theory for the simple case of bias terms in Section 3, it is revealed that its effect formally resembles a change in input correlation. This theoretical insight that the linear network approach enabled is then expanded on in the second part of the paper.

Phrases like "might hence be ubiquitous beyond these datasets" arouse suspicion. Also, how can I believe that early-time learning of OCS is "universal" when it vanishes in even linear networks if you ablate the biases?

We now realise how our wording can suggest that OCS learning will necessarily be observed in every learning setting, which we do not believe to be true. As you correctly point out, the controlled linear network setting exactly identifies such a counter-example: If both bias and input correlations are completely removed, our theory predicts that OCS learning indeed vanishes.
However, we believe that this is not the usual case in real-world learning settings. Instead, architectural bias and data correlations are mild assumptions. As far as we know, biases are a common way to shift data and the default in deep learning libraries.
As for data structure, only the statistics of the leading eigenmode are relevant for OCS learning. We present a theoretical argument in Apx. A.5.6 why a weak symmetry in the data will already satisfy this premise on the eigenmode. We provide evidence in Fig. 6 (of the revised paper) that such correlations are indeed approximately present in the image datasets we tested.

I put very little weight on experiments that work with MNIST and CIFAR (...)

We share the view that the core contribution of our work lies in developing a theoretical model of OCS learning. In the response above, we have argued that our theory shows that data correlation and architecture are robust factors to induce OCS learning, but strictly applies to linear networks only. To be considered a candidate to explain observations in nonlinear learners, we needed to test and ablate these factors also in the nonlinear setting. We agree that such evidence can only be suggestive, and have made this limitation more clear in the discussion section. We hence cannot conclude that OCS learning is inevitably present, as the premises we identify – although plausibly true in many cases – are not necessarily fulfilled.
While we do ablations for artificial learners, for humans we cannot claim what is the mechanism behind OCS learning, but only point to candidate factors. We had addressed this in our discussion: “ $...$ the ambiguity between architecture and data in driving the OCS does not allow us to determine the underlying mechanism in human learners.” Still, we aim to disentangle this question in future experiments via the systematic variation of input correlations between stimuli presented to human learners which could point to the factors we identified, or possibly rule them out altogether.

You remark that "characteristic learning signatures" arising from early time learning of the OCS involve a sort of step-wise learning structure. While I'm aware a variety of toy synthetic settings studied by theorists exhibit this behavior (see $7$ , for instance), no non-trivial data (ImageNet, C4) gives rise to dynamics that look like this (learning curves smoothly go to low loss rather than being step-wise). So certainly the word "universal" does not describe these types of learning curves that are "characteristic" of early-time learning of OCS.

We agree that our results do not imply that early, step-wise OCS learning will necessarily be observed in any dataset. While the step-wise structure of early bias learning is reported in some cases $1,2$ , it may not be visible in the loss curve in all cases: We observe this in the shallow linear setting in Fig. 2B, and also in our simulations with CelebA seen in Figure 15 in Appendix A.9. We were intrigued by your point and added the loss to the Figure 15 for completeness. There, although the loss curve is not stepwise, the output still is transiently driven towards the OCS, which is the core of our claim. While such evidence is suggestive of our claims being more broadly applicable, we still see the key contribution of our work to understand elementary components of learning, such as in the simple settings we consider.

$1$ Karpathy, A. (2019) A Recipe for Training Neural Networks.

$2$ Ye, H. J., Zhan, D. C., & Chao, W. L. (2021). Procrustean training for imbalanced deep learning. ICCV.

Thank you very much for your thorough review and the detailed and constructive feedback provided. We are encouraged by your assessment that our work spans a rare breadth from ML theory and experiments to human behaviour. We are also glad that you found the bridge to human experiments we develop to be a particular strength of our paper.

Weaknesses

The paper is quite confusing in terms of how it is laid out, and what the scope of which claims are. In my opinion, it tries do a bit too much. For instance, tensor trickery related to the NTK feels out of place/a low order concern when you didn't even take the time to define the hierarchical learning task clearly. I had to look through the Appendix as well as (Saxe, 2019) to get a sense of the task. Some specific weaknesses:

Thank you for this valuable information on the readability of our work. We now clearly realise that a reference to the previous literature alone is not enough for a self-contained presentation of the task and we have updated the PDF with a dedicated section on the hierarchical task. In an effort to improve linearity we have also now moved the discussion of the NTK to the Appendix.

It's not clear the scope of the claims being made. Either you are 1) remarking that the learning curves between linear, nonlinear networks, and humans on a certain class of task look similar but the mechanistic reasons are not well understood, and this is interesting and warrants further study or 2) making a claim that ALL neural networks are marked by this early-time behavior, and they are driven by mechanisms you identify. I would be more sympathetic to the contributions of this paper if you were making claim (1), but since you use the words "universal" and "mechanistic" in your abstract, I take it you are making claim (2). If this is the main claim, then the evidence presented is inadequate as justification. This is for the reasons that follow.

While we agree that we don’t have evidence to make claim (2) in a general sense, we do believe that our findings extend beyond (1) by manipulating the same factors which drive OCS learning across several artificial learners as well as providing general theoretical arguments. We now include an additional setting which directly demonstrates how OCS learning in linear networks can be purely driven by input correlations. In the answers below we detail these points. We have rephrased relevant sections in our paper to better reflect the precise nature of our claims.

We thank you for your thoughtful comments and constructive feedback. We are delighted that you found our work a significant contribution to our understanding of particular types of simplicity biases. We are also encouraged by the fact that you found our connections to human learning and nonlinear models interesting and unique.

Weaknesses

It would be helpful to include a brief description of the hierarchical learning task in addition to Fig 1 for those unfamiliar with the setup

Thank you for pointing this out. We agree with you on this point and have added a clear exposition of the hierarchical learning task to the main text of the paper that makes explicit reference to Figure 1. We hope that this change will increase the presentation of the paper.

Questions

Does the theoretical argument nicely extend to linear networks with more than two layers? In the multi-layer case, would a bias term present in an intermediate layer result in a bias towards OCS?

While the two-layer case captures the essential factors for OCS learning, our analysis generalises to the case of multiple layers in a notationally involved, but conceptually straightforward manner. The starting point to see this is the NTK in Proposition 3, which readily generalises to the multi-layer case $1$ . The general pattern of Proposition 3 extends: Bias terms in all layers affect OCS learning. The closer to the output a bias is located in the architecture, the more immediate it will affect the output statistics. Adding more layers will attenuate the effect of bias terms in earlier layers towards the output through downscaling by the initialization scale of deeper layers (the terms \\sigma\_{\\mathbf{W}^\\ell}^2 in Proposition 3).

$1$ Shi, J., Shea-Brown, E. and Buice, M., 2022. Learning dynamics of deep linear networks with multiple pathways. Advances in neural information processing systems, 35, pp.34064-34076.

In what ways do the linear network solutions extend Saxe et al. (2014) precisely? Through the inclusion of bias terms, and a relaxed assumption (Proposition 1)?

We here discuss how we differ from Saxe et al. (2014) in particular. For a discussion of our general contribution, please refer to the global response above.
Indeed, we needed to generalise their assumptions to make progress on our results. Saxe et al. had mostly considered whitened inputs, as have several subsequent works like Braun et al. (2022), Kunin et al. (2024), or Shi et al. (2022). Proposition 1 answers this question and is a simple condition that allows us to derive learning dynamics for several interesting datasets which contain input correlations.

Second, to the best of our knowledge, we for the first time extend predictions made by the formalism directly to human learners, where prior work had conjectured similarities $1, 2$ . We were able to directly demonstrate some behavioural similarities in terms of OCS learning. We argue in Appendix A.5.6 that this is likely due to symmetries in the data that make the theory’s premises approximately apply in many settings.

References

$1$ Rogers, T.T. and McClelland, J.L., 2005. A parallel distributed processing approach to semantic cognition: Applications to conceptual development. In Building object categories in developmental time (pp. 353-406). Psychology Press.

$2$ Saxe, A.M., McClelland, J.L. and Ganguli, S., 2019. A mathematical theory of semantic development in deep neural networks. Proceedings of the National Academy of Sciences, 116(23), pp.11537-11546.

We thank you for your thoughtful comments and feedback. We are glad that you found our work to be a useful and theoretically sound contribution to our understanding of simplicity biases in initial training phases. Furthermore, we are delighted that our approach which bridges human experiments, ML experiments, and theory was received as a strength of our work.

Weaknesses

The most significant weakness of this work is the lack of a definition and clear description of the studied “hierarchical task”. While the authors describe it as a "hierarchical category learning task" with hierarchical levels, there’s no clear explanation of the precise structure or production rules, such as the nature of the relationships between input categories and output labels. Without knowledge of the distribution of X,Y , and precise definitions of the hierarchical levels, I find it practically impossible to fully understand the study, particularly the empirical aspects.

The hierarchical task requires learning a mapping from one-hot, input vectors to output vectors that are depicted in Fig. 1B. Hereby each output vector is “three-hot”, i.e. the vector has three entries/labels. The hierarchical structure arises from the similarity between output vectors where some labels $y^m(\\mathbf{x}\_i)$ are more general and correspond to more than one input $\\mathbf{x}\_i$, while labels corresponding to the bottom of the hierarchy are specific to a single input vector $\\mathbf{x}\_i$. The task is motivated in the literature on semantic cognition and leverages the fact that semantic information is usually hierarchically structured $1, 2$ .

We now realise that a reference to the literature alone as in our initial submission is insufficient for a self-contained presentation of our approach. We have revised the paper and have added a dedicated section that introduces the task.

The link between the model of data and the OCS is also not discussed.

Our work identifies two complementary factors driving OCS learning: Bias terms in the architecture, and a specific property of correlations in the input data. The latter factor describes the interaction between data structure and OCS learning. We argue in Appendix A.5.6 on theoretical grounds that the necessary correlation property (the presence of the $\\mathbf{1}\_N$ eigenvector) requires mild symmetry that is satisfied for the hierarchical task in particular, but extends to a more general class of datasets. This argument is empirically supported by Figures 5 and 6 (in the revised PDF).

The continuous correct-rejection scores, which are crucially used to infer OCS learning, are not defined within the main text.

Thank you for highlighting this. We have now moved the definition of the metric and the associated intuition from the appendix into the main text under Section 3 Setup, which should help to increase the interpretability of results.

The figures and the corresponding captions are poor and challenging to interpret. Figure 1 contains 6 panels that lack sufficient standalone explanations. In Figure 2, the reader does not know what different colors represent.

We have updated all figure captions with the intention of making them self-contained. In particular, we have updated the caption of Figure 1 with more extensive explanations in our revised paper, we hope that the revision is sufficiently clear for standalone interpretation.

The different colours in Figure 2 represent the asymptotic magnitude of the effective singular values $a\_i(t)$ and $b\_i(t)$ for the deep and shallow linear networks respectively. Darker colours represent higher asymptotic values. We have addressed this concern in a revised draft that contains this description in the figure caption.

Proposition 3 must assume some initialization statistics for the weights that are not specified.

Thank you, we overlooked this point. We now state the initialisation statistics as an assumption in the revised paper.

Dear VeTy

Many apologies. It appears that our last upload did not go through. Thank you for pointing this out. The new pdf with all modifications should now be up.

Best wishes,

The authors

Thank you again for your feedback. We would like to emphasize that the author-reviewer discussion period ends soon. Based on your comments we have made substantial revisions to the paper and have provided a detailed response. Hence, we kindly request you to please consider a reevaluation of your score if your concerns are addressed. Please let us know if you have any additional questions or concerns!

Thank you for your review and thoughtful comments. We are delighted that you thought of our results as intriguing and we are encouraged by your assessment that our work is of interest for researchers beyond the machine learning community.

Weaknesses

My major concern about the paper is about its clarity:

- Plots are not explained well, e.g., what does "output activation top" represent in Figure 3?

Thank you for highlighting this point. We understand that our figure captions were not sufficiently self-contained and we have now revised all figure captions to include more detail. For the particular case of Figure 3, “Output activation top” denotes the activation of a single output unit of the model corresponding to the top level of the output hierarchy. We plot the response to all inputs $\\mathbf{x}\_i$. While we do explain the figure in the main text in Section 4.1, Indifference, we have clarified this in the figure caption.

- Even after reading Appendix A.4, I'm still a bit confused about the definition of TNR rate, in equation 7, the level index appears on the left side of the equation but not on the right, should the term in right hand side be to denote the target label for a specific level? Besides, this definition seems important and should potentially be moved to the main paper.

We can see where the confusion comes from. The metric is slightly unusual but was chosen as it allowed us to directly compare network and human responses for signatures of OCS learning. Key is the mapping from level $k$ to the relevant start and end indices in the output vector. We have made this mapping more clear in a revised version of the metric. Following your suggestion, we have now also moved this metric to the main text of the paper.

- The hierarchical MNIST dataset is only explained in the Appendix, and it's not quite clear how the hierarchical label structure is established for the dataset.

We develop our theory starting from a task which considers perfect within-class similarity, i.e. each output label is only associated to a single (one-hot) input. For the experiments with humans and CNNs, we consider the more realistic case where there is natural variation in each input class.
For the "hierarchical MNIST" task we used the ten digit classes provided by MNIST and then sampled 8 classes randomly. For each image in each class, we then replaced the default one-hot label corresponding to each class $i$ with the corresponding hierarchical, “three-hot” label $\\mathbf{y}\_i$ seen in Fig. 1B. E.g., all images corresponding to MNIST digit “1” might be assigned some randomly chosen “three-hot” output vector $\\mathbf{y}\_i$. We added this description to the revision of the manuscript in Section 5, Setup.

We would like to emphasis that the discussion period ends quite soon. We have made major revisions based on your comments and provided detailed responses to your questions. Hence, we kindly request you to please check the revisions and rebuttal and consider adjusting your score if your concerns are addressed. Please let us know if you have any additional questions!

Scope of claims (Rev. oCMD)

Reviewer oCMD in particular expressed concerns about the scope and support of our claims as a primary weakness of our work. We appreciate that being precise here is crucial and would like to delineate the exact scope of our claims. We have revised our language throughout in the updated PDF to clarify the exact scope of our claims.

What we do not claim:
We do not claim that OCS learning will inevitably be observed regardless of possible architectures and datasets. For example, in the controlled setting of linear networks, we identify cases in which OCS learning will not occur: linear networks will not display OCS learning in the absence of input correlations and architectural bias terms. In nonlinear networks we demonstrate the same phenomenon empirically, i.e. the absence of OCS learning in absence of bias terms and input correlations.

What we do claim:

1. Linear networks
   We do claim that OCS learning in linear networks emerges via two distinct factors entering the output dynamics in the same way:
   a. Bias terms in the model architecture
   b. Particular input correlations, i.e. whenever $X^TX$ contains a constant eigenvector.
   To clearly highlight point b we have now isolated a setting with solvable learning dynamics in which only input correlations drive OCS learning and added it Appendix A.11 of the revised paper.

2. Nonlinear networks
   We observe that OCS learning occurs in nonlinear networks, and that the same factors as in linear networks causally control OCS learning in the settings we consider (see Figure 5 and Figure 6 of the revised paper). Two findings support these observations: First, the results from Kang et al. $1$ show that for various network architectures and datasets, networks revert to the OCS in generalisation settings. This reversion implies that the OCS is a component of the learned network function.
   Second, our theory does not depend on details of the dataset, rather, it only relies on a very coarse property of the dataset correlation statistics, its leading eigenmode. We argue in Appendix A.5.6 that if mild symmetries are present, this will hold true for many datasets.
   Together, this empirical and theoretical evidence suggests that OCS learning and the factors behind it are promising candidates to explain observations beyond linear networks.

3. Human learners
   For human learners, we state in our discussion that we cannot make any specific mechanistic claim beyond the empirical observation that they do display OCS learning. However, we think that future experiments which leverage the recording of neural data or the targeted manipulation of input correlations in experimental stimuli could disentangle between the driving factors we identified, or possibly rule them out altogether.

References

$1$ Kang, Katie, et al. Deep neural networks tend to extrapolate predictably. arXiv preprint arXiv:2310.00873 (2023).