Paper: Understanding Sparse Feature Updates in Deep Networks using Iterative Linearisation (ICLR-2024)

Reviewer comment (Experiments):
The paper performs a very limited analysis of *iterative linearization*. The only contribution that is convincingly supported is the claim that a few $8-12$ updates are sufficient to achieve performance comparable to full training with $K=1$. However, in its current form, this statement does not go much beyond a simple expectation of how interpolation between $K=1$ and $K=\infty$ should look like. An example of a more insightful result would be the whole curve of dependence of a given performance metric on $K$ in the region of its most significant change (probably $K\lesssim 100$ given the provided results). While formally equivalent, the number of feature updates $T/K$ seems to be a more convenient variable than frequency $K$.

Author response:
To start with, the idea of a phase diagram here is very interesting and not something we looked into in detail, we are running experiments now to investigate this. However, most likely based on what we have seen, it will show that the improvement “per feature update” increases as K increases but the final result remains very similar (the training graph gets closer to going straight up then across). While interesting and a useful graph to include, I don't know that it sheds any new light onto general behaviour.

While we do list the contributions in the introduction, they were in paragraph form. We have split them out into dot points in the updated version and hopefully clarified it. We list these below now to also push back on the claim that our only contribution is the small number of feature updates necessary (which we already believe is a significant contribution for values of K=10^6 as that is not “a simple expectation of the interpolation”).
* The iterative linearisation algorithm itself is new (to our knowledge), and is definitely new from the context of understanding generalisation and feature learning in neural networks. This opens up an important tool to understand feature learning through better experimentation.
* We show mathematically that iterative linearisation is equivalent to the Gauss-Newton algorithm with a very large K and very small learning rate limit (when each step is equivalent to running gradient flow to convergence).
* Generalisation impact of damping in small neural networks: we show empirically that by adjusting damping in Gauss-Newton over multiple orders of magnitude, there is a point where we get a large improvement in generalisation. This is already of independent interest to training neural networks, however combined with our proposed connection to iterative linearisation it provides an important hypothesis connecting two independent areas - namely feature learning and damping in second order optimisation.
* We show an example with a similar number of feature updates where it does not perform as well due to the training loss getting to zero before there is a chance to learn features.

Reviewer comment (Presentation):
The paper does not have the *contributions*, making it difficult to distill the main claims and how they are supported.

Author response:
While we do list the contributions in the introduction, they were in paragraph form. We have split them out into dot points in the updated version and hopefully clarified it. We list these below now to also push back on the claim that our only contribution is the small number of feature updates necessary (which we already believe is a significant contribution for values of K=10^6 as that is not “a simple expectation of the interpolation”).
* The iterative linearisation algorithm itself is new (to our knowledge), and is definitely new from the context of understanding generalisation and feature learning in neural networks. This opens up an important tool to understand feature learning through better experimentation.
* We show mathematically that iterative linearisation is equivalent to the Gauss-Newton algorithm with a very large K and very small learning rate limit (when each step is equivalent to running gradient flow to convergence).
* Generalisation impact of damping in small neural networks: we show empirically that by adjusting damping in Gauss-Newton over multiple orders of magnitude, there is a point where we get a large improvement in generalisation. This is already of independent interest to training neural networks, however combined with our proposed connection to iterative linearisation it provides an important hypothesis connecting two independent areas - namely feature learning and damping in second order optimisation.
* We show an example with a similar number of feature updates where it does not perform as well due to the training loss getting to zero before there is a chance to learn features.

Reviewer comment (Writing):
The discussion of feature learning via pruning in the paragraph after eq. (2) is not clear in its current form.

Author response:
Feature learning vs pruning: We agree that this is a bit of limited discussion. It was added due to many questions about why this is distinct from ideas about the lottery ticket hypothesis. We have rephrased to remove any claim of the amount of possible options in place of the simpler observation that hidden representations must change under pruning. While this doesn’t give as much intuition, we hope this removes your issue with an unsupported claim.

