Title: ADASCALE SGD: A SCALE-INVARIANT ALGORITHM FOR DISTRIBUTED TRAINING

Abstract: When using distributed training to speed up stochastic gradient descent, learning rates must adapt to new scales in order to maintain training effectiveness. Re-tuning these parameters is resource intensive, while fixed scaling rules often degrade model quality. We propose AdaScale SGD, a practical and principled algorithm that is approximately scale invariant. By continually adapting to the gradient's variance, AdaScale often trains at a wide range of scales with nearly identical results. We describe this invariance formally through AdaScale's convergence bounds. As the batch size increases, the bounds maintain final objective values, while smoothly transitioning away from linear speed-ups. In empirical comparisons, AdaScale trains well beyond the batch size limits of popular "linear learning rate scaling" rules. This includes large-scale training without model degradation for machine translation, image classification, object detection, and speech recognition tasks. The algorithm introduces negligible computational overhead and no tuning parameters, making AdaScale an attractive choice for large-scale training.

1. INTRODUCTION:

Large datasets and large models underlie much of the recent success of machine learning. Training such models is time consuming, however, as stochastic gradient descent algorithms can require days or weeks to train effectively. Thus, procedures that speed up SGD are valuable. Faster training enables consideration of more data and models, which expands the capabilities of machine learning.

To speed up SGD, distributed systems can process thousands of training examples per iteration. But training at large scales also creates a major algorithmic challenge. Specifically, learning rates must adapt to each scale. Without choosing these training parameters carefully, scaled SGD frequently trains low-quality models, producing a waste of resources rather than a useful model. To adapt learning rates, "fixed scaling rules" are standard but unreliable strategies. Goyal et al. (2017) popularized "linear learning rate scaling," which can work well, especially for computer vision tasks (Krizhevsky, 2014; Devarakonda et al., 2017; Jastrzębski et al., 2018; Smith et al., 2018; Lin et al., 2019) . For other problems or larger scales, however, linear scaling often fails. This fact is well-known in theory (Yin et al., 2018; Jain et al., 2018; Ma et al., 2018) and in practice (Goyal et al., 2017) . Other fixed scaling rules are also undependable. Golmant et al. (2018) test three rules-linear, root, and identity-and conclude that each one often degrades model quality. Shallue et al. (2019) compute near-optimal parameters for many tasks and scales, and the results do not align with any fixed rule. To ensure effective training, the authors recommend avoiding such rules and re-tuning parameters for each new scale-an inconvenient and resource-intensive solution.

We propose AdaScale SGD. A practical but principled algorithm, AdaScale more reliably scales training by adapting to the gradient's variance. Decreased gradient variance is the fundamental impact of large batch sizes. Thus, scaling provides little gain if the variance is already "small" at small scales. In such cases, AdaScale increases the learning rate conservatively, and large-scale training progresses similarly to the small-batch setting. For iterations with "large" gradient variance, Ada-Scale increases the learning rate aggressively, and the per-iteration progress dramatically increases.

AdaScale is approximately scale invariant, a quality that simplifies large-batch training. With no changes to learning rates or other inputs, AdaScale can train at many scales with similar results. This leads to two important innovations: (i) AdaScale improves the translation of training configurations between scales, which is useful for scaling up tasks or adapting to dynamic resource availability; and (ii) AdaScale works at scale with simple learning rate schedules, which eliminates the need for "warm-up" heuristics (Goyal et al., 2017) . Qualitatively, AdaScale and warm-up have similar effects on learning rates, but with AdaScale, this behavior emerges from a principled and adaptive mechanism, not hand-tuned parameters.

We provide theoretical results that formalize this approximate scale invariance. Bounds for all scales converge to identical objective values. In contrast, the linear scaling rule requires fewer iterations but compromises model quality and training stability, causing divergence as the batch size increases.

We perform large-scale empirical evaluations on five training benchmarks. Tasks include image classification, machine translation, object detection, and speech recognition. The results align well with our theory, as AdaScale systematically preserves model quality across many scales. This includes training ImageNet with batch size 32k and Transformer with 262k max tokens per batch.

To provide context for our description of AdaScale, Figure 1 includes results from a simple scaling experiment using CIFAR-10 data. These results illustrate the concept of scale invariance, AdaScale's qualitative impact on learning rates, and a failure case for the linear scaling rule.

2. PROBLEM FORMULATION:

We focus on quickly computing approximate solutions to the problem

minimize w∈R d F (w) , where F (w) = E x∼X [f (w, x)] . (P1)

Here w parameterizes a machine learning model, while X denotes a distribution over batches of training data. We assume that F and f are differentiable and that

E x∼X [∇ w f (w, x)] = ∇F (w).

Stochastic gradient descent is a popular algorithm for solving (P1). Let w t denote the model parameters when iteration t begins. During this iteration, SGD samples a batch x t ∼ X and computes the gradient g t ← ∇ w f (w t , x t ). SGD then applies the update w t+1 ← w t -η t g t . Here η t is the learning rate. Given a schedule lr : Z ≥0 → R >0 , SGD defines η t = lr(t). For our experiments in §4, lr is an exponential decay or step decay function. SGD completes training after T iterations.

To speed up training, practitioners often parallelize gradient computation across multiple devices. Algorithm 1 defines a scaled SGD algorithm. At scale S, the algorithm samples S independent batches during each iteration. After computing the gradient for each batch in parallel, the algorithm applies the mean of these gradients (in place of g t ) when updating model parameters.

But scaling training in this way creates a considerable algorithmic challenge. Each new scale requires a new learning rate schedule, which is inconvenient and resource intensive to obtain. To help address this challenge, we propose a scaled SGD algorithm that is approximately scale invariant. Definition 1. Let w T denote the (possibly random) result of a scaled SGD algorithm. Fixing all algorithm inputs except scale S, the algorithm is scale invariant if w T does not depend on S.

A scale-invariant algorithm makes parallelizing training significantly easier. Such an algorithm can scale to any available amount of computational resources, and there is no need for parameter retuning, unreliable heuristics, or algorithmic expertise from users.

Algorithm 1 Scaled SGD function Scaled_SGD(S, lr, T , X , f , w 0 )

for t = 0, 1, 2, . . . , T -1 do ḡt ← compute_gradient(w t , S, X , f ) η t ← lr(t) w t+1 ← w t -η t ḡt return w T function compute_gradient(w t , S, X , f ) in parallel for i = 1, . . . , S do x (i) ← sample_batch(X ) g (i) ← ∇ w f (w t , x (i) ) return 1 S S i=1 g (i) Algorithm 2 AdaScale SGD function AdaScale(S, lr, T SI , X , f , w 0 ) initialize τ 0 ← 0; t ← 0 while τ t < T SI do ḡt ← compute_gradient(w t , S, X , f ) # Compute gain r t ∈ [1, S] (see §3.3): r t ← E[σ 2 (wt)+ ∇F (wt) 2 ] E 1 S σ 2 (wt)+ ∇F (wt) 2 η t ← r t • lr( τ t ) w t+1 ← w t -η t ḡt τ t+1 ← τ t + r t ; t ← t + 1 return w t

