Keywords: Unconstrained nonconvex optimization, first-order methods, global rate of convergence, asynchronous methods, parallel computing, heterogeneous deep learning problems.

1 Introduction↩︎

[1] [2]

We consider the problem of finding a first-order critical point for the optimization problem \[\label{problem} \min_{X\in I\hskip -2pt R^N} F(X) = \mathbb{E}\!\left[f(X,\xi) \right]\tag{1}\] where \(f\) is a smooth, possibly nonconvex function of \(X\) and \(\xi\) is a suitably defined random variable. In other words, we are seeking \(X\) such that \(\nabla_X \mathbb{E}\!\left[f(X,\xi) \right] = 0\). Motivated by the growing importance of very large models in deep learning, a substantial literature has been devoted to the study of asynchronous algorithms for solving (1 ). By far the most popular are stochastic asynchronous steepest descent methods (also called asynchronous gradient methods), whose modern efficient lock-free framework has been pioneered by the Hogwild! method [3]. This method, like most of the methods we discuss in this paper, maintains a vector \(\widehat{X}\) of values for the problem’s variables in a memory shared by a number of processes/processors, each of which sequentially

1. selects a subset or block of the variables from \(\widehat{X}\),

2. computes the gradient of \(f\) at \(\widehat{X}\) with respect to the variables of the block (this requires reading the vector \(\widehat{X}\) from shared memory),

3. computes a gradient step (with method-specific stepsize) in the subspace spanned by this block using the block gradient, resulting in new values of the block variables,

4. updates the values of \(\widehat{X}\) corresponding to the block with these new values.

The whole operation is asynchronous because each process/processor does not wait for other processes to complete their tasks (it is lock-free). Of course, this description is extremely synthetic and methods of the type just described may still differ on a number of issues.

Geometry.

The first is to define how a particular block of variables is selected in Step 1. This selection can be random or deterministic, and is characterized by the size of the allowed “overlap” between blocks selected by different processes, ranging from totally separable (the blocks associated with running processes are disjoint), partially separable (limited overlap between running blocks is allowed — for an introduction to partial separability, see [4]–[6]), or arbitrary (the size of the overlap is not constrained).

Atomicity.

The next issue is the management of the read/write operations between local processes and shared memory. Because of the asynchronous nature of the whole algorithm, it may happen that different processes attempt to read or write the same component of \(\widehat{X}\) at the same time. To cope with this problem, it is generally assumed that the reading and writing operations are atomic at some granularity level, meaning that reading/writing an atomic component of \(\widehat{X}\) by a process cannot be interrupted by another process. We will see below that this concept also applies to other quantities than \(\widehat{X}\), and that, in addition to purely hardware and message passing considerations, it also can be limited by the internal structure of the object to read/write, or by the norm used to measure it (for instance, objects like positive definite matrices need to maintain their positive definite nature).

Consistency.

In addition, it may well happen that a process is reading parts of \(\widehat{X}\) while other parts are being updated by other processes. This blurs the notion of algorithm’s iterates, because it is possible that a block gradient is computed for a value of the variables corresponding to an incomplete, ongoing update of \(\widehat{X}\). Most algorithms (but not all) allow these inconsistent read/writes.

Delays.

In addition, the asynchrony of the whole algorithm implies that when a block gradient is computed, it uses the values of \(\widehat{X}\) read from shared memory, but these values may result from writes (by other blocks) somewhat in the past. The time difference between the moment \(\widehat{X}\) is updated and the moment where it is used for computing a block gradient is called a delay. The analysis of all algorithms considered here imposes some kind of upper bound on such delays, either explicitly or in expectation.

Choice of stepsize.

Finally, once the block gradient is computed, yielding a descent direction, a stepsize along this direction must be defined. Algorithms in the literature either assume that the stepsize is a “suitably small” constant, or that it is given by a (user) predefined decreasing sequence converging to zero.

Having brushed the landscape, we now review chronologically significant proposals in the literature. The pioneering Hogwild! method proposed in [3] was instrumental in introducing lock-free variants of asynchronous steepest-descent. It applies to strongly convex functions with partially separable geometry, for which block gradients are computable. Its stochastic convergence analysis assumes that the gradient oracles are bounded and that the constant stepsize can be chosen small enough compared to the gradient’s Lipschitz constant. Inconsistent reads are allowed. This is not the case for the AsySG-con [7] which, as the end of its name suggests, insists on coherent reads, but applies to possibly nonconvex functions with full-size gradients without any separability requirement. Its analysis assumes bounded gradients and predefined diminishing stepsizes. The ADrock algorithm [8] is of a slightly different type, as it is better viewed as a fixed point method, but it faces similar issues. It uses full gradient vectors and requires that stepsizes decrease to zero. Its analysis holds for convex functions with bounded gradients but without any specific separability. The KroMagnon algorithm [9] is similar to Hogwild! in that it is a gradient-based method, allows block gradient evaluations and requires a partially separable geometry. Establishing its rate of convergence needs convex functions but no bound on the gradients is necessary. The requirement on partially separable geometry is relaxed for the versions of Hogwild! studied in [10], [11] and boundedness of the gradients is no longer requested. Instead of assuming a small enough stepsize as in the original version, the convergence analysis requires a stepsize decreasing to zero. ASAGA [12] is another gradient-based method requiring, as is the case for the original Hogwild!, partially separable geometry, strongly convex functions and a sufficiently small stepsize. However, there is no need to assume bounded gradients. The method occasionally evaluates the full gradient (in shared memory) in order to reduce the variance of the estimates. The FDG algorithm of [1] partitions the variables into blocks and interestingly proposes to exploit the structure of backpropagation in neural networks to compute the gradients with respect to the variables of a block independently from those of another, in turns leading to an asynchronous delayed-gradient method. The analysis stops with a descent lemma requiring bounded gradients and knowledge of the gradient’s Lipschitz constant. The AsyFLEXA [13] deterministic algorithm is again slightly different in that it relies on convex local models (instead of purely linear ones as is the case for most gradient-based methods) and applies to possibly nonconvex functions with block gradient evaluations without separability requirement. Once more, the analysis is carried out for small enough stepsize, but without assuming bounded gradients. The reference [14] discusses several variants of a stochastic gradient-based method with momentum, as seen in the context of continuous differential equations. We report here on PASSM, a variant for potentially nonconvex totally separable functions with block gradient evaluations, which appears to perform best. Bounded gradients are required for its analysis and its stepsizes follow a geometrically decreasing sequence. Interestingly, a lock is required for writing the updated variables in shared memory. One should also mention contributions of [15] and [16]. The first introduces the concept of “elastic consistency”, formalizing the conditions on delays for non-adaptive gradient methods and showing that such conditions are necessary for convergence of this type of methods. The second introduces general recurrence occurring in many non-adaptive asynchronous stochastic gradient methods and uses them to sharpen existing convergence results, in particular by considering average delays rather than maximum ones. Finally, the proposal of [2] considers an asynchronous method for stricly convex functions with known Lipschitz constant, putting special emphasis on communication costs.

Table 1 provides an easy-to-read summary of the above discussion, where column “stoch.” indicates whether the method is stochastic, and “cons. read” indicates whether the method requires consistent reading of shared memory. In the “method type” column, the abbreviations “convex approx”, “prec. grad.” and “mom” stand for “convex approximation”, “preconditioned gradient” and “momentum”, respectively. Finally, the abbreviation “adapt.” in the “stepsize” column refers to adaptive stepsize, the technique we will develop below.

All the algorithms we have reviewed so far exhibit two (in our view, significant) limitations. The first and most important one is that pure gradient steps can lead to very slow convergence on (even moderately) ill-conditioned problems. This issue has long been addressed by the use of steepest-descent steps (a norm-dependent concept, at variance with gradient steps) and (adaptive) preconditioning (see [19], for instance). While these techniques have successfully been applied for solving problem (1 ) in the synchronous case, for instance in methods like AdaGrad [20]–[22], Adam [23], Shampoo[24], [25] or Muon[26]–[28] or in the distributed case [29], they haven’t yet (to the authors’ knowledge) been considered for asynchronous optimization. The second limitation is that theoretical analysis of many of the published algorithms ignore the use of momentum, despite its widespread nature (see [29], for instance).

