Loss Landscape of Shallow ReLU-like Neural Networks: Stationary Points, Saddle Escape, and Network Embedding

Thank you very much for your review! We appreciate that you consider our approach interesting.

We made modifications to our script, highlighted in red, and marked with the orange-colored text 'BZd8' (your reviewer ID) to indicate changes based on your suggestions.

You mentioned a few weaknesses, which we address below.

---

"1. I am not convinced on the definition of stationary points for this paper (Definition 3.3)... The point I am not convinced is 'why are the radial terms and tangential terms considered differently?' - is it because the radial term is differentiable but the tangential term is non-differentiable?"

Yes, your understanding is correct. To facilitate the understanding of the difference, we also produce new plots in Appendix G. There, one can observe that the loss is always differentiable in the radial direction, which is not the case for the tangential direction (see the paragraph from line 1481).

At a high level, Relu-like networks can be understood as several patches of linear networks pieced together. Changing the norm of the input weights (moving along the radial direction) of a hidden neuron will not change the activation pattern of that hidden neuron. Namely, the training inputs that activate the hidden neuron through the positive/negative leg of the activation function will not switch to another leg. So, moving along the radial direction will not cross the non-differentiable boundary between the pieces of linear networks. However, moving along one tangential direction and the opposite might result in different activation patterns of a hidden neuron and thus enter a different piece of linear network. Therefore, there is a fundamental difference between these two types of directions.

"The definition of stationary points also felt more like a "local minimum" to me (this could possibly be wrong - but...)."

We point out that not all stationary points by our definition are local minima. Our definition of stationary points specifies that all the first-order directional derivatives should be non-negative, but there could be loss-decreasing direction resulting from higher-order directional derivatives. A good example of this is the origin of function $f(x) =x^3$, which is "flat" in terms of first-order directional derivatives (so it is a stationary point by our definition), but has decreasing directions given by third-order terms in Taylor expansion. A less toy, non-differentiable example is provided in Appendix C2.

"...my understanding is that the defined saddle point has zero differential at differentiable parameters and local minimum for non-differentiable parameters"

Your understanding of saddle points is correct. Non-differentiable directions must bend upward at saddle points. Otherwise, GD will slide away without significantly slowing down, disqualifying the point as a stationary point.

"Overall, the definition of stationary points in Definition 3.3 seems to be quite restricted."

We believe our definition of stationary points is not restricted because all the conventional stationary points of smooth networks are still stationary points by our definition. In other words, our Definition 3.3 extends the conventional notion of stationarity. The conventional stationary points should have zero directional derivatives toward all directions, which qualifies them as stationary points by our definition (which states that stationary points should have non-negative directional derivatives toward all directions).

---

"2. Also, it would be good to explicitly mention that the notion of stationary points this paper discusses may be different from the notion from other papers in the literature, at the introduction and in the abstract. For example, there exist papers that completely characterize Clarke stationary points of a one-hidden layer ReLU network [1], so to make the contributions more concise, describing the notion of stationarity could help."

Thank you very much for this suggestion. We have added one sentence in the abstract and one paragraph in the introduction to clarify the difference between our notion of stationarity and other previous notions. We have also included the reference you suggested. Please see line 18 and line 73 for more details.

"4. Most part of main results are discussions rather than concise theorems or statements."

We believe you are referring to the parts where we introduce the numerical experiment (Section 4.1.1) and contextualize Figure 4 (Section 4.2). These parts are relevant for presenting the implications of our main results, thus we kept them in the main text and actually deferred several theorems/lemmas/propositions regarding the rarity of local maxima and network embedding entirely to the appendix (Appendix D,M,O). Nevertheless, we agree that such an arrangement made our presentation style seem too heuristic. Thus, we make the following modifications to the script.

• To disambiguate our result that precludes other saddle types encountered during training, we compose Corollary 4.7 (see line 383) which replaces the originally loose discussions.

• As it caused confusion for other reviewers, we specify that Fact 4.6 and Fact 4.8 are rigorously proved results from a previous paper and our paper, respectively, rather than empirical observations (see footnote 4 and line 403).

• We defer the discussion regarding training dynamics from small but non-vanishing (previously called "finitely small") initialization to Appendix J. It meaningfully extends our discussion but could seem digressive.

