Feature-Learning Networks Are Consistent Across Widths At Realistic Scales

Summary

This paper studies the convergence of key quantities (like test predictions, feature kernels) in Neural Networks as width increases, in feature-learning regimes. The paper is largely an empirical study, and presents width-convergence results on both image and language tasks. The paper also uses ensembling to consider the role of finite-width variance from randomness in initialisation, and finishes with a toy model that looks into the causes behind the bias in finite-width models compared to their infinite width counterparts. **Update: after clarifications in the rebuttal and additional experiments, I am raising my score to 6**

Strengths

1. The question that this paper seeks to address, i.e. understanding finite-width NNs, is important and open. 2. The paper has some interesting conclusions e.g. the convergence of feature-learnt test predictions at realistic widths (though see weaknesses too for some concerns), or the breakdown of finite-width bias in terms of eigenfunctions rather than eigenvalues. 3. The paper is well-written and clear.

Weaknesses

1. My main concern is that some of the results are not convincing/not matched to the claims of the paper. For example, for the qualitiative results on ImageNet experiments on Figure 2b and 3b, it isn't clear to me that these curves have converged in width (say comparing N=256 or 512). For the quantitative equivalent to 3b, Figure 3e, it's a bit surprising that the initial y-intercepts on the true logits RMSE with width 512 are so large, even at width N=256 (and seemingly not consistent with the claim that these predictions have converged at these evaluated widths). This is especially confusing given that it seems one of the claims in Figure 2 is that ImageNet results have converged for earlier training times compared to later times, and suggests that the reason why the True logit RMSE to width 512 is small at later epochs is simply because different networks are converging to the same correct prediction, rather than any width-limit occuring. Similar conclusions (larger RMSE at earlier training steps) appear to be the case in Figure 3f for WikiText-103. 2. Likewise, I find the fact that we are often judging convergence in width off of 3-5 widths (e.g. Figure 1a-c) to be a bit unconvincing. I realise computational constraints are likely at play here, but it would be very good for the contribution of the paper to have more widths presented. Moreover, the argument that "a quantity's deviations between successive widths decrease as widths are increased" (line 54) is evidence for consistency is a bit weak to me: the harmonic series blows up to infinity despite satisfying this property. 3. Going back to the motivation (e.g. line 86), the extent to which establishing convergence to infinite-width feature learning limits helps the community in understanding finite-width NNs might be quite modest, as such limits are very difficult to understand (significantly harder to understand compared to say kernel regimes), and also computationally prohibitive (e.g. Yang and Hu 2020, or Bordelon and Pehlevan 2022)

Questions

Please respond to be weaknesses above, and also answer: 1. Do you set the output weights to 0 in order to reduce the variance from function initialisation in your results? If not, then it might be worth looking at e.g. for the ensembling results. 2. Do the experiments use identical mini-batch orderings across different widths? Also, please correct the following typos: 1. 'we put all of our in...' (line 37) 2. 'the training lost' (figure 2 caption)

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

2 fair

Presentation

3 good

Contribution

2 fair

Limitations

There isn't much discussion of limitations within the submission.

Summary

The authors study the so-called “feature learning” parameterization of neural networks that is given by the maximal update parameterization (muP). This is in contrast to the NTK parameterization for infinite-widths, which leads to the “lazy learning” regime which is claimed to not learn features, and also in contrast to the standard parameterization for finite-width networks, which the authors also claim (I believe in the Appendix) leads to no feature learning in the large-width limit. The central question of this paper is whether a finite-width (but sufficiently wide) network in the muP setting captures the dynamics of its infinite wide counterpart (which the authors empirically examine by proxy via the largest width network they are able to train, which I’ll refer to as the infinite-proxy). They consider this question in the online setting, where new data is presented in each batch, and the offline setting, where data is recycled over epochs. Throughout the paper, the authors study the following metrics: training loss, train / test error, logits with respect to the true class for a test sample, attention matrices, top eigenvalues, centered kernel alignment, preactivation neuron distributions, and kernel spectra. The central findings of the paper are: - Sufficiently wide networks begin to show similar loss, error, and true logit trajectories over the course of training, though this degrades over training on more difficult tasks (e.g. ImageNet vs. CIFAR tasks) - You can do a bias-variance decomposition of the finite-width networks by considering ensembles of predictors, and the paper claims that, empirically, sufficiently wide networks don’t suffer much from bias nor variance but narrow networks can suffer from both and that the bias is actually a larger negative effect than variance - In offline training, if the networks start to overfit then narrow and wide networks show that bias and variance measured here to not go to 0 (there is a claim that on test the bias and variance goes to 0 but on train it doesn’t, I will defer comments on this to later) EDIT: Thank you for your thorough response, I continue to be more on the accept side than reject and opt to keep my score.

