Gradient descent (GD) with large learning rates can successfully optimize neural networks even though classical smoothness-based learning rate conditions are violated. Yet theory for this regime largely covers linear models or linearly separable data. We analyze a minimal setting outside both regimes. Specifically, a two-neuron ReLU network trained with logistic loss on 1D non-linearly separable data. The objective is nonconvex and nonsmooth, and GD exhibits edge-of-stability (EoS) behavior. We identify an active cone in parameter space that contains every perfect classifier, where the problem becomes piecewise logistic regression with a varying effective dataset. Within this cone, we prove a finite-time loss bound for constant learning rates and show that GD converges in direction to an explicit max-margin vector, while the orthogonal component converges to the unique minimizer of a non-strongly convex potential. We further establish catapult results where exponentially small loss is reached after a single loss spike. Furthermore, we propose an adaptive schedule that lets learning rates grow to infinity without leaving the active cone. We prove that this schedule achieves a loss rate of O(1/(TlogT), which is faster than the standard O(1/T) rate. Our theoretical results extend to local convergence on more general network and are validated on MNIST dataset. Overall, we show that this setup provides a tractable neural-network model for studying large and adaptive step sizes beyond linear predictors and separable data.
Figures & tables
Figure 1: Training a simplified ReLU network with GD on a three point non-linearly separable dataset. (a) shows the loss landscape with respect to the hidden unit weight and bias w and b . The active cone C contains all perfect classifiers, the dead region D has zero gradient, while the loss on E is bounded below by a constant. Two GD trajectories with η=1 are shown. (b) Loss and (c) normalized sharpness ηλmax(∇2L(θt))/2 (above 1 means above the EoS threshold) for three runs; the loss is non-monotone and the sharpness exceeds 2/η . (d) is the evolution of the learned classifier on the data points for the initialization (0.2,0.1,0) with η=1 , which approaches (2∣x∣−1)log(ηt)+(1−∣x∣)log3 ( Theorem 4.5 ).
Figure 2: Two-step catapult. GD from θ0=(0.2,0.1,−0.1) with constant step sizes η=16,64,256 .
Figure 3: Phase diagram of constant-step GD ( b0=∣w0∣/2 , g0=0 , T=3000 ). Color: final loss; cyan: boundary of the set of runs that ever leave C . Runs above the boundary end in the dead region with L=log427 . White lines: sufficient conditions (I1) and (I2) . Arrow: catapult regime of Theorem 5.1 .
Figure 4: Late-stage ReLU neural network training on binary MNIST exhibits behaviour similar to the simplified model on the three-point dataset. We train a two-layer neural network with 16 hidden neurons. (a) and (b) use a trainable output layer while (c) and (d) and use a fixed output layer. The loss (a) is asymptotically O(1/(tlnt)) while the loss in (c) is asymptotically O(1/t) . (b) and (d) show the activation patterns for certain examples towards the end of training. Black means the hidden unit is active i.e. the ReLU is not zero. Even in the late training stage, the activation patterns continue to switch which mirrors the mechanism in our simplified model. Each example moves between different feature vectors in the transformed space, giving rise to piecewise logistic regression dynamics.
Symbol
Meaning
Value
ξt
index of the active middle feature
− if bt≤0 , + if bt>0
ne,nξ
multiplicities of the active features
2,1
γ
max-margin of the signed features
1/5
hs
component of xs orthogonal to ϑ^
he=(53,1) , h−=(−52,0) , h+=(−52,−2)
H
largest norm of an orthogonal component
maxs∥hs∥=24/5
m
margin offset ( Lemma D.1 )
2/30
Table 1: Notation. Here s∈{e,−,+} indexes the three signed features.
Appendix figures & tables5 assets
Supplementary material from the paper’s appendix.
Appendix
Figure 5: Two-step catapult from θ0=(0.2,0.1,−0.1) , for which ν≈0.375 . (a) The loss increases after the first update and drops after the second. (b) The normalized losses approach the limits in Theorem 5.1 .
Configuration
Digit
MNIST index
Patterns
New patterns
Fully trainable
3
46247
49
10
Fully trainable
5
38658
16
0
Fixed output
3
46247
56
0
Fixed output
5
20171
25
0
Appendix
Table 2: Activation-pattern measurements for the displayed examples. MNIST indices are zero-based indices in the original training split. “Patterns” counts distinct gate vectors in the final 150 iterates. “New patterns” counts distinct vectors in iterates 9001–10000 that were absent from iterates 5001–9000.
Configuration
Final loss
Train accuracy
Test accuracy
Fully trainable
4.89253×10−5
100%
94.53%
Fixed output
5.59540×10−4
100%
95.22%
Appendix
Table 3: Final performance after 104 full-batch gradient-descent updates on the MNIST 3-versus-5 training subset.
Figure 6: Implicit bias with constant steps ( θ0=(0.2,0.1,0) , η∈{0.25,0.5,1} , 105 iterations). (a) Coordinates of ϑt in the rotated basis: qt (black) and bt (gray) for η=1 , and pt−5logt for each η (colors as in (b)); dotted lines are the limits of Theorem 4.5 (a). (b) tL(θt)→5/η . (c) Endpoint and middle margins minus their predicted growth tend to zero.
We study the convergence dynamics of Gradient Descent (GD) in a minimal binary classification setting, consisting of a two-neuron ReLU network and two training instances. We prove that even under these strong simplifying assumptions, while GD successfully converges to an optimal robustness margin, effectively maximizing the distance between the decision boundary and the training points, this convergence occurs at a prohibitively slow rate, scaling strictly as Θ(1/ln(t)). To the best of our knowledge, this establishes the first explicit lower bound on the convergence rate of the robustness margin in a non-linear model. Through empirical simulations, we further demonstrate that this inherent failure mode is pervasive, exhibiting the exact same tight convergence rate across multiple natural network initializations. Our theoretical guarantees are derived via a rigorous analysis of the GD trajectories across the distinct activation patterns of the model. Specifically, we develop tight control over the system's dynamics to bound the trajectory of the decision boundary, overcoming the primary technical challenge introduced by the non-linear nature of the architecture.
Guy Smorodinsky, Sveta Gimpleson, Itay Safran
Faculty of Computer and Information Science, Ben-Gurion University of the Negev, Israel · Department of Computer Science, Reichman University, Israel
In this paper, we study the gradient descent dynamics for jointly training both layers of a one-hidden-layer ReLU network to fit a linear target function. Concretely, we consider a realizable setting where inputs are drawn i.i.d. from a Gaussian distribution and labels follow a planted linear model. This stylized framework captures salient features of end-to-end training in inverse problems and certain auto-encoder models. Despite its apparent simplicity, the dynamics remain poorly understood, in part because the loss landscape contains multiple non-strict saddle points, making it unclear why gradient descent from random initialization reliably escapes bad stationary regions. We provide a detailed characterization of the optimization landscape and prove that gradient descent from a moderately small random initialization-simultaneously training both layers-converges to a global minimizer at a linear rate with order-wise optimal sample complexity. Our analysis tracks the trajectory through three phases: an alignment phase in which hidden weights progressively align with the planted direction while the output weights maintain the correct sign pattern; a growth phase in which the norms of both layers increase while preserving alignment; and a local refinement phase in which the aligned neurons rapidly converge to the planted direction, yielding fast local convergence. To rigorously show that GD avoids non-strict saddles, we develop trajectory-level control arguments for the end-to-end dynamics. In addition, we establish novel uniform concentration results that hold along the entire trajectory, and are essential for obtaining order-wise optimal sample complexity. We corroborate our theory with extensive experiments across a range of configurations.
Berk Tinaz, Changzhi Xie, Mahdi Soltanolkotabi
Department of Electrical and Computer Engineering University of Southern California Los Angeles, CA, USA · Department of Computer Science University of Southern California Los Angeles, CA, USA
Adaptive gradient methods can favor max-margin separators that differ from gradient descent, yet a fixed positive numerical stability constant eventually changes the update geometry again. This paper studies the rate-controlled middle case for full-batch linear classification on separable data. For memoryless stability-annealed smoothed-sign descent with weighted exponential loss, we prove that the normalized iterates converge to the minimizer of a convex Burg-type barrier over a margin slice. The proof rewrites the dynamics exactly as entropic mirror ascent on a concave dual objective, controls the dual gap by a KL recursion, and yields an explicit S_t^{-1/2} normalized-iterate envelope. The static barrier geometry is fully characterized, including KKT conditions and both endpoint limits. Experiments validate the exact dual identities to floating-point error, illustrate the predicted path and rate diagram, and show an empirical fixed-epsilon crossover scaling in cumulative time. We further report robustness and boundary diagnostics for logistic tails, fixed-epsilon crossover, and adaptive-method variants, delineating the scope of the proved smoothed-sign theory.
Xiangwu Wang, Chengwei Cao, Yicheng Song +2
1The University of Hong Kong · University of California, San Diego · 3Beijing University of Aeronautics and Astronautics