"Furthermore, while mentioning saddle points, it is confusing whether the authors are referring to 'their notion of saddle points' or the notion of saddle points in the literature."

We believe such confusion is due to our discussion of training dynamics (Section 4.2), where we mentioned "saddles" in the "saddle-to-saddle dynamics" studied by [3-5] several times. Their notion of saddle points coincides with ours entirely but they failed to formalize such a notion. Defining such a notion is one of our contributions, which we explicate below.

[3-5] investigated the training dynamics of ReLU networks and found that there could exist points around which GD lingers extremely long, they loosely refer to these points as saddle points without rigorously defining the potentially non-differentiable saddle points. (Please refer to Assumption 3 of [5] and its discussion, which summarizes the saddle points studied by [3,4], for further details.) As a side note, they referred to such GD-slowing points as saddle points since the slowdown effect of such points on GD is indicative of a small "effective" gradient in the vicinity, reminiscent of the morphology of the conventional smooth saddle points.

We formally identify such GD-slowing points: First, GD being trapped near a point for a long time is equivalent to saying that there does not exist any directions with negative first-order directional derivative around that point, which is our definition of a stationary point. Second, if such a stationary point is not a local minimum or maximum, we call it a saddle point. We demonstrate the effectiveness of our definition in Section 3.2 and Appendix C. We also add a paragraph to clarify the connection between our notion of stationary point and the one in previous works [3-5] in our introduction (See line 73).

It is also noteworthy that the term "saddle-to-saddle" that [3-5] invoked (without a formal definition of saddle point) was first coined for deep linear networks [6], where the loss is smooth and the saddles are well-defined. [3-5] slightly abused the term for the nonsmooth case based on phenomenological similarity.

---

You also asked the following questions.

"1. Addressing the weaknesses would be good. If I can see that the author's definition of saddle points makes sense, and how they are related to saddle points in other papers, I would raise my score."

We believe our previous answers have addressed this question.

---

"2. Some clearer wording is needed: in Lemma 3.2, we have \delta(s_i) > 0 sufficiently small. Is there a clearer way to state this?"

We have changed the following

$\displaystyle\\mathbf d_{ji}^{\\mathbf v_i}=\\sum_{\\substack{k:\\\\ (\\mathbf w_i+\\Delta s_i\\mathbf v_i)\\cdot\\mathbf x_k>0}} \\alpha^+e_{kj}\\mathbf x_k + \\sum_{\\substack{k:\\\\ (\\mathbf w_i+\\Delta s_i\\mathbf v_i)\\cdot\\mathbf x_k<0}} \\alpha^-e_{kj}\\mathbf x_k$, where $\\Delta s_i > 0$ is sufficiently small.

to

$\displaystyle\\mathbf d_{ji}^{\\mathbf v_i}=\\lim_{\\Delta s_i \\searrow0^+}\\left(\\sum_{\\substack{k:\\\\ (\\mathbf w_i+\\Delta s_i\\mathbf v_i)\\cdot\\mathbf x_k>0}} \\alpha^+e_{kj}\\mathbf x_k + \\sum_{\\substack{k:\\\\ (\\mathbf w_i+\\Delta s_i\\mathbf v_i)\\cdot\\mathbf x_k<0}} \\alpha^-e_{kj}\\mathbf x_k\\right)$

" In pg. 8, could you be more specific what larger initialization scale 8.75 x 10^-4 means?"

It means that we initialize the parameter with $\mathcal{N}(0,(8.75 \times 10^{-4})^2)$. We have changed our script accordingly (see line 1696). This part of discussion is now deferred to Appendix J since it slightly deviates from our main theory.

I'd like to thank the authors very much for their comprehensive response, as well as accommodating most of my feedback. I wanted to ask you a few final remarks before making my final decision:

1. I now think I understand the intuition behind considering the differentiable and non-differentiable parts differently: in "saddle-to- saddle" dynamics when the loss function stagnates, the non-differentiable part of the loss function should bend upwards, as if not, gradient descent will escape the point "even though the subdifferential contains a 0". I now understand why Clarke subdifferentials cannot fully grasp the phenomenon. I think this is a very good intuition to understand saddle-to-saddle dynamics, and in fact shows that saddle-to-saddle dynamics is a misnomer - because we are not moving from a saddle point to another, rather, we are moving from a "specific point in the loss landscape where the loss function stagnates (which is not a saddle point in classical notion)" to another.