Strengths

There are a lot of experiments and perspectives here to support the claims the authors are making and I don’t think I’ve seen a finite-width study of muP yet, myself. Experiments are also done with MLPs, ResNets, and Transformers, which is an extensive enough set of model families to convince readers that these phenomena are not unique to any one model class. I think the conception of each experiment is nice, and it is clear from the explanation of each plot what it is trying to convince me of. It is clear to me that after a certain width, the finite-width networks begin to act similarly, though there are still minor differences (more on this later). The ensemble experiments give rise to a sort of bias-variance decomposition, though I think this is not formally a bias-variance decomposition because the definition of bias here given is with respect to the infinitely-wide version, considered as the largest width network they can train, but I think the infinitely-wide version of the network is not necessarily the optimal predictor (e.g. E[y|x] in a regression setup), unless it is in this setting and I have missed the explanation why? However, the variance aligns to the definition of variance I am familiar with, and in general I still think this is an interesting definition in-and-of-itself, just not perhaps a formal bias-variance decomposition. I appreciate the thoroughness of everything and overall I find this to be a useful contribution with respect to studying muP networks, something I haven’t really seen before in such a practical setting. I can see through the empirical results that this parameterization does lead to a reasonably stable behavior when the network width is sufficiently large, and it is a useful result for those who want to study and understand the muP dynamics. On one hand, I think there is enough here that it will be of use to some people in the community and is worth accepting. On the other hand, I have many comments and questions and I think a second draft and resubmitting to another conference would make this much stronger. I am giving it a weak accept because everything is technically sound and interesting, but I hope the authors consider my comments below.

Weaknesses

I think it could be misleading to say that you are studying feature-learning networks in the sense that I believe standard parameterization also does lead to “feature learning”, though admittedly I don’t think feature learning itself is a well-defined term. I understand that muP is meant to indicate some non-trivial change in the feature kernel, and in the appendix you argue that standard-parameterization leads to no changes in the feature kernel when width goes -> infinity, doesn’t this still leave room for feature learning in SP in the finite-width (in particular, the narrow-width) regime? I’m not totally sure how you could assess at what point a SP model is doing feature learning or not, but I am open to hearing thoughts on this. Basically, if the point of the paper is to study realism, like realistic scales and realistic models + data then I think it’s too strong to assume that there is no feature learning in realistic SP models. As such, I think it could be more precise to say that you are studying muP models at realistic (finite) scales, and that there are potentially feature-learning dynamics that are not covered in this paper nor encompassed by the muP parameterization. Though I admit I don’t fully know all of the details of muP, so maybe this is just my lack of understanding. To be more clear, I think it is obvious once you read the paper that you are only studying muP models, so my prior comment is mostly about the title and abstract, which does not refer to muP and could make a potential reader believe that you are studying, somehow, all feature-learning networks. Maybe the paper could make this point about all feature-learning networks if you compared these muP plots to the same for a modern SP model over a few different widths, perhaps a model that is reported from some recent paper on these datasets (they are very well studied datasets so there will be lots of options to choose from). Is muP outperforming a state-of-the-art SP model? If not, can you claim something about whether there is feature learning or not in that SP model? I understand these questions are not the central focus of your paper, and so they are not affecting my decisions on accept / reject. It is mostly a comment that when the title and abstract says you are studying feature-learning networks this is what I immediately think about and it is something I am wondering as I read your paper, and perhaps having that thought in my head undercuts the value of your paper, to me, while I read it. I think all of the figures are far too small, especially the fonts. It is hard to read without zooming in on the computer, and on printing it is too small. Especially considering that you want to convince the reader of “consistent” behaviors over increasing width and we are meant to observe this by seeing that loss curves, etc become closer on the higher widths, larger and clearer pictures seem very important to me. I think the small figures undersell the evidence of the claim a bit, it’s not always clear to me that some curves are closer while others are not. In that sense, Figure 3(d-f) gives me the clearest evidence of the claims being true, while many of the rest I really have to stare for a bit and zoom in to see the curves more clearly and convince myself that it’s true. I also feel that there is so much going on here, it’s not necessarily a bad thing but it’s a bit overwhelming as a paper. The sections just seem like they’re going so quickly through results just stating this figure shows this, now this other figure shows something else, etc etc. I feel that I’m missing intuitions, explanations, contextualization in existing literature and more. If I were deeply involved in muP research then maybe all I need is a clearly stated set of results and a wide variety of experiments to cross-reference, but I think the paper could benefit from more exposition. Please see questions below.

