Deep Learning Through A Telescoping Lens: A Simple Model Provides Empirical Insights On Grokking, Gradient Boosting & Beyond

We thank this reviewer for their in-depth review and their positive assessment of our work! We were delighted by their appreciation of the utility of the telescoping model itself, the wide range of phenomena we are able to consider using our model and the insights we provide. We respond to major points raised in the review below.

(1) Computational burden. We agree that, as highlighted upfront in Sec 3, the telescoping model is more expensive to obtain than a standard neural network and therefore may not be possible to implement for large scale architectures such as modern LLMs. However, we note that the model is sufficiently tractable that it could be applied across all three case studies on the standard benchmarks from existing literature (all of which have received substantial attention in the recent literature despite their relatively moderate scale), and we were able to utilize it to provide significant insight. We believe that exploring computationally efficient methods of extending this approach to larger scale settings provides an exciting direction for future work (as alluded to in l.123)

(2) Approximation quality. We also appreciate that, with the exception of models for which the first-order Taylor expansion holds exactly, the approximation quality will always deteriorate for a sufficiently high learning rate (note that this is effectively by definition as eqn (2) holds exactly as $\gamma \to 0$). Indeed, making this potential limitation explicit was the purpose of including the experiment with varying learning rates in Fig 1. However, we generally found the appropriate learning rates for provided experiments to achieve good approximation quality (across different optimization choices, as further investigated in Apx D.1).

(3) Increased complexity over linear models. Indeed, we sacrifice some of the analytic simplicity of the fully linearized lazy model in order to obtain a more accurate model of a neural network. Unfortunately, we believe this to be unavoidable as a linear model will generally be unable to accurately model a full neural network in standard applied settings (indeed including even the most of toy-ish settings we consider, e.g. MNIST in Fig 1). We would argue that any conclusions drawn from a proxy model can only be valid if that proxy is a good approximation. Therefore, we believe the telescoping model sacrifices only the minimal degree of simplicity necessary in order to maintain an accurate approximation of the full network.

(4) Experiments. Allow us to re-emphasize that the goal of this work was to provide a tool that can help us to better understand specific previously observed phenomena in deep learning. Therefore, in our experimental evaluation we primarily revisit the specific, well-known, canonical experimental settings in which the studied phenomena have been observed in previous work (e.g. grokking, double descent, LMC). We have rigorously pursued this approach by reproducing key experiments from several works (e.g. [BHMM19, LMT22, KBGP24, FDP+20]) and integrating aspects of the telescoping model to better understand their findings. Further, Apx D includes additional results for different settings of some of the case studies. We also appreciate that our concise presentation of results in the main text (which are reinforced by 4 further pages of experiments in Apx D) may have disguised the empirical contribution of this work, making it appear that our evaluations are less extensive on the surface. In the general comment, we do provide several additional results on different datasets and using different experimental settings, which were specifically requested in other reviews and further reinforce the findings of our paper. We hope that these additional results may also further reassure this reviewer.

Questions.

• Q1. Related to (3) above, important theoretical progress in understanding neural networks has emerged from the NTK literature. Many of these results require some assumptions e.g. infinite width limit, lazy learning, etc. The telescoping model provides a setting in which an arbitrary neural network and optimization routine can be iteratively cast using NTKs without large deviations from the exact network. We envision this can provide a bridge between theoretical and empirical research upon which future theoretical research may be able to build by applying tools from this literature to this iterative approximation instead.
• Q2. Indeed, a powerful component of the telescoping model is that it does not require us to discard complex aspects of the optimization process or architecture. As shown in Sec 3.1, important practical extensions to vanilla SGD (i.e. momentum, weight decay, Adam) can be explicitly derived such that they can be exactly integrated into the telescoping model. Complexities of the model architecture, such as nonlinearities, on the other hand are entirely captured by the tangent features $\nabla_\theta f_{\theta_t}(x)$. That is, these gradients of the model outputs with respect to the model parameters capture how the networks architecture propagates between input and output.
• Q3. We agree that it is an interesting next step to seek insights on the nature of loss landscapes through insights derived from LMC. In this case, we are hesitant to make strong general claims about loss landscapes on the basis of our findings. While it is clear from eqn (10) that constant gradients are a sufficient condition for LMC, in practice we cannot broadly claim that gradients will become exactly constant and a complete theory of LMC may need to take into account further intricacies of neural networks. We hope that a more complete theory of LMC can be developed through future research taking into account the telescoping insights we provide here, which would provide valuable insights into the nature of loss landscapes in deep learning.

