Unifying Optimization and Dynamics to Parallelize Sequential Computation: A Guide to Parallel Newton Methods for Breaking Sequential Bottlenecks

Xavier Gonzalez

math.NAMar 17, 2026

Unifying Optimization and Dynamics to Parallelize Sequential Computation: A Guide to Parallel Newton Methods for Breaking Sequential Bottlenecks

Authors: Xavier Gonzalez

Organizations: STANFORD UNIVERSITY

Abstract

Massively parallel hardware (GPUs) and long sequence data have made parallel algorithms essential for machine learning at scale. Yet dynamical systems, like recurrent neural networks and Markov chain Monte Carlo, were thought to suffer from sequential bottlenecks. Recent work showed that dynamical systems can in fact be parallelized across the sequence length by reframing their evaluation as a system of nonlinear equations, which can be solved with Newton's method using a parallel associative scan. However, these parallel Newton methods struggled with limitations, primarily inefficiency, instability, and lack of convergence guarantees. This thesis addresses these limitations with methodological and theoretical contributions, drawing particularly from optimization. Methodologically, we develop scalable and stable parallel Newton methods, based on quasi-Newton and trust-region approaches. The quasi-Newton methods are faster and more memory efficient, while the trust-region approaches are significantly more stable. Theoretically, we unify many fixed-point methods into our parallel Newton framework, including Picard and Jacobi iterations. We establish a linear convergence rate for these techniques that depends on the method's approximation accuracy and stability. Moreover, we give a precise condition, rooted in dynamical stability, that characterizes when parallelization provably accelerates a dynamical system and when it cannot. Specifically, the sign of the Largest Lyapunov Exponent of a dynamical system determines whether or not parallel Newton methods converge quickly. In sum, this thesis unlocks scalable and stable methods for parallelizing sequential computation, and provides a firm theoretical basis for when such techniques will and will not work. This thesis also serves as a guide to parallel Newton methods for researchers who want to write the next chapter in this ongoing story.

Figures & tables

Panel of Figure 1 — Figure 1: Unrolling an SSM. We shade the initial state $s_{0}$ to indicate that we know the initial condition.

Panel of Figure 2 — Figure 2: Graphical diagram showing the equivalence (based on currying) between an SSM driven by inputs and an autonomous system with time-varying transition dynamics. We shade the inputs $u_{t}$ to indicate that they are known.

SSM	State ( $s_{t}$ )	Input ( $u_{t}$ )	Transition ( $f$ )
Linear Dynamical Systems (LDS) [ 154 ]	State	Input	Linear
Deep SSMs [ 87 , 76 , 208 ]	Stack of states	Input	Linear
Recurrent Neural Networks (RNNs) [ 118 , 55 , 101 , 219 , 38 ]	Hidden state	Input	RNN cell
MCMC [ 53 , 71 , 244 ]	Current sample	Noise	Transition kernel
Sampling from diffusion models [ 214 , 201 , 137 ]	Noisy image	Noise	Denoising function
Explicit differential equation solvers [ 67 , 125 , 111 ]	Current state	N/A	Numerical integrator

Table 1: Some illustrative examples of state space models (SSMs).

Panel of Figure 3 — Figure 3: A linear Gaussian state space model (LGSSM): The LGSSM consists of latent variables $s_{t}$ and observed variables $o_{t}$ . The generative model of the LGSSM consists of dynamics $s_{t+1}\sim\mathcal{N}\left(As_{t},Q\right)$ and emissions $o_{t+1}\sim\mathcal{N}\left(Cs_{t+1},R\right)$ .

Panel of Figure 4 — Figure 4: Parallel Scan for Matrix Multiplication. We illustrate a divide-and-conquer approach to compute the product $A_{4}A_{3}A_{2}A_{1}$ . Note that this divide-and-conquer approach naturally leads to $\mathcal{O}(\log T)$ depth.