3. ADASCALE SGD ALGORITHM:

This section introduces our AdaScale algorithm. As motivation, we first consider the role of gradient variance in SGD. We later provide practical guidance for variance estimation and momentum tuning.

3.1. INTUITION: IDENTITY SCALING, LINEAR SCALING, AND GRADIENT VARIANCE:

We now consider two fixed scaling rules, which influence the design of AdaScale. One of these rules is identity scaling, which keeps the training configuration constant for all scales: Definition 2. To apply the identity scaling rule to Algorithm 1, use the same lr and T for all S.

Note that this rule has little practical appeal, since it fails to reduce the number of training iterations. A second and more popular strategy is linear learning rate scaling: Definition 3. To apply the linear learning rate scaling rule to Algorithm 1, use lr(t) = S •lr S1 (St) and T = T S1 /S , where lr S1 and T S1 denote the learning rate schedule and total steps for S = 1.

Conceptually, linear scaling treats SGD as a perfectly parallelizable algorithm. If true, applying gradients from S batches in parallel achieves the same result as doing so in sequence.

For special cases of (P1), the identity and linear rules result in scale-invariant algorithms. To show this, we first define the variance quantities Σ(w) = cov x∼X (∇ w f (w, x), ∇ w f (w, x)) , and σ 2 (w) = tr(Σ(w)) . In words, σ 2 (w) sums the variances of each entry in ∇ w f (w, x). By sampling batches independently, scaling fundamentally impacts SGD by reducing this variance. Given w t in Algorithm 1, we have cov(ḡ t , ḡt ) = 1 S Σ(w t ) and E [ḡ t ] = ∇F (w t ). Here, only the covariance depends on S. Consider the special case of zero gradient variance. In this case, identity scaling performs ideally: Proposition 1 (Scale-invariant SGD for deterministic gradients). If σ 2 (w) = 0 for all w ∈ R d , then applying identity scaling to Algorithm 1 results in a scale-invariant algorithm.

Although identity scaling does not speed up training, Proposition 1 is critical for framing the impact of large scales. If the gradient variance is "small," then we cannot expect large gains from increasing S-a larger scale has little effect on ḡt . With "large" variance, however, the opposite is true: Proposition 2 (Scale-invariant SGD for extreme stochasticity). Consider fixed covariance matrix Σ ∈ S d ++ , learning rate value η ∈ R >0 , and training duration T . For a given ν ∈ R >0 , assume ∇ w f (w, x) ∼ N (∇F (w), ν Σ), and apply linear scaling to Algorithm 1 with lr S1 (t) = ν -1 η and T S1 = ν T . The resulting scaled SGD algorithm is scale-invariant in the limit ν → +∞.

In less formal terms, linear scaling leads to scale-invariance in the case of very large gradient variance (as well as small learning rates and many iterations, to compensate for this variance). Since increasing S decreases variance, it is natural that scaling yields large speed-ups in this extreme case.

In practice, the gradient's variance is neither zero nor infinite, and both identity and linear scaling may perform poorly. Moreover, the gradient's variance does not remain constant throughout training. A scale-invariant algorithm, it seems, must continually adapt to the state of training.

3.2. ADASCALE DEFINITION:

AdaScale, defined in Algorithm 2, adaptively interpolates between identity and linear scaling, based on the expectation of σ 2 (w t ). During iteration t, AdaScale multiplies the learning rate by the "gain ratio" r t ∈ [1, S]: η t = r t • lr( τ t ). Here τ t is the "scale-invariant iteration," defined as τ t = t-1 t =0 r t . The idea is that iteration t performs the equivalent of r t single-batch iterations, and τ t accumulates this progress. AdaScale concludes when τ t ≥ T SI , where T SI is the total scaleinvariant iterations. Since r t ∈ [1, S], AdaScale requires at least T SI /S and at most T SI iterations.

The identity and linear rules correspond to two special cases of AdaScale. If r t = 1 for all t, the algorithm equates to SGD with identity scaling. Similarly, if r t = S for all t, we have linear scaling. Thus, to approximate scale-invariance, §3.1 suggests setting r t ≈ 1 when the gradient's variance is small and r t ≈ S when this variance is large. AdaScale achieves this by defining

r t = E σ 2 (w t ) + ∇F (w t ) 2 E 1 S σ 2 (w t ) + ∇F (w t ) 2 .

The expectations here are with respect to the distribution of w If S = 1 in AdaScale, then r t = 1 for all iterations. For larger scales, r t depends on E σ 2 (w t ) and E ∇F (w t ) 2 , and a practical implementation must efficiently approximate these values. Fortunately, the per-batch gradients g

(1)

t , . . . , g (S) t

and aggregated gradient ḡt are readily available in distributed SGD algorithms. This makes approximating r t straightforward. In particular, we define

σ2 t = 1 S-1 S i=1 g (i) t 2 -S S-1 ḡt 2 ,

and

μ2 t = ḡt 2 -1 S σ2 t .

Here σ2:

t and μ2 t are unbiased estimates of E σ 2 (w t ) and E ∇F (w t ) 2 . To ensure robustness to estimation variance, we estimate r t by plugging in moving averages σ2 t and μ2 t , which average σ2 t and μ2 t over prior iterations. Our implementation uses exponential moving average parameter θ = max{1 -S/1000, 0}, where θ = 0 results in no averaging. We find that AdaScale is robust to the choice of θ, and we provide evidence of this in Appendix C. To initialize, we set r 0 ← 1, and for iterations t < (1 -θ) -1 , we define σ2 t and μ2 t as the mean of past samples. Before averaging, we clip estimates so that σ2 t ≥ 10 -6 (to prevent division by zero) and μ2 t ≥ 0 (to ensure r t ∈ [1, S]). To verify these estimators, Figure 2 compares moving average estimates to offline estimates using model checkpoints. These plots also provide examples of gain ratios for practical problems. We note that numerous prior works-for example, (Schaul et al., 2013; Kingma & Ba, 2015; McCandlish et al., 2018) -have relied on similar moving averages to estimate gradient moments.

One final practical consideration is the momentum parameter ρ when using AdaScale with momentum-SGD. The performance of momentum-SGD depends less critically on ρ than the learning rate (Shallue et al., 2019) . For this reason, we find that AdaScale often performs well if ρ remains constant across scales and iterations. This approach to momentum scaling has also succeeded in prior works involving the linear scaling rule (Goyal et al., 2017; Smith et al., 2018) .

4. EMPIRICAL COMPARISONS:

We evaluate AdaScale on five practical training benchmarks. We assess scale invariance by comparing training curves across scales. We assess impact on training times by comparing total iterations. We consider a variety of tasks, models (He et al., 2016a; b; Amodei et al., 2016; Vaswani et al., 2017; Redmon & Farhadi, 2018) , and datasets (Deng et al., 2009; Krizhevsky, 2009; Everingham et al., 2010; Panayotov et al., 2015) . For each benchmark, we use one simple learning rate schedule. Specifically, lr is an exponential decay function for cifar10 and speech, and a step decay function otherwise. We use standard lr parameters for imagenet and yolo. Otherwise, we use tuned parameters that approximately maximize the validation metric (to our knowledge, there are no standard schedules for solving speech and transformer with momentum-SGD). We use momentum ρ = 0.9 except for transformer, in which case we use ρ = 0.99 for greater training stability. For S > 1, AdaScale's learning rate increases gradually during initial training, despite the fact that lr is non-increasing. Unlike warm-up heuristics (Goyal et al., 2017) , this behavior emerges naturally from a principled algorithm, not hand-tuned user input. Thus, AdaScale provides not only a compelling alternative to warm-up but also a plausible explanation for warm-up's success.