We thank the reviewer for the feedback. We are delighted that the reviewer finds our proxy model for a neural network interesting! Before we move to address the specific points raised in the review below, we would like to make two high-level comments. First, while we similarly believe that interesting theoretical insights can be derived from this model in follow-up work, we would like to note that the explicitly stated goal of the current paper was to use it as a tool for generating new empirical insight, unifying the study of a range of modern deep learning phenomena. Second, we also find it unfortunate that Secs 4.1 and 4.2 were only "skimmed" and therefore barely feature in the review, and would like to highlight that these were the parts that received particular praise in all other reviews.

1. Insights from Sec 3.1. First, we note that as stated in L131, eqs 6-8 were not only included to provide insight, but were also necessary in order to replicate later experiments from existing work that required Adam, weight decay, and momentum to be applied in the telescoping model. Further, we believe that Sec 3.1 will be of interest to anyone wishing to gain deeper intuition into the effect different optimization strategies will have on model predictions. Usually, these strategies are only discussed in terms of their effect in parameter space, and not in function space as we do here – even though function space is what ML is ultimately interested in. While we appreciate that this reviewer was not interested in the insights derived from Sec 3.1, we note that both YCvv and b22S explicitly mentioned this section as a strength of the paper. We hope that in examining the new empirical results we have now provided (see the pdf in the top-level response) that demonstrates how these insights are reflected in training a real neural network this reviewer will also find value in Sec 3.1.

2. Clarifying cited paragraph on complexity. This sentence highlights that the cited complexity metric of [CJvdS23] accounts for what the model has learned during the fitting process, and is thus able to capture the interaction of model, optimizer and a particular dataset. Many other metrics only consider what is learnable by the entire hypothesis class for arbitrary data. We note in the opening paragraph of this section that the former type of metric is preferable for understanding double descent and grokking, where the phenomena only arise due to interaction of the model, its learning process and the data. We hope that this alleviates any confusion. We would also be happy to clarify any further unclarities in the writing if more precise issues were pointed out!

3. LMC. We thank this reviewer for raising this question and are happy to clarify why there is no contradiction between these works and ours.

• Fort et al., which we cite numerous times as [FDP+20], find that “deep network training exhibits a highly chaotic rapid initial transient” for the first 2-3 epochs after which the relative change in the NTK decreases throughout training and plateaus at a low value. This is in complete agreement with our work which studies this through the lens of the telescoping model. Note that our Fig 5 and their Fig 7 - which are similar in spirit but study slightly different objects and use different metrics - observe a very similar trend. We further supplement their results by examining differences in tangent feature changes between pretrained and randomly initialized models.
• Ortiz-Jimenez et al. focus on a very different setting: they consider the context of task arithmetic, where a generalist model is being finetuned several times, each time on a new dataset, where the input space $\mathcal{X}$ of each task is disjoint from the others, with resulting task vectors then added to the base models weights and, therefore, is subtly different to weight averaging (they consider addition and subtraction). Please note that we strictly focus on finetuning routines that only differ in “batch orderings and data augmentations” (L356) followed by weight averaging between the resulting models, and do not claim that our empirical findings will transfer to their setting. Nonetheless, we agree that this reference is relevant to include and will discuss it in the updated version! Thank you for pointing this out.

4. Experiments. The goal of this work was to provide a tool that can help us to better understand specific previously observed phenomena in deep learning. Therefore, in our experimental evaluation we primarily revisit the exact well-known, canonical experimental settings in which the studied phenomena have been observed in previous work (e.g. grokking, double descent, LMC). We have rigorously pursued this approach by reproducing key experiments from several works (e.g. [BHMM19, LMT22, KBGP24, FDP+20]) and integrating aspects of the telescoping model to better understand their findings. We also verified that the telescoping approximation holds for the different implementation choices discussed in Sec 3.1 in Apx D.1.

We also appreciate that our concise presentation of results in the main text (which are reinforced by 4 further pages of experiments in the Appendix) may have disguised the empirical contribution of this work, making it appear that our evaluations are less extensive on the surface. For example, the double descent experiments alone - which only constituted a part (Fig 2) of case study 2 - required training 88 models per setting resulting in 432 hours of compute time on an A100 GPU.

Finally, in the pdf in the general comment we provide several additional results which were specifically requested in other reviews and further reinforce the findings of our paper. In particular, we also provide some empirical verification of our expressions from Sec 3.1 which we hope will further reassure this reviewer.

Thank you for the very detailed and constructive review! Limited by space constraints, we respond to major comments below.

(1) Clarity. We appreciate that our writing was somewhat densely packed in the mentioned sections, and will utilize the camera-ready extra page to decompress for improved clarity. We will particularly expand the paragraph on $p^0_s$ and will more explicitly write out $\hat{s}(x)$ in the main text (currently derived for different optimizers in Apx B.1), which can recursively be obtained from eq 5 after substituting in the squared loss gradient. We will also rework the abstract to include a more explicit summary of the finding of each case-study.

