Paper: No Wrong Turns: The Simple Geometry Of Neural Networks Optimization Paths (ICLR-2024)

Reviewer comment (Evaluation):
This paper seems to give a misleading message. The title says the neural networks have "simple geometry", and the abstract says "optimization trajectories never encounter significant obstacles, [....] maintain stable dynamics during the majority of training." I think that those properties are examined only for state-of-the-art neural networks. There are lots of techniques to alleviate optimization issues, so the overall claim of this paper severely trivializes the techniques in common practice. (On a related note: It would be interesting to investigate how much each technique contributes to making optimization benign.)

Author response:
We agree with the reviewer that architectural improvements are likely a big contributor to the phenomenon we observe. In fact, in Appendix D, we list architectural characteristics as a possible factor to the behavior we observe, and discuss it (see architectural characteristics paragraph). While we observe that the stable behavior of cosine similarity is, in fact, preserved on simple architectures (vanilla MLP with ReLU activations), our experiments cannot conclude on the contribution of e.g. batch norm, skip connections, etc. 
We agree this is an interesting line of research for future work. However, as our experiments are quite expensive (requires double computation as each training run needs to be computed twice), the range of ablations we have conducted is already quite extensive. 

Most importantly, the main goal of our work is to characterize the properties of neural loss landscapes in practical settings, i.e. using modern architectures. As we discuss, our main goal is not to explain why these properties are present, but to quantify them and their relation with various factors (such as optimizers, depth and width, seeds, task, etc...). Since these properties seem very relevant to training dynamics, better understanding them mean we can better understand training behavior, and build a more practically relevant theory. Thus, while indeed our observations are only validated on heavily engineered architectures, we argue that this instead makes our work more relevant, and validating our claims on simple architectures that are not deployed in practice would be less meaningful.

Reviewer comment (Evaluation):
I find that RSI is confusing to understand: in Figure 2, the measure for ImageNet and WikiText is initially close to zero, so does this mean the training dynamics is not ideal in the beginning? "RSI and EB follow predictable trends" do not imply that optimization dynamics have benign properties.

Author response:
The values of RSI and EB appear close to $0$ in Figure 2 due to being small compared to the late stages of training, where both RSI and EB go to infinity (see late training behavior paragraph). But this is just a matter of scale, and the cosine similarity in fact quickly stabilizes.  The speed at which the iterates $w_t$ approach $w_T$ is in fact entirely determined by RSI, EB, and the learning rate (see Equation 4), and the value of RSI and EB in the early stages, while small in comparaison to late training, are not negligible. For a fixed cosine similarity $\frac{RSI}{EB}$, observing smaller values of RSI and EB simply means the ideal learning rate should be larger. 
"RSI and EB follow predictable trends" does mean that the optimization trajectory advances toward its last convergence point at a steady pace, and that despite non-convexity and stochasticity, all iterations (with very few exceptions) contribute similarly in term of parameter space progression. The stability of the values of RSI and EB across minibatches, and its predictability across phases of training, mean that ideal values of the LR are stable across minibatches, and predictable across phases of training. This explains why simple LR schedules can be so efficient in training neural networks despite their non-convexity and the stochasticity of gradient sampling.

Reviewer comment (Experiments):
The cosine similarity is low. While I agree with the argument that there is stochasticity, I think cosine similarity should be compared with the baseline where (i) gradients are replaced by noise, or (ii) gradients are calculated based on random labels or random tokens in a minibatch. Without comparison with the pure noise baseline, we do not know if cosine similarity is significantly correlated with the signal direction.