The main contribution of the present paper is to propose and analyze a (to the authors’ knowledge, first ever) class of algorithms containing several stochastic asynchronous adaptively preconditioned gradient methods, possibly using momentum, for the nonconvex case. In particular, AdaNorm, full and diagonal AdaGradas well as adaptive variants of Shampooand Muonare covered. Like for most recent proposals, bounded gradients are not necessary for its analysis. The proposed algorithm builds on the ADPREC framework for (synchronous) stochastic first-order methods [30], which is itself based on the use of general dual norms and covers a large class of stochastic adaptive gradient-based methods. The present proposal takes advantage of its powerful generality, and is defined in the geometric setting of totally separable problems, as it is the case for PASSM. While this might be seen as theoretically restrictive, it still covers the very important case of layer-wise preconditioning in training large neural networks (see [31]–[34] for instance), hence justifying our interest. As a side-benefit, our proposal also extends the theory presented in [30] in allowing inexact step normalization strategies, a clear advantage in practice for methods like Muonwhere the normalization is performed using a possibly inexact Newton-Schultz iteration [25], [29], [35]–[37]. The proposed addition of momentum to asynchronous steepest-descent methods (the second line for our paper in Table 1) is also useful to enhance numerical performance, but it may come with a (mostly theoretical) cost. For the common momentum definition of using the gradient, assuming that the step-size is sufficiently small may be necessary to prove convergence, but the resulting degraded convergence rate may improved by adopting a suitable strategy for the momentum parameter. Remarkably, no assumption on the step-size is required for the momentum-less variant, or for a definition of the momentum using the momentum itself.

The paper is organized as follows. The momentum-less asynchronous framework is presented in Section [s:algonomom], where it is shown to be a special case of an easier-to-analyze delayed framework, whose rate of convergence is studied in Section [s:momentumless]. Two different variants of momentum are introduced and the resulting rates of convergence examined in Section [s:momentum]. Numerical experiments illustrating the use of asynchronous block-wise adaptive Muon/AdaGradare presented in Section [s:numerics]. Finally, some conclusions and perspectives are discussed in Section [s:conclusion] while Appendix 1 and Appendix 2 give the details of the convergence proofs for the methods of Sections [s:algonomom] and [s:momentumless], respectively.

Notations: In what follows, \(\|\cdot\|_E\) denotes the Euclidean norm. If \(\|\cdot\|\) is a norm on some space \({\pazocal{S}}\) endowed with an inner product \(\langle\cdot,\cdot\rangle\), its dual norm \(\|\cdot\|_D\) is defined by \(\|x\|_D = \max_{y\in {\pazocal{S}}}\langle y,x \rangle/ \|y\|\). The symbol \(I_n\) denotes the identity matrix of dimension \(n\) and \(0_n\) the zero vector of size \(n\). If \(x\) is a vector, \([x]_i\) denotes its \(i\)-th component and \([x]_{{\pazocal{I}}}\) the set of the components with indices in \({\pazocal{I}}\). If \({\pazocal{S}}\) is a set, \(|{\pazocal{S}}|\) denotes its cardinal.

2 ASADPREC: a stochastic asynchronous first-order framework↩︎

The basis of our approach is to assume that the problem variables come in disjoint blocks of homogeneous nature. In deep learning applications, this occurs for instance when some variables are layer weights and other thresholds or biases of activation nodes. The first may be scalar and the second matrices [31], [32]. In line with [30], we assume that there are \(L\) blocks of variables, and problem (1 ) is therefore cast in the Cartesian product space \[I\hskip -2pt R^N = \prod_{\ell=1}^L I\hskip -2pt R^{n_\ell\times m_\ell} = \prod_{\ell=1}^LI\hskip -2pt R^{d_\ell}.\] We denote by \({\pazocal{I}}_\ell\) the set of indices of the variables in block \(\ell\). For \(X\in I\hskip -2pt R^n\), \([X]_{{\pazocal{I}}_\ell}\) is then the subset of variables in \(X\) belonging to block \(\ell\). Following [30], we then specify, for each block, a particular norm \(\|\cdot\|_\ell\) and its dual \(\|\cdot\|_{\ell,*}\). The inner product on \(I\hskip -2pt R^N\) is then defined as \[\left\langle U,V\right\rangle = \sum_{\ell=1}^L \left\langle U_\ell,V_\ell\right\rangle_F \;\;where\;\; \left\langle A,B\right\rangle_F = \mathop{\mathrm{tr}}(A^TB)\] and the norm on \(I\hskip -2pt R^N\) given by \[\label{big-norm} \|V\|^2 = \sum_{\ell=1}^L \|V_\ell\|_\ell^2\tag{2}\] whose dual norm is \[\label{big-dual-norm} \|U\|_*^2 = \sum_{\ell=1}^L \|U_\ell\|_{\ell,*}^2.\tag{3}\]

Because of the structure of the variables’ space, it is natural to avoid mixing variables of different blocks, and this therefore defines the geometry of the asynchronous minimization method we are going to construct. Moreover, as is shown by the practical success of methods like Muon for matrix variables, different minimization methods may be appropriate for different types of variables. For each block, we therefore define a collection of methods applicable to variables in the block³. As was already the case in [30], each such method for block \(\ell\) is specified by a preconditioner’s update operator \({\pazocal{U}}: I\hskip -2pt R^{d_\ell}\rightarrowI\hskip -2pt R^{d_\ell}\) and a normalization operator \({\pazocal{N}}: I\hskip -2pt R^{d_\ell}\rightarrow I\hskip -2pt R^{d_\ell}\). The collection of methods applicable for block \(\ell\) is then given by \({\pazocal{M}}_\ell =\{({\pazocal{U}}_1,{\pazocal{N}}_1), \ldots, ({\pazocal{U}}_{a_\ell},{\pazocal{N}}_{\hat{m}_\ell})\}\).

While this way to describe first-order adaptively preconditioned gradient methods may seem abstract, it was shown in [30] that it covers a fairly representative set of popular methods, such as AdaNorm, full and diagonal Adagrad, Shampoo and Muon⁴ (see [30] for details).

We now specify the ASADPREC⁵ algorithmic framework , meaning that its description can be instantiated in a number of different specific algorithms. In what follows, we use the shorthand “ASADPRECalgorithm” to refer to any algorithm belonging to the ASADPRECframework.

Some comments on this algorithm are useful at this stage.

1. We have stated the framework in a form as general as possible. It is however important to observe that significant simplifications occur if, as is often practical, the computing architecture allocates to each processor a set of blocks disjoint from that of other processors. In that setting, the preconditioner \(\widehat{\Pi}_\ell\) may be stored locally on the processor in charge of block \(\ell\) and Steps 6 and 9 are unnecessary. Moreover, this design also ensures that not two processors can work on the same block at the same time, so that Steps 2, 3 , 4 and 11 can be replaced by a simpler Step 3 where the block selection is restricted to those assigned to the current processor. Thus all steps in brackets can be ignored and Step 3 simplified, thereby significantly reducing interprocessor communication costs and resulting in a much leaner operation.

2. Pushing asynchronicity to the component- (rather than block-) level is difficult for \(\Pi_\ell\) if one wishes to allow for non-diagonal preconditioners (as in full AdaGrador Shampoo), because inconsistently read symmetric matrices may fail to be positive-definite.

3. No writing conflict can occur in Steps 9 and 10 because we have assumed a totally separable geometry, with no overlap between blocks.

4. The computation of the block’s gradient \(\widetilde{G}_{k(\ell_t)}\) in Step 4 may or may not require the computation of the complete gradient \(\nabla f(X_{k(\ell_t)},\xi_{k(\ell_t)})\), depending on the particular form of the objective function and the computer architecture.

5. The general framework of ASADPRECis reminiscent of the block-Jacobi [38], [39] setting for solving linear systems of equations.