(2) Comments on Sec 3.1.

1. The insights on functional updates that can be derived from the telescoping sum subsume those of the lazy regime as a special case. As we highlight in the paragraphs on momentum and weight decay, assuming constant gradients will simplify some of these expressions. Some observations could thus also be made when assuming lazy learning – e.g. adaptive learning rates acting by rescaling the kernel – however, we have not seen these observations discussed explicitly anywhere. We will make even more clear which observations also hold in the lazy regime!
2. Claims in real networks. This is point is precisely why we included experiments on approximation quality – Fig 1 and Fig 6 in Apx D show that the approximation matches real networks very closely for the different implementation choices and that the insights in Sec 3.1 thus apply to real networks. We appreciate the suggestion to additionally numerically verify the tradeoff between momentum and weight-decay empirically, and have included an experiment that does so in the pdf in the top-level response!
3. Utility of insights. We believe that Sec 3.1 will be of interest to anyone wishing to gain more intuition and understanding of the effect different optimization strategies will have on model predictions. Usually, these are only discussed in terms of their effect in parameter space, and not in function space as we do here – even though function space is what ML is ultimately interested in! We comment further on the specific momentum weight decay example in the experiment in the top-level response.

(3) Comments on Sec 4.

1. Double descent. Note that Apx. D.2 already includes additional experiments, indeed demonstrating the same behavior for double descent on MNIST. Additionally, we provided a new experiment in the pdf of the top-level response.
2. Grokking. In Apx. D.2, we actually already conducted additional experiments varying two such mechanisms – the initialization scale [LMT22] and task alignment [KBG24]. As discussed in Apx D.2 we found in both cases that later generalisation (i.e. more extreme grokking) is predicted by an effective parameter gap emerging later in training. As suggested in the review, we repeated this experiment for another mechanism – the number of spurious dimensions – on the parity prediction task of Miller et al (TMLR2024). The results, included in the top-level comment, indicate that also here grokking is predicted by the effective parameter gap growing. We will better link to this in the updated main text, thank you for the suggestion!
3. GBTs. Note that results using additional datasets are included in Apx D.3 where a similar behavior is observed in the kernels.
4. LMC. We agree that unlike the preceding sections, the case study in Sec 4.3 is more exploratory in nature. The key point of interest is using the telescoping lens to come to the observation in eq 10 that constant gradients are sufficient for LMC. However, in practice gradients will not reach perfect stabilization and therefore this cannot provide a complete explanation for this phenomenon. We were therefore deliberate in not claiming any overly strong conclusions, and only include preliminary evidence of a specific factor contributing to LMC emerging by empirically showing that the model gradients indeed become closer to stable throughout training, and that – like the loss barrier – the change in model gradients is lower in pretrained than in randomly initialized models. We think there is significant opportunity to further investigate this perspective, and will update the section to be more explicit on scope and potential future avenues.

(4) Minor points.

1. Approximation quality section (Fig 1). Building on the answer to point 2 on Sec 3.1 above, we found this section essential as it demonstrates that (and when) the telescoping sum is indeed a reasonable model of a neural network in a practical setting [note that the neural networks used here are the same as in Sec 4.1] and that therefore the derived insights hold, while those of the lazy model may not. Further, the main purpose of including the lazy model here was to provide a reference point for the scale of the approximation error, allowing us to better evaluate how good an approximation the telescoping model is. We will make the purpose of this section more clear in the updated manuscript.
2. Limitations. We found it important to be upfront about potential computational limitations. The comment on training is in reference to the lazy model, which is sometimes used as an alternative way of training (as a linear regression in the tangent features). We will make this more clear!
3. The term pedagogical. We are in full agreement that the goal of a NeurIPS paper should be an ``improvement to the understanding of the field''. Our intention was to imply the telescoping model to be a pedagogical device in the sense that it is a tool through which researchers can better understand aspects of deep learning. However, we are happy to substitute this term as this was obviously not clear.
4. Additional relevant references. Thank you for these, we will include them in our literature review!

[Continued from above]

As derivation of $\mathbf{s}\_{\theta\_T}(x)$ requires only straightforward yet somewhat tedious algebra, we had relegated it to the Appendix to ensure readability of Section 4.1. However, we now see that the reader could benefit from some of the above discussion in the main text to make it more self-contained. As indicated in our rebuttal, we are very grateful for the suggestion and will use some of the additional space in the camera-ready to expand the corresponding paragraph on page 5.

2. Approximation performance.