Reviewer comment (Writing):
The description of iterative linearization in sec. 3 is confusing as it is different from the one described in the introduction (which was quite clear). In particular, algorithm 1 does not fully describe the iterative linearization procedure and appears to be incomplete. Eq. (4) is referenced in sec. 3.2 as the *feature learning step*, but it only states that the features of the linearized network are gradients.

Author response:
Iterative linearisation description: The description in section 3 is equivalent to the one in the introduction but was phrased to make the updating of the features as explicit as possible as that was the common point which was hard to get across. We have rephrased it to be more consistent with the usage in the introduction in the hope that that will help here. Similarly, algorithm 1 was exactly equivalent but we have made the dependence on theta explicit in the closure in the hope that that will improve clarity. Regarding the feature learning step, please notice the subscript s in equation (4) - this is a step to update the features of the network. With the rephrasing we used above, we have now called attention to this in the text as it is no longer an equation to update the features directly. This is a key point and we want to get this idea across as clearly as possible so any more information about which parts you found difficult are helpful.

Reviewer comment (Theory):
*A proxy measure of feature learning* defined in sec. 3.2 is simply a hyperparameter of the proposed iterative linearization scheme, and therefore cannot measure in any sense the actual feature learning happening during optimization. Also, the discussion in the paragraph after Def. 1 is not clear, for example, it is unclear what is meant by *absolute amount of feature learning*.

Author response:
Proxy measure: Perhaps the term 'measure' is inappropriate here and has caused some confusion - we have removed the reference to measure there as it seems to be off putting. Yes this is a hyperparameter, but as we explain, it can control the frequency (and amount - see Figure 6 & 7) of feature learning so is a useful proxy as to the amount of feature learning which occurred. Similarly the term “absolute amount of feature learning” has been replaced by “true amount of feature learning” - the latent true value for which we use our proxy,

Reviewer comment (Writing):
It is not clear which linear model is considered in sec. 4.3.

Author response:
Linear model in 4.3: the only linear model we consider in this paper is the linearised version of the neural network being trained, however we have made this explicit here to improve clarity

Reviewer comment (Novelty):
Sec. 4.2 seems absolutely redundant, as its only conclusion — that linearized models have poor generalization — is well known in the literature and already discussed by the authors in the introduction.

Author response:
Section 4.2 redundancy: we have removed this and covered the details in 4.1

Reviewer comment (Presentation):
In figure 3, the neural network and its linearization seem to be swapped in the legend (a small typo not affecting the quality of the paper).

Author response:
Typo in Figure 3: Good catch, we’ve fixed that now :)

Reviewer comment (Theory):
The main weakness of this paper is that most of the claims are made by numerical experiments without rigorous theoretical justification.

Author response:
Theoretical justifications: This is an empirical paper, while we could make some theoretical justification, when the assumptions of such a justification will clearly not be met, we believe that this is not necessary in order to show our results. We do believe that future work investigating iterative linearisation from a more theoretical direction could provide useful insights but it is not in the scope of this work, nor necessary for what we show.

Reviewer comment (Theory):
The authors assume that the gradient descent step is to perform feature learning but do not give detailed justifications. In fact, gradient descent may also memorize the noise to fit training data points. In some special settings (e.g., very wide neural network, specifically designed initialization), gradient descent can also behave similarly to only learning random features.

Author response:
GD step performing feature learning: we completely agree that gradient descent can do memorisation, we only make a claim about when feature learning is happening - as stated in 3.2 “We want to point out that feature learning only happens in Equation (4), but not in Equation (3)”. During equation (3) we are simply using the features that are already there so we can safely say that no feature learning is happening, so all feature learning must be in equation (4). We acknowledge that memorisation and many other phenomena may also occur in the same step but this does not affect any of the later claims or conclusions as we specifically are looking at limiting feature learning - we are ambivalent as to what else happens.