One thing I would like to know (which I believe makes the paper much stronger) is rigorously checking if this is always the case (i.e. for non-differentiable cases, if the loss is not bent upwards, gradient descent escapes too easily) - at least in the simplest case, i.e. for scalar functions. However, without such rigorous analysis, I am now quite convinced that the notion of the author's stationary point is an important object in understanding the saddle-to-saddle dynamics.

2. Some minor confusions were:

1. In Line 73-74, the sentence "where first-order directional derivatives toward all directions are nonnegative" could be made clearer if we write "first-order one-sided directional derivatives".
2. In Line 1098 - 1101, I think the sentences are vague. Also, I am curious if the authors don't have to mention anything about coordinate alignment, which (if I understood correctly) is important in guaranteeing the existence of derivatives.
3. I didn't understand why "all higher order derivatives are null" in Line 1105-1106.
4. I see that you changed the definition of $d_{ji}^{v_i}$ in equation (3). However, I am confused: is this same as just saying

$$\sum_{k: w_i \cdot x_k \geq 0} \alpha^{+}e_{kj}x_k + \sum_{k: w_i \cdot x_k \leq 0} \alpha^{-}e_{kj}x_k,$$

sending $\Delta s_i \rightarrow 0$. Or something similar?

I would like to thank the authors again for their sincere response. It helped me a lot in understanding the value of the work.

Thank you very much for your review! We appreciate that you noted how our method tame the non-differentiability in ReLU-like networks and liked our effort to make our result understandable.

We made modifications to our script, highlighted in red, and marked with the green-colored text '8foJ' (your reviewer ID) to indicate changes based on your suggestions.

You mentioned a few weaknesses, which we address below:

---

"There are constraints for the theoretical analysis"

We base our theories on minimal assumptions (line 131): (1) $\alpha^-\neq\alpha^+$, so that the activation function is non-differentiable, and (2) $d>1$ so that we can avoid lengthy discussions for degenerate cases. Our proof of the main result (Theorem 4.2) relies on a rigorous examination of Taylor expansion in all directions. In this sense, we believe our theoretical analysis is not subject to much constraint.

Later in the weaknesses section, you also mentioned our characterization of stationarity is "counter-intuitive", and might not identify local maxima as stationary points. We figure this could be the "constraint" that you mention here. However, we believe our non-conventional definition is a strength and exactly suits our needs, which we will explain when answering to that weakness.

---

"the experiments are performed on synthesized data instead of real data."

The experiment is an illustrative example of Theorem 4.2 and Fact 4.8 (originally Consequence 4.7). These theoretical results are justified by proofs rather than experiments. The simplicity of experiments is for better clarity. We need to track the angles and norms of all the weights during all saddle escaping. To clearly present this, we can only use a few training samples in low dimensions. In more complex setups, it will be hard to fully present the training dynamics since there will be much more neuron groups and escaped saddles. Moreover, in ReLU training dynamics papers [1-3], demonstrating theories by thoroughly probing toy examples is a common practice.

To enhance the credibility of our empirical results, we added two more numerical experiments in Appendix I, investigating the training dynamics of networks with larger input/output dimensions.

---

"I have the impression that the definition of stationary point in this work is not exactly the extension of stationary point of smooth function, which is where f′(x)=0 regardless of trapping GD or not."

First, we think that our definition of stationary point should be regarded as an extension of the stationary points of smooth functions since all the stationary points of smooth networks are still stationary points by our definition. The stationary points of smooth functions have zero directional derivatives towards all directions and thus qualify as stationary points by our definition (which states that stationary points should have non-negative directional derivatives toward all directions).

Second, a definition for stationarity that can differentiate between whether gradient descent (GD) is trapped or not is more relevant for understanding the optimization process of neural networks. Other notions of stationarity that do not have implications for training dynamics might not be as useful in this context. Several previous works [1,2,4] investigated the training dynamics and noted the existence of GD-trapping points on the ReLU network landscape but were not able to systematically define and characterize such points. Our definition of stationarity fills in this gap in theory. We now also mention this to contextualize our contribution in line 73.

