Convergence and Dynamical Behavior of the ADAM Algorithm for Nonconvex Stochastic Optimization 2: The Continuous-Time Adam Trajectory Converges to the Critical Points of FResearch Paper
Motivation
Adam (Kingma and Ba, 2015) is the default optimizer for training neural networks. At step it keeps an exponential moving average of stochastic gradients and an exponential moving average of their squared coordinates, corrects both for their initialization at zero (the bias correction or debiasing step), and moves the iterate by coordinatewise. Despite its use, its convergence on nonconvex objectives was poorly understood when Barakat and Bianchi wrote their paper: Reddi, Kale and Kumar (2018) exhibited convex online problems on which constant-parameter Adam fails to converge, and earlier analyses either dropped the bias correction or replaced Adam by a modified algorithm.
Barakat and Bianchi (arXiv:1810.02263v4, SIAM J. Math. Data Sci. 2021) study Adam through the ODE method: as the step size with and , the interpolated iterates follow a non-autonomous ordinary differential equation, the continuous-time version of Adam. This mission formalizes their convergence theorem for that ODE (Theorem 3.2): trajectories started as Adam is started, with zero moment estimates, approach the critical points of the objective.
Setting
Let be continuously differentiable with locally Lipschitz gradient , and let be locally Lipschitz. In the paper is the objective and the coordinatewise second moment of the stochastic gradient; §7 only uses the two properties above. Assume is coercive ( as ) and coordinatewise for every . Fix constants with .
The state is , and is the part where . All vector operations are coordinatewise. The Adam field (3.3) is, for ,
and (ODE) is . The factors and are the continuous-time bias corrections; they blow up as . A global solution with initial condition is a continuous , continuously differentiable on , satisfying (ODE) for every , with . The critical set is .
The proof passes through the autonomous field , its semiflow on , and the notions of equilibrium point, strict Lyapunov function and asymptotic pseudotrajectory (APT) of a semiflow.
Formalization targets
Goal: Theorem 3.2
Assume in addition that has empty interior in . For every global solution with initial condition , the set is non-empty and
The theorem does not claim that converges; convergence to a single critical point needs a Łojasiewicz condition (Theorem 3.4, a separate mission).
Milestones
The milestones are the numbered results of §7.1–7.3 used in the proof, in the order the proof uses them: positivity of along solutions (Lemma 7.4); decrease of along (Lemma 7.5, autonomous part); compactness of trajectories (Propositions 7.6, 7.7) and lower bounds on (Lemma 7.8); well-posedness of the autonomous system and its semiflow (Proposition 7.13); the abstract limit-set theorem for APTs (Proposition 7.14, quoted from Benaïm); the strict Lyapunov function (Proposition 7.15); and the APT property of the (ODE) solution (Proposition 7.17).
Significance
Theorem 3.2 is the deterministic core of the paper's analysis of Adam. Combined with the paper's tracking result (constant-step Adam converges in probability to the ODE solution as , Theorem 4.3) it says that for small steps Adam's iterates spend their time near critical points of ; it is also the input of the decreasing-step almost-sure convergence theorem (Theorem 5.2) and of the convergence-rate theorem (Theorem 3.4), which both rely on the same Lyapunov function and limit-set argument. The condition under which the Lyapunov function decreases is satisfied by the default parameters , .
The result is proved in the paper; to the best of current knowledge it has not been machine-checked. The mission produces a formal proof of the theorem, and reusable formal statements about semiflows on metric spaces: strict Lyapunov functions, and the limit set of a relatively compact asymptotic pseudotrajectory.
Difficulty
The obvious route is to show that decreases and apply LaSalle's invariance principle. This fails on two counts. First, (ODE) is non-autonomous, and LaSalle's principle is a statement about autonomous systems: the derivative of along the flow vanishes only on a set that depends on . Second, the field is singular: blows up at and is not Lipschitz near , because of , so neither existence and uniqueness nor continuous dependence on initial data come from standard theorems. The energy of the autonomous limit is a Lyapunov function but not a strict one: it is constant along some non-equilibrium orbits, so non-increase of an energy alone does not identify the limit points.
Formalization scope
Each block lives in EuclideanSpace ℝ (Fin d) and is their product; Mathlib's norm on the product is the maximum of the three Euclidean norms, which is equivalent to the Euclidean norm of and changes no statement. and are abstract, with the hypotheses of the paper's §7 (Assumptions 7.1, 7.2, 2.3, 2.4); this generality implies the printed theorem for the , of (2.2). Solutions are maps constrained only on (or ), and the equation is required at only. The field is extended off by , as on p. 12 of the paper. A semiflow is Mathlib's Flow ℝ≥0 M, on the subtype or on a compact subset of it. The semiflow in the milestones is always assumed to be the flow of , never an arbitrary flow. The APT and limit-set definitions are the published StochApproxDyn.LimitSet ones.
The goal is stated for every global solution of the specific field (3.3), and keeps " is non-empty" in the conclusion. A statement about some solution of an unspecified field, or with " has empty interior" strengthened to " is finite", is a different theorem and does not count.
A complete development needs: Grönwall-type comparison and chain-rule arguments along trajectories; well-posedness of the autonomous system on and continuity of its flow; and the limit-set theory of asymptotic pseudotrajectories (Benaïm 1999), adapted to strict Lyapunov functions in the paper's sense. The last part is independent of Adam and reusable in stochastic approximation. Contributions to any milestone are welcome; Propositions 7.13 and 7.14 are the largest pieces of infrastructure.
Selected references
- A. Barakat, P. Bianchi, Convergence and Dynamical Behavior of the ADAM Algorithm for Nonconvex Stochastic Optimization, SIAM J. Math. Data Sci. 3(1), 2021; arXiv:1810.02263v4. https://arxiv.org/abs/1810.02263v4
- D. P. Kingma, J. Ba, Adam: A Method for Stochastic Optimization, ICLR 2015. https://arxiv.org/abs/1412.6980
- S. J. Reddi, S. Kale, S. Kumar, On the Convergence of Adam and Beyond, ICLR 2018. https://arxiv.org/abs/1904.09237
- M. Benaïm, Dynamics of stochastic approximation algorithms, Séminaire de Probabilités XXXIII, Lecture Notes in Math. 1709, Springer, 1999. https://doi.org/10.1007/BFb0096509