Pos. 1	Pos. 2	Pos. 3	Pos. 4	Pos. 5	Pos. 6	Pos. 7	Pos. 8
${\color[rgb]{0.5,0.5,0.5}A_{1}}$	${\color[rgb]{0.5,0.5,0.5}A_{2}}$	${\color[rgb]{0.5,0.5,0.5}A_{3}}$	${\color[rgb]{0.5,0.5,0.5}A_{4}}$	${\color[rgb]{0.5,0.5,0.5}A_{5}}$	${\color[rgb]{0.5,0.5,0.5}A_{6}}$	${\color[rgb]{0.5,0.5,0.5}A_{7}}$	${\color[rgb]{0.5,0.5,0.5}A_{8}}$
$A_{1}$	${\color[rgb]{0,0,1}A_{1:2}}$	$A_{3}$	${\color[rgb]{0,0,1}A_{3:4}}$	$A_{5}$	${\color[rgb]{0,0,1}A_{5:6}}$	$A_{7}$	${\color[rgb]{0,0,1}A_{7:8}}$
$A_{1}$	$A_{1:2}$	$A_{3}$	${\color[rgb]{0,0,1}A_{1:4}}$	$A_{5}$	$A_{5:6}$	$A_{7}$	${\color[rgb]{0,0,1}A_{5:8}}$
$A_{1}$	$A_{1:2}$	$A_{3}$	$A_{1:4}$	$A_{5}$	$A_{5:6}$	$A_{7}$	${\color[rgb]{0,0,1}A_{1:8}}$

Table 2: Up-sweep for multiplying

A_{1},A_{2},\ldots A_{8}

Pos. 1	Pos. 2	Pos. 3	Pos. 4	Pos. 5	Pos. 6	Pos. 7	Pos. 8
$A_{1}$	$A_{1:2}$	$A_{3}$	$A_{1:4}$	$A_{5}$	${\color[rgb]{0,0,1}A_{1:6}}$	$A_{7}$	$A_{1:8}$
$A_{1}$	$A_{1:2}$	${\color[rgb]{0,0,1}A_{1:3}}$	$A_{1:4}$	${\color[rgb]{0,0,1}A_{1:5}}$	$A_{1:6}$	${\color[rgb]{0,0,1}A_{1:7}}$	$A_{1:8}$

Table 3: Down-sweep for multiplying

A_{1},A_{2},\ldots A_{8}

Figure 5: Newton’s method for root-finding. Here we illustrate 3 iterations of Newton’s method for root-finding on the one-dimensional cubic function

\mathbf{r}(\mathbf{s})=(\mathbf{s}-0.4)^{3}+0.45(\mathbf{s}-0.4)

. We observe that each iteration of Newton’s method involves linearizing the function to obtain

\mathbf{\hat{r}}^{(i)}(\cdot)

(shown in color) and then finding the zero of this linearization to obtain our next guess.

Figure 6: Newton’s method can globally diverge. A graphical depiction showing how Newton’s method for root-finding can globally diverge for a simple function like

\mathbf{r}(\mathbf{s})=\mathbf{s}^{1/3}

Figure 7: Comparison of standard sequential evaluation of an SSM (left) with parallel iterative evaluation of an SSM (right) . In the parallel iterative paradigm, we make a guess over the entire sequence, as indicated by the top right row labeled $s_{1}^{(0)},s_{2}^{(0)},\ldots$ . Using parallel computation over the sequence length, we find an update $\Delta\mathbf{s}^{(i)}$ to go from our current guess $\mathbf{s}^{(i)}$ for the entire trajectory to our next guess $\mathbf{s}^{(i+1)}$ . Adapted from Figure 1 of [ 142 ] .

Figure 8: The accuracy of evaluating with parallelized methods (DEER and quasi-DEER) as opposed to sequential evaluation. The parallelized methods converge to the correct trace within numerical precision. The hidden state size is

D=4

and the sequence length is

T=10,000

