A Minimal ReLU Classifier on Non-Linearly Separable Data: Convergence and Implicit Bias at and beyond the Edge of Stability

Quan Nguyen; Robin Yadav; Param Mody; Mark Schmidt

cs.LGOct 16, 2026

A Minimal ReLU Classifier on Non-Linearly Separable Data: Convergence and Implicit Bias at and beyond the Edge of Stability

Authors: Quan Nguyen, Robin Yadav, Param Mody, Mark Schmidt

Organizations: University of British Columbia · Canada CIFAR AI Chair (Amii)

Abstract

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/(T\log{T})$ , 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 $\mathcal{C}$ contains all perfect classifiers, the dead region $\mathcal{D}$ has zero gradient, while the loss on ${\mathcal{E}}$ is bounded below by a constant. Two GD trajectories with $\eta=1$ are shown. (b) Loss and (c) normalized sharpness $\eta\lambda_{\max}(\nabla^{2}\mathcal{L}(\theta_{t}))/2$ (above $1$ means above the EoS threshold) for three runs; the loss is non-monotone and the sharpness exceeds $2/\eta$ . (d) is the evolution of the learned classifier on the data points for the initialization $(0.2,0.1,0)$ with $\eta=1$ , which approaches $(2|x|-1)\log(\eta t)+(1-|x|)\log 3$ ( Theorem 4.5 ).

Figure 2: Two-step catapult. GD from

\theta_{0}=(0.2,0.1,-0.1)

with constant step sizes

\eta=16,64,256

Figure 3: Phase diagram of constant-step GD ( $b_{0}=|w_{0}|/2$ , $g_{0}=0$ , $T=3000$ ). Color: final loss; cyan: boundary of the set of runs that ever leave $\mathcal{C}$ . Runs above the boundary end in the dead region with $\mathcal{L}=\log\frac{27}{4}$ . White lines: sufficient conditions (I1) and (I2) . Arrow: catapult regime of Theorem 5.1 .

Panel of Figure 4 — 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 $\mathcal{O}(1/(t\ln{t}))$ while the loss in (c) is asymptotically $\mathcal{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
$\xi_{t}$	index of the active middle feature	$-$ if $b_{t}\leq 0$ , $+$ if $b_{t}>0$
$n_{\mathrm{e}},\,n_{\xi}$	multiplicities of the active features	$2,\ 1$
$\gamma$	max-margin of the signed features	$1/\sqrt{5}$
$h_{s}$	component of $x_{s}$ orthogonal to $\hat{\vartheta}$	$h_{\mathrm{e}}=\big(\tfrac{3}{\sqrt{5}},1\big)$ , $h_{-}=\big(-\tfrac{2}{\sqrt{5}},0\big)$ , $h_{+}=\big(-\tfrac{2}{\sqrt{5}},-2\big)$
$H$	largest norm of an orthogonal component	$\max_{s}\\|h_{s}\\|=\sqrt{24/5}$
$m$	margin offset ( Lemma D.1 )	$2/\sqrt{30}$

Table 1: Notation. Here

s\in\{\mathrm{e},-,+\}

indexes the three signed features.

Appendix figures & tables5 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 5: Two-step catapult from

\theta_{0}=(0.2,0.1,-0.1)

, for which

\nu\approx 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\times 10^{-5}$	$100\%$	$94.53\%$
Fixed output	$5.59540\times 10^{-4}$	$100\%$	$95.22\%$

Appendix

Table 3: Final performance after

10^{4}

full-batch gradient-descent updates on the MNIST 3-versus-5 training subset.

Figure 6: Implicit bias with constant steps (

\theta_{0}=(0.2,0.1,0)

\eta\in\{0.25,0.5,1\}

10^{5}

iterations). (a) Coordinates of

\vartheta_{t}

in the rotated basis:

q_{t}

(black) and

b_{t}

(gray) for

\eta=1

, and

p_{t}-\sqrt{5}\log t

for each

\eta

(colors as in (b)); dotted lines are the limits of Theorem 4.5 (a). (b)

t\,\mathcal{L}(\theta_{t})\to 5/\eta

. (c) Endpoint and middle margins minus their predicted growth tend to zero.

schedule	$T$	stays in $\mathcal{C}$	median $\log_{10}\mathcal{L}(\theta_{T})$	max $\log_{10}\mathcal{L}(\theta_{T})$
$\eta_{t}=\|w_{t}\|/2$	$2\cdot 10^{4}$	$20/20$	$-4.63$	$-4.63$
$\eta_{t}=1/\mathcal{L}(\theta_{t})$	$300$	$18/20$	$-25.16$	$\phantom{-}0.28$
$\eta_{t}=\|w_{t}\|^{2}/2$	$2000$	$0/20$	$\phantom{-}1.05$	$\phantom{-}2.21$

Appendix

Table 4: Twenty random initializations in

\mathcal{C}

Neural Network Optimization

Neural network optimization studies algorithms and training dynamics that help models learn efficiently, reliably, and generalize well.

See all 321 papers

On the Rate of Convergence of GD in Non-linear Neural Networks: An Adversarial Robustness Perspective
Mar 2, 2026Guy Smorodinsky, Sveta Gimpleson, Itay SafranNeural Network Training Dynamics Gradient Descent Dynamics
Faculty of Computer and Information Science, Ben-Gurion University of the Negev, Israel · Department of Computer Science, Reichman University, Israel
When Both Layers Learn: Training Dynamics of Representing Linear Models via ReLU Networks
Jun 3, 2026Berk Tinaz, Changzhi Xie, Mahdi SoltanolkotabiReLU Neural Networks Neural Network Optimization
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
Stability Annealing Selects the Implicit Bias of Smoothed Sign Descent: A Rate-Indexed Barrier Path on Separable Data
Jul 7, 2026Xiangwu Wang, Chengwei Cao, Yicheng Song +2Sign-Based Optimization Implicit Regularization
1The University of Hong Kong · University of California, San Diego · 3Beijing University of Aeronautics and Astronautics

Mar 2, 2026J/K move · Enter open · S save

On the Rate of Convergence of GD in Non-linear Neural Networks: An Adversarial Robustness Perspective

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

Neural Network Training Dynamics Gradient Descent Dynamics Optimization Convergence Analysis Adversarial Robustness

Open paper AskPDF arXiv

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.

Abstract

Neural Network Optimization

Explore similar work

On the Rate of Convergence of GD in Non-linear Neural Networks: An Adversarial Robustness Perspective

When Both Layers Learn: Training Dynamics of Representing Linear Models via ReLU Networks

Stability Annealing Selects the Implicit Bias of Smoothed Sign Descent: A Rate-Indexed Barrier Path on Separable Data