Our analysis of ASADPREC’s global rate of convergence is based on the view that, from an analytical standpoint, it can be viewed as a synchronous algorithmic framework with delays. To establish this interpretation, we define a global iterate \(\widehat{X}_t\) \((t=0,\ldots)\) as the state of the shared memory \(\widehat{X}\) at each moment \(t\) where \(\widehat{X}\) is modified by a block-atomic update. Because of the read operation of Step 1 is atomic, we have that \(X = \widehat{X}_t\) and, after Step 5, \(\widetilde{G}_\ell=[\nabla_X f(\widehat{X}_t,\xi_\ell)]_{{\pazocal{I}}_{\ell_t}}\). If \(\ell_t\) is the index then selected at Step 3, the algorithm’s loop on block \(\ell_t\) continues, starting from \((\widehat{X}_t,[\nabla_X f(\widehat{X}_t,\xi_{k(\ell_t)})]_{{\pazocal{I}}_{\ell_t}})\). Now consider the interval \(t\) to \(t+1\) and all steps being computed in parallel in this interval. Ignore all such steps whose computation does not terminate at \(t+1\), and focus on the (unique, with probability one) block \(q_t\) which terminates at \(t+1\). The computation of this step was started at some time \(t-\tau_{t,q_t}\) (where the nonnegative integer \(\tau_{t,q_t}\) is called a delay) using \(\widehat{X}_{t-\tau_{t,q_t}}\) and the approximate block gradient \([\nabla_X f(\widehat{X}_{t-\tau_{t,q_t}},\xi_{t-\tau_{t,q_t}})]_{{\pazocal{I}}_{q_t}}\). Moreover, because the selection process at Step 3 prevents two different processors to work on block \(q_t\) at the same time, we have that \([\widehat{X}_t]_{{\pazocal{I}}_{q_t}}=[\widehat{X}_{t-\tau_{t,q_t}}]_{{\pazocal{I}}_{q_t}}\). We may then (formally) consider that this step on block \(q_t\) was started at time \(t\) from \(\widehat{X}_t\), but using the “delayed approximate block gradient” \([\nabla_Xf(\widehat{X}_{t-\tau_{t,q_t}},\xi_{k(q_{t-\tau_{t,q_t}})})]_{{\pazocal{I}}_{q_t}}\). Thus the sequence of overlapping step computations for different blocks is equivalent to a sequence where, in each interval from \(t\) to \(t+1\), only the block steps terminating at the end of the interval are considered and are deemed to have started at \(t\) using a delayed approximate block gradient. This is exactly how the DADPRECalgorithm described proceeds, at each iteration selecting a subset of blocks for which a complete step is computed using a approximate, possibly delayed, block gradient. We may thus see the sequence of iterates \(\widehat{X}_t\) generated by an ASADPRECalgorithm as being instead generated by a DADPRECalgorithm, thereby identifying \(\widehat{X}_t\) for ASADPRECwith \(X_t\) for DADPREC. We also identify \(\widetilde{G}_{\ell_t}\), \(\Pi_{\ell_t}\) and \(Z_{\ell_t}\) in the former with \(\widetilde{G}_{t,\ell}\), \(\Pi_{t,\ell}\) and \(Z_{t,\ell}\) in the latter.

Note that Step 4 of DADPREC is purely formal, because that it does not involve any computation. Note also that \(\widetilde{G}_{t,\ell}\) in DADPRECis a fairly general approximation of the true block gradient, as it may include delays and other errors, such as resulting from sampling.

The reason for introducing the DADPRECframework is that its convergence analysis can be derived, with moderate effort, from that of the proposal in [30]. We describe the necessary assumptions and the results obtained (for DADPRECand thus for ASADPREC) in the next section, and postpone to the Appendix the detailed description of the modifications of the analysis of [30].

As is clear from the algorithms’ statements, this analysis does not include momentum. Variants including momentum will be considered in Section [s:momentum].

3 Convergence of the momentum-less DADPREC↩︎

Our analysis of the DADPRECframework starts by formulating standard conditions on the problem (1 ).

Defining \(G(X) = \nabla_X F(X)\), we then require the following smoothness assumption.

We now introduce conditions that define the class of methods \(({\pazocal{U}}_\ell,{\pazocal{N}}_\ell)\) for which we develop our theory. These are defined for a given block \(\ell\in\{ 1, \ldots, L \}\) and depends on the choice of dual norm \(\|\cdot\|_{\ell,*}\) for this block. Again, we stress that these abstract conditions do hold for several popular methods for specific choices of the methods \(({\pazocal{U}},{\pazocal{N}})\in{\pazocal{M}}_\ell\), including AdaNorm, full and diagonal Adagrad, as well as adaptive variants of Shampoo and Muon [30].

We also need to consider stochastic conditions on the quality of the gradient oracle.

Observe that Assumption 5 is made at the block level, and is weaker than assuming the more standard condition that \(\mathbb{E}_{t}\!\left[\widetilde{G}_t \right] = G_t= G(X_t)\). Similarly, Assumption 6 is also made at the block level. Moreover (?? ) requires a bound on the cumulative variance of all block gradient oracles over all past iterates, an approach also more general than assuming conditional variance at every iteration. In particular, it allows large variance at early iterations provided later iterations compensate. Trade-offs between different realizations and/or different blocKs are also theoretically possible.

We finally impose a minimal condition on the asynchronous nature of the algorithm.

For the purpose of analysis, we also define, for \(t\geq 0\) and \(\ell\in\{ 1, \ldots, L \}\), \[{\pazocal{J}}_{t,\ell} = \{j\in\{ 0, \ldots, t \} \mid \ell \in {\pazocal{A}}_j\},\] the set of indices of iterations at which block \(\ell\) is effectively updated, and immediately note that the summations \(\sum_{j=0}^t\sum_{\ell\in{\pazocal{A}}_j}\) and \(\sum_{\ell=1}^L\sum_{j\in{\pazocal{J}}_{t,\ell}}\) are interchangeable (for instance in (?? )).

A general theorem stating the global rate of convergence of any DADPREC(and thus ASADPREC) algorithm may now be stated under these assumptions.

Note that, since we prove (in Appendix) that each preconditioner \(\Pi_{t,\ell}\), is bounded by the (in expectation) very slowly growing \(\Theta_t\) (see ([Thetak-def])), and since \(Z_t\) is the preconditioned approximate gradient, condition (?? ) is akin to a (undelayed) cumulative affine variance condition of the form \[\sum_{j=0}^s\sum_{\ell=1}^L \mathbb{E}_{j}\!\left[\|\widetilde{G}_{j,\ell}-G_{j,\ell}\|^2 \right] \le \nu_t^2 +\omega^2\sum_{j=0}^s\sum_{\ell=1}^L\mathbb{E}_{j}\!\left[\|G_{j,\ell}\|_{*,\ell}^2 \right],\] whose iteration-wise variant has already been used in analysis of first-order methods (see [40], [41] or, for the stronger ” affine-\(^*\) ” version, [42]). In particular, the “strong growth” assumption motivated by over-parametrized problems and used in [41] to derive an improved convergence rate for AdaGradis essentially subsumed (at the cumulative level) for “preconditioned cumulative strong growth” (that is if \(\nu_t=0\) is assumed in (?? )).

As stated, the rate of convergence of DADPREC is dominated by the term in \(\kappa_c\Theta_t/\sqrt{t+1}\) in ([thetruerate]), where \(\Theta_t\) depends on \(\nu_t\), the bound on the cumulative variance of the block gradient oracles, which may itself depend on \(t\). The next corollary describes what can be said if one makes a more specific assumption on the oracle (block) variance at each iteration.

From the ASADPRECpoint of view, \(|{\pazocal{J}}_{t,\ell}|\) corresponds to \(k(\ell_t)\), and therefore counts the number of updates to block \(t\), while \(\psi_t\) is their maximum taken on all blocks. We have that \(\psi_tc\le t+1\) in general, but if, at each iteration \(t\), \(|{\pazocal{A}}_t| \ll L\) (a not unusual situation), then \(\psi_t \ll t+1\), suggesting an acceleration of the convergence rate compared with the undelayed algorithm (where \(\psi_k = t+1\)).

Observe that the continuity of the bound expressed in Theorem [thm:dadprec-convergence] as a function of \(\nu_t\) is lost in the statement of Corollary [therate2], because the constants hidden in the \({\pazocal{O}}(\cdot)\) notation depend on \(\alpha\). In particular, this formulation does not support taking the limit for \(\alpha\) tending to one. Observe also that we could weaken ([betaisktogamma]) by requiring that \[\label{alt-var-cond} \mathbb{E}_{t}\!\left[\|\widetilde{G}_{t,\ell}-G_{t,\ell}\|^2 \right] \le \frac{\sigma_\ell^2}{t^\alpha}+\omega^2\|Z_{t-1,\ell}\|_*^2\tag{4}\] for \(\ell\in{\pazocal{A}}_t\) with the convention that \(Z_{-1,\ell}=0\), since, obviously, \[\sum_{\ell=1}^L\sum_{j=0}^{t-1}\mathbb{E}_{j}\!\left[\|Z_{j,\ell}\|_*^2 \right] \le \sum_{\ell=1}^L\sum_{j=0}^t\mathbb{E}_{j}\!\left[\|Z_{j,\ell}\|_*^2 \right].\] The advantage of (4 ) over ([betaisktogamma]) is that the right-hand-side of this new condition is now measurable at iteration \(t\).

As was noted in [30], although the rates of convergence obtained in Corollary [therate2] under the bias and variance Assumptions 5 and 6 are reasonable, they do not quite match, for high-variance regimes, the best obtained so far for (momentum-less, synchronous) AdaGrad and AdaNorm [41]. However they recover (in order) the best \({\pazocal{O}}(\sqrt{\log(k+1)\log(\log(k+1))/(k+1)})\) rate in the preconditioned cumulative strong growth context⁶. Importantly, they do so for a large class of asynchronous algorithms.

4 Convergence of DADPRECwith momentum↩︎

Adding momentum to first-order methods has often been instrumental in improving numerical performance⁷. It is therefore of interest to consider asynchronous versions of such methods. We investigate two versions of this idea in this section, which share the common framework ASADPREC.Mstated below for some given sequences \(\{\mu_k\}\) of momentum parameters with \(0\le \mu_k\le \mu_{\max} < 1\) for all \(k\ge 0\), with \(k\) indexing the number of updates to block \(\ell\).