Figure 9: Evaluating an untrained GRU. Relative performance of sequential, DEER and quasi-DEER for evaluating a randomly initialized (and untrained) GRU on ( Top Row ) wall-clock time, averaged over 20 random seeds and ( Bottom Row ) memory, averaged over 3 random seeds. All experiments use a 16GB V100 SMX2 (memory capacity indicated by the black dashed line) and Newton methods were run to convergence. Missing points in each series indicate the GPU ran out of memory. In these settings, quasi-DEER has a runtime commensurate with DEER, but with lower memory consumption. Therefore, quasi-DEER can work at scales where DEER cannot.

Figure 10: Evaluating an untrained GRU. Sublinear and linear timing regimes for parallelized algorithms. The above experiments were run on a 32 GB V100 with a batch size of 1. As in Figure 9 , we use 20 seeds for timing, 3 seeds for memory, and the dashed black line indicates the memory capacity of the GPU (32 GB). We observe that in smaller regimes in

D

and

T

that the wall-clock time shows sublinear scaling indicative of the use of parallel algorithms. However, when the GPU becomes saturated, the benefits of parallelization are reduced and we begin to see linear scaling in wall-clock time with

T

Figure 11: Training a GRU with DEER. Comparison of DEER and quasi-DEER during GRU training for the C. elegans time-series classification task (Section 9.2 ). Each time series has length

T=17,984

. We show the median, and 5-95% interval across a rolling window of 20 training steps. (Left) DEER and quasi-DEER have the similar validation accuracy trajectories, indicating similar training dynamics. The sequential trace shown is for 24 hours of training (compared to 11 and 4 hours for the whole DEER and quasi-DEER traces). (Center) Each quasi training iteration is 2.5 times faster than each DEER training iteration. Sequential training steps took more than 6 seconds each (not pictured). (Right) Each quasi training iteration requires (approximately) 2 times more Newton iterations to converge, indicating that each quasi Newton step is approximately 5 times faster than the corresponding DEER Newton step.

Figure 12: Graphical Depiction of Trust-Region Methods. We show both an undamped Gauss-Newton step (red) and a stabilized trust-region step (blue). The solid lines indicate the contours of the merit function $L$ we want to minimize. The dashed lines indicate the contours of the quadratic surrogate that Gauss-Newton is minimizing on this iteration. The dotted lines indicate the trust region around $\mathbf{s}^{(i)}$ ; trust-region methods restrict the update to this ball, resulting in this case in an update that reduces the objective. Figure adapted from [ 178 , Figure 4.1] .

Panel of Figure 13 — Figure 13: Graphical Diagram of the ELK LGSSM. We provide a graphical diagram illustrating how the LM update in the context of parallelizing nSSMs is the MAP solution to posterior inference in an appropriately constructed LGSSM. Without any observations (i.e. $\lambda=0$ , or equivalently observations with infinite variance), we simply recover the DEER update. However, by using our previous state $\mathbf{s}_{1:T}^{(i)}$ as our observations, we restrict the dynamics to a trust region.

Figure 14: ELK stabilizes parallel evaluation of an AR GRU. (Top Left) The mean absolute difference (MAD) evaluated on the outputs converges rapidly for all four methods on a sequence generated by an untrained AR GRU. (Top Right) The MAD for evaluating a trained AR GRU. Undamped DEER variants are unstable and converge slowly (using the reset heuristic). ELK stabilizes and accelerates convergence. (Bottom) The output after 1, 100, 1000, and 2000 Newton iterations. The black dotted line is the true trace. ELK and quasi-ELK converge rapidly, but DEER and quasi-DEER are unstable. The lines where DEER and quasi-DEER are zero depict the zeroing heuristic.

Algorithm	Time per Newton step (ms, mean $\pm$ std)	Newton steps to convergence	Total time to convergence (ms)
Sequential Evaluation
Sequential	N/A	N/A	$96$
Parallelized Methods
DEER	$0.282\pm 0.0005$	$4449$	$1255$
Quasi-DEER	$0.087\pm 0.0002$	$7383$	$642$
ELK	$3.600\pm 0.0670$	$172$	$619$

Table 4: Time to evaluate a length

T=10,000

trained AR GRU using sequential vs parallelized methods. We note the dynamax package [ 144 ] we used for the parallel Kalman filter implementation in ELK is not optimized for speed, and hence these run times could be further improved.