For imagenet, we also consider elastic scaling. Here, the only change to AdaScale is that S changes abruptly after some iterations. We consider two cases: (i) S increases from 32 to 64 at τ t = T SI /4 and from 64 to 128 at τ t = T SI /2, and (ii) the scale decreases at the same points, from 128 to 64 to 32. In Figure 4 , we include training curves from this setting. AdaScale remains approximately scale invariant, highlighting AdaScale's value for the common scenario of dynamic resource availability. As a baseline for all benchmarks, we also evaluate linear scaling with warm-up (LSW). As inputs, LSW takes single-batch schedule lr S1 = lr and single-batch steps T S1 = T SI , where lr and T SI are the inputs to AdaScale. Our warm-up implementation closely follows Goyal et al. (2017) . LSW trains for T S1 /S iterations, applying warm-up to the first 5.5% of iterations. During warm-up, the learning rate increases linearly from lr S1 (0) to S • lr S1 (0).

Table 2 compares results for AdaScale and LSW. LSW consistently trains for fewer steps, but doing so comes at a cost. As S grows larger, LSW consistently degrades model quality and sometimes diverges. For these divergent cases, we also tested doubling the warm-up duration to 11% of iterations, and training still diverged. In contrast, AdaScale preserves model quality for nearly all cases.

As a final comparison, Figure 5 demonstrates AdaScale's performance on cifar10 with many different lr schedules. We consider a 13×13 grid of exponential decay schedules and plot contours of resulting validation accuracy. At scale 16, AdaScale results align with accuracies for single-batch training, illustrating that AdaScale is approximately scale-invariant for many schedules. Moreover, AdaScale convincingly outperforms direct search over exponential decay schedules for scaled SGD at S=16. For training at scale, AdaScale provides a more natural learning rate parameterization.

5. SCALE-INVARIANT CONVERGENCE BOUND:

We now present convergence bounds that formalize the approximate scale invariance of AdaScale.

The bounds provide identical convergence guarantees for all scales, meaning that in terms of upper bounds on training loss, AdaScale is scale invariant. For comparison, we include an analogous bound for the linear scaling rule. Qualitatively, the bounds agree closely with our empirical results.

Let us define F * = min w F (w). Our analysis requires a few assumptions that are typical of SGD analysis of non-convex problems (see, for example, (Lei et al., 2017; Yuan et al., 2019) ):

Assumption 1 (α-Polyak-Łojasiewicz). For some α > 0, F (w) -F * ≤ 1 2α ∇F (w) 2 for all w.

Assumption 2 (β-smooth). For some β > 0, ∇F (w) -∇F (w ) ≤ β ww for all w, w .

Assumption 3 (Bounded variance).:

There exists a V ≥ 0 such that σ 2 (w) ≤ V for all w.

We emphasize that we do not assume convexity. The PL condition, which is perhaps our strongest assumption, is proven to hold for some nonlinear neural networks (Charles & Papailiopoulos, 2018) .

We consider constant lr schedules, which result in simple and instructive bounds. To provide context for the AdaScale result, we first present a straightforward bound for single-batch training:

Theorem 1 (Single-batch SGD bound). Given Assumptions 1, 2, 3 and η ∈ (0, 2β -1 ), consider Algorithm 1 with S = 1 and lr(t) = η. Defining γ = ηα(2 -ηβ) and ∆ = 1 2γ η 2 βV , we have

E [F (w T ) -F * ] ≤ (1 -γ) T [F (w 0 ) -F * ] + ∆ .

The bound describes two important characteristics of the single-batch algorithm. First, the suboptimality converges in expectation to at most ∆. Second, convergence to ∆ + requires at most log((F (w 0 ) -F * ) -1 )/ log((1 -γ) -1 ) iterations. We note similar bounds exist, under a stronger variance assumption (Karimi et al., 2016; Reddi et al., 2016; De et al., 2017; Yin et al., 2018) . Importantly, our AdaScale bound converges to this same ∆ for all practical values of S: Theorem 2 (AdaScale bound). Define γ, ∆ as in Theorem 1. Given Assumptions 1, 2, 3, S ≤ γ -1 , and η ∈ (0, 2β -1 ), define w T as the result of Algorithm 2 with lr(t) = η and scale S. Then

E [F (w T ) -F * ] ≤ (1 -γ) TSI [F (w 0 ) -F * ] + ∆ .

This bound for AdaScale is scale invariant, as it does not depend on S. Like single-batch SGD, the suboptimality converges in expectation to at most ∆, but AdaScale achieves this for all scales. In addition, AdaScale speeds up training by a factor r = 1

T T -1 t=0 r t . That is, convergence to ∆ + requires at most r-1 log((F (w 0 ) -F * ) -1 )/ log((1 -γ) -1 ) iterations (since T SI ≤ τ T = rT ).

As a final comparison, we provide an analogous bound for linear scaling, which is not scale invariant: Theorem 3 (Bound for linear scaling rule). Define γ and ∆ as in Theorem 1. Given Assumptions 1, 2, 3, S ≤ γ -1 , and η ∈ (0, 2(Sβ) -1 ), consider Algorithm 1 with lr(t) = Sη. We have

E [F (w T ) -F * ] ≤ 1 -γ • 2-Sηβ 2-ηβ ST [F (w 0 ) -F * ] + 2-ηβ 2-Sηβ ∆ .

Unlike Theorem 2, this bound converges to a value that increases with S. In addition, a smaller range of learning rates guarantees convergence. In practical terms, this means that linear scaling often leads to worse model quality and greater risk of divergence, especially for large S. These differences appear throughout our empirical comparisons in §4.

Finally, we note both Theorem 2 and Theorem 3 require that S ≤ γ -1 . For practical problems, γ is small, and we can safely ignore this constraint. Otherwise single batch training would converge quickly, due to Theorem 1, and smaller scales would result in fast training.

6. RELATION TO PRIOR WORK:

While linear scaling with warm-up is perhaps the most popular scaling rule, researchers have considered a few alternative strategies. "Square root learning rate scaling" (Krizhevsky, 2014; Li et al., 2014; Hoffer et al., 2017; You et al., 2018) multiplies learning rates by the square root of the batch size increase. Across scales, this preserves the covariance of the SGD update. Establishing this invariant remains poorly justified, however, and often root scaling degrades model quality in practice (Goyal et al., 2017; Golmant et al., 2018; Jastrzębski et al., 2018) . AdaScale adapts learning rates by making η t E ḡt 2 invariant across scales, which results in our scale-invariant bound from §5. Finally, we might also consider model-specific scaling rules, such as LARS for CNNs (You et al., 2017) . AdaScale solves the general problem (P1), making AdaScale applicable to many models.