"It appears to me that Clarke subdifferential the authors mentioned could be better in this regard, in the sense that the origin of both |x| and −|x| are considered to be a stationary point."

As mentioned in our previous answer, a reasonable and useful definition of stationarity for GD should capture the trapping or slowing down of GD in the vicinity of the points of interest. As we can see, when a GD reaches near the origin of $\vert x \vert$, it effectively halts, while the origin of $-\vert x\vert$ will not slow down GD at all. Thus, the former should count as a stationary point while the latter should not, which is why Clarke Subdifferential does not suit our needs.

"With the authors’ definition of a stationary point, I can imagine that local maxima are usually not stationary points."

As mentioned in our previous answers, smooth local maxima are still characterized as stationary points by our definition. However, non-smooth local maxima, like the origin of $-\vert x \vert$, are not since they do not trap GD at all.

Thank you very much for your review! We appreciate that you consider our definition that handles stationarity novel and interesting.

We also made modifications to our script, highlighted in red, and marked with the brown-colored text 'WEjG' (your reviewer ID) to indicate changes based on your suggestions.

You mentioned a few weaknesses, which we address below.

---

"1. The current studies focus on the two-layer ReLU neural network and it is unclear how the result can be generalized to deep networks, especially since the characteristic of escape neurons can be complicated when the network is deep, it is unclear if we can draw meaningful conclusions there using the proposed notion."

We have showcased how our method can be used for deeper networks in Appendix B, where we derived directional derivatives for arbitrarily deep networks. Given our method, it should be conceptually easy to perform Taylor expansion based on directional derivatives for deeper networks to characterize and classify the stationary points. But it will be technically involved since there will be terms of much higher orders to control. Nonetheless, if one conducts Taylor expansion using our method, the equivalent concept of "escape neuron" for deeper networks should naturally emerge.

In this paper, we limit our main discussion to shallow networks since it already answers unresolved questions about non-smooth loss landscape [1-3], training dynamics [4-8], and nonsmooth network embedding [9]. These previous works all discussed shallow networks.

---

"2. The current study focuses on a non-differentiable activation function only, and the conclusion will be trivial if using a differentiable activation function. As more and more smooth activations are used such as GELU and ELU, the result will have less meaningful impact on practice."

From a practical point of view, ReLU has the advantage of low memory and computation complexity [10]. Thus, it can be advantageous for resource-constrained scenarios. The community has been revisiting ReLU [10,11] recently for its potential for efficiency.

From a theoretical point of view, in our paper, we developed tools for non-smooth, non-convex optimization landscape, which is generally hard and warrants close investigation.

---

"3. Fact 4.6 is an important conclusion to discuss the dynamics, but it was provided as an empirical observation. As this paper mainly focuses on theoretical characterization, it would increase the rigor of the analysis by providing a theoretical justification."

Fact 4.6 was proved in Section 9.5 of [8]. For clarity, we now mention this in footnote 4 in P7. Thank you for mentioning this.

---

You also raised the following questions, and we will answer them below.

"1. The escape neuron condition requires that the output weight of a specific neuron be zero. Does that imply the landscape of scalar output 2-layer neural network is rather trivial? It seems to be sufficient to just perturb the weight of this specific neuron to escape the saddle."

You are correct about the structure and the descent direction near the saddle point of such networks. This landscape, being non-convex and non-smooth, is typically challenging. There is a long active line of research dedicated to understanding it [1-9], and it was not until our paper that the relative simpleness of its structure is fully revealed. Nevertheless, we point out that the rigorous characterization of the training dynamics on such loss landscape in the general case is still an open question.

As a side note, we provide two new numerical experiments in Appendix I, one of which is for multidimensional outputs and discovers escape neurons whose output weights are not zero.

"2. Although the authors argue that the sub-differentiable definition is unsuitable for a non-convex function, does it also fail specifically in two-layer neural networks? In particular, it would be helpful to provide an example in the ReLU network to show that the sub-gradient can not properly capture the saddle point."

The definition of subgradient is the following:

Let $f: \mathbb{R}^n \to \mathbb{R}$ be a convex function. The subgradient set of $f$ at $\mathbf x_0$ is defined as:

$$\\partial f(\\mathbf x_0) = \\left\\{ \\mathbf g \\in \\mathbb{R}^n : f(\\mathbf x) \\geq f(\\mathbf x_0) + \\mathbf g^T (\\mathbf x - \\mathbf x_0), \\, \\forall \\mathbf x \\in \\mathbb{R}^n \\right\\}.$$

One can verify that the subgradient cannot be defined for the origin of the loss landscape of our setup in general. Since the directional derivatives toward all the directions are zero at the origin (by Lemma 3.2), if the subgradient $\{\mathbf g\}$ can be defined, it must mean that $\mathbf g = \mathbf 0$. Taking it to the definition of the subgradient, we find that, if the subgradient at the origin is $\{\mathbf g = \mathbf 0\}$, the origin must be a local minimum, which is generally not the case (see all of our numerical examples in Section 4.1.1 and Appendix I).

The origin is a stationary point (since the directional derivatives toward all directions are zero), but not a local minimum or a local maximum (Theorem D.1) in general. Thus, it is a saddle point where subgradients cannot be defined.

We have also added this example to line 1082.

---

Thank you again for your review! Your valuable suggestions have helped us improve our script. We believe our responses above address your questions, and we hope you consider raising your rating.

References

[1] Sahs, J., Pyle, R., Damaraju, A., Caro, J. O., Tavaslioglu, O., Lu, A., Anselmi, F., and Patel, A. B. Shallow univariate ReLU networks as splines: Initialization, loss surface, Hessian, and gradient flow dynamics, Frontiers in Artificial Intelligence, 2022.

[2] Cheridito, P., Jentzen, A., and Rossmannek, F. Landscape analysis for shallow neural networks: Complete classification of critical points for affine target functions, Journal of Nonlinear Science, 2022.

[3] Liu, B., Liu, Z., Zhang, T., and Yuan, T. Non-differentiable saddle points and sub-optimal local minima exist for deep ReLU networks, Neural Networks, 144, 2021.

[4] Chistikov, D., Englert, M., and Lazic, R. Learning a neuron by a shallow ReLU network: Dynamics and implicit bias for correlated inputs, NeurIPS '23.

[5] Kumar, A., and Haupt, J. Directional convergence near small initializations and saddles in two-homogeneous neural networks, Transactions on Machine Learning Research, 2024.

[6] Boursier, E., Pillaud-Vivien, L., and Flammarion, N. Gradient flow dynamics of shallow ReLU networks for square loss and orthogonal inputs, NeurIPS'22.

[7] Boursier, E., and Flammarion, N. Early alignment in two-layer networks training is a two-edged sword, 2024.

[8] Maennel, H., Bousquet, O., and Gelly, S. Gradient descent quantizes ReLU network features, 2018.

[9] Fukumizu, K., Yamaguchi, S., Mototake, Y., and Tanaka, M. Semi-flat minima and saddle points by embedding neural networks to overparameterization, NeurIPS'19.

[10] Tayaranian, M., Mozafari, S. H., Clark, J. J., Meyer, B., and Gross, W. Faster inference of integer SWIN transformer by removing the GELU activation, 2024.

[11] Wortsman, M., Lee, J., Gilmer, J., and Kornblith, S. Replacing softmax with ReLU in vision transformers, 2023.

Thank you again for your constructive comments! As the deadline for revising scripts draws near, we would like to know if you have further suggestions or concerns. We would be pleased to address them.

Best,

Authors

Thank you very much for your review! We appreciate that you consider our analysis to be novel and interesting and our presentation to be intuitive and insightful.

We made modifications to our script, highlighted in red, and marked with the blue-colored text 'k7xc' (your reviewer ID) to indicate changes based on your suggestions.

You also raised a few concerns, which we address below.

---

(1) You mentioned the following as one of the weaknesses:

"To understand the saddle escaping or saddle-to-saddle dynamics, it would be better to more analyses of gradient dynamics in the neighborhood of saddle points. For example, can we have the escape direction in Corollary 4.4 related to the gradient, implying that GF/GD can help escape the saddle points?"