Questions

How small of a difference on a metric like cross-entropy loss is needed in order to consider it a negligible difference, or “consistent” as defined in this paper? I imagine this is specific to dataset, model, and other choices. It would help if you gave some idea of this in the paper. You could compare the infinite-proxy network (the largest width network that is meant to be a proxy for the infinitely-wide network) to the expected regressor in a simple, synthetic setting (like mixture of gaussians or something like this) in order to see whether the bias of narrow width is like the true bias or it is some new bias-like term you’re proposing, but not exactly equivalent to the actual bias of the model. On individual model experiments and true logit experiments, are you averaging errors and losses over multiple different initializations when you plot? The ensemble experiment tells me something about averaging predictions and then computing error / loss, but I think since there is such a high variance with respect to random initialization of the network that the losses, and true logit RMSE and other measurements in most experiments should all be reported as an average over initializations (or test points when applicable), and I just haven’t seen anywhere in the paper or appendix that says that this is what you’re doing, though I may have missed it. I’m a bit confused about the claim in Section 3.2 that on train the bias/variance don’t go to zero and on test they do so this would indicate benign overfitting. I think it may come from the definition of bias-variance you have given, potentially not aligning to the notion of statistical consistency that would indicate benign overfitting. Could you talk a bit about why you think this is showing benign overfitting, and whether your bias-variance is actually the same as statistical consistency? In the first sentence of the second paragraph of Section 4, starting “Concretely, we see that although the eigenvalue spectrum of the …” I find it quite confusing to read, and I can’t help but feel there must be simpler language to get across the point you’re trying to make. Also what is “the task”? Is it fitting the data? This section reads a bit disconnected to me from the rest of the paper and introduces a lot of new terminology and ideas. While I follow the general point, it feels like you’re bringing in a whole second paper’s worth of topics and then trying to compress it all into one page, which makes it feel very rushed and somewhat out of place to me with respect to the prior experiments. I could be wrong, and maybe my fellow reviewers will disagree with my view on this. I think you’re arguing that the eigenfunctions and how the weight is placed on them is affected by narrow width more significantly than the spectrum itself, and that this would contribute to your notion of bias being large? To me, it feels that you’re throwing a lot of new ideas at the reader in the end of your paper, like “spread out target function power into slower modes” and “putting more of the task into smaller eigenmodes that take longer to be learned” and none of these ideas are motivated or explained anywhere in the paper? Maybe you are assuming that the reader is very familiar with what these phrases would mean but I’ve read a number of papers in this domain and while I can probably make some assumptions about what you’re trying to say I think it’s a lot to ask of the reader to just go with it, unless they are deeply familiar with this specific vernacular.

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

2 fair

Contribution

3 good

Limitations

N/A

Summary

The paper investigates the impact of width on the dynamics of feature-learning neural networks with both vision and language tasks. The empirical results demonstrate that during the early stage of training or during the entire training process for easy tasks, wide neural networks trained on online data not only exhibit identical loss curves but also display consistent point-wise test predictions. Nevertheless, for more challenging tasks, in the later stage of training, deviations across finite widths become more pronounced. The authors then decompose the deviation to bias and variance and show that the variance scales inversely with width. Furthermore, the authors demonstrate that ensembles of narrow networks can have inferior performances compared with a single wide network.

Strengths

1. The paper is well-organized and clearly presented. 2. All the experimental results are clearly motivated and effectively support the claims made in the paper. 3. The supplementary material is comprehensive, providing extensive details and codes for numerous experiments conducted in the study, enhancing reproducibility and transparency.

Weaknesses

he reviewer's major concern about the paper is the potential implications of the paper. While the empirical results presented in the study are well-documented, it remains challenging to discern how this work motivates or contributes to the broader research field. In the conclusion section, the authors assert that the findings "motivate the applicability of infinite-width feature-learning models in reasoning about large-scale models trained on real-world data." However, the empirical results within the paper indicate that deviations caused by varying widths amplify for more complex tasks or longer training time. Moreover, there is no apparent evidence suggesting that increasing width follows the law of diminishing marginal utility. Consequently, the reviewer struggles to grasp how the results presented in the paper effectively support the applicability of infinite-width models.

Questions

In section 3, the authors decompose the deviation into bias and variance, how did the authors quantify the two terms? It seems like bias is represented by the training loss, and variance is quantified by the improvement brought by ensemble?