Figure 15: Experiment to show how to set the hyperparameters for (quasi)-ELK on the AR GRU pre-trained to generate a noisy sine wave (Figure 14 in the main text). Top row plots Newton steps; bottom row plots wall-clock time. Lower is better for all plots. ( Left ) median steps/time to convergence over $\lambda$ over $15$ sequences. Quartiles are shaded but are very small. DEER methods are independent of $\lambda$ . ( Center ) Updated version of Figure 14 instead plotting MAD as a function of wall-clock time. ( Right ) Time to convergence is robust as a function of convergence threshold $\epsilon$ . Median and quartiles across 15 sequences are shown. DEER methods are nearly constant at the thresholds considered (very slight positive slope). Note we plot for increasing $\lambda$ corresponding to a smaller trust region, and reducing $\epsilon$ corresponding to a tighter convergence threshold.

Figure 16: Evaluating the Lorenz96 system in parallel. (Top two rows) : Same format as Figure 15 . (Bottom row) : Plot of Lorenz96 trajectory during optimization. DEER methods are noticeably more unstable than ELK methods.

Figure 17: PL constant $\mu$ captures the flatness of the merit function landscape. We provide a schematic illustrating how a smaller PL constant $\mu$ results in flatter merit function landscapes.

Figure 18: Threshold phenomenon in DEER convergence based on system predictability. In a family of RNNs, DEER has fast convergence for predictable systems and prohibitively slow convergence for chaotic systems. Left (Theory): We depict Theorem 5.3 , illustrating how the conditioning of the optimization problem degrades as $T$ and the LLE ( $\lambda$ ) increase. Center (Experiment): We vary $\lambda$ across the family of RNNs, and observe a striking concordance in the number of DEER optimization steps empirically needed for convergence with our theoretical characterization of the conditioning of the optimization problem. Right: For 20 seeds, each with 50 different values of $\lambda$ , we plot the relationship between $\lambda$ and the number of DEER steps needed for convergence for the sequence length $T=1000$ (gray line in left and center panels). We observe a sharp increase in the number of optimization steps at precisely the transition between predictability and unpredictability.

Figure 19: Robust relationship in mean field RNN between variance parameter

g

and LLE of the system. For 20 seeds, we observe a robust and non-decreasing relationship between the scalar parameter

g

and the LLE of the resulting mean-field RNN. The plot above is made for

50

different values of

g

from

0.5

2.0

(linearly spaced). We estimate the LLE over a sequence length of

T=9999

Figure 20: Convergence rates and wallclock time for many optimizers. We supplement the mean-field RNN experiment by also considering quasi-Newton and gradient descent methods (top) , and recording wallclock time, including for sequential evaluation (bottom)

Figure 21: DEER converges quickly for Langevin dynamics in a two-well potential. (Left) An illustration of the two-well potential state space in $D=2$ . We superimpose a contour plot of the potential on a color scheme showing the spectral norm of the dynamics Jacobian (blue indicates stability, red instability). (Center) A trace plot for the $y$ -coordinate. The LLE of the system is $-0.0145$ . (Right) We observe that this system, which has negative LLE, enjoys sublinear scaling in the sequence length $T$ in the number of DEER iterations needed to converge. We plot the median number of DEER steps to convergence over 20 seeds.

Figure 22: In this plot, we provide additional information about the behavior of DEER when rolling out Langevin dynamics on a two-well potential. (Left) We observe that across 20 random seeds (including different Langevin dynamics trajectories), the LLE for intermediate DEER iterations becomes negative after the first iteration. Consequently, we observe that the merit function (Center) experiences a spike on the very first DEER iteration (following initialization, which was the only trajectory with positive LLE), before trending towards convergence. As the system spends most of its time in contracting regions, we observe (Right) that the number of DEER iterations needed for convergence scales sublinearly with the sequence length

T

. We plot the min-max range for 20 seeds, and observe that even out of 20 seeds, the maximum number of DEER iterations needed to converge on a sequence length of

T=10,000

is around