The two variants are specified by detailing Step 9 either as \[\left.\begin{array}{lcl} \Pi_\ell &= &\Pi_\ell^-+{\pazocal{U}}_\ell(\widetilde{G}_\ell)^2,\\ M_\ell &= & \left\{\begin{array}{ll} \mu_k M_\ell^-+(1-\mu_k) \widetilde{G}_\ell &\;\; if\;\; k>0,\\ \widetilde{G}_\ell &\;\; if\;\; k=0, \end{array}\right. \end{array}\right\}({\small {\sf ASADPREC.MG}})\] or as \[\left.\begin{array}{lcl} M_\ell & = & \left\{\begin{array}{ll} \mu_kM_\ell^-+(1-\mu_k) \widetilde{G}_\ell &\;\; if\;\; k>0,\\ \widetilde{G}_\ell &\;\; if\;\; k=0 \end{array}\right.\\ \Pi_\ell & =& \Pi_\ell^-+{\pazocal{U}}_\ell(M_\ell)^2. \end{array}\right\}({\small {\sf ASADPREC.MM}})\] Note that, in both cases, the momentum \(\widehat{M}_\ell\) and the index \([\widehat{K}]_\ell\) must in general be read from and updated in shared memory, unless, again, specific blocks are affected to specific processors, in which case they can be maintained in local memory. All bracketed steps may thus be omitted in this context.

Fortunately, the convergence theory can be adapted to cover both variants. This is achieved again by considering DADPREC.MMand DADPREC.MG, the obvious extensions of DADPRECcorresponding to ASADPREC.MMand ASADPREC.MG, respectively, where ([Pi-def]) and ([z-def]) are now replaced by

Because the update formula for \(Z_{t,\ell}\) above now involves the momentum \(M_{t,\ell}\) instead of \(\widetilde{G}_{t,\ell}\), we reformulate Assumption 3 to reflect this change⁸ as follows.

Armed with this reformulated assumption, it is now possible to establish that the expected global rate of convergence of a DADPREC.MM algorithm is essentially identical to that of the version without momentum.

We note at this point that our momentum definition and convergence analysis for the DADPREC.MMalgorithm do not make the assumption that the stepsize parameter \(\eta\) is sufficiently small, in contrast with previous proofs of convergence where results assume that \(\eta\) is small enough, in particular smaller than a multiple of the (usually unknown) Lipschitz constant \(L_G\) [43]–[45]. Establishing convergence results for the DADPREC.MGvariant unfortunately requires additional assumptions. In particular, the stepsize \(\eta\) has to be chosen small enough. This alternative theory hinges on the fact that \({\pazocal{U}}(\widetilde{G}_{k,\ell})\) can be viewed as a perturbation of \({\pazocal{U}}(M_{k,\ell})\) , and therefore that the structural relations of Assumption 8 are themselves perturbed when the DADPREC.MGalgorithm is used. Fortunately, it remains possible to bound these perturbations by a mix of variance-related and second-order terms, the first of which potentially affecting the resulting global convergence rate. The final outcome is given by the following theorem and corollary.

As it turns out, assuming ([ass:sub-add]) is not restrictive⁹, but supposing ([small-eta]) is definitely less desirable because the value of the Lipschitz constant \(L_G\) or that of \(\kappa_Z\) (if it exists) is usually unknown. It is however common practice in the literature [43]–[45]. Note the term \(\kappa_{\nu\nu}\nu_t^2\) in the middle term of the expression of \(\Theta_t\) in ([Thetak-def-alt]), which is the first significant difference between ([Thetak-def-alt]) and ([Thetak-def]) and induces a modified convergence rate. The second difference is the presence of the factor \(\max[\mu_{\pi(j-1,\ell),\ell},\mu_{j,\ell}]^2\) in the definition of the cumulative variance, which has the opposite effect of potentially improving the convergence rate. This is expressed by the following corollary, where \(|{\pazocal{J}}_{t,\ell}|\) is the number of updates to block \(\ell\) up to iteration \(t\).

Again, be aware that the constants hidden in the \({\pazocal{O}}(\cdot)\) depend on \(\alpha\) and \(\beta\), in particular preventing taking limits for \(\alpha\) tending to one. The rate of convergence now depends on the value of \(\alpha+2\beta\), indicating how acting on the momentum parameter can partly alleviate the effect of high gradient oracle’s variance. Because the convergence rate is determined by \(\theta_t\), choosing a small momentum parameter in effect reduces the propagation of larger errors in the gradient oracle across iterations. In particular, setting \(\mu_{t,\ell}\) to a multiple of \(|{\pazocal{J}}_{t,\ell}|^{-{\scriptstyle \frac{1}{4}}\alpha}\) recovers the rate of the momentum-less variant also for the high-variance regime (\(\alpha<1\)). If however \(\mu_{t,\ell}\) is kept constant (or bounded away from zero), that is if \(\beta=0\), then ([var-cond-alt2]) is stronger than ([betaisktogamma]) and the rate of convergence for \(\alpha<1\) is now \({\pazocal{O}}\left((k+1)^{{\scriptstyle \frac{1}{2}}-\alpha}\right)\) instead of \({\pazocal{O}}\left((k+1)^{-{\scriptstyle \frac{\alpha}{2}}}\right)\), requiring in particular that \(\alpha > {\scriptstyle \frac{1}{2}}\).

5 Approximate normalization↩︎

We have assumed, in our theory, that normalization of the step is exact, in the sense that \(\|{\pazocal{N}}_{t,\ell}(Z)\|_\ell=1\) for all \(t\), all \(\ell\in{\pazocal{A}}_t\) and all \(Z\inI\hskip -2pt R^{n_\ell \times m_\ell}\). This choice has been made for simplicity, but the reader can readily convince himself/herself that it is not necessary, and that it is sufficient to assume the existence of constants \(0<\kappa_{1,{\pazocal{N}}} \le 1 \le \kappa_{2,{\pazocal{N}}}\) such that \(\|{\pazocal{N}}_{t,\ell}(Z)\|_\ell\in [\kappa_{1,{\pazocal{N}}}, \kappa_{2,{\pazocal{N}}}]\) for all \(t\), \(\ell\in{\pazocal{A}}_t\) and all \(Z\inI\hskip -2pt R^{n_\ell \times m_\ell}\). The constants \(\kappa_{1,{\pazocal{N}}}\) and \(\kappa_{2,{\pazocal{N}}}\) then percolate though all proofs, complicating the expression of relevant factors even more, but without affecting the order of global convergence. This observation is particularly useful for methods involving matrix blocks, such as Shampooor Muon, where step normalization is typically performed by a possibly inexact Newton-Schulz iteration [25], [29], [35]–[37].

6 Numerical experiments↩︎

6.1 Experimental objectives↩︎

The experiments in this section illustrate the practical behavior of the asynchronous block-parallel adaptive preconditioning method (ASADPREC) introduced in the previous sections. The convergence analysis in Section [s:momentumless] establishes that the asynchronous block-coordinate iteration converges under a bounded-delay condition, provided the block preconditioners capture sufficient local curvature to control the inter-block coupling. The experiments are designed to verify that these theoretical predictions hold on concrete machine learning problems, and to quantify the practical trade-offs involved.

The OFFO (Objective-Function-Free Optimization) framework of [30] provides the algorithmic foundation: the parameter vector is decomposed into \(L\) blocks, each equipped with its own adaptive preconditioner, and the method requires only gradient evaluations—no function value computations. In the synchronous regime, this reduces to a preconditioned block-Jacobi iteration; in the asynchronous regime, each block is updated independently using a possibly stale snapshot of the other blocks. Classical results on asynchronous chaotic relaxation [46] and asynchronous iterations [38], [39] predict convergence when the spectral radius of the “asynchronous iteration matrix” is less than one—a condition that is favored by weak inter-block coupling and strong local preconditioning. In practice, the throughput gain from asynchronous execution depends on the computational heterogeneity between blocks. When all blocks have equal per-iteration cost, asynchronous execution offers no advantage over the synchronized variant (the barrier imposes no idle time). When block costs differ, the synchronization barrier forces fast blocks to idle while waiting for the slowest block, and the throughput gain from removing the barrier is bounded by an Amdahl-type ratio: the cost of the slowest worker relative to the average. The experiments are designed to explore this relationship across a range of heterogeneity levels. More specifically, the experiments address the following questions:

1. Does asynchronous block-parallel optimization preserve convergence quality (training loss, gradient norm, generalization metric) relative to its synchronous counterpart, as predicted by the bounded-delay convergence theory?

2. What throughput gain does asynchronous execution provide, and how does this gain relate to the measured computational heterogeneity between blocks (Amdahl-type bound)?