Rating

5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

Yes.

Summary

This is a primarily empirical paper studying the agreement between real, finite-width neural networks and their limiting behavior under infinite width. Rather than using the neural tangent kernel as the descriptor of infinite-width behavior, the paper focuses on a different parameterization called feature-learning networks, which vary from standard-parameterization networks in that intermediate layer weights do not remain fixed even in the infinite-width limit. For these feature-learning networks trained in the online setting (without repeating training data) across a mixture of architectures and vision and language tasks, the paper demonstrates that large but finite-width networks tend to agree across a wide variety of metrics including overall loss, per-sample predictions, and the distribution of learned weights. However, if networks are too narrow, or are trained in the offline setting (when data is repeated during training) this agreement can break down. This gap between finite and infinite-width performance is then studied using the empirical NTK (at initialization) as well as the empirical after-kernel, and the results suggest that finite width induces deviation in the kernel eigenvectors (but not its eigenvalues) compared to the infinite-width limit.

Strengths

I really appreciate the thorough and systematic experimentation. Deep learning theory benefits from the ability to model infinite-width networks, but it’s important to know how well and under what circumstances such infinite-width limits translate into finite-width practical behaviors. In particular, I appreciate that this paper includes a mixture of convolutional and attention-based neural networks on a mixture of vision and language tasks, and studies finite and infinite-width agreement across loss, per-sample predictions, and network weights. It also finds some interesting and useful boundaries for when we can expect agreement to hold versus when finite-width networks behave differently, and takes steps to explain why.

Weaknesses