Reviewer comment (Evaluation):
Since it is difficult to exactly quantify how much feature learning is performed in gradient descent, there might be some potential concerns by only comparing the number of gradient steps. It is possible that feature learning only happens in a small number of early steps of neural network training, then the "effective number" of feature learning steps could be much smaller than the total iteration numbers.

Author response:
Does all the feature learning happen early? This appears not to be the case (e.g. see Vyas et al (2022) in our references). But perhaps it would be helpful for you for us to include a reproduction of this where we train with SGD for various lengths before one-shot linearising? This is a simple plot to make and we can include that in the final version (we have done similar experiments already but not made a clear plot of them)

Reviewer comment (Experiments):
As claimed by the authors, iterative linearization is similar to the Gauss-Newton methods, it would be good to also present the results for Gauss-Newton method in the experiments.

Author response:
Gauss-Newton on NNs: we agree that this would be an interesting result, it is unfortunately intractable computationally to run a true second order method on a network with > 450k parameters which is why we ran it on smaller networks instead (Figure 9). We do however run an analogue of this in Figure 8 where we train the linearised model to convergence using Adam (if we were using MSE instead of softmax then this would be roughly the same as the Gauss-Newton method with a slight bias due to using Adam). One experiment that we have considered that maybe would assuage this concern for you would be to include the loss term equivalent to the damping factor for a version of the experiment from Figure 8, if you feel this would satisfactorily cover this case then we can run this.

Reviewer comment (Experiments):
In deep neural network training, people typically use a larger learning rate for SGD/GD at least in the early stages, the authors may consider trying a larger learning rate in the experiments, rather than using 1e-3 for all experiments.

Author response:
Larger learning rates: Even with very small learning rates, we see the same feature learning and generalisation behaviour we wish to study. We acknowledge that this does not reflect standard practice but there is already a lot still to explain with only small learning rates and keeping these small allows us to run experiments that would diverge with larger learning rates.

Reviewer comment (Experiments):
What happens if applying Gauss-Newton to train the neural network?

Author response:
Gauss-Newton: As we covered above, we show in Figure 8 that when we do an analogue of Gauss-Newton for softmax networks and approximate the inversion using Adam, we get that it performs worse as would be expected from very little feature learning and the train loss approaching 0 very quickly. We also note that we do train neural networks using Gauss-Newton directly for Figure 9 where it isn’t computationally intractable.

Reviewer comment (Experiments):
What if using larger learning rates for gradient descent?

Author response:
Larger learning rates: The proposed algorithm will diverge if the learning rates are too large, but the effects which we intend to study already occur at small learning rates so we do not believe this invalidates our results.

Reviewer comment (Theory):
Is there any theoretical explanation to back up the observations?

Author response:
Theoretical justifications: This is an empirical paper, while we could make some theoretical justification, when the assumptions of such a justification will clearly not be met, we believe that this is not necessary in order to show our results. We do believe that future work investigating iterative linearisation from a more theoretical direction could provide useful insights but it is not in the scope of this work, nor necessary for what we show.