Many prior works have also considered the role of gradient variance in SGD. McCandlish et al. (2018) study the impact of gradient variance on scaling efficiency. These general findings also apply to AdaScale, as gradient variance similarly determines AdaScale's efficiency. Much like AdaScale, Johnson & Guestrin (2018) also adapt learning rates to lower amounts of gradient variance-in this case when using SGD with importance sampling. Because the variance reduction is relatively small in this setting, however, distributed training can have far greater impact on training times. Lastly, many algorithms also adapt to gradient moments for improved training, given a fixed amount of variance-see (Schaul et al., 2013; Kingma & Ba, 2015; Balles & Hennig, 2018) , just to name a few. AdaScale adapts learning rates across scales, which correspond to different amounts of gradient variance. Perhaps future algorithms will combine approaches in order to achieve both goals.

7. DISCUSSION:

SGD is not perfectly parallelizable. Unsurprisingly, the linear scaling rule can fail at large scales. In contrast, AdaScale accepts sublinear speedups in order to better preserve model quality. What do the speed-ups from AdaScale tell us about the scaling efficiency of SGD in general? For many problems, such as imagenet with batch size 32.8k, AdaScale establishes lower bounds on SGD's scaling efficiency. An important remaining question is whether AdaScale is optimally efficient, or if other practical algorithms can achieve similar scale invariance with fewer iterations.

AdaScale provides a useful new parameterization of learning rate schedules for large-batch SGD. We provide a simple lr schedule, which AdaScale adapts to learning rates for scaled training. From this, warm-up behavior emerges naturally, which produces quality models for many problems and scales. Even in elastic scaling settings, AdaScale adapts successfully to the state of training. Given these appealing qualities, it seems important to further study such learning rate schedules.

Based on our empirical results, as well as the algorithm's practicality and theoretical justification, we believe that AdaScale is valuable for speeding up training in practice. 

A PROOFS:

In this appendix, we prove the results from §5 and §3. We first prove a lemma in §A.1, which we apply in the proofs. We prove Theorem 1 in §A.2, Theorem 2 in §A.3, and Theorem 3 in §A.4. We also prove Proposition 1 in §A.5 and Proposition 2 in §A.6.

A.1 KEY LEMMA Lemma 1. Given Assumptions 1, 2, 3 and η ∈ (0, 2β -1 ), define γ = ηα(2-ηβ) and ∆ = 1 2γ η 2 βV . Consider Algorithm 2 with lr(t) = η. For all iterations t, we have

E [F (w t ) -F * ] ≤ [F (w 0 ) -F * ] t-1 t =0 (1 -r t γ) + ∆ .

Proof. We prove this by induction. To simplify notation, let us define F (w) = F (w) -F * . For t = 0, we have

E[ F (w 0 )] = F (w 0 ) ≤ F (w 0 ) -1 t =0

(1 -r t γ) + ∆ . Note we are using the convention -1

i=0 x i = 1. For t ≥ 1, assume the inductive hypothesis E[ F (w t-1 )] ≤ F (w 0 ) t-2 t =0

(1 -r t γ) + ∆ .

(1)

Applying Assumption 2 (smoothness) and the update equation w t = w t-1 -r t-1 ηḡ t-1 , we have

F (w t ) ≤ F (w t-1 ) + ∇F (w t-1 ), w t -w t-1 + β 2 w t -w t-1 2 = F (w t-1 ) -r t-1 η ∇F (w t-1 ), ḡt-1 + r 2 t-1 η 2 β 2 ḡt-1 2 .

Taking the expectation with respect to the S random batches from step t, we have

E F (w t ) | w t-1 ≤ F (w t-1 ) -r t-1 η ∇F (w t-1 ) 2 + r 2 t-1 η 2 β 2 E ḡt-1 2 | w t-1 .

Now taking the expectation with respect to the distribution of w t-1 , it follows that

EQUATION

)

For the last term, we have

EQUATION

Combining (3) with (2), we have

EQUATION

In this last step, we applied Assumption 1 (PL condition) and plugged in definitions for γ and ∆.

To complete the proof, we apply (1):

E[ F (w t )] ≤ (1 -r t-1 γ) F (w 0 ) t-2 t =0 (1 -r t γ) + ∆ + r t-1 γ∆ = F (w 0 ) t-1 t =0

(1 -r t γ) + ∆ .

A.2 PROOF OF THEOREM 1:

Theorem 1 (Single-batch SGD bound). Given Assumptions 1, 2, 3 and η ∈ (0, 2β -1 ), consider Algorithm 1 with S = 1 and lr(t) = η. Defining γ = ηα(2 -ηβ) and ∆ = 1 2γ η 2 βV , we have

E [F (w T ) -F * ] ≤ (1 -γ) T [F (w 0 ) -F * ] + ∆ .

Proof. The theorem is a special case of Lemma 1. In particular, Algorithm 1 with inputs lr(t) = η, S = 1, and T iterations is equivalent to Algorithm 2 with T SI = T and the same scale and learning rate inputs. This follows from the fact that r t = 1 for all iterations of AdaScale when S = 1. Thus, we can obtain the result by plugging t = T into the bound from Lemma 1.

A.3 PROOF OF THEOREM 2:

Theorem 2 (AdaScale bound). Define γ, ∆ as in Theorem 1. Given Assumptions 1, 2, 3, S ≤ γ -1 , and η ∈ (0, 2β -1 ), define w T as the result of Algorithm 2 with lr(t) = η and scale S. Then

E [F (w T ) -F * ] ≤ (1 -γ) TSI [F (w 0 ) -F * ] + ∆ .

Proof. Let T denote the total iterations for Algorithm 2. Applying Lemma 1, we have

E [F (w t ) -F * ] ≤ (F (w 0 ) -F * ) T -1 t =0

(1 -r t γ) + ∆ .

(5)

Now note that for any r ≥ 1 and x ∈ [0, 1], we have

EQUATION

This holds because for any r ≥ 1 and x ∈ [0, 1], the function (1 -x) r is convex in x, and 1 -rx is tangent to this function at x = 0. Thus,

EQUATION

Note that this requires 1 -r t γ ≥ 0 for all t, which is true because r t ≤ S ≤ γ -1 . Now plugging (7) into ( 5),

E [F (w t ) -F * ] ≤ (F (w 0 ) -F * )(1 -γ) T -1 t =0 r t + ∆ ≤ (F (w 0 ) -F * )(1 -γ) TSI + ∆ .

In the last step, we use the stopping condition of Algorithm 2 (T SI ≤ τ T = T -1 t=0 r t ).

A.4 PROOF OF THEOREM 3:

Theorem 3 (Bound for linear scaling rule). Define γ and ∆ as in Theorem 1. Given Assumptions 1, 2, 3, S ≤ γ -1 , and η ∈ (0, 2(Sβ) -1 ), consider Algorithm 1 with lr(t) = Sη. We have

E [F (w T ) -F * ] ≤ 1 -γ • 2-Sηβ 2-ηβ ST [F (w 0 ) -F * ] + 2-ηβ 2-Sηβ ∆ .

