Deep Learning Through A Telescoping Lens: A Simple Model Provides Empirical Insights On Grokking, Gradient Boosting & Beyond
Deep learning sometimes appears to work in unexpected ways. In pursuit of a deeper understanding of its surprising behaviors, we investigate the utility of a simple yet accurate model of a trained neural network consisting of a sequence of first-order approximations telescoping out into a single empirically operational tool for practical analysis. Across three case studies, we illustrate how it can be applied to derive new empirical insights on a diverse range of prominent phenomena in the literature -- including double descent, grokking, linear mode connectivity, and the challenges of applying deep learning on tabular data -- highlighting that this model allows us to construct and extract metrics that help predict and understand the a priori unexpected performance of neural networks. We also demonstrate that this model presents a pedagogical formalism allowing us to isolate components of the training process even in complex contemporary settings, providing a lens to reason about the effects of design choices such as architecture&optimization strategy, and reveals surprising parallels between neural network learning and gradient boosting.
Paper
Similar papers
Peer review
Summary
The paper proposes a new approach to understanding neural networks by examining a model consisting of a sequence of first-order approximations telescoping into an empirically operational tool. This model aims to isolate components of the training process and offers a lens to analyze the effects of design choices in neural networks. The authors illustrate the applicability of this model through case studies on double descent, linear mode connectivity etc.
Strengths
The paper introduces a novel telescoping model that bridges the gap between theoretical and empirical research in neural networks. The approach is used to derive new insights into well-known phenomena such as double descent. The paper covers a wide range of phenomena and provides a detailed analysis of each, contributing significantly to the understanding of neural network behavior. The model offers a simplified yet effective way to understand the impact of various design choices on neural network training.
Weaknesses
The telescoping model involves increased computational costs, especially for large networks and datasets, which may limit its practical applicability. The quality of the approximation depends heavily on the learning rate and the optimizer choice, which might not generalize well across different settings. While the model is intended to simplify understanding, the incremental approximation can become complex to implement and analyze for certain networks and tasks. Although the paper provides case studies, more extensive empirical validation across diverse datasets and architectures could strengthen the claims.
Questions
In what ways does this model bridge the gap between purely theoretical analysis and empirical observations in neural network research? How does the model handle non-linearities and complex interactions within neural networks? Can the authors provide a quantification of how these features affect their approximation? Can you elaborate on how the telescoping model reveals linear mode connectivity within the loss landscape? What are the implications of this connectivity for network training and generalization?
Rating
6
Confidence
2
Soundness
3
Presentation
3
Contribution
2
Limitations
None
Summary
This paper tries to understand the design of optimizer, model architectures and some deep learning phenomenon empirically. Additionally, the authors find a proxy model of neural networks other than the simply linearized model.
Strengths
1. The topic is interesting, especially the proxy model of neural networks. I would like to see more theoretical insights derived from the proxy model, but the current paper didn't put efforts on it.
Weaknesses
1. First of all, limited insights gained from the proxy model of neural networks when dealing with the modern optimizer and architectures. In section 3.1, the authors try to connect their proposed proxy model with some modern optimization techniques. However, no rigorous proof or insightful experiments are provided. Complex formulas including Eq. 6-8 only provides limited insights. 2. As for case study part, the writing confuses me a lot. For example, in Section 4.1, in line 200-201, "A candidate complexity measure that avoids the shortcomings listed above because it only considers the behavior of the final fitted model was recently used by [CJvdS23] in their study of non-deep double descent." I cannot understand what this mean. 3. Still the case study part, I strongly disagree that "LMC and train-time transition into lazy regime." (As I am not familiar with GBT and grokking, I skimmed those parts). [cite 1] and [cite 2] demonstrates that lazy regime is not sufficient to explain LMC and model merging. [cite 1] Fort, Stanislav, Gintare Karolina Dziugaite, Mansheej Paul, Sepideh Kharaghani, Daniel M. Roy, and Surya Ganguli. "Deep learning versus kernel learning: an empirical study of loss landscape geometry and the time evolution of the neural tangent kernel." Advances in Neural Information Processing Systems 33 (2020): 5850-5861. [cite 2] Ortiz-Jimenez, Guillermo, Alessandro Favero, and Pascal Frossard. "Task arithmetic in the tangent space: Improved editing of pre-trained models." Advances in Neural Information Processing Systems 36 (2024). 4. There are seeming no many experiments in both case study part and Sec 3. I would expect more experiments relating their proposed proxy model to the case studies and optimization techniques.
Questions
N/A
Rating
3
Confidence
3
Soundness
1
Presentation
1
Contribution
1
Limitations
N/A
Summary
The paper replaces the traditional "lazy learning" regime approximation with an approximation using a telescoping sum. This splits up the usual interval of approximation $\theta_T - \theta_0$ into many smaller pieces. The authors then show why this approximation is useful for explaining various interesting aspects of modern deep learning: 1. Generalisation behaviour (double descent, grokking and the effect of inductive bias on model complexity) 2. The connection between gradient boosting and modern deep neural networks (DNNs) 3. The success of weight averaging 4. Modern design choices regarding architecture and optimiser
Strengths
In general, I think the authors did a reasonably good job. Particular praise goes to the following aspects of the work: 1. The extension of the lazy regime approximation to a telescoping sum seems like a very natural extension of the current set of ideas. The use of the sum seems logical and may become a useful theoretical insight. 2. I liked the choice of case studies. Phenomena like grokking, double descent, and weight averaging are clearly of interest to the community. 3. The explanations provided seem relatively intuitive under the telescoping approximation. I will now touch on how well I think the paper does on the areas of originality, quality, clarity, and significance: 1. *Originality.* The use of a telescoping sum is clearly somewhat original. In addition, I thought the analysis of GBTs seemed fairly novel and interesting. 2. *Quality.* There are some elements of high-quality work in the paper. I think the analysis and empirical evidence for the generalisation section were fairly good. I also thought the analysis of momentum and weight decay seemed interesting. (However, I have issues/suggestions for both of these sections in weaknesses). 3. *Clarity.* I thought the authors did a fairly good job of justifying their use of a telescoping sum. However, I find it hard to come up with other strengths, as I think clarity is what this paper is most lacking. 4. *Significance.* I think the insights are of decent significance. For example, if it was demonstrated that the effective parameter count metric had further predictive power, it could be a useful tool in the analysis of grokking. (I provide suggestions as to how it could be expanded upon).
Weaknesses
I will structure this section into three subsections. The first two subsections concern general where I believe the paper may be lacking. In the third section, I list additional minor issues. I then end with a general assessment. ### (1) Lack of clarity and writing style I think the paper lacks clarity in the writing. Here are some major areas where this causes issues: 1. **Abstract.** The abstract does not describe the idea of the paper well nor provides a clear explanation of the impact of the work. I found it very hard to read at first. 2. **Explanation of insights.** Due to the density of the writing, I found it hard to understand the theoretical insights. 3. **Explanation of setup in section 4.1.** I found the explanation of the $p_{\hat{s}}^0$ measure to be quite confusing. I could not see where $\hat{s}(x)$ was clearly defined. ### (2) Importance of theoretical insights and empirical backing In section 3.1, there are several interesting observations about modern optimisation choices and how they can be more easily reconciled under the telescoping approximation. However, I am unsure of several things: 1. **Utility of telescoping sum.** Are the observations provided specific to the telescoping sum approximation or would they hold with the traditional lazy approximation? 2. **Claims in real networks.** Do the claims made hold in real networks? It would have been nice to see experiments demonstrating the claims made given that we are still working with an approximation. 3. **Utility of insights.** What should I do with the insights provided (supposing they hold in real cases)? For example, now that I know that momentum offsets the effect of weight decay on learned function updates, should I try and avoid using weight decay and momentum? In my experience, it seems that weight decay and momentum do work pretty well together in practice. In section 4, there are also many interesting observations. However, I have a few concerns: 1. I would have liked to have seen more experiments demonstrating that cases of double descent and effective parameter counts are related in the way described. I think it is likely to be true, I just don't think that one experiment is sufficient. 2. In the grokking cases, I would like to see whether the metric proposed (namely mismatch in the effective parameter count) has more predictive power. For example, in the literature, there have been ways of inducing grokking proposed. For example, in [this paper](https://openreview.net/pdf?id=ux9BrxPCl8) it is shown that adding spurious dimensions to the input reliably increases grokking. I would like to see whether cases with a greater grokking gap correlate with a greater divergence in effective parameter use. 3. In the second case study, the theory of why GBTs perform better than neural networks seems reasonable. However, I think the empirical evidence is lacking. A single plot with correlation does not seem sufficient. 4. In the third case study, I am unsure as to whether the trend in the pre-activation gradients aligning with the accuracy gap is sufficient to demonstrate the hypothesis. ### (3) Minor gripes with the paper 1. If I have understood correctly, the paragraph discussing approximation error seems unnecessary. I would place this in an appendix and spend two sentences on it in the main discussion. It seems quite obvious that the telescoping sum is going to approximate things better. 2. The practical limitations section also seems unnecessary. We know that it is a tool for analysis rather than an alternative approach for training. This does not need to be explained. 3. I'm not sure how I feel about the use of the word "pedagogical" throughout the paper. I don't know if the impact of a NIPS paper is meant to be a tool for teaching. Rather, the desired impact should be a novel result or an improvement to the understanding of the field. ### A general assessment Generally, I think the paper provides some interesting insights but is attempting to do too much. I think if it were to be rejected, I would encourage the authors to rewrite it in a much simpler style focused on far fewer topics with many more experiments. Update: I have moved my overall rating from a 5 to a 6 as I believe some of my concerns were addressed in the rebuttal. I also moved my rating of soundness from a 2 to a 3.
Questions
Some questions: 1. As someone who has read relatively few papers using tangent kernel arguments, I am curious how other groups have extended the traditional lazy -> rich approximation? 2. Could the authors provide a better explanation of $p_{\hat{s}}^0$ and associated nomenclature? What is $\mathcal{I}_0$ for example? 3. Are there meaningful experiments to run with the insights in section 3.1? For example, could the authors look at verifying the reported trade-off between momentum and weight decay? Some other suggestions: 1. A few grokking papers which are not included in the literature review are: - [Paper 1](https://openreview.net/forum?id=GH2LYb9XV0). Grokking in linear estimators without regularisation, which may go against intuitions with the measure proposed - [Paper 2](https://openreview.net/forum?id=ux9BrxPCl8). Discusses complexity as a theory of grokking. - [Paper 3](https://openreview.net/forum?id=UHjE5v5MB7). Talks about corrupted labels, might be interesting in discussion of future work where the proposed measure could be used to analyse cases with corrupted labels 2. Paper on double descent not mentioned: [Double Descent Paper](https://arxiv.org/pdf/2111.08234).
Rating
6
Confidence
3
Soundness
3
Presentation
2
Contribution
3
Limitations
I think limitations are appropriately addressed given that the paper is not focused on performance. It might be interesting to see cases where the telescoping sum does not approximate some set of phenomena well. Perhaps some discussion of this instead of the current limitations paragraph would be more meaningful.
Response [2/2]
[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!
Summary
This work proposes a new analytical model of neural networks (NNs) extending popular 'linearized' approximations over the initial parameters built from the gradient vectors. In particular, they consider approximating the full learning trajectory of NNs as a sequence of first-order approximations rather than a single one. The authors connect their model to several prominent empirical observations of the properties of neural networks, showing how it can be used to compare different optimizers and learning frameworks, and also justify phenomena such as grokking, double-descent, and linear-mode connectivity.
Strengths
1) The authors provide a comprehensive first look into the properties and implications of the proposed framework. For instance, they consider extensions to vanilla SGD updates, including momentum and weight decay - showing how they can connect quite naturally through a weight averaging perspective (start of Sec 3.1). 2) The Appendix is functionally used to relegate more comprehensive secondary details, allowing the authors to keep the text free of long proofs and extended literature reviews. This makes the main text concise and pleasant to read, while still allowing the interested readers to delve deeper into specific topics. 3) The paper does a good job at highlighting the relevance of the proposed framework, providing grounded examples that connect it with several popular understandings of neural network dynamics and behavior in Section 4. Out of the three examples, I found the discussion regarding comparing neural networks and gradient boosting through their relative approximated dynamics and kernels (4.2) of particular interest.
Weaknesses
1) The proposed telescoping approximation partially defeats the purpose of the simplified linear models, by introducing back much complexity that linearization precisely seeks to abstract away. Perhaps, some additional discussion about the limitations in terms of theoretical analysis would have been useful. 2) Also due to this increased complexity, most experiments are carried out in relatively toy-ish settings (e.g., MNIST). 3) There is some overuse of inline math notation (e.g., Section 4.1) hindering readability. While I understand the limitations due to the imposed page constraints, I would encourage the authors to try and adjust their text to lay out their equations more clearly in later revisions. Typos: "due rescaling" -> due to rescaling (117)
Questions
I would appreciate if the authors could address the criticism highlighted above in their rebuttal.
Rating
7
Confidence
3
Soundness
3
Presentation
2
Contribution
3
Limitations
The authors explicitly discuss the main limitation of the proposed methodology (increased complexity) at the end of Section 3, although mostly from a practical (and not theoretical) viewpoint. Overall, I found the work well-written, quite comprehensive, and insightful - building directly on top prior work, yet, providing some interesting new perspectives. Hence, I am leaning toward acceptance.
Response to rebuttal for b22S
I would like to thank the authors for their detailed response, it cleared up many of my questions and concerns. As a result, I will update my rating to 6 (weak accept) and my assessment of soundness to 3. A few lingering questions and comments: 1. Could the authors explain in a comment what is meant by $\hat{s}(x)$, unfortunately, I was not able to understand its meaning from Appendix B.2. 2. I still believe the section on the approximation performance is not really necessary. My original concern was that even if the approximation holds with reference to the way the learning trajectories look, the approximation might not hold when used to derive other properties with the telescoping sum. (I hope this makes sense!)
Response [1/2]
We thank the reviewer for the very quick response and for engaging with our rebuttal – we really appreciate it! We are delighted that our rebuttal alleviated most concerns and hope that our responses below address any leftover reservations. 1. **What is $\mathbf{s}\_{{\theta}\_{T}}(x)$?** _Motivation._ [CJvdS23] showed that their _effective parameter_ metric $p^0\_s$ can be used to explain double descent in some _non-deep_ models (linear regression and tree-based methods) because it allows to distinguish between train-time and test-time complexity. In section 4.1, we explore whether the telescoping model allows us to make use of $p^0\_s$ in deep models too, which is not obvious a priori as the use of $p^{0}\_s$ requires to be able to write model predictions in smoother form. *Extended background.* [CJvdS23] build on ideas from the literature on nonparametric regression smoothers, which are methods that output predictions that are a _linear combination_ of the nx1 vector of training labels $\mathbf{y}$. That is, such methods issue predictions $\hat{f}(x)=\hat{s}(x)\mathbf{y}$ where $\hat{s}(x)$ is a 1xn vector assigning a weight to each training example. Prototypical examples of (the very broad class of) smoothers are k-Nearest Neighbor (kNN) estimators, which use smoother weights $\hat{s}(x) \in \\{0, 1/k\\}^n$, and ordinary linear regression (OLS), which uses smoother weights $\hat{s}(x)=x’(X’X)^{-1}X’$. Effective parameters $p^{0}\_{s}=\frac{n}{|\mathcal{I}\_0|}\sum\_{j\in \mathcal{I}\_0} ||\hat{s}(x\_j)||^2$ then aim to make different types of smoothers comparable to each other w.r.t. the amount of smoothing they do, by measuring the average squared norm of the smoother weights when issuing predictions for a given set of inputs $\\{x_j\\}\_{j \in \mathcal{I}\_0}$. Here $\mathcal{I}\_0$ is simply used as a shorthand to collect the indices of the inputs for which the used effective parameters are measured; e.g. $\mathcal{I}\_{train}=\\{1, \ldots, n\\}$ the indices of the training data and $\mathcal{I}\_{test}=\\{n+1, \ldots, n+m\\}$ the indices of the test data. Intuitively, $p^0\_s$ provides a measure for how non-uniform and extreme the learned smoother weights are – the higher $p^{0}\_{s}$, the more complex the model. For example, OLS with $p$ covariates has $p^{0}\_{s}=p$ and a kNN estimator has $p^{0}\_{s}=\frac{n}{k}$. *How does the telescoping model enable us to use $p^0\_s$?* For the squared loss, we exploit that the SGD functional updates $\Delta \tilde{f}\_t(x) = \gamma \nabla\_{{\theta}} f\_{{\theta}\_{t-1}}(x)^\top \mathbf{T}\_t(\mathbf{y} - \tilde{\mathbf{f}}\_{\theta\_{t-1}})$ use only a linear combination of the training labels, which is how we recursively construct $\mathbf{s}\_{\theta\_t}(x)$. To see this, assume for simplicity that at initialization $f\_{\theta\_0}({x})=0$ (otherwise, the prediction at initialization is carried forward as additional constant term $c^0\_{\theta\_t}(x)$ in the paper). Then, the first functional update is $$\Delta \tilde{f}\_1(x)=\gamma \nabla\_{\theta} f\_{\theta_{0}}(x)^\top \mathbf{T}\_1(\mathbf{y} - \mathbf{0}) =\underbrace{\gamma \nabla\_{\theta} f\_{\theta\_{0}}(x)^\top \mathbf{T}\_1}\mathbf{y} = \mathbf{s}\_{\theta\_1}(x)\mathbf{y}$$ which means that $\tilde{f}\_{\theta\_1}(x)=0 +\mathbf{s}\_{\theta\_1}(x)\mathbf{y}$ is also a linear combination of the training labels. Letting $\mathbf{S}\_{\theta\_1}$ denote the nxn matrix that has as the ith row $\mathbf{s}\_{\theta\_1}(x_i)=\gamma \nabla\_{\theta} f\_{\theta\_{0}}(x_i)^\top \mathbf{T}\_1$ for the ith training example and $\mathbf{I}\_n$ denote the nxn identity matrix, we note that the second update can then be written as $$\Delta \tilde{f}\_2(x) = \gamma \nabla\_{\theta} f\_{\theta\_1}(x)^\top \mathbf{T}\_2(\mathbf{y} - \mathbf{S}\_{\theta\_1}\mathbf{y}) = \underbrace{\gamma \nabla\_{\theta} f\_{\theta\_1}(x)^\top \mathbf{T}\_2(\mathbf{I}\_n - \mathbf{S}\_{\theta\_1})}\mathbf{y}= \Delta \mathbf{s}\_{\theta\_2}(x)\mathbf{y}$$ which is also a linear combination of the training labels. Then, by adding the functional update to the prediction from the first step, we get the prediction at the second step $$\tilde{f}\_{\theta_2}(x)=0 +\mathbf{s}\_{\theta\_1}(x)\mathbf{y} +\Delta \mathbf{s}\_{\theta\_2}(x)\mathbf{y} = (\mathbf{s}\_{\theta\_1}(x) +\Delta \mathbf{s}\_{\theta\_2}(x))\mathbf{y} = \mathbf{s}\_{\theta\_2}(x) \mathbf{y}$$ which is again a linear combination of the training labels. Then, by recursion, any future update analogously adds a term $$\Delta \mathbf{s}\_{\theta\_t}(x) = \gamma \nabla\_{\theta} f\_{\theta\_{t-1}}(x)^\top \mathbf{T}\_t(\mathbf{I}\_n - \mathbf{S}\_{\theta\_{t-1}})$$ to the previous smoother weight. Finally, this then gives $$\mathbf{s}\_{\theta\_T}(x)=\sum^T\_{t=1}\Delta \mathbf{s}\_{\theta\_t}(x)$$ where the $\Delta \mathbf{s}\_{\theta\_t}(x)$ depend only on the gradient feature matrices $\\{\mathbf{T}\_{t'}\\}_{t'\leq t}$. [Continued in next comment below]
Thanks for your rebuttal
I would like to thank the authors for their concise rebuttal which managed to address my main concerns, and for providing additional experiments and analysis. I also read the other reviews, and while I do understand some of their shared concerns e.g., "limited insights [...] when dealing with the modern optimizer and architecture" (gQrQ) I believe these have been mostly addressed/should be addressable (e.g., by toning down some of the claims), and do not warrant outright rejection. I believe expecting each paper to provide definite evidence for some of the key questions in the field is unrealistic. However, I believe this paper makes a solid effort to provide insightful novel evidence and analysis, which is likely to be of interest to the wider research community. Thus, I stand by my original assessment and would argue this paper clearly meets the bar for acceptance.
Decision
Accept (poster)