Author response:
We do compare to the baseline where gradients are replaced by noise, as suggested by the referee, in Section 6 (paragraph "induced Bias"). We also discuss the low value of cosine similarity in Section 4 (paragraph "low value of cosine similarity") and remark that in fact, if it wasn't very low, we would be able to train neural networks in very few steps, which know not to be the case. The cosine similarity being low can be compared to a quadratic function being badly conditioned: optimization is difficult, but follows a simple underlying structure, which eventually drives the convergence. Interestingly, RSI and EB would also follow predictable trends in the case of a random walk due to the high dimensionality having a regularizing effect. This can be interpreted as random walks having a stationary update rule. The complex structure of neural networks suggest the training dynamics could be highly complex (see contrasting examples in Section 6). Our work finds that it is not the case.

Reviewer comment (Experiments):
I understand that this paper is mainly about investigating the empirical properties of training neural networks from the optimization perspective. But from a practical point of view, what can we learn? Do we know what part of the architecture of, say a transformer, is crucial for benign optimization? Which techniques (layer normalization, dropout, learning rate warmup, etc) are playing a role?

Author response:
As stated in our reply to Weakness 1, our work does not provide direct insights on which architectural features enable such properties. The goal of our work is to improve our understanding of the fundamental properties of neural loss landscapes that make first order optimization so surprisingly efficient, rather than exploring their origin. Once a sufficiently refined characterization of empirical neural loss landscapes is obtained, we will be able to derive theoretical results in more realistic settings than previous works, which will hopefully be able to guide the design of more efficient optimization algorithms. A few more steps in this direction are therefore necessary to lead to immediate algorithmic improvements.

Reviewer comment (Novelty):
There is a closely relevant paper titled "SGD Converges to Global Minimum in Deep Learning via Star-convex Path" which appeared in ICLR 2019, which also empirically studies the SGD trajectories and shows that the star-convexity is satisfied for most epochs. I have not checked the follow-up results carefully, so there may be more existing results that discover surprisingly benign geometric properties of neural network training. The paper should contextualize itself relative to this existing result(s).

Author response:
This is an excellent point. We were surprised to find that "SGD Converges to Global Minimum in Deep Learning via Star-convex Path" [1] is not already part of our bibliography, as it is a paper that has inspired our work and was one of the first added to our bib file. We intended to discuss differences with this seminal work in detail, and to be transparent, we forgot. We will add a paragraph contextualizing our work with respect to this paper. 

The limitations in [1] that we try to address in ours are the following:
- Their work assume interpolation regime and only evaluates on CIFAR and MNIST, without typical training components such as batch normalization, weight decay, momentum, etc. By evaluating in realistic settings, with popular architectures and optimizers, on a variety of more advanced tasks (including in the non-interpolation regime), our work is closer to empirical practice.
- The epoch-wise star-convexity is fairly limited as it only considers the average of first-order developments over a full epoch. The iteration-wise star-convexity is only verified with respect to the last iterate of the sequence (just like in our work) but their theoretical results require star-convexity with respect to *all* minimizers of the given minibatch, which seems impossible to confirm empirically (and is in fact unlikely due to permutation invariance). Our theoretical results do not have such limitations and apply directly from our observations. 
- [1] Assumes L-smoothness of each $l_i$. Even when the loss is made smooth, it typically entails an extremely high constant $L$ which is unlikely to explain the good behavior of training, and it is intractable to estimate in practice. In our work, the role of smoothness is played by EB.
- Perhaps most importantly, our work **quantifies** the properties of optimization paths on neural networks. In [1], the paths are simply observed to be star-convex, which would be akin to only observing the non-negativity of RSI (or $\gamma$) in our work. This quantification allows us to suggest LR schedules, guarantee linear convergence, etc. 

In summary, prior works ([1], Goodfellow et al. (2015) and Lucas et al. (2021)) observed that the loss landscapes is simpler than expected, whereas we attempt to formally characterize and quantify this simplicity, providing more advanced tools to build theoretical results, and making a significant step toward practical relevance in our opinion.