35

System	LLE (System)	LLE (Observer)	DEER Steps (System)	DEER Steps (Observer)
ABC	0.16	-0.08	4243	3
Chua’s Circuit	0.02	-1.37	697	14
Kawczynski-Strizhak	0.01	-3.08	29396	2
Lorenz	1.02	-6.28	30000	3
Nosé–Hoover Thermostat	0.02	-0.13	29765	3
Rössler	0.01	-0.07	29288	7

Table 5: Comparison of system and observer LLEs and number of DEER steps for

T=30,000

and Euler discretization step size

\Delta t=0.01

Figure 23: Equivalence between a contractive nSSM and an $\mathcal{O}(\log T)$ stack of linear state-space models. Contractivity implies that nonlinear dynamics can be decomposed into a hierarchy of $\mathcal{O}(\log T)$ layers of linear SSMs, each of which can be evaluated in $\mathcal{O}(\log T)$ time by a parallel scan.

Fixed-point method	Order	Transition matrix $\boldsymbol{\tilde{A}_{t+1}}$	Parallelization
Newton	first-order	$\dfrac{\partial f_{t+1}}{\partial s_{t}}(s_{t}^{(i)})$	Parallel Scan (dense matrix multiplication)
Quasi-Newton	quasi first-order	$\operatorname{diag}\!\left[\dfrac{\partial f_{t+1}}{\partial s_{t}}(s_{t}^{(i)})\right]$	Parallel Scan (elementwise vector multiplication)
Picard	zeroth-order	$I_{D}$	Prefix Sum (vector addition)
Jacobi	zeroth-order	$0$	Map (embarrassingly parallel)

Table 6: Summary of fixed-point iteration schemes as linear dynamical systems . We list the methods by the order of their approximation. While higher order methods may converge in fewer iterations, each iteration may be more costly. For example, the prefix sum and parallel scan have

\mathcal{O}(\log T)

depth, while a single Jacobi iteration has constant depth. For all the methods, each iteration is an LDS, i.e. they can be written in the form of equation 20 where

\tilde{A}_{t+1}

is the transition matrix. By Proposition 3.1 , these methods are guaranteed to converge in at most

T

iterations. "Order" refers to the highest number of derivatives taken: Newton and quasi-Newton methods use first derivatives, while Picard and Jacobi methods do not use derivatives of

f_{t}

Figure 24: Comparing Picard and Jacobi iterations on a diagonal LDS. For the underlying dynamical system $s_{t+1}=\alpha s_{t}$ , we plot the norm of the error $\mathbf{e}^{(i)}$ for Jacobi and Picard iterations. We denote the empirical slope of $\|\mathbf{e}^{(i)}\|_{2}$ for Jacobi iterations by $\gamma_{J}$ .

Figure 25: A single Newton iteration solves the $S_{5}$ group word problem, whereas the number of iterations required for the other methods increases with sequence length. We consider the task of evaluating the product of $S_{5}$ group elements. A: The group word problem can be expressed as an LDS with input-dependent state-transition matrices. B: An example input-dependent transition matrix $A_{t}$ for permutation $(1\ 5\ 2\ 4\ 3)$ , in cycle notation. C: For each fixed-point method and a range of sequence lengths, $T$ , we compute the median (over ten random seeds) number of fixed-point iterations to converge (top) and the median wall-clock time (bottom). While a single Newton iteration is sufficient to solve the $S_{5}$ problem, the number of iterations required for the other methods increases with the sequence length.

Figure 26: Picard iterations struggle to parallelize RNNs. We evaluate GRUs with random parameter initialization for different sequence lengths

T

and hidden state sizes

D

. A: The nonlinear dynamics of a GRU, following [ 64 ] , where

x_{t}

is the hidden state,

u_{t}

is the input, and the notation

\mathrm{Linear}[\cdot,\cdot]

indicates a linear readout from the concatenation of two vectors. B: A representative Jacobian matrix

\nicefrac{{\partial f_{t}}}{{\partial x}}