3. How does exponential moving average (EMA) momentum interact with the stale reads inherent to asynchronous execution?

The benchmarks are chosen to cover a range of architectural patterns (dense MLPs, sparse embedding models, hybrid architectures) and heterogeneity levels, so that the interplay between block structure, computational cost imbalance, and asynchronous scheduling can be studied in a controlled setting. The goal is not to achieve state-of-the-art predictive performance, but to verify the theoretical predictions and to identify the regimes in which asynchronous block parallelism is most beneficial. We consider three algorithmic variants. All three apply the same per-block preconditioned updates and differ only in scheduling and synchronization.

Synchronous single-gradient (SyncSingle):

A single mini-batch is sampled and the stochastic gradient \(G_k = \nabla F_{\mathcal{B}_k}(x_k)\) is computed once. All \(L\) blocks are updated from the same snapshot, using Steps 8 and 10 of the ASADPRECframework. This variant minimizes gradient computation cost but assumes that a single gradient evaluation can be shared across all block solvers.

Synchronous multi-gradient (SyncMulti):

Each of the \(L\) workers independently computes the full stochastic gradient on the same shared batch and snapshot, then extracts its own block component. Workers are synchronized by barriers: no block update is applied until every worker has completed its gradient computation. The mathematical iterate is identical to that of, SyncSingle, using Steps 8 and 10 of ASADPREC, but the wall-clock cost per iteration includes the overhead of \(L\) gradient evaluations and barrier synchronization.

Asynchronous block-owner (ASADPREC):

The synchronization barriers are removed. Each block \(X_\ell\) is permanently assigned to a dedicated worker. All workers share access to the full parameter vector \(\widehat{X}= (\widehat{X}_1, \dots, \widehat{X}_L)\). Worker \(\ell\) repeatedly:

1. reads the full vector \(\widehat{X}\) by cloning each block \(\widehat{X}_j\) under block \(j\)’s lock, ensuring that all parameters within a given block come from the same publication epoch (intra-block consistency); different blocks may come from different epochs (inter-block asynchrony);

2. samples a local mini-batch \(\mathcal{B}_{t,\ell}\), computes the full stochastic gradient \(\nabla F_{\mathcal{B}_{t,\ell}}(\widehat{X}_k)\), and extracts the block component \(G_{t,\ell} = [\nabla F_{\mathcal{B}_{t,\ell}}(\widehat{X}_k)]_\ell\);

3. applies the preconditioned update to \(\widehat{X}_\ell\) and publishes the result atomically under block \(\ell\)’s lock.

This is the asynchronous block-Jacobi model [38], [39] used by ASADPREC. A bounded-delay throttle pauses any worker whose publication counter exceeds the slowest worker’s by more than \(\tau_{\max} = 512\).

6.2 Test problems and architectures↩︎

Table 2 summarizes the five benchmarks. Each model is decomposed so that one or two blocks use the matrix-valued adaptive Muon preconditioner (Newton–Schulz polar approximation [26] for SVHN, SVD for the others) while the remaining blocks use diagonal AdaGrad.

FashionMNIST [47]:

A one-hidden-layer MLP with \(3.26 \times 10^6\) parameters decomposed into 4 blocks: \(W_1\) (Muon), \(b_1\), \(W_2\), \(b_2\) (AdaGrad).

MoE-FMNIST [47]:

A Mixture-of-Experts architecture applied to flattened 28×2828×28 Fashion-MNIST images, with a backbone MLP (\(784 \rightarrow 512784\rightarrow 512\)), a top-2/4 softmax router, four experts of varying width (\(512 \rightarrow h \rightarrow 512\)) with \(h\in\{2048,1024,256,64\}\)), and an output head (\(512 \rightarrow 10512 \rightarrow 10\)), decomposed into 6 blocks: \(W_{1,1}\) and \(W_{2,1}\) (Muon, Newton–Schulz, 5 iterations), expert 3 and expert 4 full parameters, backbone and expert 1 and 2 biases and second-layer weights, router and output head (all AdaGrad).

CovType [48]:

A two-hidden-layer MLP with \(1.61 \times 10^5\) parameters (200,000 samples, 80/20 split), decomposed into 4 blocks: \((W_1,b_1)\), \(W_2\) (Muon), \(b_2\), \((W_3,b_3)\) (AdaGrad).

MovieLens 100K [49]:

A NeuralMF model with \(4.71 \times 10^5\) parameters (100,000 ratings, 943 users, 1,682 items), decomposed into 5 blocks: \(P\), \(Q\), \((b_u, b_i)\) (sparse row-wise AdaGrad), \(W_1\) (Muon), \((c_1, W_2, c_2)\) (AdaGrad).

Criteo [50]:

A click-through-rate model with \(6.35 \times 10^5\) parameters, decomposed into 3 blocks: embedding table (AdaGrad), \(W_1\) (Muon), \((b_1, W_2, b_2)\) (AdaGrad).

SVHN [51]:

A two-hidden-layer MLP with \(14.7 \times 10^6\) parameters applied to flattened \(32\times 32\times 3\) images, decomposed into 5 blocks: \(W_1\) and \(W_2\) (Muon, Newton–Schulz, 5 iterations), \(b_1\), \(b_2\), \((W_3, b_3)\) (AdaGrad). This benchmark features two Muon blocks of different sizes (\(12.6 \times 10^6\) and \(2.1 \times 10^6\) entries), creating a hierarchy of block costs.

6.2.1 Measured block heterogeneity↩︎

The raw block cost ratio (\(t_{\max}/t_{\min}\)) exceeds \(90\) on all benchmarks, but the effective heterogeneity is considerably lower. Each worker’s per-iteration cost has two components: the gradient computation \(t_{\mathrm{grad}}\) (identical for all workers, since each computes the full gradient \(\nabla F\)) and the block update \(t_{\mathrm{block},\ell}\). When \(t_{\mathrm{grad}}\) is large relative to \(t_{\mathrm{block},\ell}\), all workers run at similar speeds regardless of block update costs. Table 3 reports per-block computational costs measured on an Apple M3 Pro processor (single-threaded, median over 50 repetitions). The effective heterogeneity ratio is defined as \((t_{\mathrm{grad}} + t_{\max}) / (t_{\mathrm{grad}} + t_{\min})\).

On FashionMNIST, the Muonpolar factorization (\(105\) ms) far exceeds the gradient (\(5.6\) ms), so the effective ratio remains high (\(19.6\)). On CovType, the gradient (\(0.8\) ms) exceeds the slowest block update (\(0.42\) ms), reducing the effective ratio to \(1.5\times\). SVHN exhibits the highest effective ratio (\(25.8\)) thanks to the Newton–Schulz iteration on \(W_1 \in \mathbb{R}^{4096 \times 3072}\) (\(580\) ms) relative to the gradient (\(23\) ms).

6.3 Implementation and experimental protocol↩︎

All experiments use PyTorch on a shared-memory CPU architecture (Apple M3 Pro, 12 cores) with torch.set_num_threads(1). The CPython GIL prevents byte-code level parallelism, but PyTorch C++ operations (matrix multiplications, SVD, tensor cloning) release the GIL, allowing concurrent execution across cores. The ASADPRECmode uses per-block locking: each block is read and written under its own lock, ensuring intra-block consistency while allowing inter-block asynchrony. A fixed wall-clock budget of 30 seconds is used for problems CovType, MovieLens and Criteo, a budget of 60 seconds for Fashion-MNIST, and MoE-FMNIST and a budget of 120 seconds for SVHN. All runs use 8 seeds (\(1234, \dots, 1241\)), and identical initial parameters per seed. Throughput is reported in model-equivalent updates per second: \(\mathrm{throughput} = N_{\mathrm{block}} / (L \cdot T)\).

6.4 Results without momentum↩︎

Tables 4 and 5 present the results obtained with the momentum-less ASADPRECin the setting just described. We start by reporting, for all benchmarks the full-dataset training loss and the Euclidean norm of the full training gradient \(\|\nabla F(X)\|\) as convergence diagnostics.

ASADPRECobtains the best gradient (in norm) for Criteo, FashionMNIST and MovieLens, while there is little difference between SyncSingleand ASADPRECfor CovType. The gradient norm is higher for SVHN and MoE-FMNIST (due to delayed gradients). The best loss value is obtained by ASADPRECfor 4/6 benchmarks. SyncSingleremains better for MovieLens and Criteo because it performs only one gradient evaluation per iteration.