Reviewer comment (Evaluation):
A recently discovered phenomenon Edge of Stability (EoS) shows that GD trajectory experiences progressive sharpening and then reaches and hovers around the stable convergence threshold $2/\text{(step size)}$. When EoS is in effect in the later phase of training, the training dynamics is rather unstable, with occasional peaks in loss curves. Similarly, when you run SGD instead of GD, we also often observe these loss peaks in the training curves. In contrast, the RSI, EB, and $\gamma$ curves look quite stable, without any significant peaks or perturbations. My question is: did you observe any peaks or unstable convergence in your experiments? Do the "fundamental properties" (Section 4.1) continue to hold when we observe EoS? If so, how can we reconcile these seemingly contradictory trends?

Author response:
**Question 1 :** since we train in realistic settings, we inevitably observe jumps in the training loss. The reason this does not contradict EoS observations is that small perturbations in high-curvature directions (like a small modification of very impactful features) can lead to large variations of the training loss, without impacting significantly the trajectory in the parameter space.

Reviewer comment (Experiments):
It seems that using large batch sizes increases $\gamma$ throughout training. What happens if we use full-batch training (i.e., GD)?

Author response:
**Question 3:** We assume RSI and EB would be even better behaved by taking out the stochasticity. We did not conduct these experiments, first because it would be very computationally expensive to train until convergence on datasets like ImageNet, and second because we believe the observations in the deterministic setting would be less significant due to being empirically unrealistic. A major surprise in our work is that RSI and EB remain remarkably well behaved despite stochastic gradients having high variance. Such observations would be much less significant in the full batch setting.

Reviewer comment (Evaluation):
**The nature of cosine similarity in high dimensions:** One of the fundamental issues about this paper is to what extent cosine similarity is a meaningful measure in high dimensions, which here are in the orders of tens of millions. The particular value of cosine similarity is somewhat difficult to grasp as a result, without the lack of a baseline that contextualizes its value. While the curves look alright and pretty stable throughout, the moment one looks at the value at the y-axis (which, e.g., is in the order of 1e-3 for ImageNet), doubts emerge if one is reading too much signal in the noise. This is in addition to the fact that the cosine similarity is between vectors that at their core are based on the same network structure will tend to have a higher cosine similarity (contrast that with arbitrary random vectors in the same dimensions). All of this makes it hard to be convinced by their argument about the stability of this cosine similarity (and optimization proceeding at a regular pace), when this value of cosine similarity almost hints at near orthogonality.

Author response:
**Nature of the cosine similarity in high dimensions (Weakness 1):**
Interpreting the significance of our observations is the most challenging component of our submission and the referee is right to point out that it is difficult to grasp. 

Regarding the low value of the cosine similarity (~1e-3 for ImageNet), we discuss it explicitly in page 6 (“Low Value of Cosine Similarity” paragraph). If the cosine similarity was stable at a high value, then SGD would converge at an unreasonably fast pace. It is thus expected that cosine similarity would be low on average. We encourage the reviewer to compare our observations to a very badly dimensioned quadratic, for which optimization is slow and difficult, but the dynamics are simple, stable, and efficiently captured by the curvature. 

The relationship between sampled gradients and $w-w^\star$ is also discussed in detail in all of page 8, where we compare to anisotropic random walks in very high dimension as the referee suggests. 

Finally, the referee seems to oppose our claim about the stability of the cosine similarity to its low value (near orthogonality). The stability of the cosine similarity is relevant to the regularity of the convergence (in the parameter space), while its value is relevant to the rate of that convergence. Thus we argue these observations are not contradictory. Optimization trajectories on neural nets effectively progress at a slow pace, but with surprising regularity, encountering no significant setbacks… or wrong turns.

Reviewer comment (Evaluation):
**The lack of clear insights that can be drawn:** Consider Figure 2, middle column, where the results are shown for ResNet50/ImageNet. It is unclear if anything much is being said by these graphs, except towards the end where it hints at being close to convergence. And that too, because the denominator in these expressions $\|w-w^\ast\|$ starts becoming small. More generally, a lot of these plots contain a sense of vagueness about the actual insights that can be drawn from them. The paper says there are "predictable" properties, but a less glorified way of phrasing would be they are somewhat banal.