from a GRU trajectory, which is not well approximated by the identity matrix. C: For each fixed-point method and a range of sequence lengths,

T

, and state sizes,

D

, we compute the median (over ten random seeds) number of fixed-point iterations to converge (top row) and the median wall-clock time (bottom row). Picard iterations take nearly

T

iterations to converge, while the other fixed point methods yield order-of-magnitude speed-ups over sequential evaluation

Figure 27: Understanding the convergence rates in Figure 26 . In the setting of the GRU experiment for

D=8

and

T=1000

, we plot relevant quantities for understanding the convergence rates of different methods over

10

random seeds. (Top left.) We plot the spectral norm of the approximate Jacobian for the quasi-Newton iterations we consider in this paper, i.e.

\mathrm{diag}[A_{t}(s_{t-1}^{\star})]

. (Top right.) For each of the

10

random seeds, we plot

\|\tilde{\mathbf{J}}_{\mathrm{QN}}(\mathbf{s}_{1:T}^{\star})^{-1}\mathbf{\|}_{2}

. We observe that they are always larger than one. (Bottom left.) We plot the difference between approximate Jacobians and true dynamics Jacobians over all time steps and seeds for quasi-Newton, Jacobi, and Picard iterations. We observe that this difference for Picard iterations is always larger than one, and so we would intuitively expect Picard iteration to be very slow for parallelizing GRUs. This behavior is precisely what we see in Figure 26 . (Bottom right.) Across the

10

random seeds, we plot the value of

\gamma

for Jacobi and quasi-Newton iterations (Picard would be

O(T)

and so is not shown). Because

\|\tilde{\mathbf{J}}_{\mathrm{J}}(\mathbf{s}_{1:T}^{\star})^{-1}\mathbf{\|}_{2}=1

, the 10

\gamma_{J}

’s are equivalent to the maximum values from the differences in (bottom left) over the 10 random seeds. However, since (top right) shows that

\|\tilde{\mathbf{J}}_{\mathrm{QN}}(\mathbf{s}_{1:T}^{\star})^{-1}\mathbf{\|}_{2}>1

, we observe that the values of

\gamma_{QN}

are larger than in (bottom left). In summary, because the values of

\gamma_{J}

are smaller than the values of

\gamma_{\mathrm{QN}}

, we would intuitively expect Jacobi to converge in fewer fixed-point iterations, which is exactly what we observe in Figure 26 .

Figure 28: Jacobi iterations struggle when the dynamics Jacobian is close to the identity. We evaluate Langevin dynamics for a potential

\phi

. A: The nonlinear dynamics of Langevin dynamics for a potential

\phi

and step size

\epsilon

, where

x_{t}

is the state and

w_{t}

is Gaussian noise. B: The Jacobian for Langevin dynamics is well-approximated by the identity matrix, especially for small step size

\epsilon=

1\text{\times}{10}^{-5}$$ . C: We evaluate Langevin dynamics for larger dimensions, plotting the median of 10 random seeds. Jacobi iteration consistently take

T

steps and are always slower than sequential, while the other fixed-point methods converge in fewer

T

steps and can be faster than sequential. The missing Newton iteration points indicate the GPU ran out of memory.

Appendix figures & tables1 asset

Supplementary material from the paper’s appendix.

Appendix

Figure 29: Illustration of Theorem 3.6 of [ 221 ] . In this illustration,

j=3

. The portion shaded red must be zero for the proof by induction to work, showing that

\mathbf{G}

cannot be an arbitrary matrix. The portion shaded blue can be nonzero, and the proof by induction will still hold.

Lyapunov Exponent Estimation

See all 8 papers

May 12, 2026

Parallel-in-Time Training of Recurrent Neural Networks for Dynamical Systems Reconstruction

Reconstructing nonlinear dynamical systems (DS) from data (DSR) is a fundamental challenge in science and engineering, but it inherently relies on sequential models. Recent breakthroughs for sequential models have produced algorithms that parallelize computation along sequence length

T

, achieving logarithmic time complexity,

\mathcal{O}(\log T)

. Since sequence lengths have been practically limited due to the linear runtime complexity