We recognize the first point being made by the reviewer - the telescoping model is strictly a better approximation of a neural network than the lazy model by design as the difference is simply taking smaller approximation terms. This point does not require major discussion in the main text. The more important aspect of this section is demonstrating that, at least on the tasks evaluated, the difference in approximation quality results in a categorical improvement from a model that certainly does not track the performance of a neural network to one that does. The point is cemented further in Fig 7 where we observe the lazy model catastrophically diverging from the neural network’s prediction performance early in training while the telescoping model matches performance throughout. We would argue that this categorical difference in the two models ability to emulate a neural network is vital to demonstrate as it provides the essential justification for accepting the added complexity of the telescoping over the lazy model (e.g. one could conceivably compute effective parameters in case study 1 using the lazy approximation instead (as it is a special case of the telescoping approximation), but these results demonstrate that it would be a mistake to expect conclusions drawn in that setting to transfer to a fully trained neural network).

We appreciate this reviewer also has a subtle point that matching performance is necessary but not sufficient to know that a model is generally reliable as a proxy for analysis. One could imagine designing a model that matches a neural network relatively well in terms of predictions but that are achieved by a very different mechanism (e.g. using AdaBoost with decision trees). In that case, findings about the inner workings of the proxy model could plausibly not also hold for the neural network it models. In some sense, this is an unavoidable fact of modelling any process (particularly relevant in the social sciences where natural data-generating processes are often modelled using generalized linear models). However, we would point to two strong pieces of evidence that suggest the telescoping model only induces minimal bias in its approximation. (1) Unlike in e.g. the social sciences, we do have access to the exact mathematical form of the natural phenomenon (a neural network being trained on data) and are taking a direct approximation of that neural network. (2) The empirical predictions derived from the telescoping model have been reflected in the behavior of the full neural network repeatedly throughout our empirical investigations (e.g. (a) the divergence of full neural networks and gradient boosting is well modelled by the kernel of the telescoping model or (b) the predicted counterbalancing effect of weight decay and momentum from the telescoping model was demonstrated to exist exactly in a full neural network in our new results in the pdf).

Once more, we would like to thank this reviewer for their detailed engagement with our submission. We will integrate aspects of this extended discussion into the main text as we believe they are indeed valuable to expand upon. We hope that we have provided sufficient clarification on these final concerns but would be happy to discuss further if not!

We thank this reviewer for their very constructive review and the positive assessment of our work! We were especially delighted to read that the reviewer found our uncovered connection between the telescoping model and gradient boosting to be particularly instructive. We respond to the comments raised in the review below.

(1) Increased complexity over linear models. Indeed, we sacrifice some of the analytic simplicity of the fully linearized lazy model in order to obtain a more accurate model of a neural network. Unfortunately, we believe this to be unavoidable as a linear model will typically be unable to accurately model a full neural network in standard applied settings (indeed including even the most of toy-ish settings we consider, e.g. MNIST in Fig 1). We would argue that any conclusions drawn from a proxy model can only be valid if that proxy is a good approximation. Therefore, we believe the telescoping model sacrifices only the minimal degree of simplicity necessary in order to maintain an accurate approximation of the full network. We appreciate the reviewer's suggestion to extend the limitations section beyond practical considerations by highlighting that this additional complexity can come at the cost of complicating potential future theoretical analyses, and will do so in the updated manuscript!

(2) Scale of experimental settings. While we agree that our experiments were generally conducted in relatively moderate-scale settings compared to e.g. much of the recent LLM focused empirical work, we would like to emphasize that we were able to operate at a sufficiently large scale to satisfy the three provided case studies, all of which have received substantial attention in the recent literature (despite their relatively moderate scale). For case study 2 we operate on typical size tabular datasets from a standard benchmark from the literature. For case studies 1 & 3, we primarily revisit the exact well-known, canonical experimental settings in which the studied phenomena have been observed in previous work (e.g. grokking, double descent, LMC from [BHMM19, LMT22, KBGP24, FDP+20]). While larger scale experiments were thus not necessary in order to address our selected research questions here, we certainly agree that considering how to apply the proposed telescoping model to larger scale networks would enable many interesting research directions in future work (as we allude to in l.123).

(3) Use of in-line maths. We were delighted to read that this reviewer appreciated our efforts to make the reading experience as natural as possible by relegating extended details to the appendix for interested readers. We also appreciate that the quantity of content and breadth of phenomena covered in this work required certain parts of the paper to be written quite densely. This was a fair point raised by several reviewers. We will therefore put considerable effort into using the extra page of a camera-ready version to decompress the text (e.g. less in-line maths particularly in Sec 4.1) and expand on details where there is an opportunity to improve clarity.

Typo: Thank you, fixed!