Proof. We reduce the theorem to a special case of Theorem 1. Define x = (x (1) , . . . , x(S) ), where x(i) ∼ X for each i ∈ [S], and x(1) , . . . , x(S) are jointly independent. Denote by X the distribution of x. Also define

f (w, x) = 1 S S i=1 f (w, x(i) ) .

It follows that for any w,

E x ∇ f (w, x) -∇F (w) 2 = 1 S σ 2 (w) ≤ V S .

The algorithm described in Theorem 3 is identical to running Algorithm 1 with scale 1, batch distribution X , loss f , learning rate lr(t) = Sη, and variance upper bound V S . Plugging these values into Theorem 1, we have

E [F (w T ) -F * ] ≤ (1 -Sηα(2 -Sηβ)) T [F (w 0 ) -F * ] + SηβV S -1 2α (2 -Sηβ) = 1 -Sγ • 2-Sηβ 2-ηβ T [F (w 0 ) -F * ] + 2-ηβ 2-Sηβ ∆ ≤ 1 -γ • 2-Sηβ 2-ηβ ST [F (w 0 ) -F * ] + 2-ηβ 2-Sηβ ∆ .

The last step follows from (6).

A.5 PROOF OF PROPOSITION 1:

Proposition 1 (Scale-invariant SGD for deterministic gradients). If σ 2 (w) = 0 for all w ∈ R d , then applying identity scaling to Algorithm 1 results in a scale-invariant algorithm.

Proof. Since the gradient variance is zero, the compute_gradient function returns ∇F (w t ), which does not depend on S. Thus, the algorithm does not depend on S in this case, which implies that it is scale-invariant.

A.6 PROOF OF PROPOSITION 2

Proposition 2 (Scale-invariant SGD for extreme stochasticity). Consider fixed covariance matrix Σ ∈ S d ++ , learning rate value η ∈ R >0 , and training duration T . For a given ν ∈ R >0 , assume ∇ w f (w, x) ∼ N (∇F (w), ν Σ), and apply linear scaling to Algorithm 1 with lr S1 (t) = ν -1 η and T S1 = ν T . The resulting scaled SGD algorithm is scale-invariant in the limit ν → +∞.

Proof. The scaled SGD algorithm runs for ν T /S iterations and follows the update rule

w t+1 = w t -S η ν ∇F (w t ) + S η ν ξ t .

Here ξ t is normally distributed with E [ξ t ] = 0 and cov(ξ t , ξ t ) = ν S Σ. In the limit ν → +∞, this difference equation converges to a stochastic differential equation on the interval [0, η T ] (Kloeden & Platen, 1992, Chapter 9): dw = -∇F (w)dt + (η Σ) 1/2 dW(t) , where W(t) ∼ N (0, I) .

Since this SDE does not depend on S, the algorithm is scale-invariant in this limit.

B ADDITIONAL DETAILS ON EMPIRICAL COMPARISONS:

This appendix provides additional details of our experiment set-up.

B.1 LEARNING RATE SCHEDULES:

We describe the lr schedules for each training benchmark in Table 3 . We use two learning rate families: exponential decay and step decay. Using parameters η 0 , d, and w i , we define lr(t) = η 0 d (t/T S1 ) for exponential decay families and lr(t) = η 0 d n i 1[t>wi] for step decay families. Here T S1 denotes the total iterations for scale S = 1. Note that in all cases, we use simple schedules and no warm-up. For imagenet and yolo, we used standard learning rate schedules from (Goyal et al., 2017) and (Zhang et al., 2019) . For cifar10, speech, and transformer, we chose learning rate parameters, via hand-tuning, that approximately maximized model quality. This was necessary for speech and transformer, since our reference implementations train with the Adam optimizer (Kingma & Ba, 2015) , and momentum-SGD requires different learning rate values.

B.2 WARM-UP IMPLEMENTATION:

Our warm-up procedure closely follows the strategy of Goyal et al. (2017) . We apply warm-up for the first 5.5% of training iterations-we denote this number by W S . During warm-up, the learning rate increases linearly, starting at the initial learning rate for single-batch training and finishing at S times this value. After warm-up, we apply linear scaling to the single-batch schedule. Following Goyal et al. (2017) , we modify this scaled schedule so that the total iterations, including warm-up, is proportional to S -1 . For step-decay schedules, we omit the first W S iterations after warm-up.

For exponential decay schedules, we compress the scaled schedule by W S iterations, using slightly faster decay.

B.3 BENCHMARK-SPECIFIC IMPLEMENTATION DETAILS:

Here we describe implementation details that are specific to each benchmark task.

B.3.1 cifar10:

We train ResNet-18 (preactivation) models (He et al., 2016b) , using the standard training data split for CIFAR-10 (Krizhevsky, 2009) . We use weight decay = 5 × 10 -4 . For batch normalization, we use parameters momentum = 0.995 and = 2 × 10 -5 , and we do not train the batch normalization scaling parameters. We apply standard data augmentation during training. Specifically, we pad images to 40 × 40 and random crop to 32 × 32, and we also apply random horizontal reflections.

B.3.2 imagenet:

For ImageNet classification (Deng et al., 2009) , we train ResNet-50 models (He et al., 2016a) .

Our implementation closely follows the implementation of Goyal et al. (2017) . We use stride-2 convolutions on 3 × 3 layers. For each block's final batch normalization layer, we initialize the batch norm scaling parameters to 0 (and we initialize to 1 everywhere else). We use weight decay parameter 10 -4 . Since each GPU processes 128 examples per batch, we use ghost batch normalization (Hoffer et al., 2017) with ghost batch size 32. We resize input images to 224×224×3.

For data augmentation, we apply random cropping and left-right mirroring during training.

B.3.3 speech:

We use Amodei et al. (2016) 's Deep Speech 2 model architecture. The model consists of two 2D convolutional input layers, five bidirectional RNN layers, one fully connected layer, and softmax outputs. Each convolutional layer has 32 filters. The RNN layers use GRU cells with hidden size 800. We apply batch normalization to the inputs of each layer. The batch norm parameters are momentum = 0.997 and = 10 -5 . The loss is CTC loss. The inputs to the network are log spectrograms, which we compute using 20ms windows from audio waveforms sampled at 16 kHz.

The training data is the train-clean-100 and train-clean-360 partitions of the OpenSLR LibriSpeech Corpus, which amounts to 460 hours of recorded speech. We evaluate models on the dev-clean partition.

B.3.4 transformer:

We train Transformer base models (Vaswani et al., 2017) . We use dynamic batching with at most 256 tokens per example. In Table 2 , the "batch size" is the maximum number of tokens processed per iteration. Our implementation closely follows that of Vaswani et al. (2017) . Unlike Vaswani et al., we use only the final model for evaluation instead of the average of the last five checkpoints. We train on the WMT 2014 English-German dataset and evaluate on the newstest2014 test set.

B.3.5 yolo:

We train YOLOv3 models (Redmon & Farhadi, 2018) . To achieve high mAP scores, we also apply mixup (Zhang et al., 2018) and class label smoothing, following (Zhang et al., 2019) . We also use focal loss. We use batch normalization momentum= 0.9 and weight decay = 5 × 10 -4 . We resize input images to 416×416 (for both training and validation). We report mAP values at IOU threshold 0.5. We use the Pascal VOC 2007 trainval and 2012 trainval datasets for training and the 2007 test set for validation (Everingham et al., 2010) . During training, we initialize the darknet-53 convolutional layers with weights trained on ImageNet.

B.4 MISCELLANEOUS:

In practice, wall time speed-ups also depend on system scaling efficiency. Since most aspects of system scaling relate orthogonally to the training algorithm, we limit our scope to algorithmic aspects of training.

For Figure 5 , one dimension defines initial value lr(0), and the second dimension specifies total decrease lr(T SI )/lr(0). For single-batch training, we use T = 39.1 × 10 3 steps. We run AdaScale and the LW baseline at S = 16, and we compare the final validation accuracies.

C ROBUSTNESS TO AVERAGING PARAMETER:

In this appendix, we test the robustness of AdaScale to the averaging parameter θ for estimating gain ratios (see §3.3). When θ = 0, AdaScale does not average estimates of gradient moments. The closer θ is to 1, the more that AdaScale averages across iterations.

Using the cifar10 benchmark, we compare four values of θ at scales S = 8 and S = 32. The case θ = 1 -S/1000 corresponds to the cifar10 experiment for Figure 1 . We average the resulting metrics over five trials. Figure 6 contains the training curves.

We also include final metric values in Table 4 .

For the three smaller settings of θ, the results align very closely. This suggests that AdaScale is robust to the choice of θ. When θ = 1 -S/10000, we see that smoothing more significantly biases gain ratio estimates, which leads to more contrasting results.

D ADDITIONAL EMPIRICAL RESULTS:

This appendix provides additional empirical results. As can be seen from empricial results in §4, AdaScale takes more steps than LSW for all scales. In order to understand the contribution of increased step number to the improved performance of Ada-Scale, we compare AdaScale to another algorithm called LSW+. LSW+ runs for the same number of steps as AdaScale. It scales the steps axis of the LSW learning rate schedule while keeping the learning rate axis the same. Specifically, it takes single-batch schedule lr S1 = lr and steps T S1 = T AS as inputs, where T AS is the average number of iterations (over five trials) taken by AdaScale. It applys warm-up to the first 5.5% of the T AS iterations. During warm-up, the learning rate increases linearly from lr S1 (0) to S • lr S1 (0).

As can be seen from Table 5 , the behavior of LSW+ is generally similar to that of LSW. As expected, LSW+ improves upon LSW, but LSW+ still degrades model quality at larger scales for all problems.

For speech, transformer, and yolo, LSW+ diverges at the largest scales.

We also note that LSW+ is not a practical algorithm, because it requires either (i) first running AdaScale to determine the number of iterations; or (ii) tuning the number of iterations. The second path is inconvenient in practice. Moreover, it seems for a fair comparison, we would also need to consider AdaScale with tuning. Thus, even if LSW+ matched AdaScale, which it does not, AdaScale would still be preferable to LSW+. 

:



Figures and Tables:

Figure 1: Motivating results. For cifar10, AdaScale preserves model quality for many scales S. When plotted in terms of scale-invariant iterations, training curves align closely. With AdaScale, "warm-up" behavior emerges from adapting a simple learning rate schedule (exponential decay) to scale S (learning rate plot cropped to show behavior). Meanwhile, linear scaling (with warm-up heuristic) degrades model quality as S increases.

Figure 3: AdaScale training curves. For many scales and benchmarks, AdaScale trains quality models.Training curves align closely in terms of τt. In all cases, ηt warms up gradually at the start of training, even though all lr schedules are simple exponential or step decay functions (which are non-increasing in t).

Figure 3 (and Figure 1) contains AdaScale training curves for the benchmarks and many scales. Each curve plots the mean of five distributed training runs with varying random seeds. As S increases, AdaScale trains for fewer iterations but consistently preserves model quality. Illustrating AdaScale's approximate scale invariance, the training curves align closely when plotted in terms of scale-invariant iterations.

Figure5: Scale invariance for many learning rate schedules. Heat maps cover the space of exponential decay lr schedules for cifar10. At scale 16, validation accuracies for AdaScale align closely with results for single-batch training, with the space of 94+% schedules growing moderately with AdaScale. With LSW, no schedule achieves 94% accuracy. On the right, direct lr search at scale 16 produces inferior results to AdaScale (here the total iterations, 3.28k, is the average total iterations among 94+% AdaScale trials). Thus, AdaScale induces a superior family of schedules for scaled training. The white '×' indicates the lr used for Figure1.

Figure 6: AdaScale training curves with varying moving average parameter.

Figure 7: Gain ratios for transformer. Plots compare moving average rt estimates to values computed offline (using 1000 batches).

t , and we must approximate r t in practice (see§3.3). This definition of r t ensures that as S increases, E[ w t+1 -w t , ∇F (w t ) ] and E[ w t+1 -w t 2 ] increase multiplicatively by r t . This leads to our scale-invariant bound in §5.

Table content: <table><tr><td>3.3 PRACTICAL CONSIDERATIONS</td><td/><td/><td/></tr><tr><td/><td colspan="2">imagenet, S=16</td><td colspan="2">imagenet, S=128</td></tr><tr><td>Gain ratio rt</td><td>0 4 8 12 16</td><td/><td>0 32 64 96 128</td></tr><tr><td/><td>0</td><td>10k 20k 30k</td><td>0</td><td>2k 4k 6k</td></tr><tr><td/><td/><td>Iteration t</td><td/><td>Iteration t</td></tr><tr><td/><td/><td>Offline evaluation</td><td/><td>Online estimate</td></tr></table>

Table 1 summarizes our training benchmarks. Due to space limitations, we provide additional implementation details in Appendix B.

Table content: <table/>

Overview of training benchmarks. For imagenet, AdaScale is approximately scale invariant, even if S changes abruptly (at τt = 133k, 225k). Unlike AdaScale, LSW degrades model quality in this setting (see Table2). Elastic scaling comparisons consider one random trial; future versions of this work will include five trials.

Table content: <table><tr><td>Name</td><td>Task</td><td>Model</td><td>Dataset</td><td>Metric</td></tr><tr><td>cifar10</td><td colspan="2">Image classification ResNet-18 (v2)</td><td>CIFAR-10</td><td>Top-1 accuracy (%)</td></tr><tr><td>imagenet</td><td colspan="2">Image classification ResNet-50 (v1)</td><td>ImageNet</td><td>Top-1 accuracy (%)</td></tr><tr><td>speech</td><td>Speech recognition</td><td>Deep speech 2</td><td>LibriSpeech</td><td>Word accuracy (%)</td></tr><tr><td colspan="4">transformer Machine translation Transformer base WMT-2014</td><td>BLEU</td></tr><tr><td>yolo</td><td>Object detection</td><td>YOLOv3</td><td colspan="2">PASCAL VOC mAP (%)</td></tr></table>