\mathcal{O}(T)

of classical backpropagation through time, this opens new avenues for DSR. This paper studies two prominent classes of parallel-in-time algorithms for this task, both of which leverage parallel associative scans as their core computational primitive. The first class comprises models with linear yet non-autonomous dynamics and a nonlinear readout, such as modern State Space Models (SSMs), while the second consists of general nonlinear models which can be parallelized using the DEER framework. We find that the linear training-time recurrence of the first class of models imposes limitations that often hinder learning of accurate nonlinear dynamics. To address this, we augment DEER with Generalized Teacher Forcing (GTF), a novel variant within the more general nonlinear framework that ensures stable and effective learning of nonlinear dynamics across arbitrary sequence lengths. Using GTF-DEER, we investigate the benefits of training on extremely long sequences (

T>10^4

) for DSR. Our results show that access to such long trajectories significantly improves DSR if the data features long time scales. This work establishes GTF-DEER as a robust tool for data-driven discovery and underscores the largely untapped potential of long-sequence learning in modeling complex DS.

Florian Hess, Florian Götz, Daniel Durstewitz

1Dept. of Theoretical Neuroscience, Central Institute of Mental Health,2026 Mannheim, Germany · Faculty of Mathematics and Computer Science, Heidelberg University, Germany · Faculty of Physics and Astronomy, Heidelberg University, Germany +1

Nonlinear System Identification Recurrent Neural Networks State Space Models

Ask PDF

Jun 1, 2026

Stochastic convergence of parallel asynchronous adaptive first-order methods

A new class of asynchronous adaptive first-order optimization methods is introduced, comprising asynchronous variants of several popular algorithms. Versions of these methods using momentum and/or inexact normalization are also considered. The convergence of methods in the class on non-convex functions is analyzed in a fully stochastic setting, and is shown to be (up to logarithmic factors) of order O(1/sqrt{t}) under reasonable assumptions. Numerical experiments suggest that such asynchronous adaptive algorithms are very relevant in heterogeneous large-scale machine learning systems.

Serge Gratton, Philippe L. Toint

Universit´e de Toulouse, INP, IRIT, Toulouse, France · Universit´e de Toulouse, INP, IRIT, Toulouse, France, and NAXYS, University of Namur, Namur, Belgium

Stochastic Optimization Nonconvex Stochastic Optimization Stochastic Optimization Convergence Adaptive Gradient Methods

Ask PDF

May 18, 2026

Perfect Parallelization in Mini-Batch SGD with Classical Momentum Acceleration

Accelerating stochastic gradient methods with classical momentum schemes, such as Polyak's heavy ball, has proven highly successful in training large-scale machine learning models, particularly when combined with the hardware acceleration of large mini-batch computations. Yet, the effect of classical momentum on stochastic mini-batch optimization has been poorly understood theoretically, with prior works requiring strong noise assumptions and extremely large mini-batches. In this work, we develop a general theory of stochastic momentum acceleration for optimizing over quadratics in the interpolation regime, a popular abstraction for studying deep learning dynamics which also includes classical methods such as randomized Kaczmarz and coordinate descent. Our framework encompasses both heavy ball and Nesterov-style momentum, allows for arbitrary mini-batch sizes, and makes minimal assumptions on the stochastic noise. In particular, we show that acceleration from classical momentum is directly proportional to the gradient mini-batch size (up to a natural saturation point), thereby enabling perfect parallelization of mini-batch computations. Our theory also provides a simple choice for the momentum parameter, which is shown to be effective empirically.

Sachin Garg, Michał Dereziński

University of Michigan

Momentum Methods Stochastic Optimization Convergence Accelerated Gradient Methods Stochastic Gradient Descent

Ask PDF

Abstract

Lyapunov Exponent Estimation

Explore similar work

Parallel-in-Time Training of Recurrent Neural Networks for Dynamical Systems Reconstruction

Stochastic convergence of parallel asynchronous adaptive first-order methods

Perfect Parallelization in Mini-Batch SGD with Classical Momentum Acceleration