You also asked the following question:

"Is it possible to come up with a more rigorous example (with gradient dynamics analysis) for Figure 4 to understand the saddle-to-saddle dynamics?"

Dynamics analysis around saddle points is a contribution of previous works [1-3], which our work extends.

Previous works, such as [1,2], gave rigorous examples with gradient dynamics analysis. Concretely, [1] studied a case where the training inputs are all orthogonal, and [2] studied a case where there is only one effective neuron learned. But as [3] pointed out, these previous works could not preclude other types of saddle points, as the techniques they developed for the dynamics in the vicinity of certain saddle points cannot offer insight into the global structure of the loss landscape. We fill this gap in the theory by devising a method to systematically characterize and classify the stationary points on the entire loss landscape.

To further contextualize our contribution, we add the following sentence to line 440:

To describe such a process, [1,2] rigorously characterized the gradient flow for toy examples, revealing trajectories consistent with Figure 4.

---

(2) You suggested:

"It is better to present those implications or corollaries of the stationary point characterization more formally instead of descriptions like Consequence 4.7."

Thank you for this suggestion. We will use more formal explanations.

• Consequence 4.7 (now Fact 4.8) is, in fact, rigorously proved with Taylor expansion, as is stated in Remark F.1. Nevertheless, to clarify the rigorousness of this result, we will mention in line 403 that this consequence is actually proved. We chose not to call it a corollary as it is a partial restatement of Theorem 4.2. But, to avoid confusion, we will now call it a "Fact" instead of a "Consequence" and stress that it is a partial restatement of Theorem 4.2.

• Moreover, to avoid ambiguity, we compile our results that preclude other types of saddle points into the following corollary.

Corollary 4.7: Following the setup and notation introduced in Section 2, let $\\mathbf P(t) = (\\mathbf w\_i(t), h\_{j\_0 i}(t))\_{i\\in I}\\in \\mathbb{R}^D$ be the parameter trajectory of a one-hidden-layer scalar-output ReLU network trained with the empirical squared loss $\\mathcal{L}$, where $t \\geq 0$ denotes time. Suppose the trajectory satisfies $\\frac{\\mathrm d\\mathbf P}{\\mathrm d t} = -\\frac{\\partial \\mathcal{L}}{\\partial \\mathbf P}$, in which we specify $\\frac{\\partial \\mathrm{ReLU}(x)}{\\partial x} = \\mathbb{1}\_{\\{x> 0\\}}$ in the chain rule. Let $\mathbf A\in \mathbb{R}^D$ be an arbitrary vector. We initialize the network with $\mathbf{P}(0)=\sigma \mathbf A$. Under these conditions, if $\lim\_{\sigma \searrow 0}\left(\inf\_t\Vert\mathbf P(t) - \overline{\mathbf P} \Vert\right) = 0$, where $\overline{\mathbf P}= (\overline{\mathbf w}\_i, \overline{ h}\_{j\_0 i})\_{i\in I}$ is a non-minimum stationary point, we must have $\overline{\mathbf w}\_i = \mathbf{0}$ and $\overline{ h}\_{j\_0 i} = 0$ for some $i\in I$.

Please refer to line 383 for more details.

---

(3)

"the toy examples only have one-dimensional inputs (with bias)... More empirical results need to be shown to corroborate the theoretical results' implications."

We supplement our discussion with two new numerical experiments in Appendix I, one with 3-dimensional input and scalar output, the other with 2-dimensional input and output. Both are consistent with our theories.

---

Thank you again for your review! Your valuable suggestions have helped us improve our script. We believe our responses above address your questions, and we hope you consider raising your rating.

References

[1] Boursier, E., Pillaud-Vivien, L., and Flammarion, N. Gradient flow dynamics of shallow ReLU networks for square loss and orthogonal inputs, NeurIPS'22.

[2] Chistikov, D., Englert, M., and Lazic, R. Learning a neuron by a shallow ReLU network: Dynamics and implicit bias for correlated inputs, NeurIPS '23.

[3] Kumar, A., and Haupt, J. Directional convergence near small initializations and saddles in two-homogeneous neural networks, Transactions on Machine Learning Research, 2024.