Author response:
**Lack of clear insights that can be drawn (Weakness 2):**
Considering the well known complexity of neural loss landscapes, observing such “banal” (as the referee puts it) patterns for fundamental quantities is, paradoxically, surprising. We refer to Figure 5 for contrasting examples of what the cosine similarities could look like for trajectories computed on functions that are either non-convex or stochastic. The significance of our observations is discussed in the “significance” paragraph of page 6. Moreover, the stepwise optimal learning rate is directly tied to the local value of RSI and EB. This banal behavior of RSI and EB (which we denote as “predictable”) is thus what allows banal learning rate schedules to work so well. Unpredictable variations, as exhibited in Figure 5, would make it impossible to design a near-optimal learning rate schedule.

Reviewer comment (Theory):
The paper does not provide any further theoretical analysis based on their empirical findings. Although existing works have demonstrated that well-behaved RSI and EB can lead to the convergence of first-order algorithms, it is not clear whether trajectory-dependent RSI and EB properties can lead to similar results.

Author response:
**Theoretical guarantees for local versions of RSI and EB (Weakness 1 and Question 2):**
Interestingly, the linear convergence guarantees induced by RSI and EB only hinges on the local values of RSI and EB in the iterates that are explored. While previous theoretical works indeed only considered RSI and EB as global properties, Eq. (4) only depends on the local values of RSI and EB. The optimal learning rate $\eta^\star$ in Eq. (5) [which we will rename $\eta^\star(w)$ to clarify that it is dependent on $w$], and the corresponding step-wise guarantee in Eq. (6), only depend on the local value of RSI and EB. To obtain the guarantee of Eq. (7), it is thus sufficient for RSI and EB to be bounded on explored iterates. 

Alternatively, the interpolation conditions derived in “Gradient Descent is Optimal Under Lower Restricted Secant Inequality and Upper Error Bound” (NeurIPS 2022) state that as long as (local) RSI and (local) EB (wrt to some $w^\star$) are bounded by $\mu$ and $L$ on a finite set of points and gradients (e.g. the iterates), then there exists a function $f$ for which RSI and EB are bounded globally by $\mu$ and $L$, that exactly interpolates these points and gradients. Our trajectory is thus also a valid trajectory on $f$, and convergence guarantees obtained from global properties naturally convert to local properties for iterates.

Considering the referee’s remark, we agree that the above points are not sufficiently clear. We propose to add a short paragraph at the end of Section 3 to emphasize that bounds on (local) RSI and (local) EB for iterates provide the same guarantees as global bounds. (Same for bounds on cosine similarity with adequate lr schedules).

Reviewer comment (Reproducibility):
The paper does not provide practical guidance on network training. Most conclusions are explanatory.

Author response:
**Lack of practical guidance (Weakness 2):**

We absolutely agree with the referee. We believe that optimization of neural networks has been largely guided by intuitive engineering instead of theory in recent times, mostly because the classes of functions involved are insufficiently understood for relevant theory to emerge. We believe our work to be a promising – but incomplete – step towards understanding the structure of neural loss landscapes most relevant to 1st order optimization. Our work is indeed explanatory, but we hope that after further contributions to our understanding of empirical neural loss properties, we will eventually be able to develop a more adequate theory leading to useful algorithmic prescriptions ! 
Perhaps the most immediate impact on theory is that we believe our work supports that RSI and EB are much more relevant assumptions to use in optimization theory for neural networks than e.g. smoothness and strong convexity.

Reviewer comment (Writing):
In contrasting examples, what is w*? The (known) global optimum or the converged point?

Author response:
**$w^\star$ in contrasting examples (Question 1):**

We use the converged point to be in the same setting as our experiments. We will specify this explicitly.