We now turn to primary generalization metrics, depending on the benchmark. For classification problems (FashionMNIST, CovType, SVHN), we report the test accuracy (fraction of correctly classified test samples). For MovieLens, we report the root mean square error (RMSE), defined as \(\mathrm{RMSE} = \sqrt{\frac{1}{n}\sum_{i=1}^n (\hat{r}_i - r_i)^2}\), where \(\hat{r}_i\) and \(r_i\) are the predicted and true ratings. For Criteo, we report the area under the ROC curve (AUC), which measures the probability that a randomly chosen positive example (clicked ad) is ranked higher than a randomly chosen negative example by the model’s predicted score; AUC \(= 0.5\) corresponds to random ranking and AUC \(= 1\) to perfect ranking. Generalization metrics are reported in Table 5, along with the throughput of ASADPRECand the speedup of ASADPRECcompared with SyncMulti.

ASADPRECis best on all benchmarks except Criteo.

Globally, Tables 4 and 5 suggest that stale (delayed) gradients do not degrade the quality of the training and that the additional iterations made possible by asynchrony translate into better generalization.

6.5 Results with momentum↩︎

The approaches yielding DADPREC.MM/ASADPREC.MMand DADPREC.MG/ASADPREC.MGnot only differ in theory, but their practical requirements are also significantly different, in particular for the Shampooand Muoncase. Indeed, if the SVD of \(M_\ell\) is given by \(M_\ell = U_\ell \Sigma_\ell V_\ell^T\), the Muonrecursions¹⁰ are written, in this case, as

*[1ex] The main difference is that Muon.MM only requires \(\|M_\ell\|_{*,\ell}\), which can be computed with a single polar decomposition (using SVD or Newton-Schulz iterations), while Muon.MG also requires \(\|\widetilde{G}_\ell\|_{*,\ell}\). This potentially requires another decomposition, which can significantly slow down the algorithm. It is therefore of interest to consider faster techniques for computing \(\|\widetilde{G}_\ell\|_{*,\ell}\). We have considered both the stochastic Lanczos quadrature [52] and a warm-started power iteration for this task. As it turns out, the first, although theoretically suitable, appears to be too slow for practical use. By contrast, the second, which can only provide a deterministic lower bound, nevertheless works remarkably well. The details of this technique are given in Appendix 3.

When applied to ASADPREC, these considerations lead us to a family of momentum variants, of which we have retained six, whose characteristics are given in Table 6. In this table, \(\beta\) is the rate of momentum decay appearing in ([mutell]) (with \(\mu_{\max}=0.9\)) and the last column indicates whether a warm-started power iteration is used to approximate \(\|\widetilde{G}_\ell\|_{*,\ell}\) in the Muonupdates.

To keep the paper as concise as possible, we only report in Table 7 the values of the final gradients (in norm) for these six variants, and refer the reader to Appendix 3 for more results.

As can be seen from this table, no specific momentum variant appears to dominate, but the momentum-less variant (first column) is never the best.

6.6 Discussion↩︎

Globally, our results suggest the following conclusions.

1. Asynchronous parallelism preserves and improves convergence quality. DADPRECis the best momentum-less variant for 5/6 benchmarks, and is also best in many case where momentum is used.

2. Heterogeneity drives the speedup. The per-block locking protocol—guaranteeing intra-block consistency while allowing inter-block asynchrony—matches the classical asynchronous block-Jacobi model [38], [39]. The effective heterogeneity ratio, which accounts for the shared gradient computation cost, is a better predictor of the asynchronous speedup than the raw block cost ratio.

3. The cost of momentum is critical for matrix updates like Muon. This makes the MM approach attractive from this point of view, because its lower cost per iteration allows more iterations in the time budget, resulting in gains in accuracy between +1.8 and +3.6 points compared with the MG approach.

4. The estimation of the nuclear norm of \(\widetilde{G}_\ell\) using the warm-started power iteration is an efficient proxy, and does produce the best results for some problems.

5. As expected, the effect of decaying momentum (i.e.\(\beta>0\) in ([mutell])) is problem-dependent, because of its direct interaction with the variance of the problem’s oracle.

6. Although globally promising, our tests did not allow us to select a clear winner among the studied variants.

It should however be stressed that all our experiments were conducted assuming the “full gradient evaluation model”, that is the case where \(G_{t,\ell} = [\nabla F_{\mathcal{B}_{t,\ell}}(\widehat{X}_k)]_\ell\) cannot be computed on its own, but requires the evaluation of \(\nabla F_{\mathcal{B}_{t,\ell}}(\widehat{X}_k)\). Should the evaluation of \(G_{t,\ell}\) be possible at a lower cost¹¹, the heterogeneity ratios would then be dramatically increased, which is likely to give the asynchronous option used in ASADPRECa substantial advantage.

7 Conclusions and perspectives↩︎

We have presented a new class of asynchronous adaptive first-order methods including asynchronous variants of AdaNorm, AdaGrad, Shampooand Muon, as well as variants thereof using momentum and/or inexact step normalization. We have shown that, under reasonable assumptions, all methods in the class have an global convergence rate (essentially) of order \({\pazocal{O}}(1/\sqrt{t})\) in expectation.

We also described numerical experiments with these methods, providing motivation for asynchronous adaptive optimization where different blocks use different geometries, different blocks have drastically different computational costs, and synchronization barriers become prohibitively expensive. Such situations naturally appear in modern large-scale machine learning systems mixing diagonal adaptive methods, Our results, although limited, provide empirical support for the claim that asynchronous adaptive block-coordinate optimization methods may be relevant in heterogeneous large-scale machine learning systems.

Appendix 1: Detailed analysis for momentum-less DADPREC↩︎

The proof of Theorem [thm:dadprec-convergence] follows the broad lines of [30]. In order to keep our argument as compact as possible, we refer, in what follows, to that reference for the proofs of results directly extracted from it¹².

Our first step is to analyze the (expected) descent at iteration \(t\).

The following trace inequalities may then be proved.

Because not every block is necessarily updated at iteration \(t\) in the DADPRECalgorithm, we now need to re-prove the suitable variant of the classical telescoping argument.

We now recall from [30] a crucial consequence of the previous theorem, which provides a bound \(\Theta_t\) on the expected trace of all preconditioners at iteration \(t\).

This last theorem finally gives us all the ingredients to derive the convergence and complexity results of Theorem [thm:dadprec-convergence]. The bound ([Thetak-def-app]), together with Assumption 4, indeed implies convergence of the criticality measure on each block with a uniform rate-of-convergence bound. As shown by the next theorem, the expected global rate of convergence is (in order) proportional to the value of \(\Theta_t\). The proof of this statement is substantially more involved than [30] because of the introduction of delays.

Proof of Theorem [thm:dadprec-convergence]↩︎

