arXiv:2605.13434v1 [cs.LG] 13 May 2026
Rescaled Asynchronous SGD: Optimal Distributed Optimization under Data and System Heterogeneity
Ammar Mahran KAUST [email protected]
Artavazd Maranjyan KAUST [email protected]
Peter Richtárik KAUST [email protected]
Abstract Asynchronous stochastic gradient descent (ASGD) is a standard way to exploit heterogeneous compute resources in distributed learning: instead of forcing fast workers to wait for slow ones, the server updates the model whenever a gradient arrives. Vanilla ASGD applies each arriving gradient with the same weight. When local data distributions are heterogeneous, this becomes problematic: faster workers contribute more updates, and we show theoretically that the method is biased toward a frequency-weighted average of the local objectives rather than the desired global objective. Existing remedies typically move away from the simple ASGD template by introducing gathering phases, buffering, or extra memory. We show that this is unnecessary. Keeping the standard ASGD mechanism, we recover the correct objective by rescaling worker-specific stepsizes in proportion to their computation times, so that each worker contributes the same aggregate learning rate over a cycle. In the non-convex setting, under smoothness and bounded heterogeneity assumptions, we prove that the resulting method, Rescaled ASGD, converges to stationary points of the correct global objective in the fixed-computation model. Its time complexity matches the known lower bound in the leading term, while the effects of staleness and data heterogeneity appear only in lower-order terms. Experiments confirm that the method converges to the correct objective and is competitive with state-of-the-art baselines.
1
Introduction
We consider the non-convex distributed optimization problem ( ) n 1X min F (x) := Fi (x) , n i=1 x∈Rd
(1)
where Fi : Rd → R, x 7→ Eξ∼Di [fi (x; ξ)] denotes the local objective function of worker i = 1, . . . , n, which is the expected value of the sample loss fi (x; ξ) over data points ξ drawn from their local distribution Di . In parallel stochastic gradient descent methods, workers collaboratively solve this problem by computing stochastic gradients of their local objectives, which are then processed by a central server to find an ε-stationary point, that is, a random vector x̄ ∈ Rd such that 2 E ∥∇F (x̄)∥ ≤ ε. Preprint.
Table 1: Comparison of worst-case wall-clock time complexities for parallel stochastic first-order methods in the fixed-computation time model with data and system heterogeneity to achieve εstationarity in (1). We denote the arithmetic mean and maximum, respectively, of the workers’ computation times by τA , τmax . Problem parameters include the initial optimality gap ∆ := F (x0 ) − F ∗ and global smoothness constant L (Assumption 2.3), the target stationarity ε > 0, and a bound on the stochastic gradient variance σ 2 (Assumption 2.2). Asymptotic Optimality: time complexity matches the lower bound established by Tyurin and Richtárik [2023] in the leading term and achieves ε-stationarity for arbitrarily small ε. No Idle Workers: all workers remain busy and computational resources are fully utilized. No Memory Overhead: no gradients or model iterates need to be stored for later use. Method (Reference)
Time Complexity (Leading Term)
Asymptotic Optimality
No Idle Workers
No Memory Overhead
Naive Minibatch SGD
∆Lσ 2 nε2 τmax
✘
✘
✔
∆Lσ 2 nε2 τA
✔
✔
✘
∆L′ σ 2 (†) nε2 τA
✘
✔
✘
∆L(σ 2 +ζ 2 ) τmax (‡) nε2
✘
✘
✘
∆Lσ 2 nε2 τA
✘(‡)
✔
✔
∆Lσ 2 nε2 τA
✔
✔
✔
[Cotter et al., 2011, Dekel et al., 2012] Malenia SGD
[Tyurin and Richtárik, 2023] Ringleader ASGD
[Maranjyan and Richtárik, 2026] Concurrent ASGD
[Koloskova et al., 2022] Delay-Adaptive ASGD
[Mishchenko et al., 2022a] Rescaled ASGD
Theorem 4.1, Corollary 4.2
(†) Maranjyan and Richtárik [2026] rely on a constant L′ ≥ L to bound data heterogeneity. In certain settings, L′ = O(L), and
Ringleader ASGD achieves optimal time complexity.
(‡) Koloskova et al. [2022], Mishchenko et al. [2022a] rely on a stricter assumption on the data heterogeneity, ∥∇F (x) − ∇F (x)∥2 ≤ i ζ 2 . In Concurrent ASGD, the constant ζ 2 affects the leading term of the time complexity. In Delay-Adaptive ASGD, the gradient norm
can be bounded only up to this constant, therefore not achieving arbitrary ε-stationarity, and thus no asymptotic optimality. If the computation times are equal across workers, Delay-Adaptive ASGD achieves arbitrary ε-stationarity with the leading term of the time complexity scaling with τA . See Section D for details.
In this model, workers differ across two dimensions. First, each worker has their own local objective function Fi which may differ from one another and the global objective function F . These objective functions vary across workers, typically due to highly heterogeneous local data distributions Di , as is common in federated learning [McMahan et al., 2017, Caldas et al., 2019, Kairouz et al., 2021, Wang et al., 2021], for example. While early asynchronous literature primarily focused on the homogeneous data setting [e.g., Tyurin and Richtárik, 2023, Maranjyan et al., 2025d], recent extensions to the heterogeneous setting [Mishchenko et al., 2022a, Koloskova et al., 2022, Nguyen et al., 2022] often require restrictive similarity assumptions or algorithmic modifications. In the present work, our focus is on the general case allowing for data heterogeneity (Fi ̸= Fj ) under a relaxed similarity assumption (Assumption 2.5). Second, workers may also differ in the time it takes them to compute stochastic gradients, due to, for example, hardware differences [Dutta et al., 2018, Li et al., 2020, Maranjyan et al., 2025a,b,c]. We denote the time needed for worker i to compute a stochastic gradient by τi throughout the text and explicitly allow for these to differ across workers. Naive Distributed Optimization In Naive Minibatch SGD [Cotter et al., 2011, Dekel et al., 2012], each worker i computes a stochastic gradient in τi units of time and sends it to the central server. Once the server has received gradients from all workers, gradients are averaged, and the model is updated and sent back to the workers for computation of another round of gradients. Each round thus takes τmax := maxi τi units of time, leaving faster workers idle after delivering their gradients while they wait for the global model update. While this fully synchronized procedure simplifies the theoretical analysis to that of standard SGD, it severely underutilizes fast workers whose idle time could be spent more productively. 2
Algorithm 1 Asynchronous SGD 1: Input: initial point x0 ∈ Rd , stepsizes γk > 0 2: Set y0,i = x0 , ∀i ▷ yk,i = worker i’s model before update k. 3: Each worker i begins calculation of ∇fi (y0,i ) 4: for k = 0, 1, . . . do 5: Some worker ik delivers stochastic gradient ∇fik (yk,ik ) 6: Server updates xk+1 = xk − γk ∇fik (yk,ik ) ▷ Model update with stale gradient. 7: Update worker model:
yk+1,ik = xk+1 yk+1,j = yk,j , ∀j ̸= ik 8: Worker ik begins calculating ∇fik (xk+1 ) 9: end for
▷ Update model held by ik . ▷ Models held by other workers unchanged.
Leveraging Faster Workers Tyurin and Richtárik [2023] introduce an improved variant, Malenia SGD, that addresses worker idleness. In this algorithm, workers compute a number of stochastic gradients at the current model until the server receives at least one stochastic gradient from each. Thus, faster workers do not idly wait after computation of a single stochastic gradient, but continue computing stochastic gradients at the same model iterate. These gradients are then averaged per worker, so that each worker effectively delivers one stochastic gradient with reduced variance. Once the server has obtained at least one gradient from each worker, the model is updated and sent back to all workers for the next round. While the gradient gathering phase is asynchronous, the model update happens synchronously ensuring all workers compute gradients at the same model iterate in each round. Building on this, Maranjyan and Richtárik [2026] introduce an asynchronous algorithm, Ringleader ASGD, in which the number of model updates is greater, allowing for greater exploration of the search space. Here, the gradient gathering phase is followed by n asynchronous updates, one from each worker. Notably, Ringleader ASGD utilizes a gradient table in which gradients from all workers are stored to ensure that once an asynchronous update is taken, the step direction approximates the gradient of the global objective function rather than that of the workers’ local objective functions. As in Naive Minibatch SGD and Malenia SGD, the gathering phase causes the slowest worker to dictate when updates from all workers can be applied. Nevertheless, analysis shows that the leading term of the wall-clock time complexity of Malenia SGD and Ringleader ASGD scales with the arithmetic mean of the computation times, not their maximum. This speed-up stems from the variance reduction of having faster workers compute multiple stochastic gradients per round. Asynchronous SGD In contrast to these methods, Asynchronous SGD never synchronizes workers nor delays the application of updates from any workers. Instead, the server updates the model whenever it receives a gradient from a worker, and immediately forwards the updated model to the same worker for computation of their next stochastic gradient. This method is described in Algorithm 1.1 Asynchronous SGD avoids the problem of idle workers, but introduces two novel problems: staleness
and objective inconsistency. Because the server updates its model whenever a gradient arrives from any worker, the gradient applied in a given iteration typically does not correspond to the model of the previous iteration; in other words, the server performs model updates with stale gradients. Moreover, as faster workers contribute more gradient updates than slower ones, the search trajectory may be biased towards faster workers’ local objective functions, resulting in the wrong objective function being targeted. Koloskova et al. [2022] analyze a modified version of Algorithm 1. Here, workers compute stochastic gradients in parallel and the server updates the model whenever an update arrives, just like in Algorithm 1. Unlike Asynchronous SGD, the server then selects a worker uniformly at random to compute a stochastic gradient at the updated model. Crucially, workers are sampled uniformly at random, and even workers who are still busy computing a stochastic gradient may be sampled for computation of another gradient at this updated iterate. Thus, the total number of gradients in flight is 1 Henceforth, Asynchronous SGD refers to Algorithm 1 specifically, while "asynchronous SGD" refers to the wider class of
asynchronous SGD methods.
3
kept constant throughout this algorithm. While theoretical analysis shows this method to converge to a stationary point, it does not leverage faster workers to compute more gradients as the shares of model updates by all workers are, by construction, equal in the long run. Mishchenko et al. [2022a] offer the latest analysis of Algorithm 1 under arbitrary delays. In the heterogeneous data setting, the authors can only establish a bound on the gradient norm with a non-vanishing term capturing the degree of data heterogeneity, i.e., a measure of how much the gradients of the local objective functions differ from that of the global objective function. As the authors rightly point out, "we cannot really expect good performance for arbitrarily heterogeneous losses under arbitrary delays." From Gradient to Gradient Step Fortunately, there is often more temporal structure in practice than assumed in the analysis of Mishchenko et al. [2022a] that we can exploit. To a first-order approximation, computation times may, in many settings, be assumed to be fixed over time. To see how this added structure can be leveraged, let us first consider what happens when Pn Naive Minibatch SGD takes a step of stepsize α: The model x0 is moved in direction −α · n1 i=1 ∇fi (x0 ), so that the contribution of worker i to this update is −α/n∇fi (x0 ). In expectation, this becomes −α/n∇Fi (x0 ). Next, consider what Asynchronous SGD can do in the same time, and let us assume, for simplicity, 0 that worker i computes Ki stochastic gradients at points y00 , . . . , yK in the time needed for i −1 one Naive Minibatch SGD update. If each step of worker i has stepsize γi , then the contribution PKi −1 0 of worker i over this time is − k=0 γi ∇fi (yk0 ). Provided that the Ki points y00 , . . . , yK are i −1 0 close to x and the evaluated gradients do not change abruptly, this contribution is approximately equal to −Ki γi ∇Fi (x0 ) in expectation. Comparison with the contribution made by worker i in Naive Minibatch SGD suggests choosing a worker-specific stepsize γi = α/nKi to approximate the ideal gradient step −α∇F (x0 ) through the sum of all workers’ contributions. The sequence of stochastic gradient steps thus computed by Asynchronous SGD is, in general, not an unbiased estimator of the true gradient step −α∇F (x0 ). However, our analysis shows that this bias can be controlled efficiently. Moreover, by letting faster workers take more steps we achieve greater variance reduction and faster convergence. Contributions and Insights Our primary conceptual insight is a shift in perspective: asynchronous methods do not need to approximate the true global gradient direction in every single iteration. Instead, we show that by approximating the global gradient step—split across workers and over time—we can establish convergence to an ε-stationary point. We show that rescaling worker-specific stepsizes (Rescaled ASGD) neutralizes objective inconsistency under data and system heterogeneity. In the fixed-computation model, Rescaled ASGD achieves a near-optimal wall-clock time complexity to reach an ε-stationary point, matching known lower bounds in the leading term (cf. Table 1). As an additional result of our analysis, we show that Vanilla ASGD, which employs the same stepsizes for all workers, targets a frequency-weighted average of the local objectives. Rescaled ASGD maintains the standard ASGD template (Algorithm 1) without introducing memory
overhead, gathering phases, or worker idleness. Through proof-of-concept experiments on heterogeneous data (cf. Section 5), we confirm that Rescaled ASGD accurately targets the global objective and remains competitive with memory- and synchronization-heavy baselines.
2
Problem Setup
Let ξk denote the random sample drawn by worker ik for the stochastic gradient delivered in iteration k. As is standard, we assume the collection of samples {ξk }k∈N to be mutually independent. Moreover, we make the following assumptions about the stochastic gradients: Assumption 2.1 (Unbiased Gradients). The stochastic gradients are unbiased, that is ∀x ∈ Rd ,
Eξ∼Di [∇fi (x, ξ)] = ∇Fi (x),
∀i ∈ {1, . . . , n} .
Assumption 2.2 (Bounded Variance). The stochastic gradients have bounded variance σ 2 ≥ 0, that is, 2 Eξ∼Di ∥∇fi (x, ξ) − ∇Fi (x)∥ ≤ σ 2 , ∀x ∈ Rd , ∀i ∈ {1, . . . , n} . 4
Recall the standard definition of smoothness. Definition 2.1 (Smoothness). A differentiable function F : Rd → R is called LF -smooth if ∥∇F (x) − ∇F (y)∥ ≤ LF ∥x − y∥ ,
∀x, y ∈ Rd .
By convention, LF denotes the smallest such constant. For our analysis, we require the functions involved to be smooth. Assumption 2.3 (Global Objective Function). The global objective function F is differentiable and L-smooth, and bounded below by F ∗ := inf x∈Rd F (x) > −∞. We denote the initial optimality gap by ∆ := F (x0 ) − F ∗ . Assumption 2.4 (Local Objective Functions). Each local objective function Fi is differentiable and Li -smooth. We denote by Lmax := maxi Li the maximum of these local smoothness constants. Assumption 2.5 (Bounded Heterogeneity). There exist constants ζ, ρ ≥ 0 such that 2
2
∥∇Fi (x)∥ ≤ ζ 2 + ρ2 ∥∇F (x)∥ ,
∀x ∈ Rd ,
∀i ∈ {1, . . . , n} .
Assumptions 2.1 to 2.3 are standard in first-order methods [Ghadimi and Lan, 2013, Bubeck, 2015, Bottou et al., 2018]. Limiting function heterogeneity as in Assumption 2.5 is necessary to establish convergence for Asynchronous SGD and similar methods that take steps based on gradients from a single worker.
While early asynchronous literature often relied on the overly restrictive assumption of bounded gradients (ρ = 0) [Bertsekas, 2011, Recht et al., 2011], recent analyses [Mishchenko et al., 2022a, Koloskova et al., 2022] relax this but still constrain the gradient scaling. Our assumption is even weaker, allowing us to encompass a significantly broader class of functions. In Section C, we offer a deeper discussion on these bounded heterogeneity assumptions.
3
Asynchronous SGD with Cyclic Update Schedule and Rescaled ASGD
Following Mishchenko et al. [2022a], Tyurin and Richtárik [2023], we assume worker i to take a fixed amount of time, τi , to compute one stochastic gradient. These times may differ across workers, but do not vary over time for a given worker. To facilitate the theoretical analysis and derive the convergence guarantees in Section 4, we must ensure that workers deliver a fixed number of updates over a given time horizon. To that end, we make the following assumption: Assumption 3.1 (Harmonic Periods). The set of computation times {τ1 , . . . , τn } satisfies τi/τj ∈ N or τj/τi ∈ N for all i, j ∈ {1, . . . , n}. To the best of our knowledge, this structural assumption is a novelty in the optimization literature and merits some discussion. First, while it imposes structure on the computation times, it preserves the primary characteristic of interest: worker heterogeneity. Second, unlike the functional assumptions introduced in Section 2, this constraint can be strictly enforced in practice. Real-time systems scheduling often uses harmonic task sets to maximize resource utilization in exactly this manner [Han and Tyan, 1997, Sudvarg et al., 2024]. For example, a small delay can be added to round the computation times to the next-largest power of two, slowing down any given worker by a factor of two at worst. Third, this structure renders the theoretical analysis tractable and allows us to derive the interpretable bounds on the wall-clock time complexity presented in Section 4. The performance of Rescaled ASGD remains robust in practice even when this assumption is relaxed. In Section 5, we demonstrate that our method remains robust in settings with non-harmonic periods and stochastically fluctuating computation times. This suggests that while our fixed-computation model imposes strict theoretical structure, the conceptual insights generalize to much broader, unstructured settings. Cyclic Update Schedule Under Assumption 3.1, Asynchronous SGD as presented in Algorithm 1 becomes more structured. Assuming, for ease of exposition, that the server processes updates 5
delivered by multiple workers at the same time in a fixed order, the sequence of worker indices ik is fully deterministic and follows a cyclic update schedule. After τmax time units, each worker will have sent Ki := τmax τi updates in a fixed order that will be repeated exactly over the next τmax time units. We Pnrefer to one such pass as a cycle, to τmax as the (wall-clock) cycle duration, and define K := i=1 Ki as the total number of updates received from all workers over the course of one cycle. We may now reformulate Algorithm 1 in this more structured form as shown in Algorithm 2. Here, m denotes the cycle and k the iteration within a cycle. Algorithm 2 Asynchronous SGD with Cyclic Update Schedule 1: Input: initial point x0 ∈ Rd , stepsizes γkm > 0 2: Set x00 = x0 ▷ xm k = server model before update k in cycle m. 0 0 m 3: Set y0,i = x , ∀i ▷ yk,i = worker i’s model before update k in cycle m.
0 4: Each worker i begins calculation of ∇fi (y0,i ) 5: for m = 0, 1, 2, . . . do 6: for k = 0, . . . , K − 1 do m 7: Worker ik delivers stochastic gradient ∇fik (yk,i ) k
m m m Server updates xm ▷ Model update with stale gradient. k+1 = xk − γk ∇fik (yk,ik ) Update worker model: m yk+1,i = xm ▷ Update model held by ik . k+1 k m m yk+1,j = yk,j , ∀j ̸= ik ▷ Models held by other workers unchanged. 10: Worker ik begins calculating ∇fik (xm k+1 ) 11: end for m+1 m 12: Set y0,i = yK,i , ∀i ▷ Last point seen in current cycle. 13: Set xm+1 = xm ▷ End-of-cycle update. K m+1 m ▷ First point of new cycle. 14: Set x0 = xK 15: end for
8: 9:
Algorithm 2 initializes a shared point x0 , after which all workers proceed asynchronously without any synchronization or idle periods. We refer to {xm }m∈N as the cycle iterates and to {xm k }k=1,...,K as Pn P n the inner iterates within cycle m. Note that each cycle comprises K = i=1 Ki = i=1 τmax/τi = nτmax τH updates, where τH is the harmonic mean of the computation times. m Let yk,i be the local model held by worker i at the start of inner iteration k in cycle m. This model ′ ′ corresponds to a past server iterate xm k′ , originating either from the previous cycle (m = m − 1, ′ ′ k ≥ k) or the current one (m = m, k ′ < k). As a result, the gradient staleness, or delays—that is, the number of server updates between a worker reading the model and applying its gradient—is bounded by K. Bounded staleness of this form is a well-established setting in the asynchronous optimization literature [Agarwal and Duchi, 2011, Recht et al., 2011, Lian et al., 2015].
By definition, cycle iterates xm capture the server state at the start of cycle m. Thus, each worker i contributes exactly Ki updates between xm and xm+1 . This additional structure allows us to analyze the sequence of cycle iterates. Rescaled ASGD We aim to find a stationary point of the objective function (1). Our proposed method sets the stepsizes in Algorithm 2 to γkm = γik ∝ τik . By doing so, even though the number of gradients delivered varies across workers, the aggregate stepsizes taken along each worker’s descent direction over a cycle, γi Ki = γi · τmax/τi , are equal. The accumulated gradient update taken by Rescaled ASGD over the course of a cycle m can then be decomposed (cf. Lemma B.7) into the exact scaled global gradient α∇F (xm ), a bias term bm , and a noise term νm : K−1 X k=0
where α :=
PK−1 k=0
m γik ∇fik (yk,i ) = α∇F (xm ) + bm + νm , k
γik is the cycle stepsize.
The bias term bm , unlike the noise term νm , does not vanish in expectation. As we will see next, however, the impact of this bias can be efficiently controlled by scaling down the stepsizes. 6
4
Convergence Guarantees
Our main result establishes that Rescaled ASGD with properly chosen, worker-specific stepsizes targets the equal-weighted average F defined in (1). For a chosen stepsize parameter γ > 0, we set the worker-specific stepsizes as τH γi := γ · τi · , (2) nτmax where the last factor serves as a normalizing constant. Under this rescaling, we obtain the following convergence guarantee, the proof of which, as well as those of all other results in this section, is deferred to Section B. Theorem 4.1 (Convergence to the Equal-Weighted Average). Let the worker-specific stepsizes be chosen according to (2) and 2.1 to 2.5 and 3.1 hold. If the stepsize parameter n suppose Assumptions o 1 1 satisfies 0 < γ ≤ min 6LτH , 5Lmax ρτmax , then, after M ≥ 1 cycles, the cycle iterates xm generated by Algorithm 2 satisfy M −1 i 1 X h ∆ γτA Lσ 2 2 E ∥∇F (xm )∥ ≤ c0 · + c1 · + c2 · γ 2 τA2 L2max σ 2 + ζ 2 M m=0 γτH M K
(3)
for some absolute constants c0 , c1 , c2 > 0. Note that the bound on the expected gradient norm is established for the cycle iterates xm only. Over the course of a cycle, the inner iterates are pulled toward the workers’ local objectives, whose stationary points do not generally coincide with those of the global objective. The presence of K, the total number of updates per cycle, in the second term reflects this cycle-level perspective. The last term in (3) reflects the effect of the cycle bias and scales with γ 2 . By scaling down the stepsize parameter γ, we can shrink this term and achieve ε-stationarity for arbitrarily small ε > 0. The following corollary provides a bound on the worst-case time complexity to reach such an ε-stationary point. Corollary 4.2 (Time Complexity for the Equal-Weighted Average). In the setting of Theorem 4.1, the worst-case wall-clock time complexity to find an ε-stationary point of F is ! p ∆Lσ 2 ∆Lmax σ 2 + ζ 2 τmax ∆Lmax ρ τmax ∆L O τA + τmax . (4) τA + τmax + nε2 ε1.5 τH ε τH ε 2
This time complexity is notable for several reasons. First, its first and last terms, ∆Lσ nε2 τA and ∆L τ , match the theoretical lower bound established by Tyurin and Richtárik [2023]. The middle ε max terms reflect penalties incurred from the combination of data heterogeneity and staleness. For small ε, the leading term dominates the complexity bound. This scales efficiently with the arithmetic mean of computation times, τA , rather than their maximum, τmax . The local function parameters (Lmax ) and data heterogeneity measures (ζ, ρ) affect only lower-order terms, showing that the slowdown caused by function heterogeneity and gradient staleness is not of first-order concern. The τmax/τH multiplier in these terms represents the average number of gradients delivered per worker during a cycle and captures the impact of staleness. Objective Inconsistency under Equal Stepsizes Our analysis allows us to precisely quantify the objective inconsistency for other stepsize choices (see Theorem B.9 for a more general result). In Vanilla ASGD, all workers use the same stepsize γi = γ/K , so that γi/τi ∝ 1/τi . Consequently, Algorithm 2 now targets a frequency-weighted average of the local functions, F̃ (x) :=
n X i=1
w̃i Fi (x),
τ −1 w̃i := Pn i −1 ∝ τi−1 . j=1 τj
(5)
Our next theorem makes this precise. Theorem 4.3 (Convergence to the Frequency-Weighted Average). Let the workers’ stepsizes be equal, γi = γ/K , and suppose Assumptions 2.1 to 2.5 and 3.1 hold for F̃ .2 If the stepsize parameter 2 That is, we replace F, ∆, L, ζ, ρ by F̃ , ∆, ˜ L̃, ζ̃, ρ̃ in Assumptions 2.3 and 2.5.
7
n o 1 m satisfies 0 < γ ≤ min 61L̃ , 5Lmax ρ̃ , then, after M ≥ 1 cycles, the cycle iterates x generated by Algorithm 2 satisfy M −1 ˜ 2 ∆ γ L̃σ 2 1 X m ≤ c0 · E ∇F̃ (x ) + c1 · + c2 · γ 2 L2max σ 2 + ζ̃ 2 M m=0 γM K for some absolute constants c0 , c1 , c2 > 0. As in Theorem 4.1, the bias term on the right-hand side scales with γ 2 , so that Algorithm 2 converges to an ε-stationary point of the frequency-weighted average F̃ if the stepsize is sufficiently small. To better understand the workings of Algorithm 2 with equal stepsizes, we must consider its wall-clock time complexity. Corollary 4.4 (Time Complexity for the Frequency-Weighted Average). In the setting of Theorem 4.3, the worst-case wall-clock time complexity to find an ε-stationary point of F̃ is q ˜ max σ 2 + ζ̃ 2 ˜ max ρ̃ ˜ L̃σ 2 ˜ L̃ ∆L ∆L ∆ ∆ O τH + τmax + τmax + τmax . (6) nε2 ε1.5 ε ε Note that the leading term of this complexity bound now scales with the harmonic mean τH . This is in contrast to Corollary 4.2, where the arithmetic mean τA appears instead. This captures the price of neutralizing objective inconsistency: In Vanilla ASGD, fast workers take unscaled steps, driving the global model forward at a higher rate governed by the harmonic mean of the workers’ computation speeds. With rescaled stepsizes, on the other hand, stepsizes of fast workers are shrunk (γi ∝ τi ) to prevent them from dominating the optimization trajectory. While this ensures convergence to the equal-weighted average, scaling down the updates from faster workers slows the global learning progress, shifting the time complexity bottleneck from the smaller harmonic mean τH to the larger arithmetic mean τA .
5
Experiments
We compare Rescaled ASGD against Malenia SGD and Ringleader ASGD, two state-of-the-art methods with optimal or near-optimal theoretical wall-clock time complexities, in a data-heterogeneous setup. We train a two-layer neural network on MNIST [LeCun et al., 2010]. To enforce maximal heterogeneity, we partition the data by label [Hsu et al., 2019] across n = 10 workers, such that each worker holds images from exactly one class. See Section E for setup and methodology details.
10
Cross-entropy loss
Cross-entropy loss
10
Rescaled ASGD Malenia SGD Ringleader ASGD
0
10−1
0
5000
10000
15000
20000
25000
30000
Rescaled ASGD Malenia SGD Ringleader ASGD
0
10−1
0
Simulated Wall-Clock Time
5000
10000
15000
20000
25000
30000
Simulated Wall-Clock Time
(a) Fixed Times.
(b) Fluctuating Times.
Figure 1: Solid lines denote the median loss across five random seeds, with shaded regions indicating the minimum and maximum. Malenia SGD and Ringleader ASGD are slowed down under fluctuating computation times as the gradient gathering phase takes longer. Rescaled ASGD shows virtually identical performance in both settings. We evaluate two computation time settings. First, fixed harmonic periods (Assumption 3.1) matching our theoretical setup. We assign τi ∈ {1, 2, 4, 8, 16} to two workers each, making the fastest workers 8
16 times faster than the slowest. Second, a fluctuating setting where computation times of worker i are sampled from an exponential distribution with mean τi . While outside our theoretical framework, the underlying mechanism of equalized long-run learning progress remains. Figure 1 shows the loss trajectory. Under fluctuating times, Malenia SGD and Ringleader ASGD degrade slightly due to stragglers slowing the gathering phase, whereas Rescaled ASGD remains unaffected. To illustrate algorithmic differences, Figure 2 plots the cumulative stepsizes over time. For a fair comparison, worker i takes a step of ατi/nτmax in Rescaled ASGD, ensuring the total stepsize per cycle, taking τmax units of time, is α. Ringleader ASGD workers take steps of α/n to match this sum, while Malenia SGD requires no rescaling as it takes one step per cycle. In the fixed setup (a), all methods share an average progress rate of α per τmax units of time, with Rescaled ASGD and Ringleader ASGD yielding smoother trajectories. Under fluctuating times (b), Malenia SGD and Ringleader ASGD are bottlenecked by the slowest worker due to their gathering phases. Although faster workers compute continuously, the server model stalls until the straggler finishes, making the expected gathering duration greater than τmax . Rescaled ASGD continuously applies updates, exploring the search space unimpeded. 16 Rescaled ASGD Malenia SGD Ringleader ASGD
14 12
Cumulative Step Taken (×α)
Cumulative Step Taken (×α)
16
10 8 6 4 2 0
0
2
4
6
8
10
12
12 10 8 6 4 2 0
14
Wall-Clock Time (×τmax )
Rescaled ASGD Malenia SGD Ringleader ASGD
14
0
2
4
6
8
10
12
14
Wall-Clock Time (×τmax )
(a) Fixed.
(b) Fluctuating.
Figure 2: Cumulative stepsize taken over wall-clock time. The gathering phases dilate under fluctuating computation times, causing Malenia SGD and Ringleader ASGD to stall in the search space. Rescaled ASGD takes steps at an unimpeded rate of α/τmax .
6
Conclusion
In this work, we introduced Rescaled ASGD to address the problem of objective inconsistency in asynchronous optimization caused by the interplay of data and system heterogeneity. By proportionally rescaling worker-specific stepsizes, we proved that Asynchronous SGD converges to the true global objective without introducing memory overhead, gathering phases, or worker idle times. Furthermore, our theoretical wall-clock time complexity matches known lower bounds in the leading term for the fixed-computation model considered here. A primary limitation of our theoretical analysis is the reliance on the fixed-computation model with harmonic periods (Assumption 3.1). While our empirical results demonstrate that the method remains robust under stochastically fluctuating computation times, formally relaxing this structural assumption remains an open problem, likely requiring a different analytical approach. Future work may also explore extending this stepsize-rescaling principle to time-varying system heterogeneity. Acknowledgments The research reported in this publication was supported by funding from King Abdullah University of Science and Technology (KAUST): i) KAUST Baseline Research Scheme, ii) CRG Grant ORFSCRG12-2024-6460, and iii) Center of Excellence for Generative AI, under award number 5940. 9
References Alekh Agarwal and John C Duchi. Distributed delayed stochastic optimization. Advances in Neural Information Processing Systems, 24, 2011. (Cited on pages 6 and 14.) Yossi Arjevani, Ohad Shamir, and Nathan Srebro. A tight convergence analysis for stochastic gradient descent with delayed updates. In Algorithmic Learning Theory, pages 111–132. PMLR, 2020. (Cited on page 14.) Gerard M Baudet. Asynchronous iterative methods for multiprocessors. Journal of the ACM (JACM), 25(2): 226–244, 1978. (Cited on page 14.) Dimitri Bertsekas and John Tsitsiklis. Parallel and distributed computation: numerical methods. Athena Scientific, 2015. (Cited on page 14.) Dimitri P. Bertsekas. Incremental proximal methods for large scale convex optimization. Mathematical Programming, 129(2):163–195, October 2011. ISSN 1436-4646. doi: 10.1007/s10107-011-0472-0. URL https://doi.org/10.1007/s10107-011-0472-0. (Cited on page 5.) Léon Bottou, Frank E. Curtis, and Jorge Nocedal. Optimization Methods for Large-Scale Machine Learning. SIAM Review, 60(2):223–311, 2018. doi: 10.1137/16M1080173. URL https://doi.org/10.1137/ 16M1080173. _eprint: https://doi.org/10.1137/16M1080173. (Cited on page 5.) Sébastien Bubeck. Convex Optimization: Algorithms and Complexity, November 2015. URL http://arxiv. org/abs/1405.4980. arXiv:1405.4980. (Cited on page 5.) Sebastian Caldas, Sai Meher Karthik Duddu, Peter Wu, Tian Li, Jakub Konečný, H. Brendan McMahan, Virginia Smith, and Ameet Talwalkar. LEAF: A Benchmark for Federated Settings, December 2019. URL http://arxiv.org/abs/1812.01097. arXiv:1812.01097. (Cited on page 2.) Daniel Chazan and Willard Miranker. Chaotic relaxation. Linear algebra and its applications, 2(2):199–222, 1969. (Cited on page 14.) Andrew Cotter, Ohad Shamir, Nati Srebro, and Karthik Sridharan. Better mini-batch algorithms via accelerated gradient methods. Advances in Neural Information Processing Systems, 24, 2011. (Cited on page 2.) Ofer Dekel, Ran Gilad-Bachrach, Ohad Shamir, and Lin Xiao. Optimal Distributed Online Prediction using Mini-Batches, January 2012. URL http://arxiv.org/abs/1012.1367. arXiv:1012.1367. (Cited on page 2.) Sanghamitra Dutta, Gauri Joshi, Soumyadip Ghosh, Parijat Dube, and Priya Nagpurkar. Slow and stale gradients can win the race: Error-runtime trade-offs in distributed SGD. In International Conference on Artificial Intelligence and Statistics, pages 803–812. PMLR, 2018. (Cited on page 2.) Hamid Reza Feyzmahdavian and Mikael Johansson. Asynchronous iterations in optimization: New sequence results and sharper algorithmic guarantees. Journal of Machine Learning Research, 24(158):1–75, 2023. (Cited on page 14.) Saeed Ghadimi and Guanghui Lan. Stochastic first-and zeroth-order methods for nonconvex stochastic programming. SIAM Journal on Optimization, 23(4):2341–2368, 2013. (Cited on page 5.) Eduard Gorbunov, Filip Hanzely, and Peter Richtárik. Local SGD: Unified Theory and New Efficient Methods, November 2020. URL http://arxiv.org/abs/2011.02828. arXiv:2011.02828. (Cited on page 14.) C-C Han and H-Y Tyan. A better polynomial-time schedulability test for real-time fixed-priority scheduling algorithms. In Proceedings Real-Time Systems Symposium, pages 36–45. IEEE, 1997. (Cited on page 5.) Tzu-Ming Harry Hsu, Hang Qi, and Matthew Brown. Measuring the Effects of Non-Identical Data Distribution for Federated Visual Classification, September 2019. URL http://arxiv.org/abs/1909.06335. arXiv:1909.06335. (Cited on pages 8 and 30.) Rustem Islamov, Mher Safaryan, and Dan Alistarh. AsGrad: A sharp unified analysis of asynchronous-SGD algorithms. In International Conference on Artificial Intelligence and Statistics, pages 649–657. PMLR, 2024. (Cited on page 14.) Sashank J Reddi, Ahmed Hefny, Suvrit Sra, Barnabas Poczos, and Alexander J Smola. On variance reduction in stochastic gradient descent and its asynchronous variants. Advances in Neural Information Processing Systems, 28, 2015. (Cited on page 14.)
10
Peter Kairouz, H Brendan McMahan, Brendan Avent, Aurélien Bellet, Mehdi Bennis, Arjun Nitin Bhagoji, Kallista Bonawitz, Zachary Charles, Graham Cormode, Rachel Cummings, and others. Advances and open problems in federated learning. Foundations and Trends® in Machine Learning, 14(1–2):1–210, 2021. (Cited on page 2.) Sai Praneeth Karimireddy, Satyen Kale, Mehryar Mohri, Sashank J. Reddi, Sebastian U. Stich, and Ananda Theertha Suresh. SCAFFOLD: Stochastic Controlled Averaging for Federated Learning, April 2021. URL http://arxiv.org/abs/1910.06378. arXiv:1910.06378. (Cited on page 14.) Anastasiia Koloskova, Sebastian U Stich, and Martin Jaggi. Sharper convergence guarantees for asynchronous SGD for distributed and federated learning. Advances in Neural Information Processing Systems, 35: 17202–17215, 2022. (Cited on pages 2, 3, 5, 14, 15, 28, and 29.) Remi Leblond, Fabian Pedregosa, and Simon Lacoste-Julien. Improved Asynchronous Parallel Optimization Analysis for Stochastic Incremental Methods. Journal of Machine Learning Research, 19(81):1–68, 2018. URL http://jmlr.org/papers/v19/17-650.html. (Cited on page 14.) Rémi Leblond, Fabian Pedregosa, and Simon Lacoste-Julien. ASAGA: Asynchronous Parallel SAGA. In Aarti Singh and Jerry Zhu, editors, Proceedings of the 20th International Conference on Artificial Intelligence and Statistics, volume 54 of Proceedings of Machine Learning Research, pages 46–54. PMLR, April 2017. URL https://proceedings.mlr.press/v54/leblond17a.html. (Cited on page 14.) Yann LeCun, Corinna Cortes, and CJ Burges. MNIST handwritten digit database. ATT Labs [Online]. Available: http://yann.lecun.com/exdb/mnist, 2, 2010. (Cited on pages 8 and 30.) Tian Li, Anit Kumar Sahu, Manzil Zaheer, Maziar Sanjabi, Ameet Talwalkar, and Virginia Smith. Federated optimization in heterogeneous networks. Proceedings of Machine Learning and Systems, 2:429–450, 2020. (Cited on pages 2 and 14.) Xiangru Lian, Yijun Huang, Yuncheng Li, and Ji Liu. Asynchronous parallel stochastic gradient for nonconvex optimization. Advances in Neural Information Processing Systems, 28, 2015. (Cited on pages 6 and 14.) LO Mangasarian. Parallel gradient distribution in unconstrained optimization. SIAM Journal on Control and Optimization, 33(6):1916–1925, 1995. (Cited on page 14.) Horia Mania, Xinghao Pan, Dimitris Papailiopoulos, Benjamin Recht, Kannan Ramchandran, and Michael I Jordan. Perturbed iterate analysis for asynchronous stochastic optimization. SIAM Journal on Optimization, 27(4):2202–2229, 2017. (Cited on page 14.) Artavazd Maranjyan. First Provably Optimal Asynchronous SGD for Homogeneous and Heterogeneous Data. PhD Thesis, KAUST Research Repository, 2025. (Cited on page 15.) Artavazd Maranjyan and Peter Richtárik. Ringleader ASGD: The First Asynchronous SGD with Optimal Time Complexity under Data Heterogeneity. In The Fourteenth International Conference on Learning Representations, 2026. URL https://openreview.net/forum?id=5wqTal0EuC. (Cited on pages 2, 3, and 15.) Artavazd Maranjyan, Omar Shaikh Omar, and Peter Richtárik. MindFlayer SGD: Efficient Parallel SGD in the Presence of Heterogeneous and Random Worker Compute Times. In The 41st Conference on Uncertainty in Artificial Intelligence, 2025a. (Cited on page 2.) Artavazd Maranjyan, El Mehdi Saad, Peter Richtárik, and Francesco Orabona. ATA: Adaptive Task Allocation for Efficient Resource Management in Distributed Machine Learning. In International Conference on Machine Learning, 2025b. (Cited on page 2.) Artavazd Maranjyan, Mher Safaryan, and Peter Richtárik. GradSkip: Communication-Accelerated Local Gradient Methods with Better Computational Complexity. Transactions on Machine Learning Research, 2025c. ISSN 2835-8856. URL https://openreview.net/forum?id=6R3fRqFfhn. (Cited on page 2.) Artavazd Maranjyan, Alexander Tyurin, and Peter Richtárik. Ringmaster ASGD: The First Asynchronous SGD with Optimal Time Complexity. In Proceedings of the 42nd International Conference on Machine Learning, pages 43120–43139. PMLR, October 2025d. URL https://proceedings.mlr.press/v267/ maranjyan25b.html. (Cited on pages 2 and 15.) Brendan McMahan, Eider Moore, Daniel Ramage, Seth Hampson, and Blaise Aguera y Arcas. Communicationefficient learning of deep networks from decentralized data. In Artificial Intelligence and Statistics, pages 1273–1282. PMLR, 2017. (Cited on page 2.)
11
Konstantin Mishchenko, Francis Bach, Mathieu Even, and Blake E Woodworth. Asynchronous SGD beats minibatch SGD under arbitrary delays. Advances in Neural Information Processing Systems, 35:420–433, 2022a. (Cited on pages 2, 4, 5, 14, 28, and 31.) Konstantin Mishchenko, Grigory Malinovsky, Sebastian Stich, and Peter Richtárik. ProxSkip: Yes! Local Gradient Steps Provably Lead to Communication Acceleration! Finally! In Kamalika Chaudhuri, Stefanie Jegelka, Le Song, Csaba Szepesvari, Gang Niu, and Sivan Sabato, editors, Proceedings of the 39th International Conference on Machine Learning, volume 162 of Proceedings of Machine Learning Research, pages 15750–15769. PMLR, July 2022b. URL https://proceedings.mlr.press/v162/mishchenko22b.html. (Cited on page 14.) Yurii Nesterov. Lectures on Convex Optimization, volume 137. Springer, 2018. (Cited on page 23.) John Nguyen, Kshitiz Malik, Hongyuan Zhan, Ashkan Yousefpour, Mike Rabbat, Mani Malek, and Dzmitry Huba. Federated learning with buffered asynchronous aggregation. In International Conference on Artificial Intelligence and Statistics, pages 3581–3607. PMLR, 2022. (Cited on page 2.) Benjamin Recht, Christopher Re, Stephen Wright, and Feng Niu. Hogwild!: A Lock-Free Approach to Parallelizing Stochastic Gradient Descent. In Advances in Neural Information Processing Systems, volume 24. Curran Associates, Inc., 2011. URL https://papers.nips.cc/paper_files/paper/2011/ hash/218a0aefd1d1a4be65601cc6ddc1520e-Abstract.html. (Cited on pages 5, 6, and 14.) Sebastian U. Stich. Local SGD Converges Fast and Communicates Little, May 2019. URL http://arxiv. org/abs/1805.09767. arXiv:1805.09767. (Cited on page 14.) Marion Sudvarg, Ao Li, Daisy Wang, Sanjoy Baruah, Jeremy Buhler, Chris Gill, Ning Zhang, and Pontus Ekberg. Elastic scheduling for harmonic task systems. In 2024 IEEE 30th Real-Time and Embedded Technology and Applications Symposium (RTAS), pages 334–347, 2024. doi: 10.1109/RTAS61025.2024.00034. (Cited on page 5.) Zhirayr Tovmasyan, Artavazd Maranjyan, and Peter Richtárik. Rennala MVR: Improved time complexity for parallel stochastic optimization via momentum-based variance reduction, 2026. URL https://arxiv.org/ abs/2605.08871. (Cited on page 15.) Alexander Tyurin. Tight time complexities in parallel stochastic optimization with arbitrary computation dynamics. arXiv preprint arXiv:2408.04929, 2024. (Cited on page 15.) Alexander Tyurin and Peter Richtárik. Optimal Time Complexities of Parallel Stochastic Optimization Methods Under a Fixed Computation Model. In Advances in Neural Information Processing Systems, volume 36, pages 16515–16577, 2023. (Cited on pages 2, 3, 5, 7, and 15.) Jianyu Wang, Qinghua Liu, Hao Liang, Gauri Joshi, and H. Vincent Poor. Tackling the objective inconsistency problem in heterogeneous federated optimization. In Proceedings of the 34th International Conference on Neural Information Processing Systems, NIPS ’20, pages 7611–7623, Red Hook, NY, USA, December 2020. Curran Associates Inc. ISBN 978-1-7138-2954-6. URL https://dl.acm.org/doi/10.5555/3495724. 3496362. (Cited on page 14.) Jianyu Wang, Zachary Charles, Zheng Xu, Gauri Joshi, H Brendan McMahan, Maruan Al-Shedivat, Galen Andrew, Salman Avestimehr, Katharine Daly, Deepesh Data, and others. A field guide to federated optimization. arXiv preprint arXiv:2107.06917, 2021. (Cited on page 2.) Blake Woodworth, Kumar Kshitij Patel, Sebastian Stich, Zhen Dai, Brian Bullins, Brendan Mcmahan, Ohad Shamir, and Nathan Srebro. Is local SGD better than minibatch SGD? In International Conference on Machine Learning, pages 10334–10343. PMLR, 2020. (Cited on page 14.) Cong Xie, Sanmi Koyejo, and Indranil Gupta. arXiv:1903.03934, 2019. (Cited on page 14.)
Asynchronous federated optimization.
arXiv preprint
Mikhail Yurochkin, Mayank Agarwal, Soumya Ghosh, Kristjan Greenewald, Nghia Hoang, and Yasaman Khazaeni. Bayesian Nonparametric Federated Learning of Neural Networks. In Proceedings of the 36th International Conference on Machine Learning, pages 7252–7261. PMLR, May 2019. URL https:// proceedings.mlr.press/v97/yurochkin19a.html. (Cited on page 30.)
12
Contents 1
Introduction
1
2
Problem Setup
4
3 Asynchronous SGD with Cyclic Update Schedule and Rescaled ASGD
5
4
Convergence Guarantees
7
5
Experiments
8
6
Conclusion
9
A Related Work
14
B Proofs
15
B.1 Preliminaries . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
15
B.2 Auxiliary Results . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
17
B.3 Proof of Theorem 4.1 and Corollary 4.2 . . . . . . . . . . . . . . . . . . . . . . .
25
B.4 Proof of Theorem 4.3 and Corollary 4.4 . . . . . . . . . . . . . . . . . . . . . . .
26
C On the Necessity of Bounding Data Heterogeneity
28
D Wall-Clock Time Complexities
29
D.1 Concurrent ASGD . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
29
D.2 Delay-Adaptive ASGD . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
29
E Experimental Details
30
F Further Experiments
31
F.1
Vanilla ASGD Targets the Frequency-Weighted Average . . . . . . . . . . . . . .
31
F.2
Delay-Adaptive SGD Exhibits Objective Inconsistency . . . . . . . . . . . . . . .
31
13
A
Related Work
Asynchronous Optimization Asynchronous methods have been a cornerstone of parallel computing for over half a century [Chazan and Miranker, 1969, Baudet, 1978, Bertsekas and Tsitsiklis, 2015]. Early research focused primarily on the homogeneous data setting, where the central challenge is managing stale gradients to ensure they do not derail the learning trajectory. In these settings, the assumption of bounded delays has proven a fruitful abstraction for analysis [Agarwal and Duchi, 2011, Recht et al., 2011], a property that is naturally satisfied within the structured computation model we employ. This line of work includes asynchronous nonconvex SGD, perturbation-based analyses, tight delay-dependent bounds, and asynchronous variance-reduced methods [Lian et al., 2015, Mania et al., 2017, Arjevani et al., 2020, J Reddi et al., 2015, Leblond et al., 2017, 2018, Feyzmahdavian and Johansson, 2023, Islamov et al., 2024]. The more complex interplay between asynchronous updates and heterogeneous data has received comparatively less attention. Mishchenko et al. [2022a] examine a framework with arbitrary, unbounded delays. While their results for homogeneous data yield convergence rates to ε-stationary points comparable to ours, the heterogeneous case presents a fundamental challenge. Specifically, if worker computation times can vary without structure, the optimization trajectory becomes biased toward faster workers. Consequently, their guarantees in this setting necessitate a strong similarity assumption between local objectives and do not allow for convergence to an arbitrary ε-stationary point. By contrast, the structured nature of our computation model allows us to neutralize this bias through stepsize rescaling, enabling stronger guarantees without requiring local objectives to be nearly identical. Koloskova et al. [2022] also explore asynchronous SGD under heterogeneous data but utilize a distinct mechanism that deviates from the standard Asynchronous SGD formulation (Algorithm 1). In their model, the server assigns gradient computation tasks to workers via uniform random sampling. This process allows for multiple gradient tasks to be queued for a single worker. While this ensures that all workers contribute an equal number of gradients in the long run–matching the core motivation of our approach–it does so without utilizing faster workers’ speed to increase throughput. In their model, faster workers are characterized by lower average delays rather than higher update frequencies. While this enables convergence to arbitrary ε-stationarity, the queuing mechanism does not reflect the common physical reality of asynchronous systems where workers deliver results as soon as they are computed and the server applies them as soon as they arrive. Our approach, Rescaled ASGD, allows workers to operate at their maximum physical frequency while using rescaled stepsizes to ensure that the resulting trajectory targets the correct global objective. Objective Inconsistency The phenomenon wherein distributed methods fail to minimize the global objective (1) when workers have different computation times is a well-documented challenge in distributed optimization [Wang et al., 2020, Islamov et al., 2024]. This inconsistency arises because the stochastic process underlying the optimization becomes implicitly weighted by the relative frequencies of worker updates, effectively targeting a surrogate objective F̃ rather than the equalweighted average F . Wang et al. [2020] address this challenge within the context of local SGD, a framework where workers perform multiple local updates before periodically synchronizing with a central server [Mangasarian, 1995, Stich, 2019, Gorbunov et al., 2020, Woodworth et al., 2020]. In their analysis, the number of local steps taken by a worker is intrinsically linked to its computation speed. To neutralize the resulting bias, they propose a mechanism conceptually similar to ours: rescaling the aggregate model difference from each worker by the reciprocal of the number of local steps performed. This normalization ensures that fast workers do not pull the global model disproportionately toward their local optima, thereby enabling the algorithm to reach arbitrary ε-stationarity for the true objective. Other strategies designed to mitigate bias from heterogeneous worker participation include the use of proximal regularization to constrain local drift [Li et al., 2020], control variates to correct for client-server residuals [Karimireddy et al., 2021, Mishchenko et al., 2022b], and staleness-aware mixing coefficients that de-weight delayed updates [Xie et al., 2019]. While effective at improving stability, these methods typically focus on reducing the variance or the impact of stale information rather than enforcing the structural update parity required to eliminate objective inconsistency. In the absence of explicit rescaling or contribution equalization, the underlying optimization target remains skewed toward more active participants. 14
Our work extends the rescaling principle to the asynchronous, single-update regime, providing a memory-efficient solution. Optimal Methods in the Fixed-Computation Model The fixed-computation model offers a benchmark for evaluating the efficiency of parallel optimization algorithms. Tyurin and Richtárik [2023] derive a lower bound for the wall-clock time complexity of first-order stochastic gradient methods when worker computation times τi are constant yet heterogeneous. While they propose synchronous methods that achieve this bound in both homogeneous and heterogeneous settings, their approach inherently sacrifices the continuous throughput of asynchronous systems to maintain synchronization. Notably, in the homogeneous data setting, Vanilla ASGD can match this lower bound in regimes with high stochastic gradient variance σ 2 relative to the target precision ε. Subsequent work extended the corresponding lower-bound framework beyond fixed worker computation times to arbitrary computation dynamics [Tyurin, 2024] and to variance-reduced parallel stochastic methods [Tovmasyan et al., 2026]. Complementary to these lower-bound results, Maranjyan [2025] developed asynchronous methods that match the lower bounds on time complexities. Maranjyan et al. [2025d] improve upon these results for the homogeneous setting by introducing a variant of Asynchronous SGD that employs delay-adaptive learning rates and a hard-thresholding rule. By discarding gradients with staleness exceeding a specific threshold, they demonstrate that staleness bias can be actively managed to achieve optimal convergence regardless of the gradient variance. However, such a thresholding mechanism is fundamentally incompatible with the heterogeneous data setting, where it could systematically exclude updates from the slowest workers, thereby introducing the very objective inconsistency our work seeks to resolve. Maranjyan and Richtárik [2026] propose an asynchronous algorithm, Ringleader ASGD, designed specifically for the heterogeneous regime. While their method approaches the optimal time complexity, its leading term is still constrained by a constant related to local objective similarity, similar to the limitations observed in the work of Koloskova et al. [2022] (cf. Table 1). In contrast, Rescaled ASGD targets the equal-weighted average through stepsize rescaling rather than delay manipulation, gradient tables, or worker selection. This allows us to exploit the higher update frequency of faster workers, a benefit also noted in the empirical performance of Maranjyan and Richtárik [2026]. By ensuring that every worker’s contribution is appropriately weighted in the model space, we maintain the exploration advantages of frequent asynchronous updates.
B
Proofs
B.1
Preliminaries
m m Notation In our model, randomness enters through the stochastic gradients ∇fik (yk,i , ξk,i ) alone. k k m m For ease of notation, we drop the dependence on the samples ξk,ik ∼ Dik and write ∇fik (yk,i ) k m instead. We denote by F the sigma field generated by the stochastic gradients delivered up to the beginning of cycle m. Likewise, Fkm denotes the sigma field generated by the stochastic gradients m := delivered up to the beginning of iteration k in cycle m, and define FK F m+1 . The iterates xm k m m m m are Fk -measurable, as are yk , while the stochastic gradients ∇fik (yk,ik ) are Fk+1 -measurable as these are resolved only at the end of iteration k in cycle m.
E [x̄] denotes the unconditional expectation of a random vector x̄, Em [x̄] the conditional expectation conditional on all information available at the beginning of cycle m, i.e., conditioned on the sigma field F m , Em,k [x̄] the conditional expectation conditional on all information available at the beginning of iteration k within cycle m, i.e., conditioned on the sigma field Fkm . Similarly, 2 2 2 Vm (x̄) := Em [∥x̄ − Em [x̄]∥ ] = Em [∥x̄∥ ] − ∥Em [x̄]∥ denotes the conditional variance of a random vector x̄ conditioned on all information available at the beginning of cycle m, and Cm (x̄, ȳ) := Em [⟨x̄ − Em [x̄] , ȳ − Em [ȳ]⟩] = Em [⟨x̄, ȳ⟩] − ⟨Em [x̄] , Em [ȳ]⟩ denotes the conditional covariance of two random vectors x̄, ȳ.
We define Lmax = max Li , i=1,...,n
γmax := max γi , i=1,...,n
15
τmax := max τi . i=1,...,n
Moreover,
n
τA :=
1X τi n i=1
τH := Pn
and
n
−1 i=1 τi
denote the arithmetic and harmonic mean of computation times, respectively. Frequently used Assumptions, Definitions, and Standard Results Below, we state a series of standard results, definitions, and the assumptions laid out in Section 2 in simplified form for reference. These will be used extensively in the following proofs. Unbiased Gradients: Eξ∼Di [∇fi (x, ξ)] = ∇Fi (x)
(UG)
h i 2 Eξ∼Di ∥∇fi (x, ξ) − ∇Fi (x)∥ ≤ σ 2
(BV)
∥∇Fi (x) − ∇Fi (y)∥ ≤ Li ∥x − y∥ , ∥∇F (x) − ∇F (y)∥ ≤ L ∥x − y∥
(LS)
Bounded Gradient Variance:
Smoothness:
Lower Bound:
F (x0 ) − F (xM ) ≥ F (x0 ) − F ∗ = ∆
(LB)
Bounded Data Heterogeneity: 2
2
∥∇Fi (x)∥ ≤ ζ 2 + ρ2 ∥∇F (x)∥
(BH)
Cycle Stepsize: α :=
K−1 X
γ ik
(α)
γi2k
(A)
k=0
Sum of Squared Stepsizes: A :=
K−1 X k=0
Cycle Noise: νm :=
K−1 X k=0
m m γik ∇fik (yk,i ) − ∇Fik (yk,i ) k k
(N)
m γik ∇Fik (yk,i ) − ∇Fik (xm ) k
(B)
Cycle Bias: bm :=
K−1 X k=0
Triangle Inequality: n X i=1
vi ≤
n X i=1
∥vi ∥
(T)
Cauchy-Schwarz Inequality: d X i=1
!2 ui vi
≤
d X i=1
! u2i
d X
! vi2
(CS)
i=1
Tower Property: Em [Em,k [x̄]] = Em [x̄] 16
(TP)
Descent Lemma: F (y) ≤ F (x) + ⟨∇F (x), y − x⟩ +
L 2 ∥y − x∥ 2
(DL)
Jensen’s Inequality: h i 2 2 ∥Em [x̄]∥ ≤ Em ∥x̄∥
(JN)
Young’s Inequality: ∥u∥ ∥v∥ ≤
1 s 2 2 ∥u∥ + ∥v∥ , 2 2s
∀s > 0
(YN)
Squared Sum Inequality: n X
2
vi
i=1
Proof Outline
≤n
n X i=1
2
∥vi ∥ .
(SS)
Theorem B.9 states our main result in a more general form, showing that
Rescaled ASGD can target any convex combination (7) of the local objective functions, and es-
tablishes a bound on the expected average squared gradient norm of the cycle iterates. To that end, Lemma B.8 establishes a bound in terms of the squared norm of the cycle bias utilizing a standard descent lemma. A bound on the norm of the expected cumulative cycle bias is derived in Lemmas B.2 to B.6. Similarly, Lemma B.1 establishes a bound on the squared norm of the cycle noise. Theorems 4.1 and 4.3 restate Theorem B.9 with the appropriately chosen stepsizes. Corollaries 4.2 and 4.4 then follow immediately. B.2
Auxiliary Results
Throughout this subsection, we denote by F (x) :=
n X
wi Fi (x)
(7)
i=1
an arbitrary convex combination of the local objective functions Fi . The weights wi ≥ 0 sum to unity and are understood to be fixed constants. Assumptions 2.3 and 2.5 are understood to refer to the objective function defined in (7). Note that we recover the equal-weighted average (1) by setting wi = 1/n. The more general case presented in this subsection allows us to derive the results from the main text (Theorems 4.1 and 4.3) as special cases of the more general Theorem B.9. Lemma B.1 (Cycle Noise). Under assumptions 2.1 and 2.2, the cycle noise νm :=
K−1 X k=0
m m γik ∇fik (yk,i ) − ∇Fik (yk,i ) k k
satisfies Em [νm ] = 0,
h i 2 Em ∥νm ∥ ≤ Aσ 2 .
m Proof. Recall that yk,i denotes the model held by worker ik at the beginning of iteration k of cycle k m. Thus, "K−1 # X (N) m m Em [νm ] = Em γik ∇fik (yk,ik ) − ∇Fik (yk,ik ) k=0
(TP)
K−1 X
(UG)
K−1 X
=
k=0
=
m m γik Em Em,k ∇fik (yk,i ) − ∇Fik (yk,i ) k k γik Em [0]
k=0
17
= 0. Further, h i 2 2 Em ∥νm ∥ = Vm (νm ) + ∥Em [νm ]∥ = Vm
K−1 X
! γ ik
k=0
=
K−1 X K−1 X k=0 l=0
m m ∇fik (yk,i ) − ∇Fik (yk,i ) k k
m m m m γik γil Cm ∇fik (yk,i ) − ∇Fik (yk,i ), ∇fil (yl,i ) − ∇Fil (yl,i ) . k k l l
m m Let em k := ∇fik (yk,ik ) − ∇Fik (yk,ik ). Unbiasedness of the stochastic gradients (Assumption 2.1) m m ensures Em [ek ] = Em [Em,k [ek ]] = 0. Now consider a cross-term with k < l: m m m m m Cm (em k , el ) = Em [⟨ek , el ⟩] − ⟨Em [ek ] , Em [el ]⟩ (TP)
m = Em [Em,l [⟨em k , el ⟩]] m = Em [⟨ek , Em,l [em l ]⟩] = 0,
m where we used that em k is Fl -measurable, i.e., its randomness has been resolved by the beginning of iteration l > k.
With Assumption 2.2, we now find h i 2 Em ∥νm ∥ = Vm (νm ) =
K−1 X
m γi2k Cm (em k , ek )
k=0
=
K−1 X k=0
K−1 (BV) X
≤
h i 2 γi2k Em ∥em k ∥ γi2k σ 2
k=0
= Aσ 2 .
Lemma B.2. Under Assumption 2.4, the cycle bias bm :=
K−1 X k=0
m ) − ∇Fik (xm ) γik ∇Fik (yk,i k
satisfies 2
∥bm ∥ ≤ AL2max
K−1 X k=0
2
m yk,i − xm k
.
(8)
Proof. Using our definitions of the cycle bias bm and cycle stepsize A, we find 2 (B)
∥bm ∥ =
(T)
≤
K−1 X
2
γ ik
k=0 K−1 X k=0
m ∇Fik (yk,i ) − ∇Fik (xm ) k
!2
γ ik
m ∇Fik (yk,i ) − ∇Fik (xm ) k
18
(CS)
K−1 X
(A)
K−1 X
(LS)
K−1 X
≤
! γi2k
k=0
= A
k=0
2 m ∇Fik (yk,i ) − ∇Fik (xm ) k
k=0
m ∇Fik (yk,i ) − ∇Fik (xm ) k
k=0
≤ A
!
K−1 X
2
m − xm L2ik yk,i k
≤ AL2max
K−1 X k=0
2
m yk,i − xm k
2
,
where we used the smoothness of the local functions (Assumption 2.4) in the penultimate line. Lemma B.3. Let k ′ ∈ {0, . . . , K − 1} , m′ ∈ {m − 1, m} denote the iteration and cycle index corresponding to the last model seen by worker ik at the beginning of iteration k in cycle m, i.e., ′ m m m yk,i = xm k′ . If we denote the set of iterations between yk,ik and x by k ′ {k , . . . , K − 1} × {m − 1} , m′ = m − 1 m := Ik , {0, . . . , k ′ − 1} × {m} , m′ = m where {0, . . . , k ′ − 1} is understood as the empty set when k ′ = 0, then X 2 m m 2 c γi2l ∇fil (yl,i ) . yk,i − x ≤ K l k (l,c)∈Ikm
Proof. Note that |Ikm | ≤ K as the cyclic update schedule ensures that each worker receives a model ′ m update after at most K iterations. The vector xm k′ − x is, up to sign, the sum of the updates indexed m ′ by Ik ; the sign depends on whether m = m − 1 or m′ = m and disappears after taking the norm. The claim then follows from the squared-sum inequality: m yk,i − xm k
2
=
′
m xm k′ − x
2
2
X
=
(l,c)∈Ikm (SS)
≤ |Ikm | ≤ K
c γil ∇fil (yl,i ) l
X (l,c)∈Ikm
X (l,c)∈Ikm
2
c ) γi2l ∇fil (yl,i l
c ) γi2l ∇fil (yl,i l
2
.
Lemma B.4. Under assumptions 2.1 to 2.3 and 2.5, h i h i h i 2 2 2 m 2 m 2 2 2 m 2 m E ∇fik (yk,i ) ≤ σ + ζ + 2ρ E ∥∇F (x )∥ + 2ρ L E y − x . k,i k k
(9)
Proof. By assumptions 2.1 and 2.2, i h i h i h 2 (UG) 2 2 m m m m ) = E ) − ∇F (y ) + E ∇F (y ) Em,k ∇fik (yk,i ∇f (y m,k i i m,k i k,i k,i k,i k k k k k k k (BV)
m ≤ σ 2 + ∇Fik (yk,i ) k
2
.
For the latter term, using assumptions 2.3 and 2.5, we derive the bound m ∇Fik (yk,i ) k
2 (BH)
m ≤ ζ 2 + ρ2 ∇F (yk,i ) k
19
2
(SS)
2
m ) − ∇F (xm ) ≤ ζ 2 + 2ρ2 ∥∇F (xm )∥ + 2ρ2 ∇F (yk,i k
(LS)
2
m ≤ ζ 2 + 2ρ2 ∥∇F (xm )∥ + 2ρ2 L2 yk,i − xm k
2
2
.
Taking the unconditional expectation then gives h i i h i h 2 m 2 2 2 m 2 2 2 m m 2 E ∇fik (yk,i . ) ≤ σ + ζ + 2ρ E ∥∇F (x )∥ + 2ρ L E y − x k,ik k
Lemma B.5. Let QM :=
M −1 K−1 X X
γi2k E
h
m ∇fik (yk,i ) k
m=0 k=0
2
i
.
1 , we have For γmax ≤ 2KLρ 2
QM ≤ 2AM σ + ζ
2
+ 4Aρ
2
M −1 X m=0
h i 2 E ∥∇F (xm )∥ .
Proof. Multiplying the bound derived in Lemma B.3 by γi2k and summing over all cycles m = 0, . . . , M − 1 and iterations k = 0, . . . , K − 1, we find M −1 K−1 X X
M −1 K−1 X X X m m 2 yk,ik − x ≤K γi2k γi2l m m=0 k=0 (l,c)∈Ik
γi2k
m=0 k=0
M −1 K−1 X X
=K
m′ =0 k′ =0
2
c ∇fil (yl,i ) l
′
γi2k′ ∇fik′ (ykm′ ,ik′ )
2
×
X
γi2l .
(l,c):(k′ ,m′ )∈Ilc
m In the first line, we sum over the paths Ikm connecting yk,i and xm , which we bound by considering k c the norms of the gradients computed along these paths, ∇fil (yl,i ). To get to the second line, we l ′ ′ m′ instead consider the gradient in an iteration (k , m ), ∇fik′ (yk′ ,ik′ ), and sum up the contribution of the paths on which it appears. Since the update schedule is cyclic and the stepsizes depend only on the worker index, the sum of squared stepsizes over any block of K consecutive updates is A. For a fixed update (k ′ , m′ ), all destination updates (l, c) whose paths contain (k ′ , m′ ) lie in such a block of at most one full cycle. We can therefore bound this factor uniformly by
X (l,c):(k′ ,m′ )∈Ilc
γi2l ≤
K−1 X
γi2k = A .
k=0
Hence, M −1 K−1 X X
γi2k E
h
2 m yk,i − xm k
m=0 k=0
i
≤ AK
M −1 K−1 X X
γi2k E
h
m=0 k=0
m ∇fik (yk,i ) k
2
i
= AKQM . (10)
Now, using Lemma B.4, we have QM =
M −1 K−1 X X
γi2k E
h
m=0 k=0 (9)
≤
M −1 K−1 X X m=0 k=0 2
m ∇fik (yk,i ) k
2
i
i h i h 2 m 2 m γi2k σ 2 + ζ 2 + 2ρ2 E ∥∇F (xm )∥ + 2L2 ρ2 E yk,i − x k
= AM σ + ζ
2
+ 2Aρ
2
M −1 X m=0
h
m
E ∥∇F (x )∥ 20
2
i
2 2
+ 2L ρ
M −1 K−1 X X m=0 k=0
γi2k E
h
m yk,i − xm k
2
i
M −1 h i X 2 ≤ AM σ 2 + ζ 2 + 2Aρ2 E ∥∇F (xm )∥ + 2AKL2 ρ2 QM
(10)
m=0
≤ AM σ 2 + ζ 2 + 2Aρ2 where we use A =
PK−1 k=0
M −1 X
h i 2 2 E ∥∇F (xm )∥ + 2γmax K 2 L2 ρ2 QM ,
m=0
γi2k ≤ γmax
PK−1 k=0
Gathering QM -terms on the left side,
2 γik ≤ γmax K to obtain the last inequality.
M −1 h i X 2 2 1 − 2γmax K 2 L2 ρ2 QM ≤ AM σ 2 + ζ 2 + 2Aρ2 E ∥∇F (xm )∥ , m=0
1 , then we see that, if γmax ≤ 2KLρ M −1 h i X 2 QM ≤ 2AM σ 2 + ζ 2 + 4Aρ2 E ∥∇F (xm )∥ . m=0
1 Lemma B.6 (Bias). Under Assumptions 2.1 to 2.5, and if γmax ≤ 2KLρ , then M −1 X m=0
M −1 h h i i X 2 2 E ∥bm ∥ ≤ 2A2 K 2 L2max M σ 2 + ζ 2 + 4A2 K 2 L2max ρ2 E ∥∇F (xm )∥ .
(11)
m=0
Proof. Similar to the proof of Lemma B.5, we first bound M −1 K−1 X X m=0 k=0
m yk,i − xm k
2
≤K =K
M −1 K−1 X X m=0 k=0 M −1 K−1 X X m′ =0 k′ =0
1×
X (l,c)∈Ikm
c γi2l ∇fil (yl,i ) l
2
′
γi2k′ ∇fik′ (ykm′ ,ik′ )
2
×
X
1,
(l,c):(k′ ,m′ )∈Ilc
|
{z
≤K
}
where the last factor is bounded by K because, for a fixed update (k ′ , m′ ), all destination updates (l, c) whose paths contain (k ′ , m′ ) must lie within at most one full cycle after (k ′ , m′ ). After K updates, every worker has received a model update newer than (k ′ , m′ ), so no later path can contain this fixed update. So M −1 K−1 i X X h 2 m E yk,i − xm ≤ K 2 QM . (12) k m=0 k=0
Now M −1 X m=0
M −1 K−1 h i (8) i X X h 2 2 m E ∥bm ∥ ≤ AL2max E yk,i − xm k m=0 k=0
(12)
≤ AK 2 L2max QM .
1 , we can apply Lemma B.5 to obtain the desired bound: If γmax ≤ 2KLρ M −1 X m=0
M −1 h i h i X 2 2 2 2 2 2 2 2 2 2 2 E ∥bm ∥ ≤ 2A K Lmax M σ + ζ + 4A K Lmax ρ E ∥∇F (xm )∥ . m=0
21
Lemma B.7 (Cycle Step Decomposition). Let the worker-specific stepsizes be chosen such that γi ∝ wi τi . Then, K−1 X
Sm :=
k=0
where α :=
PK−1 k=0
m γik ∇fik (yk,i ) = α∇F (xm ) + bm + νm , k
(13)
γik is the cycle stepsize.
Proof. Let γi = cwi τi for some c > 0. Then α= = =
K−1 X k=0 n X i=1 n X i=1
γ ik
γi Ki cwi τi ·
τmax τi
= cτmax , and consequently K−1 X k=0
γik ∇Fik (xm ) = =
n X i=1 n X i=1
γi Ki ∇Fi (xm ) cwi τi ·
= cτmax
n X i=1
τmax · ∇Fi (xm ) τi
wi ∇Fi (xm )
= α∇F (xm ) .
(14)
For the cycle step, we now have Sm
=
K−1 X k=0
=
K−1 X k=0
+
m γik ∇fik (yk,i ) k
γik ∇Fik (xm )
K−1 X k=0
+
m ) − ∇Fik (xm ) γik ∇Fik (yk,i k
K−1 X k=0
(14),(B),(N)
=
m m γik ∇fik (yk,i ) − ∇Fik (yk,i ) k k
α∇F (xm ) + bm + νm .
Pn Lemma B.8. Let the target objective be F (x) = i=1 wi Fi (x), and suppose Assumptions 2.1 to 2.3 and 3.1 hold. Assume the worker-specific stepsizes are chosen such that γi ∝ wi τi . If the cycle PK−1 1 stepsize α := k=0 γik satisfies 0 < α ≤ 6L , then, after M ≥ 1 cycles, the cycle iterates xm generated by Algorithm 2 satisfy M −1 M −1 i i 1 X h 2∆ 3ALσ 2 2 3L 1 X h 2 m 2 E ∥∇F (x )∥ ≤ + + + E ∥b ∥ . m M m=0 αM α α2 α M m=0 22
Proof. In Lemma B.7, we defined the cycle step Sm , so that for the cycle iterates of Algorithm 2 Sm = xm − xm+1 holds.
By L-smoothness of the objective function F (Assumption 2.3), we obtain the standard bound [Nesterov, 2018] F (xm+1 ) ≤ F (xm ) + ∇F (xm ), xm+1 − xm + = F (xm ) − ⟨∇F (xm ), Sm ⟩ +
L m+1 2 x − xm 2
L 2 ∥Sm ∥ . 2
Taking expectations conditional on information available at the beginning of cycle m, we find h i L 2 Em F (xm+1 ) ≤ F (xm ) − ⟨∇F (xm ), Em [Sm ]⟩ + Em ∥Sm ∥ 2 h i L (13) 2 m m m = F (x ) − ⟨∇F (x ), α∇F (x ) + Em [bm ] + Em [νm ]⟩ + Em ∥Sm ∥ 2 h i L 2 m m 2 m = F (x ) − α ∥∇F (x )∥ − ⟨∇F (x ), Em [bm ]⟩ + Em ∥Sm ∥ , (15) 2 where we used Lemma B.1 to drop Em [νm ] = 0. We bound the inner product term (CS)
− ⟨∇F (xm ), Em [bm ]⟩ ≤ ∥∇F (xm )∥ ∥Em [bm ]∥ (YN) α 1 2 2 ∥∇F (xm )∥ + ∥Em [bm ]∥ ≤ 4 α h i (JN) α 1 2 2 ≤ ∥∇F (xm )∥ + Em ∥bm ∥ . 4 α
(16)
Now, 2 (13)
∥Sm ∥ = ∥α∇F (xm ) + bm + νm ∥ (SS)
2
2 2
2
≤ 3α2 ∥∇F (xm )∥ + 3 ∥bm ∥ + 3 ∥νm ∥ ,
and thus
h i h i 2 2 2 Em ∥Sm ∥ ≤ 3α2 ∥∇F (xm )∥ + 3Em ∥bm ∥ + 3Aσ 2 ,
(17)
utilizing the second part of Lemma B.1. Using (16) and (17) to further bound (15), we obtain (15)
Em F (xm+1 )
≤
(16),(17)
≤
(JN)
≤
1 For α ≤ 6L , we have
2
F (xm ) − α ∥∇F (xm )∥ − ⟨∇F (xm ), Em [bm ]⟩ +
h i L 2 Em ∥Sm ∥ 2
2
F (xm ) − α ∥∇F (xm )∥ α 1 2 2 + ∥∇F (xm )∥ + ∥Em [bm ]∥ 4 α i h L 2 2 2 + 3α ∥∇F (xm )∥ + 3Em ∥bm ∥ + 3Aσ 2 2 3 3L 2 2 m F (x ) − α− α ∥∇F (xm )∥ 4 2 i h 3ALσ 2 1 3L 2 + + + Em ∥bm ∥ . 2 α 2
3 3L 2 4α − 2 α
≥ α2 and thus
α 3ALσ 2 2 Em F (xm+1 ) ≤ F (xm ) − ∥∇F (xm )∥ + + 2 2 23
1 3L + α 2
h i 2 Em ∥bm ∥ .
Rearranging gives h i 2 F (xm ) − Em F (xm+1 ) 3L 3ALσ 2 2 2 + ∥∇F (x )∥ ≤ + + E ∥b ∥ . m m α α α2 α m
2
Taking unconditional expectation, averaging, and applying the bound F (x0 ) − E F (xM ) ≤ F (x0 ) − F ∗ = ∆ (Assumption 2.3) gives M −1 i 2 PM −1 E F (xm ) − F (xm+1 ) 1 X h m=0 m 2 E ∥∇F (x )∥ ≤ M m=0 αM M −1 h ii 3ALσ 2 3L 1 X h 2 2 + + + E E ∥b ∥ m m α α2 α M m=0 M −1 i 3L 1 X h 3ALσ 2 2 2∆ 2 + E ∥b ∥ . + + ≤ m αM α α2 α M m=0
Theorem B.9 (Convergence to a Weighted Average). Let the target objective be F (x) = Pn i=1 wi Fi (x), and suppose Assumptions 2.1 to 2.5 and 3.1 hold. Assume the worker-specific stepsizes are chosen such that γi ∝ wi τi . If the maximum stepsize satisfies 0 < γmax ≤ 5KL1max ρ PK−1 1 and the cycle stepsize α := k=0 γik satisfies 0 < α ≤ 6L , then, after M ≥ 1 cycles, the cycle m iterates x generated by Algorithm 2 satisfy 2 M −1 i ∆ A 2 A 1 X h m 2 E ∥∇F (x )∥ ≤ c0 · + c1 · Lσ + c2 · K 2 L2max σ 2 + ζ 2 , M m=0 αM α α for some constants c0 , c1 , c2 > 0, where A :=
PK−1 k=0
γi2k .
Proof. Plug the bound on the cycle bias from Lemma B.6 into Lemma B.8 to obtain M −1 i 2∆ 1 X h 2 E ∥∇F (xm )∥ ≤ M m=0 αM
+
1 Now, if α ≤ 6L , then also gives
10A2 K 2 L2max ρ2 1− α2
3ALσ 2 α 2 3L + + 2A2 K 2 L2max σ 2 + ζ 2 α2 α M −1 h i X 2 3L 2 2 2 2 1 m 2 + + 4A K L ρ E ∥∇F (x )∥ . max α2 α M m=0
2 3L α2 + α
≤ 2.5 α2 . Grouping terms involving the gradient norm on the left
M −1 i 5A2 K 2 L2max σ 2 + ζ 2 1 X h 2∆ 3ALσ 2 m 2 E ∥∇F (x )∥ ≤ + + M m=0 αM α α2
We require the first factor on the left-hand side to exceed 1/2, or, equivalently, Noting that K−1 K−1 X X 2 A= γik ≤ γmax γik = γmax α, k=0
k=0
24
10A2 K 2 L2max ρ2 ≤ 12 . α2
2 this requirement is met if 10γmax K 2 L2max ρ2 ≤ 12 , or, if
γmax ≤
1 1 ≤√ . 5KLmax ρ 20KLmax ρ
Thus, for γmax ≤ 5KL1max ρ , we find M −1 i 1 X h ∆ A A2 2 E ∥∇F (xm )∥ ≤ 4 + 6 Lσ 2 + 10 2 K 2 L2max σ 2 + ζ 2 . M m=0 αM α α
B.3
(18)
Proof of Theorem 4.1 and Corollary 4.2
Theorem 4.1 (Convergence to the Equal-Weighted Average). Let the worker-specific stepsizes be chosen according to (2) and 2.1 to 2.5 and 3.1 hold. If the stepsize parameter o n suppose Assumptions 1 1 satisfies 0 < γ ≤ min 6LτH , 5Lmax ρτmax , then, after M ≥ 1 cycles, the cycle iterates xm generated by Algorithm 2 satisfy M −1 i ∆ 1 X h γτA Lσ 2 2 E ∥∇F (xm )∥ ≤ c0 · + c1 · + c2 · γ 2 τA2 L2max σ 2 + ζ 2 M m=0 γτH M K
for some absolute constants c0 , c1 , c2 > 0. Proof. This is a special case of Theorem B.9 with stepsizes chosen according to (2). The cycle stepsize now becomes α=
K−1 X k=0
γ ik =
n X i=1
γi · K i =
n X i=1
γτi
τH τmax · = γτH , nτmax τi
the sum of squared stepsizes A=
K−1 X k=0
γi2k =
n X i=1
γi2 · Ki =
n X
γ 2 τi2
i=1
2 2 τmax τH τA τH 2 τH τA · , = γ = γ2 2 2 n τmax τi nτmax K
and the maximum stepsize γmax = γτmax
τmax τH τH =γ =γ , nτmax n K
where we used the identity K/n = τmax/τH (Assumption 3.1). γτA Therefore, A α = K . Plugging this into (18) gives M −1 i ∆ γτA Lσ 2 1 X h 2 E ∥∇F (xm )∥ ≤ 4 +6 + 10γ 2 τA2 L2max σ 2 + ζ 2 . M m=0 αM K
The bound on the cycle stepsize translates to α≤
1 1 ⇔γ≤ , 6L 6LτH
and that on the maximum stepsize to γmax ≤
1 1 ⇔γ≤ . 5KLmax ρ 5Lmax ρτmax
25
(3)
Corollary 4.2 (Time Complexity for the Equal-Weighted Average). In the setting of Theorem 4.1, the worst-case wall-clock time complexity to find an ε-stationary point of F is ! p ∆Lσ 2 ∆Lmax σ 2 + ζ 2 τmax ∆Lmax ρ τmax ∆L O τA + τA + τmax + τmax . (4) nε2 ε1.5 τH ε τH ε Proof. Under the assumptions of Theorem 4.1, we have M −1 i γτA Lσ 2 1 X h ∆ 2 + c1 · + c2 · γ 2 τA2 L2max σ 2 + ζ 2 . E ∥∇F (xm )∥ ≤ c0 · M m=0 γτ M {z } {zK } | | {zH } | T2 T1
T0
To achieve T1 ≤ 4ε , we require γ≤
1 εK . 4c1 Lσ 2 τA
To achieve T2 ≤ 4ε , we likewise require
√ ε 1 p γ≤√ . 4c2 Lmax τA σ 2 + ζ 2
The largest γ satisfying these constraints, as well as those in Theorem 4.1, is ( ) √ 1 εK ε 1 1 1 p γ̂ := min , ,√ , . 4c1 Lσ 2 τA 4c2 Lmax τA σ 2 + ζ 2 5Lmax ρτmax 6LτH To achieve T0 ≤ 2ε and bound the right-hand side of Theorem 4.1 by ε, the cycle count must satisfy 2c0 ∆ M ≥ εγ̂τ . Ignoring constant factors, this cycle bound can be expressed as H p ∆Lσ 2 ∆Lmax σ 2 + ζ 2 ∆Lmax ρ ∆L . M≳ 2 τA + τA + τmax + ε KτH ε1.5 τH ετH ε To obtain the wall-clock time complexity, we multiply the required number of cycles by the wall-clock duration of a single cycle. Under Assumption 3.1, a cycle is completed after τmax = KτnH time units. Thus, after ! p ∆Lmax σ 2 + ζ 2 τmax ∆Lmax ρ τmax ∆L ∆Lσ 2 O τA + τA + τmax + τmax nε2 ε1.5 τH ε τH ε time units, we achieve the target accuracy. B.4
Proof of Theorem 4.3 and Corollary 4.4
Theorem 4.3 (Convergence to the Frequency-Weighted Average). Let the workers’ stepsizes be equal, γi = γ/K , and suppose Assumptions 2.1 to 2.5 and 3.1 hold for F̃ .3 If the stepsize parameter n o
satisfies 0 < γ ≤ min Algorithm 2 satisfy M −1
1 X E M m=0
1 , 1 6L̃ 5Lmax ρ̃
m
∇F̃ (x )
2
, then, after M ≥ 1 cycles, the cycle iterates xm generated by
≤ c0 ·
˜ ∆ γ L̃σ 2 + c1 · + c2 · γ 2 L2max σ 2 + ζ̃ 2 γM K
for some absolute constants c0 , c1 , c2 > 0. 3 That is, we replace F, ∆, L, ζ, ρ by F̃ , ∆, ˜ L̃, ζ̃, ρ̃ in Assumptions 2.3 and 2.5.
26
Proof. This is a special case of Theorem B.9 with equal stepsizes γi = γ/K . The cycle stepsize now becomes α=
K−1 X k=0
γ = γ, K
γ ik = K ·
the sum of squared stepsizes A=
K−1 X k=0
γi2k = K ·
and the maximum stepsize γmax =
γ 2 K
=
γ2 , K
γ . K
γ Therefore, A α = K . Plugging this into (18) gives M −1 ˜ 2 ∆ 1 X γ L̃σ 2 ≤4 E ∇F̃ (xm ) +6 + 10γ 2 L2max σ 2 + ζ̃ 2 . M m=0 γM K
The bound on the cycle stepsize translates to α≤
1 1 ⇔γ≤ , 6L̃ 6L̃
and that on the maximum stepsize γmax ≤
1 1 ⇔γ≤ . 5KLmax ρ̃ 5Lmax ρ̃
Corollary 4.4 (Time Complexity for the Frequency-Weighted Average). In the setting of Theorem 4.3, the worst-case wall-clock time complexity to find an ε-stationary point of F̃ is q ˜ max σ 2 + ζ̃ 2 ˜ L̃σ 2 ˜ max ρ̃ ˜ L̃ ∆L ∆ ∆L ∆ O τH + τmax + τmax + τmax . (6) nε2 ε1.5 ε ε Proof. Under the assumptions of Theorem 4.3, we have M −1 ˜ 2 ∆ γ L̃σ 2 1 X m ≤ c0 · + c1 · E ∇F̃ (x ) + c2 · γ 2 L2max σ 2 + ζ̃ 2 . M m=0 γM K {z } | {z } | {z } | T1
T0
T2
εK To achieve T1 ≤ 4ε , we require γ ≤ 4c11 L̃σ . To achieve T2 ≤ 4ε , we likewise require γ ≤ 2 √ √1 √ε . 4c2 Lmax
σ 2 +ζ̃ 2
The largest γ satisfying these constraints, as well as those in Theorem 4.3, is √ 1 Kε 1 ε 1 1 q γ̂ := min ,√ , , . 4c1 L̃σ 2 5Lmax ρ̃ 6L̃ 4c2 L σ 2 + ζ̃ 2 max
To achieve T0 ≤ 2ε and bound the right-hand side of Theorem 4.3 by ε, the cycle count must satisfy ˜
0∆ M ≥ 2cεγ̂ . Ignoring constant factors, this cycle bound can be expressed as q ˜ max σ 2 + ζ̃ 2 ˜ L̃σ 2 ˜ max ρ̃ ∆ ˜ L̃ ∆L ∆ ∆L M≳ 2 + + + . 1.5 ε K ε ε ε
27
To obtain the wall-clock time complexity, we multiply the required number of cycles by the wall-clock duration of a single cycle. Under Assumption 3.1, a cycle is completed after τmax = KτnH time units. Thus, after q ˜ 2 ˜ ˜ max ρ̃ ˜ L̃ ∆L σ 2 + ζ̃ 2 max ∆ L̃σ ∆L ∆ O τH + τmax + τmax + τmax nε2 ε1.5 ε ε time units, we achieve the target accuracy.
C
On the Necessity of Bounding Data Heterogeneity
In this section, we present a simple counterexample to demonstrate that without a bounded heterogeneity assumption (such as Assumption 2.5), Rescaled ASGD with worker-wise constant stepsizes cannot guarantee convergence to an ε-stationary point. Consider a two-worker setup with local objective functions F1 (x) =
x2 − cx, 2
F2 (x) =
x2 + cx . 2 2
The global objective is the equal-weighted average F (x) = x2 , with ∇F (x) = x and a unique stationary point at x∗ = 0. The constant c > 0 captures the degree of data heterogeneity, c = ∥∇Fi (x) − ∇F (x)∥.
Suppose τ1 = 1, τ2 = 2, so worker 1 takes two steps of size γ per cycle while worker 2 takes only one of size 2γ. Assume the cycle begins at a point |x0 | < 1 so that ∥∇F (x0 )∥ < 1. First, worker 1 delivers the gradient computed at the initial model, ∇F1 (x0 ) = x0 − c, and the server updates the model: x01 = x0 − γ(x0 − c) = x0 (1 − γ) + γc .
Worker 1 delivers their second update, ∇F1 (x01 ) = x01 − c, and the server updates the model again: x02 = x01 − γ(x01 − c) = x01 (1 − γ) + γc = x0 (1 − γ)2 + c(2γ − γ 2 ) . Worker 2 delivers their stale gradient, ∇F2 (x0 ) = x0 + c, and the server updates the model: x1 = x03 = x02 − 2γ(x0 + c) = x0 (1 − 4γ + γ 2 ) − γ 2 c . Thus, for any fixed γ > 0, the gradient norm at the end of the cycle can be arbitrarily large if there is no restrictions on c. In particular, if c>
1 + x0 (1 − 4γ + γ 2 ) , γ2
then ∇F (x1 ) > 1. Relaxing the Assumption Koloskova et al. [2022], Mishchenko et al. [2022a] impose the constraint 2
∥∇Fi (x) − ∇F (x)∥ ≤ ξ 2 ,
∀x ∈ Rd ,
∀i .
(19)
This assumption implies our Assumption 2.5 and is therefore more restrictive. Suppose (19) holds, then 2
(SS)
2
2
∥∇Fi (x)∥ ≤ 2 ∥∇Fi (x) − ∇F (x)∥ + 2 ∥∇F (x)∥ (19)
2
≤ 2ξ 2 + 2 ∥∇F (x)∥ ,
so Assumption 2.5 also holds for ζ 2 = 2ξ 2 and ρ2 = 2. 28
Note that (19) is quite restrictive as it does not even allow for quadratics with different Hessians and smoothness constants. To illustrate, consider two simple quadratics, F1 (x) = a2 x2 and F2 (x) = 2b x2 . The deviation for the first worker is ∥∇F1 (x) − ∇F (x)∥ = | a−b 2 ||x|. For this term to remain bounded by a constant ξ 2 for all x ∈ R, we must have a = b. Our relaxed Assumption 2.5 accommodates such cases by allowing the local gradient bounds to scale proportionally with the global gradient norm.
D
Wall-Clock Time Complexities
D.1
Concurrent ASGD
Koloskova et al. [2022] give a bound on the iteration complexity to achieve ε-stationarity. The leading 2 2 term of this bound is ∆L(σε2+ζ ) , where ζ 2 is a constant bounding data heterogeneity (cf. Section C). Because Concurrent ASGD employs uniform worker-sampling and keeps the number of gradients in-flight constant throughout, each worker will contribute M/n gradients in the long run, where M denotes the number of iterations. The time needed to compute these gradients is bottlenecked by the slowest worker and thus scales with the maximum of the computation times, τmax , not their arithmetic mean, τA . In other words, the expected time to complete a large number M of iterations of Concurrent ASGD will be performed in M τmax/n units of time. The leading term of the time 2 +ζ 2 ) complexity bound is therefore ∆L(σ τmax as stated in Table 1. nε2 D.2
Delay-Adaptive ASGD
In Delay-Adaptive ASGD, stepsizes are scaled down in proportion to the gradient staleness. In our model with a cyclic update schedule (cf. Section 3), these stepsizes are therefore deterministic and do not change per within-cycle iteration after the second cycle. Let δk denote the number of model updates applied between worker ik receiving the model and the server applying the gradient in γ iteration k. Then the stepsize applied is γk = 1+δ . The total aggregate stepsize of steps taken by k worker i over one cycle is then X X γ Γi := γk = , (20) 1 + δk k:ik =i
k:ik =i
and we can decompose the cycle step as before (cf. Lemma B.7). The objective function targeted by Pn Delay-Adaptive ASGD is a weighted-average, F̂ (x) = i=1 ŵi Fi (x), with weights ŵi ∝ Γi . If the local objective functions are linearly independent, then F̂ coincides with the equal-weighted average F only if all workers have the same computation time, τi = τA . In this case, Delay-Adaptive ASGD coincides with Rescaled ASGD and Vanilla ASGD. In the general case, slow workers are penalized twice—due to the lower number of updates per cycle and the smaller stepsizes. By letting faster workers take larger steps, Delay-Adaptive ASGD may converge faster, albeit to the wrong objective function F̂ . Under Assumption 3.1, if simultaneous updates are processed in ascending order of worker indices, the iteration-specific delay δk for every update delivered by worker i remains constant across all iterations k with ik = i. With a slight abuse of notation, the worker-specific delay is then X τi n δi = = τi − 1. (21) τj τH j̸=i
τH 1 Substituting (21) into (20), we find Γi = γ τmax . With this, the cycle stepsize and sum of n τi2 squared stepsizes become n
τmax τH X −2 α=γ τi , n i=1
n 2 X 2 τmax τH A=γ τ −3 . n2 i=1 i
29
Following the same steps as in Sections B.3 and B.4, we can now derive the cycle and wall-clock time complexities. Under the same assumptions as in Theorem B.9 (imposed on F̂ ), the leading term in the cycle complexity becomes Pn −3 ˆ L̂σ 2 ∆ i=1 τi 2 , P n ε2 τmax τ −2 i=1 i
and that in the time complexity bound Pn −3 ˆ L̂σ 2 ˆ L̂σ 2 ∆ ∆ i=1 τi = τDA , P n −2 2 ε2 nε2 i=1 τi where τDA
:=
Pn
τ −3
i n Pni=1 −2 That is, the leading term of the time complexity of 2. ( i=1 τi ) Delay-Adaptive ASGD scales with τDA , which collapses to the arithmetic mean if all workers have the same computation speed. We can bound
τH ≤ τDA ≤ nτmin , which shows that Delay-Adaptive ASGD may converge faster than Rescaled ASGD—e.g., if there is one fast worker with τ1 = 1 and one slow worker with τ2 = 10, so τDA ≈ 1.96 < 5.5 < τA —or slower—e.g., under near-homogeneity with a slight skew, τ1 = 1, τ2 = · · · = τ10 , so τDA ≈ 2.01 > 1.9 = τA —but the objective function targeted in such cases is, in general, not the equal-weighted average (1). The leading term scaling with τA reported in Table 1 refers to the setting with homogeneous worker computation times under which Delay-Adaptive ASGD can guarantee convergence to an ε-stationary point for arbitrarily small ε.
E
Experimental Details
For our experiments in Section 5, we consider an image classification problem on MNIST [LeCun et al., 2010] with standard normalization. To enforce a high degree of data heterogeneity, we assign each one of the ten classes to one of ten workers. This is a special case of Dirichlet partitioning [Hsu et al., 2019, Yurochkin et al., 2019] with concentration parameter set to 0. We trim the workers’ data to ensure equal sample sizes, resulting in 5, 421 examples per client. Our model is a two-layer MLP with ReLU activations in the hidden layer of dimension 128. The workers compute stochastic gradients using minibatches of their local data of size 64. The model is trained by minimizing cross-entropy loss. We assign τi ∈ {1, 2, 4, 8, 16} to two workers each, making the fastest workers 16 times faster than the slowest. In the fixed-computation setting, worker i takes τi units of time to compute a stochastic gradient. In the fluctuation computation-time setting, worker i’s computation times are sampled from an exponential distribution with mean τi instead. For each method, we tune the stepsize parameter γ within a fixed wall-clock budget of 30, 000 time units to minimize loss, and sweep the set {10x : x = −4, . . . , 0}. The selected stepsizes all fall in the interior of this grid. We report the training loss versus wall-clock time in Figure 1 for the three methods under consideration. Each method is run five times over different seeds. The solid lines show the median over these runs and the shaded region corresponds to the minimum and maximum. To enable these computations when computation times fluctuate, the loss trajectories have been linearly interpolated over a grid of size 200 covering the simulated time horizon. We find that Rescaled ASGD outperforms Malenia SGD and Ringleader ASGD in this setup. Although our theoretical analysis suggests all three methods to have similar performance in the fixedcomputation model, it appears that Rescaled ASGD benefits from the larger number of model updates. This is consistent with the finding that, in the setting with fluctuating computation-times, the performance of Rescaled ASGD is not affected as it takes roughly the same number as updates as before, whereas Malenia SGD and Ringleader ASGD are slowed down due to, in expectation, longer gathering phases. 30
F
Further Experiments
F.1
Vanilla ASGD Targets the Frequency-Weighted Average
To illustrate that Rescaled ASGD targets the right objective function (1), while Vanilla ASGD instead targets the frequency-weighted average of the local objective (5), we consider a simple example with two workers with computation times τ1 = 1, τ2 = 2, and local objective functions F1 (x) = (x − 4)2 and F2 (x) = 2(x + 3)2 . The unique minimizer of the equal-weighted average is x∗ = −2/3, while that of the frequency-weighted average is x̃∗ = +1/2. We run both algorithms with the same cycle stepsize α = 0.01, setting gradient stochasticity to zero to isolate the effects of staleness. Their trajectories in the search space are shown in Figure 3. Panel (a) shows the full sequence of iterates, including the within-cycle iterates, while panel (b) shows the cycle iterates only. We see that Rescaled ASGD converges to the minimizer of the equal-weighted average F , while Vanilla ASGD converges to the minimizer of the frequency-weighted average F̃ . Moreover, Vanilla ASGD converges faster for the same cycle stepsize α, consistent with our discussion of Corollary 4.4 where we have seen that the leading term scales with the harmonic mean of workers’ computation times. Note that the oscillations seen in the trajectories covering all iterates in panel (a) are absent when shifting to the cycle-iterate trajectories in (b). Lastly, a close look at the trajectory of Vanilla ASGD in panel (b) shows that the convergence is not exactly to the stationary point. This, too, is consistent with our theory: Even in the absence of gradient stochasticity (σ 2 = 0), a neighborhood term reflecting the staleness bias persists. As seen in Theorems 4.1 and 4.3, this scales with γ 2 and can thus be efficiently controlled by choosing a smaller stepsize. 5
5
Rescaled ASGD Vanilla ASGD x∗ = − 2 / 3
4
Rescaled ASGD Vanilla ASGD x∗ = − 2 / 3
4
x̃∗ = +1 /2
x̃∗ = +1 /2
3
x
x
3 2
2
1
1
0
0
0
50
100
150
200
250
300
0
50
100
150
200
Simulated Wall-Clock Time
Simulated Wall-Clock Time
(a) All iterates xm k .
(b) Cycle iterates xm .
250
300
Figure 3: Rescaled ASGD targets the equal-weighted average F , Vanilla ASGD the frequencyweighted average F̃ .
F.2
Delay-Adaptive SGD Exhibits Objective Inconsistency
Delay-Adaptive ASGD [Mishchenko et al., 2022a] scales stepsizes down in proportion to the gradient staleness, i.e., the number of iterations that have passed between the worker receiving the model from the server and them delivering the gradient back. In the fixed-computation model, this approach does the opposite of Rescaled ASGD.
To illustrate, consider another two-worker setting with one fast worker (τ1 = 1) and one slow worker (τ2 = 100). 99 out of 100 gradients delivered by the fast worker will exhibit no staleness, so Delay-Adaptive ASGD does not scale the stepsizes down. Merely one gradient of the fast worker in each cycle will be shrunk by a factor of two due to the slow worker delivering theirs. The one gradient delivered by the slow worker, in contrast, will always be stale by 100 iterations, and thus the stepsize scaled down accordingly. The search trajectory of Delay-Adaptive ASGD will therefore be heavily biased towards the local objective function of the fast worker, even more so than under Vanilla ASGD. 31
Figure 4 shows a simulation for F1 (x) = (x − 1)2 , F2 (x) = (x + 1)2 . As expected, Delay-Adaptive ASGD converges to a point close to x∗1 = 1, the minimizer of worker 1’s objective function, whereas Rescaled ASGD converges to a small neighborhood around the minimizer, x∗ = 0, of the equal-weighted average. 3.0
Rescaled ASGD Delay-Adaptive ASGD x∗ = 0 x∗1 = 1
2.5
x
2.0 1.5 1.0 0.5 0.0 0
200
400
600
800
1000
1200
1400
Simulated Wall-Clock Time Figure 4: Delay-Adaptive ASGD is heavily biased towards the local objective of the faster worker. Rescaled ASGD converges to a small neighborhood around the minimizer of the equal-weighted average.
32