First-Order and Stochastic Optimization Methods for Machine Learning V: Nonconvex Stochastic Mirror DescentTextbook
Motivation
Most machine learning training objectives — deep network losses, matrix factorization, regularized empirical risk with a nonconvex loss — are not convex, yet the great majority of convergence theory available before Ghadimi and Lan's 2013 work applied only to convex problems or gave no non-asymptotic rate at all. Ghadimi and Lan (2013) established the first non-asymptotic complexity bounds for stochastic first-order methods on smooth nonconvex problems, using the norm of a gradient mapping (rather than function-value suboptimality, which is meaningless without convexity) as the convergence measure, together with a randomized stopping rule that removes the need to know in advance which iterate will be best. This mission formalizes the constrained, composite generalization of that theory — Lan's own extension (2020) to problems with a nonsmooth term and a general Bregman geometry rather than the Euclidean norm — culminating in the stochastic complexity bound for the randomized stochastic mirror descent (RSMD) algorithm.
Setting
Fix a nonempty closed convex , a continuously differentiable (possibly nonconvex) with -Lipschitz gradient, and a simple convex (possibly nonsmooth) (e.g. or ); write , (assumed finite). For a distance-generating function with modulus 1 and its prox-function , the generalized projection at with gradient-like input and stepsize is
which reduces to itself when and : is a generalized projected gradient (or gradient mapping) of at , and its norm going to zero is the right notion of "approximately stationary" for the composite, possibly-nonconvex problem .
The randomized stochastic mirror descent (RSMD) algorithm, given only a stochastic first-order oracle returning with and (Assumption 13), forms a mini-batch average of oracle calls at each step , updates via the generalized projection with , and stops at a randomly chosen index (drawn from a prescribed pmf , independently of the optimization process) rather than a deterministic final iterate.
Formalization targets
Goal — Theorem 6.6(a), RSMD complexity
for (strict for at least one ) and chosen as in (6.2.30), the expectation over both and the oracle randomness .
Supporting milestones, in attack order
- Lemma 6.4: — the bound that lets a smoothness inequality on become a descent inequality on the whole composite .
- Lemma 6.6: the three-point characterization of , the composite-problem analogue of Chapter 3's Lemma 3.4.
- Theorem 6.5 (deterministic ancestor): for the exact-gradient nonconvex MD algorithm.
- Corollary 6.4: the constant-stepsize instantiation .
Every result states its constants exactly as the book derives them; no milestone or the goal hides a rate behind an unspecified .
Significance
The goal theorem gives the complexity of the RSMD algorithm in terms of a squared generalized gradient-mapping norm — the correct convergence criterion for constrained, composite, possibly nonconvex stochastic optimization, since function-value suboptimality is not controllable without convexity and unconstrained gradient norms are meaningless once or is nonsmooth. Choosing and appropriately (a corollary this mission does not formalize) turns this bound into the celebrated total-oracle-call complexity for finding an -stationary point in expectation — the standard benchmark every later stochastic nonconvex method (variance-reduced SGD, SPIDER, and their composite/constrained variants) is compared against.
No result in this mission has a machine-checked proof on Prove2Me under this exact hypothesis set.
The two closest platform results, both from lean-optrates (Shi), are genuinely different
objects: ShiOptRates.gd_exact_rate is plain, unconstrained, deterministic gradient descent
(, no set , no composite , no generalized projection), and
ShiOptRates.Stochastic.sgd_rate is plain SGD under the same unconstrained, non-composite
setup — its filtration/conditional-expectation formalization pattern (a Filtration ℕ, μ[·|ℱ k] for the unbiasedness and variance-bound hypotheses) is the same one this mission's goal
theorem uses, confirming it as the platform's established idiom for this class of result, but the
mathematical content (plain gradient step vs. generalized-projection/mirror-descent step, no
or ) is different. Neither is reused; both are noted as the platform's nearest existing work.
Difficulty
The generalized projection replaces the Euclidean projection with an arbitrary Bregman-based prox-mapping and absorbs the nonsmooth term directly into the subproblem — a formalization that quietly assumes or would collapse every milestone here into the special case and prove nothing about the constrained composite problem the chapter is actually about. The harder difficulty is in the goal theorem's own randomness: the book's proof does not use an unconditional (marginal) form of Assumption 13, because from step 2 onward is itself a random variable (a function of the history ), so the cross-term the proof needs to vanish requires a conditional statement — "" is the book's own phrasing. A formalization using only marginal moment bounds would either be unprovable as stated or, worse, would misstate the theorem by using hypotheses too weak for the claimed conclusion.
Formalization scope
generalized_projection_gradient_bound, generalized_projection_characterization,
nonconvex_md_bound and nonconvex_md_rate are stated over a real inner product space (Chapter
6's own generality — unlike Chapter 3, §6.2.3 explicitly restricts to "the norm associated with
the inner product"), with every argmin-defined point (, and the iterate sequence )
represented by its pointwise minimality property rather than an argmin term, consistent with
this series' convention. The goal theorem, rsmd_complexity_bound, additionally introduces a
probability space (Ω,P) and a Mathlib Filtration ℕ 𝒢, with x k/G k required
𝒢(k-1)-strongly-measurable and Assumption 13 stated via MeasureTheory.condExp (𝒢 (k-1))
(conditional mean 0, conditional second moment ≤ σ²/m_k) — the conditional form the book's own
proof actually needs, not a weaker marginal substitute. The σ²/m_k bound is (6.2.40)'s
conclusion for the m_k-sample batch average, taken as a hypothesis on the already-averaged G k
directly rather than re-derived from m_k raw i.i.d. calls (that derivation is not itself a
numbered result of the book). 's independence from the process is stated via
ProbabilityTheory.IndepFun; every integrability side condition the conclusion's Bochner
integral needs to be non-vacuous is stated explicitly, guarding against the well-known trap of an
uninhabited/non-integrable hypothesis silently defaulting condExp/the integral to 0 and making
the theorem trivially true.
A trivializing formalization this mission rules out: taking and
throughout would make every generalized projection collapse to the ordinary gradient, reducing
this entire mission to a restatement of plain (stochastic) gradient descent — exactly the
ShiOptRates results already on the platform — rather than the constrained composite theory the
chapter develops; , and are kept as genuine free parameters in every milestone and the
goal.
Left out of scope, for time: Theorem 6.6(b) (the convex-case corollary on , requiring the nondecreasing/nonincreasing stepsize side-conditions of (6.2.33)/(6.2.35)); the raw-sample derivation of (6.2.40); Lemma 6.3 (the stationarity consequence of a small gradient mapping, using and the normal cone ); the 2-RSMD algorithm and its large-deviation improvement; and the gradient-free (RSMDF) variant.
Selected references
- G. Lan, First-Order and Stochastic Optimization Methods for Machine Learning, Springer Series in the Data Sciences, Springer 2020, §6.2. https://doi.org/10.1007/978-3-030-39568-1
- S. Ghadimi and G. Lan, "Stochastic First- and Zeroth-Order Methods for Nonconvex Stochastic Programming," SIAM Journal on Optimization, 23(4), 2013, pp. 2341–2368.
- S. Ghadimi, G. Lan and H. Zhang, "Mini-batch Stochastic Approximation Methods for Nonconvex Stochastic Composite Optimization," Mathematical Programming, 155(1–2), 2016, pp. 267–305 (the RSMD algorithm's original source).