Reviewer comment (Novelty):
The idea of less frequent feature learning steps is not new. It has now become usual to consider layer-wise training of a two-layer neural network, where the first layer is trained with one step gradient and a large step size, followed by the linear regression of the second-layer parameters. The examples include [Damian et al. (2022)](https://arxiv.org/abs/2206.15144) and [Ba et al. (2022)](https://proceedings.neurips.cc/paper_files/paper/2022/hash/f7e7fabd73b3df96c54a320862afcb78-Abstract-Conference.html). Therefore, in theory, it is not surprising that less frequent updates of the feature suffices to achieve high test accuracy.

Author response:
* Novelty of less frequent feature updates: Thank you for the links to papers,we had not come across the first paper and are looking forward to reading it in more detail, but I do not believe that the results imply what we have shown here. We had, however, come across the Ba et al paper, it is limited to a single layer as you say which is not the setting which we look at. Additionally, we note that one of our results is to say that when fully solving the linearised problem we get much worse generalisation implying that we are seeing something beyond what has already been studied as that is significantly more than the one step of feature learning in the Ba et al paper.

Reviewer comment (Theory):
No convergence guarantee. If an algorithm claims itself as a proxy of (S)GD training, it is necessary to have a global convergence guarantee. This paper does not provide any convergence guarantee. Because of its connection to the Gauss-Newton algorithm, it could be possible to derive local convergence, but not global convergence as NTK and the mean-field analysis do. Note that, in order to theoretically justify their argument, we need to show not only the global convergence, but also that an increase in \( K \) does not slow down the convergence under certain conditions.

Author response:
* Convergence of iterative linearisation: While we agree that this would be interesting to show, we question why it is necessary. Convergence proofs can be incredibly helpful, however when the goal is to empirically show a phenomena then we argue that they are unnecessary. For the purposes of our results it suffices that it converged in those cases.

Reviewer comment (Theory):
Validity of the proxy measure (Definition 1). When \( K \) is not large enough, the next linearization step comes before the linearized network is fully optimized. Thus, I am not sure whether this value is a number of different features that must be passed through along the way.

Author response:
* Proxy measure of feature learning: We agree that this measure of feature learning is not exact - if the linearised model isn’t fully optimised then this will be different. We treat it as a rough proxy for feature learning later in cases where it makes sense. We know of no way to truly measure feature learning accurately in these cases however this allows us to still see correlations.

Reviewer comment (Theory):
Discussion on the mean-field neural network. Although the network depth is limited to two-layer, the mean-field neural network is an important paradigm in the feature learning analysis. It attributes the optimization of neural networks to the optimization in the measure space and explains the feature learning dynamics [Chizat & Bach (2018)](https://arxiv.org/abs/1805.09545); [Mei et al. (2018)](https://arxiv.org/abs/1804.06561). Recently, several papers [Abbe et al. (2022)](https://arxiv.org/abs/2202.08658); [Abbe et al. (2023)](https://arxiv.org/abs/2302.11055); [Suzuki et al. (2023)](https://openreview.net/forum?id=tj86aGVNb3&referrer=%5BAuthor%20Console%5D(%2Fgroup%3Fid%3DNeurIPS.cc%2F2023%2FConference%2FAuthors%23your-submissions); [Bietti et al. (2023)](https://arxiv.org/abs/2310.19793)) gave convergence and generalization guarantees on specific types of problems (especially polynomials) using the mean-field neural network.

Author response:
* Relation to mean field theory: We agree that the mean field formulation can give us useful insights on feature learning and thank you for some links to newer papers which we hadn’t read previously. Due to the empirical nature of this paper we didn’t attempt to prove what iterative linearisation would accomplish in a mean field setting, nor are we aware of any works which investigate multi-step or second order methods in a mean-field setting. We have now included a discussion in the related work section

Reviewer comment (Theory):
Do you think it is possible to give some convergence guarantees on the proposed algorithm? This does not necessarily need to be for the general functions.

Author response:
* Convergence guarantees: we believe it would be possible, but beyond the scope of this work (see above for details). Specifically the form would likely have to be, for a given K and problem, there exists a learning rate such that convergence is guaranteed.

Reviewer comment (Evaluation):
How can we explicitly know that the feature is learned by the proposed algorithm? Can you evaluate the parameter alignment?

Author response:
* Regarding both the second and third questions, these both assume some kind of ground truth for measuring features and how far we are from some kind of “optimal features”. The second question wasn't completely clear to us, but there is very little way to know anything for sure about these features. This is why iterative linearisation and our proxy measure are useful - they are not exact but we should expect them to be correlated such that we can draw useful conclusions.

Reviewer comment (Theory):
Is it possible to precisely evaluate the minimum required number of feature updates for some specific problems?

Author response:
* Regarding the third question, we can possibly do something to investigate the minimum number of features needed to learn, however due to not having any ground truth about how much features are updated or what the “optimal features” are, decisions about the K value to use become very important (e.g. see Figure 8 for why we cannot just use extremely large K) and have no obvious solution