Assumption  4 gives that, for each \(\ell\in\{ 1, \ldots, L \}\), \[\sum_{j\in{\pazocal{J}}_{t,\ell}} \|\widetilde{G}_{j,\ell}\|_{*,\ell}^2 \le \kappa_\circ^2\sum_{j\in{\pazocal{J}}_{t,\ell}}\!\mathop{\mathrm{tr}}({\pazocal{U}}_{j,\ell}(\widetilde{G}_{j,\ell})^2) = \kappa_\circ^2\,\mathop{\mathrm{tr}}\!\Big(\!\!\sum_{j\in{\pazocal{J}}_{t,\ell}}{\pazocal{U}}_{j,\ell}(\widetilde{G}_{j,\ell})^2\Big) \le \kappa_\circ^2 \,\mathop{\mathrm{tr}}(\Pi_{t,\ell})\] But \[\mathop{\mathrm{tr}}(\Pi_{t,\ell}) =\sum_{i=1}^{d_\ell}\left(\sqrt{\lambda_i[\Pi_{t,\ell}]}\right)^2 \le \left(\sum_{i=1}^{d_\ell}\sqrt{\lambda_i[\Pi_{t,\ell}]}\right)^2 =\mathop{\mathrm{tr}}(\Pi_{t,\ell}^{1/2})^2,\] and thus, summing over \(\ell\in\{ 1, \ldots, L \}\) and using the fact that \(\sum_i a_i^2 \le (\sum_i a_i )^2\), \[\sqrt{\sum_{\ell=1}^L\sum_{j\in{\pazocal{J}}_{t,\ell}} \|\widetilde{G}_{j,\ell}\|_{*,\ell}^2} \le \kappa_\circ \sqrt{ \sum_{\ell=1}^L\,\mathop{\mathrm{tr}}(\Pi_{t,\ell}^{1/2})^2} \le \kappa_\circ \sqrt{\left(\sum_{\ell=1}^L\,\mathop{\mathrm{tr}}(\Pi_{t,\ell}^{1/2})\right)^2} = \kappa_\circ\sum_{\ell=1}^L\,\mathop{\mathrm{tr}}(\Pi_{t,\ell}^{1/2}).\] Using now the Cauchy-Schwarz and Jensen’s inequalities with this last bound, we deduce that \[\label{therate} \begin{align} \sum_{\ell=1}^L\sum_{j\in{\pazocal{J}}_{t,\ell}} \mathbb{E}\!\left[\|\widetilde{G}_{j,\ell}\|_{*,\ell} \right] &=\mathbb{E}\!\left[\sum_{\ell=1}^L\sum_{j\in{\pazocal{J}}_{t,\ell}}\|\widetilde{G}_{j,\ell}\|_{*,\ell} \right]\\ &\le \sqrt{L\max_{\ell\in\{ 1, \ldots, L \}}|{\pazocal{J}}_{t,\ell}|}\,\mathbb{E}\!\left[\sqrt{\sum_{\ell=1}^L\sum_{j\in{\pazocal{J}}_{t,\ell}}\|\widetilde{G}_{j,\ell}\|_{*,\ell}^2} \right]\\ &\le \kappa_\circ\sqrt{L(t+1)} \,\mathbb{E}\!\left[\sum_{\ell=1}^L\,\mathop{\mathrm{tr}}(\Pi_{t,\ell}^{1/2}) \right]\\ &\le \kappa_\circ\sqrt{L(t+1)} \,\Theta_t \end{align}\tag{13}\] where the last inequality results from ([Thetak-def-app]). Now let \(H_j = ( G_{\pi(j,1),1}, \ldots, G_{\pi(j,L),L}).\) Then \[\label{twoterms} \sum_{j=0}^t \mathbb{E}\!\left[\|G_j\|_* \right] \le \sum_{j=0}^t \mathbb{E}\!\left[\|G_j-H_j\|_* \right] + \sum_{j=0}^t \mathbb{E}\!\left[\|H_j\|_* \right].\tag{14}\] Consider the two terms of the right-hand side separately. To bound the first, we apply Assumption 2 and obtain that \[\|G_j-G_{\pi(j,\ell)}\|_*^2 \le L_G^2 \|X_j-X_{\pi(j,\ell)}\|^2 \le L_G^2\eta^2 \sum_{i=\pi(j,\ell)}^{j-1} \|Z_i\|_*^2.\] Assumption 7 then guarantees that each \(\|Z_i\|_*^2\) can only appear \(\tau\) times when summing over \(j\), and thus that \[\sum_{j=0}^t\sum_{\ell=1}^L\mathbb{E}\!\left[\|G_j-G_{\pi(j,\ell)}\|_*^2 \right] \le \tau L_G^2 \eta^2 \sum_{j=0}^t\sum_{\ell=1}^L \mathbb{E}\!\left[\|Z_j\|_*^2 \right] = \tau L_G^2 \eta^2 L \sum_{j=0}^t \mathbb{E}\!\left[\|Z_j\|_*^2 \right].\] But (8 ), (12 ) and ([Thetak-def-app]) imply that \[\label{bdDelta} \sum_{j=0}^t\sum_{\ell=1}^L \mathbb{E}\!\left[\|Z_{j,\ell}\|_{*,\ell}^2 \right] \le \Delta_t \le \kappa_0+2N\log\!\left(\mathbb{E}\!\left[\sum_{\ell=1}^L\mathop{\mathrm{tr}}(\Pi_{t,\ell}^{1/2}) \right]\right) \le \kappa_0+2N\log(\Theta_t).\tag{15}\] and thus, using (3 ), that \[\sum_{j=0}^t\mathbb{E}\!\left[\|G_j-H_j\|_*^2 \right] =\sum_{j=0}^t\sum_{\ell=1}^L\mathbb{E}\!\left[\|G_{j,\ell}-G_{\pi(j,\ell),\ell}\|_{*,\ell}^2 \right] \le \tau L_G^2 \eta^2 L \left[\kappa_0+2N\log(\Theta_t)\right].\] The Cauchy-Schwarz inequality then gives that \[\label{bdt1} \sum_{j=0}^t\mathbb{E}\!\left[\|G_j-H_j\|_* \right] \le \sqrt{t+1}\sqrt{\sum_{j=0}^t\mathbb{E}\!\left[\|G_j-H_j\|_*^2 \right]} \le \sqrt{t+1}\,\sqrt{\tau L_G^2\eta^2 L \left[\kappa_0+2N\log(\Theta_t)\right]}.\tag{16}\] Consider now the second term in the right-hand side of (14 ). We have, again using (3 ), that \[\sum_{j=0}^t \|H_j\|_* = \sum_{j=0}^t \sqrt{\sum_{\ell=1}^L\|G_{\pi(j,\ell),\ell}\|_{*,\ell}^2} \le \sum_{j=0}^t\sum_{\ell=1}^L\|G_{\pi(j,\ell),\ell}\|_{*,\ell}\] Now Assumption 7 ensures each \(\|G_{\pi(j,\ell),\ell}\|_{*,\ell}^2\) appears at most \(\tau\) times in the sum over \(j\). Thus, \[\label{aa1} \sum_{j=0}^t \|H_j\|_* \leq \sum_{j=0}^t\sum_{\ell=1}^L\|G_{\pi(j,\ell),\ell}\|_{*,\ell} \le \tau \sum_{\ell=1}^L\sum_{j\in{\pazocal{J}}_{t,\ell}}\|G_{j,\ell}\|_{*,\ell}.\tag{17}\] But Assumption 5, Jensen’s inequality and the law of total expectation ensure that \[\mathbb{E}\!\left[\|G_{j,\ell} \|_{*,\ell} \right] = \mathbb{E}\!\left[\|\mathbb{E}_{j}\!\left[\widetilde{G}_{j,\ell} \right]\|_{*,\ell} \right] \le \mathbb{E}\!\left[\mathbb{E}_{j}\!\left[\|\widetilde{G}_{j,\ell}\|_{*,\ell} \right] \right] = \mathbb{E}\!\left[\|\widetilde{G}_{j,\ell} \|_{*,\ell} \right].\] Substituting this identity in (17 ) and using (13 ), we then obtain that \[\label{bdt2} \sum_{j=0}^t \mathbb{E}\!\left[\|H_j\|_* \right] \le \tau \sum_{\ell=1}^L\sum_{j\in{\pazocal{J}}_{t,\ell}}\mathbb{E}\!\left[\|\widetilde{G}_{j,\ell}\|_{*,\ell} \right] \le \tau \kappa_\circ\sqrt{L(t+1)} \,\Theta_t.\tag{18}\] Combining (16 ) and (18 ) and dividing by \(t+1\) finally yields ([thetruerate]).

Proof of Corollary [therate2]↩︎