Table content: <table><tr><td>Task</td><td>S</td><td>Total batch size</td><td colspan="2">Validation metric AS LSW</td><td colspan="2">Training loss AS LSW</td><td colspan="2">Total iterations AS LSW</td></tr><tr><td>cifar10</td><td>1</td><td>128</td><td>94.1</td><td>94.1</td><td>0.157</td><td>0.157</td><td>39.1k</td><td>39.1k</td></tr><tr><td/><td>8</td><td>1.02k</td><td>94.1</td><td>94.0</td><td>0.153</td><td>0.161</td><td>5.85k</td><td>4.88k</td></tr><tr><td/><td>16</td><td>2.05k</td><td>94.1</td><td>93.6</td><td>0.150</td><td>0.163</td><td>3.36k</td><td>2.44k</td></tr><tr><td/><td>32</td><td>4.10k</td><td>94.1</td><td>92.8</td><td>0.145</td><td>0.177</td><td>2.08k</td><td>1.22k</td></tr><tr><td/><td>64</td><td>8.19k</td><td>93.9</td><td>76.6</td><td>0.140</td><td>0.272</td><td>1.41k</td><td>611</td></tr><tr><td>imagenet</td><td>1</td><td>256</td><td>76.4</td><td>76.4</td><td>1.30</td><td>1.30</td><td>451k</td><td>451k</td></tr><tr><td/><td>16</td><td>4.10k</td><td>76.5</td><td>76.3</td><td>1.26</td><td>1.31</td><td>33.2k</td><td>28.2k</td></tr><tr><td/><td>32</td><td>8.19k</td><td>76.6</td><td>76.1</td><td>1.23</td><td>1.33</td><td>18.7k</td><td>14.1k</td></tr><tr><td/><td>64</td><td>16.4k</td><td>76.5</td><td>75.6</td><td>1.19</td><td>1.35</td><td>11.2k</td><td>7.04k</td></tr><tr><td/><td>128</td><td>32.8k</td><td>76.5</td><td>73.3</td><td>1.14</td><td>1.51</td><td>7.29k</td><td>3.52k</td></tr><tr><td/><td>Elastic↑</td><td>various</td><td>76.8</td><td>75.7</td><td>1.15</td><td>1.36</td><td>11.6k</td><td>7.04k</td></tr><tr><td/><td>Elastic↓</td><td>various</td><td>76.6</td><td>73.8</td><td>1.23</td><td>1.46</td><td>13.7k</td><td>9.68k</td></tr><tr><td>speech</td><td>1</td><td>32</td><td>79.6</td><td>79.6</td><td>2.03</td><td>2.03</td><td>84.8k</td><td>84.8k</td></tr><tr><td/><td>4</td><td>128</td><td>81.0</td><td>80.9</td><td>5.21</td><td>4.66</td><td>22.5k</td><td>21.2k</td></tr><tr><td/><td>8</td><td>256</td><td>80.7</td><td>80.2</td><td>6.74</td><td>6.81</td><td>12.1k</td><td>10.6k</td></tr><tr><td/><td>16</td><td>512</td><td>80.6</td><td>N/A</td><td>7.33</td><td>N/A</td><td>6.95k</td><td>5.30k</td></tr><tr><td/><td>32</td><td>1.02k</td><td>80.3</td><td>N/A</td><td>8.43</td><td>N/A</td><td>4.29k</td><td>2.65k</td></tr><tr><td>transformer</td><td>1</td><td>2.05k</td><td>27.2</td><td>27.2</td><td>1.60</td><td>1.60</td><td>1.55M</td><td>1.55M</td></tr><tr><td/><td>16</td><td>32.8k</td><td>27.4</td><td>27.3</td><td>1.60</td><td>1.60</td><td>108k</td><td>99.0k</td></tr><tr><td/><td>32</td><td>65.5k</td><td>27.3</td><td>27.0</td><td>1.59</td><td>1.61</td><td>58.9k</td><td>49.5k</td></tr><tr><td/><td>64</td><td>131k</td><td>27.6</td><td>26.7</td><td>1.59</td><td>1.63</td><td>33.9k</td><td>24.8k</td></tr><tr><td/><td>128</td><td>262k</td><td>27.4</td><td>N/A</td><td>1.59</td><td>N/A</td><td>21.4k</td><td>12.1k</td></tr><tr><td>yolo</td><td>1</td><td>16</td><td>80.2</td><td>80.2</td><td>2.65</td><td>2.65</td><td>207k</td><td>207k</td></tr><tr><td/><td>16</td><td>256</td><td>81.5</td><td>81.4</td><td>2.63</td><td>2.66</td><td>15.9k</td><td>12.9k</td></tr><tr><td/><td>32</td><td>512</td><td>81.3</td><td>80.5</td><td>2.61</td><td>2.81</td><td>9.27k</td><td>6.47k</td></tr><tr><td/><td>64</td><td>1.02k</td><td>81.3</td><td>70.1</td><td>2.60</td><td>4.02</td><td>5.75k</td><td>3.23k</td></tr><tr><td/><td>128</td><td>2.05k</td><td>81.4</td><td>N/A</td><td>2.57</td><td>N/A</td><td>4.07k</td><td>1.62k</td></tr></table>

Z. Yuan, Y. Yan, R. Jin, and T. Yang. Stagewise training accelerates convergence of testing error over sgd. In Advances in Neural Information Processing Systems 32, 2019. H. Zhang, M. Cisse, Y. N. Dauphin, and D. Lopez-Paz. Mixup: Beyond empirical risk minimization. In International Conference on Learning Representations, 2018.

Table content: <table><tr><td>Z. Zhang, T. He, H. Zhang, Z. Zhang, J. Xie, and M. Li. Bag of freebies for training object detection</td></tr><tr><td>neural networks. arXiv:1902.04103, 2019.</td></tr></table>

Learning rate schedules for training benchmarks.

Table content: <table><tr><td>Benchmark</td><td colspan="2">Learning rate famliy η0</td><td>d</td><td>wi</td></tr><tr><td>cifar10</td><td>Exponential decay</td><td>0.08</td><td colspan="2">0.0133 N/A</td></tr><tr><td>imagenet</td><td>Step decay</td><td>0.1</td><td>0.1</td><td>150,240, 300,480, 400,640</td></tr><tr><td>speech</td><td>Exponential decay</td><td colspan="2">1.4 ×10 -3 0.05</td><td>N/A</td></tr><tr><td colspan="2">transformer Step decay</td><td>0.01</td><td>0.1</td><td>1,440,000</td></tr><tr><td>yolo</td><td>Step decay</td><td colspan="2">2.5 ×10 -4 0.1</td><td>160,000, 180,000</td></tr></table>

AdaScale final metrics with varying moving average parameter.

