Stochastic Subgradient Methods with Guaranteed Global Stability in Nonsmooth Nonconvex Optimization
Nachuan Xiao, Xiaoyin Hu, Kim-Chuan Toh
The Chinese University of Hong Kong, Shenzhen · Shenzhen University · National University of Singapore
math.OC, cs.AI, cs.LG, stat.ML
Submitted: 2026-08-08
Updated: 2026-08-11
Comments: 50 pages
Code: https://github.com/google/automl
License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/
Importance score: 66/100
The gist: This paper focuses on "providing convergence guarantees for stochastic subgradient methods in minimizing nonsmooth nonconvex functions." The authors address the challenge that "most existing works on
Terminology
Summary
This paper focuses on providing convergence guarantees for stochastic subgradient methods in minimizing nonsmooth nonconvex functions.
The authors address the challenge that most existing works on convergence analysis only focus on cases where f is differentiable or weakly convex,
which excludes many applications in training neural networks with nonsmooth activation functions like ReLU and leaky ReLU.
The authors propose a general framework for stochastic subgradient methods
(SGM) defined by the following update scheme:
x k+1 in x k - eta k H delta k(x k) + xi k+1
In this framework, H: R n so R n is a set-valued mapping that is graph-closed, convex-valued, and locally bounded,
eta k denotes the stepsizes, delta k refers to the approximation parameters characterizing inexact evaluations of H,
and xi k characterizes the evaluation noises.
The paper's main contributions are divided into three primary areas:
1. Guaranteed global stability of the general framework (SGM)
The authors establish convergence guarantees for (SGM) by proving that, with a coercive Lyapunov function, the sequence of iterates generated by (SGM) is uniformly bounded under mild conditions.
They prove that for any epsilon > 0, by choosing sufficiently small (possibly non-diminishing) stepsizes eta k and approximation parameters delta k, coupled with sufficiently controlled evaluation noises xi k, the generated iterates x k will
eventually stabilize within an epsilon-neighborhood of the stable set of the Lyapunov function."
The stability results vary based on the sampling technique:
-
Random Reshuffling (RR): When evaluation noises are introduced by RR,
the iterates x k of our proposed framework (SGM) stabilize around the stable set almost surely with non-diminishing but sufficiently small stepsizes.
-
With-Replacement Sampling (WRS): When noises are introduced by WRS and stepsizes diminish at the rate of o(1/ (k)), the authors
prove the high-probability convergence properties for the iterates x k generated by (SGM).
2. Convergence properties of (GSGD)
The authors develop a scheme for SGD-type methods
(GSGD), which includes:
g k in D f i(x k), m k+1 = m k + tau eta k (g k - m k), x k+1 in x k - eta k (d phi(m k+1) + rho g k).
This scheme encompasses a wide range of SGD-type methods, including heavy-ball SGD, SignSGD, and normalized SGD.
The authors prove that the coercivity of f guarantees the coercivity of the corresponding Lyapunov function,
allowing the global stability of (SGM) to be applied. Furthermore, they show that with random initialization, almost surely, (GSGD) can avoid the spurious stationary points introduced by the conservative fields D f i,
and the sequence x k asymptotically converges towards the Clarke critical points of f almost surely, under mild conditions with random initialization.
3. Convergence properties of (ADM)
The authors introduce a scheme for ADAM-family methods
(ADM), which encompasses various variants of ADAM-family methods, such as the ADAM, AdaBelief, and NADAM.
The (ADM) scheme is defined as:
g k in D f i(x k), m k+1 = m k + tau 1 eta k (g k - m k), v k+1 in v k + tau 2 eta k (V(x k, m k+1) - v k), x k+1 in x k - eta k (P+(v k+1) + epsilon 0)-2 (m k+1 + rho g k).
Because the coercivity of f cannot guarantee the coercivity of its corresponding Lyapunov function
for (ADM), the authors "introduce an auxiliary update scheme to (ADM) parameterized by K > 0, which corresponds to a coercive Lyapunov function for all positive K. By proving that this auxiliary scheme
coincides with (ADM) for a specific choice of K > 0, they establish that
the global stability and avoidance of spurious stationary points of (ADM) directly follow from those properties of the proposed auxiliary scheme. This provides convergence guarantees for
several representative popular (time independent) stochastic subgradient methods... including heavy-ball SGD, SignSGD, normalized SGD, ClipSGD, ADAM, and AdaBelief."
Improvements for AI systems
1. Mathematically Grounded Optimizer Architectures
-
Improvement: Integrate the proposed GSGD (General SGD-type) and ADM (ADAM-family) update schemes directly into deep learning optimization libraries, replacing heuristic-based updates with the specific mathematical structures defined in the paper (incorporating the auxiliary coercive Lyapunov function for ADAM-family methods).
-
What the improved system can do: It can train deep neural networks using nonsmooth activation functions (like ReLU, Leaky ReLU, or Maxout) with a mathematical guarantee that the weight iterates will remain bounded and will not diverge to infinity, even when using non-diminishing stepsizes.
2. Optimized Sampling and Learning Rate Scheduling
-
Improvement: Implement a training controller that switches between Random Reshuffling (RR) and With-Replacement Sampling (WRS) based on the convergence requirements, utilizing the specific stepsize eta k and approximation parameter delta k constraints identified in the stability proofs.
-
What the improved system can do: When using Random Reshuffling, the system can maintain higher, non-diminishing learning rates to accelerate training without risking instability; when using With-Replacement Sampling, it can automatically apply the o(1/ (k)) decay rate to ensure high-probability convergence to a critical point.
3. Spurious Stationary Point Avoidance in Non-convex Landscapes
-
Improvement: Deploy the GSGD/ADM frameworks specifically in high-dimensional, nonsmooth, nonconvex loss landscapes, utilizing the
random initialization
property proven in the paper. -
What the improved system can do: It can navigate complex loss surfaces to bypass
spurious
(fake) stationary points that typically trap standard SGD, instead asymptotically converging to true Clarke critical points, leading to higher model accuracy and better generalization.
4. Robustness to Gradient Noise and Inexact Evaluations
-
Improvement: Incorporate the delta k (approximation parameter) and xi k (evaluation noise) control mechanisms into the training loop to manage precision in gradient computations (e.g., during mixed-precision training or quantized training).
-
What the improved system can do: It can maintain global stability and convergence guarantees even when using low-precision gradients or noisy hardware-level approximations, allowing for more efficient and energy-conscious AI training on edge devices.
Sources
- A closed-measure approach to stochastic approximation
- Symbolic Discovery of Optimization Algorithms
- Convergence guarantees for RMSProp and ADAM in non-convex optimization and an empirical comparison to Nesterov acceleration
- Adam-family Methods with Decoupled Weight Decay in Deep Learning
- Hamiltonian Descent Methods
- Provable Adaptivity of Adam under Non-uniform Smoothness
- Frictionless Hamiltonian Descent and Coordinate Hamiltonian Descent for Strongly Convex Quadratic Problems
- Large Batch Training of Convolutional Networks
- Large Batch Optimization for Deep Learning: Training BERT in 76 minutes
Related papers
- Lions and Muons: Optimization via Stochastic Frank-Wolfe under Heavy-Tailed Noise
- Adam-HNAG: A Convergent Reformulation of Adam with Accelerated Rate
- Incremental Learning in Mirror Flows
- Online Control via Counterfactual Tracking
- Asynchronous Replanning in Two Population Linear Quadratic Mean Field Games: Information Requirements and Stability
- Petrov-Galerkin operator inference with application to stability-encouraging identification