Most of the weaknesses I’ll describe should be viewed as suggestions to improve the exposition of the paper, and in some cases suggestions to make the experiments clearer or more compelling. I tried to separate suggestions/complaints into this section and clarification questions into the “questions” section, but there isn’t always a perfect distinction. - The abstract uses the term “feature-learning neural networks” without definition. Thought it’s defined early on in the introduction, many potential readers may be confused by the abstract before they see the definition in the introduction. - There are a decent number of typos and small grammatical mistakes, though they don’t impede comprehehension. For example, Figure 1 caption says “Loss curves…are nearly to identical at large widths.” Line 37 says “We put all of our in mu-parameterization”—similar missing-word type typos appear throughout the text. - Essentially all of the figures require the reader to zoom in to understand what is being shown. Figures should be fully interpretable without zooming in, e.g. for readers who use a printed copy. Axis labels, legends, tick labels, and line widths all need to be substantially larger. - Lines 122-125 define a bunch of notation that is (1) confusing, especially for readers unfamiliar with Transformers, and (2) is never referenced again. Instead, figure 2 uses “d” without any subscript, which is not defined. It would be preferable to define (only) the necessary notation, and push other details to the supplement. - Figure 3 is a bit confusing. It’s not clear if the first row is literally showing results from a single test point, or if it’s an average over test points. I’m also not sure why these figures only consider the correct logit value as a metric of agreement; I would have thought a more natural comparison to be some metric of model agreement across predictions, such as overlap in where mistakes are made (used in https://arxiv.org/pdf/1905.12580.pdf). - Figures 4c and 4f and 8c and 8f all plot 1-{some quantity} on the y axis, whereas in the textual description trends are described with respect to {some quantity}. This is needlessly confusing. - Line 157 describes figure 5a as growing “steadily to a final value that it then fluctuates around” but in the figure I don’t really see a stable final value with fluctuations; rather it seems the trend continues to increase even after fluctuations arise. This is a minor point, but was a bit confusing since it otherwise seems the main purpose of the plot is to visualize when the fluctuations begin. - Figure 7: Please label the x axis with epochs either instead of or ideally in addition to training steps. The purpose of the figure is largely to investigate the effects of reusing training data from one epoch to the next, so it’s important to see when data is first reused. For example, in subfigures (a) and (b) I suspect this happens around step 10000 when the curves diverge, but if this is the case I wonder why not include more training steps so that we can see the deviations more clearly? - Equation 1 seems to be describing the NTK. I wonder why the preceding paragraph that describes its derivation doesn’t mention NTK? As I read that paragraph I kept thinking “this sounds just like NTK” but I wasn’t sure why the definitions were present unless either it’s different in some way from NTK (in which case, please specify the difference directly) or it’s a refresher for readers who aren’t familiar with NTK (in which case, please say so to avoid ambiguity).

Questions

- Section 3 describes ensembling of finite-width models as a way to reduce variance, but shows that ensembled smaller models are still worse than a single larger model. How did you decide how many smaller models to ensemble? In particular, I wonder if there is some combinatorial flavor here (inspired by https://arxiv.org/abs/1803.03635) where an exponential number of smaller models is needed. - In Section 3.2, what motivates the use of these two datasets? In particular, I would think comparisons of the effect of label noise would be easier to see if using the same dataset with vs without label noise, rather than two different datasets where one of them happens to be noisy. - In Section 4, why use NTK and after-kernel? I presume this is because there are not similarly computable limits for feature-learning networks? In general I wonder why the paper uses large but finite models as a proxy for infinite-width networks, if feature-learning limits are well understood. An explanation of this reasoning and more detail on related work for what is known about the infinite-width limits of feature learning networks would be valuable.

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

Limitations of the experiments are not directly discussed, and the paper would benefit from such a discussion, perhaps in the conclusion. One limitation would be the lack of theoretical or direct modeling of the infinite-width limit, and the use of large-width networks as a proxy.

Summary

In this work, the authors conducted an empirical investigation into the consistency across widths of the training dynamics of neural network models in the maximum-update scaling. Specifically, they demonstrated through experiments on both vision and language tasks that many behaviors in the online learning dynamics are consistent at large widths, including the loss values, single predictions, feature kernels and large-learning-rate phenomena, and that ensembling improves the performance. Deviations from the large-width behavior are reported in the offline learning setting. Finally, the authors argued through heuristics and experiments that the finite-width bias is largely due to the deviation in the tangent kernel's eigenvectors.

Strengths

Given the number of theoretical results concerning over-parameterized neural networks and their infinite-width limits, the question of the finite-width deviation from the infinite-width limit is an important one, especially under the feature-learning regimes. To my knowledge, this is the first work to put forth an empirical investigation of the large-width consistency of neural networks in the maximum-update parameterization. The authors of the paper performed a wide range of numerical experiments on various tasks with different neural network models that cover several relevant aspects of the model training behavior.

Weaknesses

Note that for feature-learning neural networks in an alternative (mean-field) parameterization, both empirical and theoretical studies of similar purposes already exist, with examples including Nguyen (2019), Chen et al. (2020) and Geiger et al. (2020). More comments below regarding parameterizations. The authors argued based on Figures 2 and 7 that, as the widths of the model increase, the training performance of the single network converges in the online learning setting but not in the offline setting, which I don't find fully convincing. The y-axis uses a linear scale in Figure 2 and logarithmic scale in Figure 7, and hence the discrepancy of the losses across different widths in Figure 7 (which appears to be quite large late in training under the logarithmic scale) may not actually be much more significant than in Figure 2. In addition, it is also hard to rule out the possibility that the width being used is not yet large enough to approximate the large-width limit (after all, 512 is still small compared to the size of the datasets). The original proposal of the maximum-update parameterization does not include normalization layers, and given their role in preserving the scale of the hidden-layer features, one would expect their effect to interact with the initialization scale of the parameters. The authors only mentioned in Appendix A.2.1 that LayerNorm was used in the experiments on CIFAR-5m and did not comment on whether normalization layers are used in the other experiments. I think the authors should discuss in detail whether normalization layers were included in the other experiments and, if so, why the phenomena observed in this work can be truly associated with the maximum-update parameterization. References: 1. Nguyen, "Mean Field Limit of the Learning Dynamics of Multilayer Neural Networks", arXiv 2019 2. Chen et al., "A Dynamical Central Limit Theorem for Shallow Neural Networks", NeurIPS 2020 3. Geiger et al., "Scaling description of generalization with number of parameters in deep learning." Journal of Statistical Mechanics: Theory and Experiment, 2020

Questions

(Please see the comments above.)

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

2 fair

Presentation

2 fair

Contribution

3 good

Limitations

N/A

Thanks

Thanks for the rebuttal and new experiments, as well as for the clarifications. I do think the true logits RMSE plots in (3d-3f) are quite misleading in the original submission, and contradict some of the arguments that the paper is trying to make (i.e. feature learning nets converge in width better at earlier training times). I would potentially suggest something like plotting the non-normalised RMSE in the predicted probabilities (post softmax) instead, which would get rid of this dependence on the norms of the output logits which vary throughout training, at the expense of being invariant to additive scalars in logit space. I think Figure 3 of the attached pdf is more convincing of the arguments of the paper compared to the plots in the submission. I would be interested to see Figure 3 plotted at various points of training to see how the convergence in loss changes as the loss becomes better. I have a further question: in Figure 1 of the attached pdf how do you set the learning rates of the different networks? Thanks!