Observe that, when \(\Theta_t\) (as defined in Theorem [thm:convergence-theta]) grows with \(t\), the bound ([thetruerate]) is dominated by the term \(\kappa_c\Theta_t/\sqrt{t+1}\). Since Corollary [therate2] solely aims at describing dominating terms for growing \(t\), we may ignore the other terms and merely consider that ([thetruerate]) reduces to \[\label{simple-conv} \frac{1}{t+1}\sum_{j=0}^t\mathbb{E}\!\left[\|G_j\|_* \right] \le \frac{\kappa_c \Theta_t}{\sqrt{t+1}},\tag{19}\] which is identical to the result (equation (3.42)) of Theorem 3.10 in [30]. The rest of the proof then follows that of [30] with \(\nu_k^2\) replaced by \[\mu_{\max}\max_{\ell\in\{ 1, \ldots, L \}}\left[\sum_{j\in{\pazocal{J}}_{t,\ell}}\frac{1}{\min[|{\pazocal{J}}_{j,\ell}|]}\right] \le \left\{\begin{array}{ll} \frac{\displaystyle \mu_{\max} \sigma_{\rm tot}^2}{\displaystyle 1-\alpha}\psi_t^{1-\alpha} & \;\; if\;\; \alpha < 1,\\ \mu_{\max} \sigma_{\rm tot}^2 \log(\psi_t) & \;\; if\;\; \alpha = 1,\\ \mu_{\max} \sigma_{\rm tot}^2\zeta(\alpha) & \;\; if\;\; \alpha > 1, \end{array}\right.\] yielding ([therate2-psi]).

Appendix 2: Detailed analysis for DADPRECwith momentum↩︎

We start by establishing a bound on the norm of the difference between the block-wise momentum \(M_{t,\ell}\) and the approximate gradient \(\widetilde{G}_{t,\ell}\).

Proof of Theorem [thm:momMM-convergence]↩︎

Observe now that the DADPREC.MMalgorithm is nothing but Algorithm DADPRECwhere the momentum \(M_{t,\ell}\) plays the role of the approximate gradient \(\widetilde{G}_{t,\ell}\). Moreover Assumption 8 exactly reflects this change of perspective. Also note that, for all \(\ell\in\{ 1, \ldots, L \}\) and \(j\in{\pazocal{J}}_{t,\ell}\), \[\|M_{j,\ell} - G_{j,l}\|_{*,\ell}^2 \le 2 \|E_{j,\ell}\|_{*,\ell}^2 + 2 \|\widetilde{G}_{j,\ell} - G_{j,l}\|_{*,\ell}^2.\] Hence, using Lemma [lem:E-form] and given that \(\bar{\mu}_{j,\ell}<1\), we have that \[\sum_{\ell=1}^L\sum_{j\in{\pazocal{J}}_{t,\ell}}\mathbb{E}\!\left[\|M_{j,\ell}-G_{j,l}\|_{*,\ell}^2 \right] <\left(\frac{3}{(1\!-\!\mu_{\max})^2}+2\right)\sum_{j=0}^t\mathbb{E}\!\left[\|\widetilde{G}_j-G_j\|_*^2 \right] +\frac{2[\tau L L_G^2\eta^2}{(1\!-\!\mu_{\max})^2}\sum_{j=0}^t\mathbb{E}\!\left[\|Z_j\|_*^2 \right].\] This provides an adapted formulation to replace (?? ) in Assumption 6, which thus continues to hold with \[\label{nuknown} \nu_t^2 = \left(\frac{3}{(1-\mu_{\max})^2}+2\right)\sum_{j=0}^t\mathbb{E}\!\left[\|\widetilde{G}_j-G_j\|_*^2 \right] \;\;and\;\; \omega^2 = \frac{2\tau L L_G^2\eta^2}{(1-\mu_{\max})^2}.\tag{21}\] As a consequence, the theory developed for the DADPRECalgorithm also holds for its DADPREC.MMvariant and this proves Theorem [thm:momMM-convergence].

Proof of Theorem [thm:convergence-alt]↩︎

This section considers the convergence of the DADPREC.MGalgorithm. We first recall a result on the impact of using this algorithm on the structural identities of Assumption 8.

From this point on, the analysis follows the lines of the argument of Appendix 1 with Lemma [lem:pert-bounds2] providing an alternate set of structural inequalities. Lemma [lem:sqrt-op] is unchanged. The modifications to Theorem [th:master-op] are minor. It is restated as follows.

Note that the form of bound ([master-mom]) differs very little from that of ([master]): besides using different constants, ([master-mom]) now involves a term in \(\kappa_{\nu\nu}\theta_t^2\). This term then percolates through the proof of Theorem [thm:convergence-theta], so that ([Thetak-def-app]) now becomes \[\label{Thetak-def-app-MG} \mathbb{E}\!\left[\sum_{\ell=1}^L \mathop{\mathrm{tr}}(\Pi_{t,\ell}^{1/2}) \right] \le \Theta_t \stackrel{\rm def}{=} \max\left[\,\kappa_\Theta,\,\frac{3\kappa_{\nu\nu}\theta_t^2}{\eta},\,T_t\,\right].\tag{22}\] The proof of Theorem [thm:dadprec-convergence] then stands without modification other than using this updated value of \(\Theta_t\), finally yielding Theorem [thm:convergence-alt].

Proof of Corollary [therate2-alt]↩︎

Finally, the proof of Corollary [therate2-alt] is similar in spirit to that of Corollary 3.11 in [30], with two twists. The first is that the formula for the block-wise variance is now given by ([var-cond-alt2]), which, together with ([mutell]), implies¹³ that \[\begin{align} \sum_{\ell=1}^L\sum_{j\in{\pazocal{J}}_{t,\ell}}\|M_{j,\ell}-G_{j,\ell}\|_{*,\ell}^2 & \le \mu_{\max}\,\sum_{\ell=1}^L \sigma_\ell^2\, \sum_{j\in{\pazocal{J}}_{t,\ell}}\frac{1}{\min[|{\pazocal{J}}_{j,\ell}|,|{\pazocal{J}}_{j-1,\ell}|]^{\alpha+2\beta}} +\omega^2\sum_{j=0}^t\mathbb{E}\!\left[\|Z_{j,\ell}\|_{*,\ell}^2 \right]\\ & \le \mu_{\max} \sigma_{\rm tot}^2\, \sum_{j=1}^t\frac{1}{|{\pazocal{J}}_{j-1,\ell}|^{\alpha+2\beta}}+\omega^2\sum_{j=0}^t\mathbb{E}\!\left[\|Z_{j,\ell}\|_{*,\ell}^2 \right]. \end{align}\] Therefore, using \(\psi_{t-1}\le\psi_t\), we may set the bound \[\label{thetaMG} \theta_t^2= \left\{\begin{array}{ll} \frac{\sigma_{\rm tot}^2}{1-\alpha-2\beta}\psi_t^{1-\alpha-2\beta} & \;\; if\;\; \alpha+2\beta < 1,\\ \sigma_{\rm tot}^2 \log(\psi_t) & \;\; if\;\; \alpha+2\beta = 1,\\ \sigma_{\rm tot}^2\zeta(\alpha+2\beta) & \;\; if\;\; \alpha+2\beta > 1.\\ \end{array}\right.\tag{23}\] The second is the presence of term \(3\kappa_{\nu\nu}\theta_t^2\) in ([Thetak-def-alt]). Since the dominant term in the convergence bound ([thetruerate-alt]) is the last, the final order of convergence is determined by the maximal power of \(\theta_t\) occuring in the expression of \(\Theta_t\), which is now \(\theta_t^2\) because of this new term. Thus substituting (23 ) in ([thetruerate-alt]) and only keeping the dominant term in \(t\) gives ([therateMG]). \(\Box\)

Appendix 3: Implementation and additional results for asynchronous momentum variants↩︎

We first describe our warm-started power iteration to compute an approximation of the nuclear norm of \(G\). It is based on the proxy \[\|G\|_* \approx \frac{\|G\|_F^2}{\sigma_{\max}[G]}\] where \(\sigma_{\max}[G]\) is the largest singular value of \(G\). Note that \(\|G\|_F^2\) is easily computable, but we estimate \(\sigma_{\max}[G]\) using two iterations of the power iterations on \(G\). This does not ensure a convergent estimator, but gives a deterministic lower bound \[\frac{\|G\|_F^2}{\sigma_{\max}[G]} = \frac{\sum_{i=1}^n \sigma_i[G]^2}{\sigma_{\max}[G]} \le \sum_{i=1}^n \sigma_i[G] = \|G\|_*,\] where equality holds if and only if all singular values are identical. Our algorithm is warm-started: we initialize the power iteration with the result obtained at the previous optimisation iteration. As the gradient \(G\) changes progressively, so does its dominant singular vector, and two iterations are typically enough to obtain an accurate estimation est\((\|G\|_*)\). The detailed algorithm is given .

The cost of using WSPOWER for an \(m \times n\) matrix amounts to \({\pazocal{O}}(4mn)\) flops for the four matrix-vector products, plus \({\pazocal{O}}(mn)\) flops for \(\|G\|_F^2\), giving a total of \({\pazocal{O}}(5mn)\) flops, compared to \({\pazocal{O}}(mn\min(m,n))\) flops for the SVD. In our experiments with \(256\times 2048\) matrices, using WSPOWER is \(\min(m,n)/5 \approx 50\) times faster.

Although a mere proxy of \(\|G\|_*\), WSPOWER works remarkably well, but its warm-start feature is crucial. Without it, accuracy results diverge on larger models (SVHN: 56%, MoE-FMNIST: 77%, to compare with SVHN 70.44% ± 0.43, MoE 88.63% ± 0.34 with warm start, this last result being the best achieved for this benchmark in all our runs) The throughput of our asynchronous MG optimization variants using WSPOWER then become comparable that of MM variants, as shown in Table 10 below.

We now complete the results mentioned in Section 6.5. Table 8 gives the final loss values obtained when using different asynchronous momentum variants.

We see that the MMconst variant dominates on CovType (0.149), Criteo (0.008!), Fashion-MNIST (0.239)and MovieLens (0.146), while MGconstA is best for MoE-FMNIST and MGdecay for SVHN.

The generalization metrics for the different variants with momentum are given in Table 9.

Care should however be exercised: compared to the momentum-less option, the loss for MovieLens improves from 0.353 to 0.008, but the accuracy worsens from 0.916 to 0.992, illustrating a possible overfitting in that case.

The throughputs per problem are given in Table 10, which shows that the introduction of the MM momentum has a limited effect on throughput, but that the MG approach has a clear negative impact, due to the double decomposition that we pinpointed in Section 6.5. Fortunately, the use of the warm-started power iteration to approximate the nuclear norm of \(\widetilde{G}_\ell\) essentially corrects this deficiency.

Finally focussing on the classification problems, we report, in Table 11, the variance of the accuracy obtained for the 8 independent runs. This table suggests that the variance on accuracy is nearly always better for the MM momentum than for the  MG approach. In particular, a comparison of MMconst with MGconst shows that the former improves the accuracy of all classification problems by as much as 3.58% (for SVHN).

References↩︎