Table content: <table><tr><td>S</td><td>θ</td><td>Final val. accuracy (%)</td><td>Final train objective</td><td>Total iterations</td></tr><tr><td>8</td><td>1 -S/10</td><td>94.0</td><td>0.153</td><td>5.75k</td></tr><tr><td/><td>1 -S/100</td><td>94.1</td><td>0.154</td><td>5.78k</td></tr><tr><td/><td>1 -S/1000</td><td>94.1</td><td>0.153</td><td>5.85k</td></tr><tr><td/><td>1 -S/10000</td><td>94.1</td><td>0.147</td><td>6.45k</td></tr><tr><td>32</td><td>0</td><td>94.0</td><td>0.145</td><td>2.02k</td></tr><tr><td/><td>1 -S/100</td><td>94.1</td><td>0.147</td><td>2.03k</td></tr><tr><td/><td>1 -S/1000</td><td>94.1</td><td>0.145</td><td>2.08k</td></tr><tr><td/><td>1 -S/10000</td><td>94.1</td><td>0.136</td><td>2.46k</td></tr><tr><td colspan="5">D.1 EFFECT OF NUMBER OF STEPS ON LINEAR SCALING WITH WARMUP</td></tr></table>

Comparison of AS and LSW+. Shorthand: AS=AdaScale, LSW+=Stretched linear scaling rule with warm-up, takes the same number of steps as AS gray=model quality significantly worse than for S = 1 (5 trials, 0.95 significance), N/A=training diverges

Table content: <table><tr><td>Task</td><td>S</td><td>Total batch size</td><td colspan="2">Validation metric AS LSW+</td><td colspan="2">Training loss AS LSW+</td><td colspan="2">Total iterations AS LSW+</td></tr><tr><td>cifar10</td><td>1</td><td>128</td><td>94.1</td><td>94.1</td><td>0.157</td><td>0.157</td><td>39.1k</td><td>39.1k</td></tr><tr><td/><td>8</td><td>1.02k</td><td>94.1</td><td>94.0</td><td>0.153</td><td>0.145</td><td>5.85k</td><td>5.85k</td></tr><tr><td/><td>16</td><td>2.05k</td><td>94.1</td><td>94.1</td><td>0.150</td><td>0.136</td><td>3.36k</td><td>3.36k</td></tr><tr><td/><td>32</td><td>4.10k</td><td>94.1</td><td>94.0</td><td>0.145</td><td>0.128</td><td>2.08k</td><td>2.08k</td></tr><tr><td/><td>64</td><td>8.19k</td><td>93.9</td><td>93.0</td><td>0.140</td><td>0.128</td><td>1.41k</td><td>1.41k</td></tr><tr><td>imagenet</td><td>1</td><td>256</td><td>76.4</td><td>76.4</td><td>1.30</td><td>1.30</td><td>451k</td><td>451k</td></tr><tr><td/><td>16</td><td>4.10k</td><td>76.5</td><td>76.5</td><td>1.26</td><td>1.27</td><td>33.2k</td><td>33.2k</td></tr><tr><td/><td>32</td><td>8.19k</td><td>76.6</td><td>76.4</td><td>1.23</td><td>1.24</td><td>18.7k</td><td>18.7k</td></tr><tr><td/><td>64</td><td>16.4k</td><td>76.5</td><td>76.5</td><td>1.19</td><td>1.20</td><td>11.2k</td><td>11.2k</td></tr><tr><td/><td>128</td><td>32.8k</td><td>76.5</td><td>75.5</td><td>1.14</td><td>1.20</td><td>7.29k</td><td>7.29k</td></tr><tr><td>speech</td><td>1</td><td>32</td><td>79.6</td><td>79.6</td><td>2.03</td><td>2.03</td><td>84.8k</td><td>84.8k</td></tr><tr><td/><td>4</td><td>128</td><td>81.0</td><td>81.0</td><td>5.21</td><td>4.22</td><td>22.5k</td><td>22.5k</td></tr><tr><td/><td>8</td><td>256</td><td>80.7</td><td>80.7</td><td>6.74</td><td>6.61</td><td>12.1k</td><td>12.1k</td></tr><tr><td/><td>16</td><td>512</td><td>80.6</td><td>N/A</td><td>7.33</td><td>N/A</td><td>6.95k</td><td>6.95k</td></tr><tr><td/><td>32</td><td>1.02k</td><td>80.3</td><td>N/A</td><td>8.43</td><td>N/A</td><td>4.29k</td><td>4.29k</td></tr><tr><td>transformer</td><td>1</td><td>2.05k</td><td>27.2</td><td>27.2</td><td>1.60</td><td>1.60</td><td>1.55M</td><td>1.55M</td></tr><tr><td/><td>16</td><td>32.8k</td><td>27.4</td><td>27.4</td><td>1.60</td><td>1.59</td><td>108k</td><td>108k</td></tr><tr><td/><td>32</td><td>65.5k</td><td>27.3</td><td>27.3</td><td>1.59</td><td>1.59</td><td>58.9k</td><td>58.9k</td></tr><tr><td/><td>64</td><td>131k</td><td>27.6</td><td>27.1</td><td>1.59</td><td>1.60</td><td>33.9k</td><td>33.9k</td></tr><tr><td/><td>128</td><td>262k</td><td>27.4</td><td>N/A</td><td>1.59</td><td>N/A</td><td>21.4k</td><td>21.4k</td></tr><tr><td>yolo</td><td>1</td><td>16</td><td>80.2</td><td>80.2</td><td>2.65</td><td>2.65</td><td>207k</td><td>207k</td></tr><tr><td/><td>16</td><td>256</td><td>81.5</td><td>81.9</td><td>2.63</td><td>2.47</td><td>15.9k</td><td>15.9k</td></tr><tr><td/><td>32</td><td>512</td><td>81.3</td><td>81.7</td><td>2.61</td><td>2.42</td><td>9.27k</td><td>9.27k</td></tr><tr><td/><td>64</td><td>1.02k</td><td>81.3</td><td>80.6</td><td>2.60</td><td>2.51</td><td>5.75k</td><td>5.75k</td></tr><tr><td/><td>128</td><td>2.05k</td><td>81.4</td><td>N/A</td><td>2.57</td><td>N/A</td><td>4.07k</td><td>4.07k</td></tr></table>

Peer review (iclr2020):
Authors use the PL condition to motivate a formula for learning rate selection. (they bound their loss by a quadratic with fixed curvature alpha, and use this to get learning rate). Their analysis very closely mirrors one presented in *Gradient Diversity* paper, it uses the same assumptions on the loss. IE compare A.1 in the paper to the B.5 in *Gradient Diversity* Gradient Diversity solves for batch size B after which linear learning rate scaling starts to break down, while this paper instead fixes B and solves for learning rate. Two results are comparable, if you take their learning rate formula in section 3.2 (need formula numbers) and solve for B which gives learning rate halfway between B=1 and B=inf, you get the same expression as *critical batch size* in equation 5 of gradient diversity paper. It*s not immediately obvious how to invert formula in Gradient Diversity paper to solve for learning rate, so I would consider their learning rate formula an interesting result. I also appreciate the breadth of experiments used for evaluation. The biggest issue I have with the paper is that I can*t tell if it*s better of worse than linear learning rate scaling from their experiment section. All of their experiments use more iterations for AS evaluation uses than for LSW evaluation. They demonstrate better training (and test) losses for AS, but because of extra iterations, I can*t tell that the improvement in training loss is due to number of iterations, or due to AS scaling.