Year: 2019

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.

heading: 1 INTRODUCTION
text: 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; Jastrzę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, AdaScale 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.

heading: 2 PROBLEM FORMULATION
text: We focus on quickly computing approximate solutions to the problem
minimizew∈Rd F (w) , where F (w) = Ex∼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 Ex∼X [∇wf(w,x)] = ∇F (w). Stochastic gradient descent is a popular algorithm for solving (P1). Let wt denote the model parameters when iteration t begins. During this iteration, SGD samples a batch xt ∼ X and computes the gradient gt ← ∇wf(wt,xt). SGD then applies the update wt+1 ← wt − ηtgt. 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 gt) 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 wT denote the (possibly random) result of a scaled SGD algorithm. Fixing all algorithm inputs except scale S, the algorithm is scale invariant if wT 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 , w0)
for t = 0, 1, 2, . . . , T − 1 do ḡt ← compute_gradient(wt, S,X , f ) ηt ← lr(t) wt+1 ← wt − ηtḡt
return wT function compute_gradient(wt, S, X , f )
in parallel for i = 1, . . . , S do x(i) ← sample_batch(X ) g(i) ← ∇wf(wt,x(i))
return 1S ∑S i=1 g (i)
Algorithm 2 AdaScale SGD function AdaScale(S, lr, TSI, X , f , w0)
initialize τ0 ← 0; t← 0 while τt < TSI do
ḡt ← compute_gradient(wt, S,X , f )
# Compute gain rt ∈ [1, S] (see §3.3): rt ← E[σ2(wt)+‖∇F (wt)‖2]
E [ 1 S σ 2(wt)+‖∇F (wt)‖2 ]
ηt ← rt · lr(bτtc) wt+1 ← wt − ηtḡt τt+1 ← τt + rt; t← t+ 1
return wt

heading: 3 ADASCALE SGD ALGORITHM
text: 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.

heading: 3.1 INTUITION: IDENTITY SCALING, LINEAR SCALING, AND GRADIENT VARIANCE
text: 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 ·lrS1(St) and T = dTS1/Se, where lrS1 and TS1 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) = covx∼X (∇wf(w,x),∇wf(w,x)) , and σ2(w) = tr(Σ(w)) . In words, σ2(w) sums the variances of each entry in ∇wf(w,x). By sampling batches independently, scaling fundamentally impacts SGD by reducing this variance. Given wt in Algorithm 1, we have cov(ḡt, ḡt) = 1SΣ(wt) and E [ḡt] = ∇F (wt). 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 ∈ Rd, 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 ḡt. With “large” variance, however, the opposite is true: Proposition 2 (Scale-invariant SGD for extreme stochasticity). Consider fixed covariance matrix Σ̃ ∈ Sd++, learning rate value η̃ ∈ R>0, and training duration T̃ . For a given ν ∈ R>0, assume ∇wf(w,x) ∼ N (∇F (w), νΣ̃), and apply linear scaling to Algorithm 1 with lrS1(t) = ν−1η̃ and TS1 = ν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.

heading: 3.2 ADASCALE DEFINITION
text: AdaScale, defined in Algorithm 2, adaptively interpolates between identity and linear scaling, based on the expectation of σ2(wt). During iteration t, AdaScale multiplies the learning rate by the “gain ratio” rt ∈ [1, S]: ηt = rt · lr(bτtc). Here τt is the “scale-invariant iteration,” defined as τt = ∑t−1 t′=0 rt′ . The idea is that iteration t performs the equivalent of rt single-batch iterations, and τt accumulates this progress. AdaScale concludes when τt ≥ TSI, where TSI is the total scaleinvariant iterations. Since rt ∈ [1, S], AdaScale requires at least dTSI/Se and at most TSI iterations. The identity and linear rules correspond to two special cases of AdaScale. If rt = 1 for all t, the algorithm equates to SGD with identity scaling. Similarly, if rt = S for all t, we have linear scaling. Thus, to approximate scale-invariance, §3.1 suggests setting rt ≈ 1 when the gradient’s variance is small and rt ≈ S when this variance is large. AdaScale achieves this by defining
rt = E [ σ2(wt) + ‖∇F (wt)‖2 ]/ E [ 1 Sσ 2(wt) + ‖∇F (wt)‖2 ] .
The expectations here are with respect to the distribution of wt, and we must approximate rt in practice (see §3.3). This definition of rt ensures that as S increases, E[〈wt+1 −wt,∇F (wt)〉] and E[‖wt+1 −wt‖2] increase multiplicatively by rt. This leads to our scale-invariant bound in §5.
3.3 PRACTICAL CONSIDERATIONS imagenet, S=16 imagenet, S=128
If S = 1 in AdaScale, then rt = 1 for all iterations. For larger scales, rt depends on E [ σ2(wt) ] and E [ ‖∇F (wt)‖2 ] , and a practical implementation must efficiently approximate these values. Fortunately, the per-batch gradients g(1)t , . . . ,g (S) t and aggregated gradient ḡt are readily available in distributed SGD algorithms. This makes approximating rt straightforward. In particular, we define
σ̂2t = 1 S−1 ∑S i=1 ‖g (i) t ‖2 − SS−1 ‖ḡt‖ 2 ,
and µ̂2t = ‖ḡt‖ 2 − 1S σ̂ 2 t .
Here σ̂2t and µ̂ 2 t are unbiased estimates of E [
σ2(wt)
]
and E
[
‖∇F (wt)‖2
]
. To ensure robustness to estimation variance, we estimate rt by plugging in moving averages σ̄2t and µ̄ 2 t , which average σ̂2t 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 r0 ← 1, and for iterations t < (1 − θ)−1, we define σ̄2t and µ̄2t as the mean of past samples. Before averaging, we clip estimates so that σ̂2t ≥ 10−6 (to prevent division by zero) and µ̂2t ≥ 0 (to ensure rt ∈ [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).

heading: 4 EMPIRICAL COMPARISONS
text: 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). Table 1 summarizes our training benchmarks. Due to space limitations, we provide additional implementation details in Appendix B.
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.
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.
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 = TSI/4 and from 64 to 128 at τt = TSI/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 lrS1 = lr and single-batch steps TS1 = TSI, where lr and TSI are the inputs to AdaScale. Our warm-up implementation closely follows Goyal et al. (2017). LSW trains for dTS1/Se iterations, applying warm-up to the first 5.5% of iterations. During warm-up, the learning rate increases linearly from lrS1(0) to S · lrS1(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.

heading: 5 SCALE-INVARIANT CONVERGENCE BOUND
text: 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 ∗ = minw 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 ∗ ≤ 12α ‖∇F (w)‖ 2 for all w. Assumption 2 (β-smooth). For some β > 0, ‖∇F (w)−∇F (w′)‖ ≤ β‖w −w′‖ 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 ∆ = 12γ η 2βV , we have
E [F (wT )− F ∗] ≤ (1− γ)T [F (w0)− 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
dlog((F (w0)−F ∗) −1)/ log((1−γ)−1)e 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 wT as the result of Algorithm 2 with lr(t) = η and scale S. Then
E [F (wT )− F ∗] ≤ (1− γ)TSI [F (w0)− 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̄ = 1T ∑T−1 t=0 rt. That is, convergence to ∆ + requires at most dr̄−1 log((F (w0)− F ∗) −1)/ log((1− γ)−1)e iterations (since TSI ≤ τT = r̄T ). 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 (wT )− F ∗] ≤ ( 1− γ · ( 2−Sηβ 2−ηβ ))ST [F (w0)− 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.

heading: 6 RELATION TO PRIOR WORK
text: 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; Jastrzębski et al., 2018). AdaScale adapts learning rates by making ηtE [ ‖ḡ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.

heading: 7 DISCUSSION
text: 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.

heading: A PROOFS
text: 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.

heading: A.1 KEY LEMMA
text: Lemma 1. Given Assumptions 1, 2, 3 and η ∈ (0, 2β−1), define γ = ηα(2−ηβ) and ∆ = 12γ η 2βV . Consider Algorithm 2 with lr(t) = η. For all iterations t, we have
E [F (wt)− F ∗] ≤ [F (w0)− F ∗] ∏t−1 t′=0(1− rt′γ) + ∆ .
Proof. We prove this by induction. To simplify notation, let us define F̃ (w) = F (w) − F ∗. For t = 0, we have
E[F̃ (w0)] = F̃ (w0) ≤ F̃ (w0) ∏−1 t′=0(1− rt′γ) + ∆ .
Note we are using the convention ∏−1 i=0 xi = 1. For t ≥ 1, assume the inductive hypothesis
E[F̃ (wt−1)] ≤ F̃ (w0) t−2∏ t′=0 (1− rt′γ) + ∆ . (1)
Applying Assumption 2 (smoothness) and the update equation wt = wt−1 − rt−1ηḡt−1, we have
F̃ (wt) ≤ F̃ (wt−1) + 〈∇F (wt−1),wt −wt−1〉+ β2 ‖wt −wt−1‖ 2
= F̃ (wt−1)− rt−1η〈∇F (wt−1), ḡt−1〉+ r2t−1η2 β 2 ‖ḡt−1‖ 2 .
Taking the expectation with respect to the S random batches from step t, we have E [ F̃ (wt) | wt−1 ] ≤ F̃ (wt−1)− rt−1η ‖∇F (wt−1)‖2 + r2t−1η2 β 2E [ ‖ḡt−1‖2 | wt−1 ] .
Now taking the expectation with respect to the distribution of wt−1, it follows that E [ F̃ (wt) ] ≤ E [ F̃ (wt−1) ] − rt−1ηE [ ‖∇F (wt−1)‖2 ] + r2t−1η 2 β 2E [ ‖ḡt−1‖2 ] . (2)
For the last term, we have E [ ‖ḡt−1‖2 ] = E [ ‖(ḡt−1 −∇F (wt−1)) +∇F (wt−1)‖2 ] = E [ ‖ḡt−1 −∇F (wt−1)‖2 + ‖∇F (wt−1)‖2
] = E [ 1 Sσ 2(wt−1) + ‖∇F (wt−1)‖2 ]
= 1 rt−1 E [ σ2(wt−1) + ‖∇F (wt−1)‖2 ] ≤ 1 rt−1 ( E [ ‖∇F (wt−1)‖2 ] + V ) . (3)
Combining (3) with (2), we have E [ F̃ (wt) ] ≤ E [ F̃ (wt−1) ] − rt−1η(1− η β2 )E [ ‖∇F (wt−1)‖2 ] + rt−1η 2 β 2V
≤ (1− rt−1γ)E [ F̃ (wt−1) ] + rt−1γ∆ . (4)
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̃ (wt)] ≤ (1− rt−1γ) ( F̃ (w0)
t−2∏ t′=0 (1− rt′γ) + ∆
) + rt−1γ∆
= F̃ (w0) t−1∏ t′=0 (1− rt′γ) + ∆ .

heading: A.2 PROOF OF THEOREM 1
text: 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 ∆ = 12γ η 2βV , we have
E [F (wT )− F ∗] ≤ (1− γ)T [F (w0)− 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 TSI = T and the same scale and learning rate inputs. This follows from the fact that rt = 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.

heading: A.3 PROOF OF THEOREM 2
text: Theorem 2 (AdaScale bound). Define γ, ∆ as in Theorem 1. Given Assumptions 1, 2, 3, S ≤ γ−1, and η ∈ (0, 2β−1), define wT as the result of Algorithm 2 with lr(t) = η and scale S. Then
E [F (wT )− F ∗] ≤ (1− γ)TSI [F (w0)− F ∗] + ∆ .
Proof. Let T denote the total iterations for Algorithm 2. Applying Lemma 1, we have
E [F (wt)− F ∗] ≤ (F (w0)− F ∗) T−1∏ t′=0 (1− rt′γ) + ∆ . (5)
Now note that for any r ≥ 1 and x ∈ [0, 1], we have
1− rx ≤ (1− x)r. (6)
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,
T−1∏ t′=0 (1− rt′γ) ≤ (1− γ) ∑T−1 t′=0 rt′ . (7)
Note that this requires 1 − rtγ ≥ 0 for all t, which is true because rt ≤ S ≤ γ−1. Now plugging (7) into (5),
E [F (wt)− F ∗] ≤ (F (w0)− F ∗)(1− γ) ∑T−1 t′=0 rt′ + ∆
≤ (F (w0)− F ∗)(1− γ)TSI + ∆ .
In the last step, we use the stopping condition of Algorithm 2 (TSI ≤ τT = ∑T−1 t=0 rt).

heading: A.4 PROOF OF THEOREM 3
text: 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 (wT )− F ∗] ≤ ( 1− γ · ( 2−Sηβ 2−ηβ ))ST [F (w0)− 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,
Ex̃ [ ‖∇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 VS . Plugging these values into Theorem 1, we have
E [F (wT )− F ∗] ≤ (1− Sηα(2− Sηβ))T [F (w0)− F ∗] + SηβV S−1
2α (2− Sηβ) = ( 1− Sγ · (
2−Sηβ 2−ηβ ))T [F (w0)− F ∗] + ( 2−ηβ 2−Sηβ ) ∆
≤ ( 1− γ · (
2−Sηβ 2−ηβ ))ST [F (w0)− F ∗] + ( 2−ηβ 2−Sηβ ) ∆ .
The last step follows from (6).

heading: A.5 PROOF OF PROPOSITION 1
text: Proposition 1 (Scale-invariant SGD for deterministic gradients). If σ2(w) = 0 for all w ∈ Rd, 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 (wt), which does not depend on S. Thus, the algorithm does not depend on S in this case, which implies that it is scale-invariant.

heading: A.6 PROOF OF PROPOSITION 2
text: Proposition 2 (Scale-invariant SGD for extreme stochasticity). Consider fixed covariance matrix Σ̃ ∈ Sd++, learning rate value η̃ ∈ R>0, and training duration T̃ . For a given ν ∈ R>0, assume ∇wf(w,x) ∼ N (∇F (w), νΣ̃), and apply linear scaling to Algorithm 1 with lrS1(t) = ν−1η̃ and TS1 = νT̃ . The resulting scaled SGD algorithm is scale-invariant in the limit ν → +∞.
Proof. The scaled SGD algorithm runs for dνT̃ /Se iterations and follows the update rule
wt+1 = wt − Sη̃ν ∇F (wt) + 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/2dW(t) , where W(t) ∼ N (0, I) .
Since this SDE does not depend on S, the algorithm is scale-invariant in this limit.

heading: B ADDITIONAL DETAILS ON EMPIRICAL COMPARISONS
text: This appendix provides additional details of our experiment set-up.

heading: B.1 LEARNING RATE SCHEDULES
text: 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 wi, we define lr(t) = η0d(t/TS1) for exponential decay families and lr(t) = η0d ∑n i 1[t>wi] for step decay families. Here TS1 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.

heading: B.2 WARM-UP IMPLEMENTATION
text: 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 WS . 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 WS iterations after warm-up. For exponential decay schedules, we compress the scaled schedule by WS iterations, using slightly faster decay.

heading: B.3 BENCHMARK-SPECIFIC IMPLEMENTATION DETAILS
text: 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 stride2 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.

heading: B.4 MISCELLANEOUS
text: 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(TSI)/lr(0). For single-batch training, we use T = 39.1× 103 steps. We run AdaScale and the LW baseline at S = 16, and we compare the final validation accuracies.

heading: C ROBUSTNESS TO AVERAGING PARAMETER
text: 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.

heading: D ADDITIONAL EMPIRICAL RESULTS
text: This appendix provides additional empirical results.

heading: D.1 EFFECT OF NUMBER OF STEPS ON LINEAR SCALING WITH WARMUP
text: 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 AdaScale, 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 lrS1 = lr and steps TS1 = TAS as inputs, where TAS is the average number of iterations (over five trials) taken by AdaScale. It applys warm-up to the first 5.5% of the TAS iterations. During warm-up, the learning rate increases linearly from lrS1(0) to S · lrS1(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+.

heading: D.2 GAIN RATIO ESTIMATION
text: Our online gain ratio estimates align closely with offline estimates (computed by averaging over 1000 batches). Figure 7 demonstrates this for the transformer task.
D.3 cifar10 SCALE INVARIANCE CURVES
Figure 8 shows additional plots for the cifar10 task. Notably, training loss curves at various scales and full view of the learning rate curves are shown.

heading: D.4 ELASTIC SCALING
text: Learning rate and gain ratio curves for the two dynamic scaling scenarios we consider (discussed in §4) align surprisingly well with the corresponding curves for the scenarios where the scale is kept constant throughout the training. This is shown in Figure 9. The abrupt change in scale causes the gain ratio to change quickly which in turn leads to an almost immediate change in learning rate. This allows the algorithm to quickly adapt to varying scales.

Review summaries:
This paper proposes a novel rule for scaling the learning rate, called the gain ratio, for when the effective batch size (induced by synchronous distributed SGD) is increased. This rule for the gain ratio uses estimates of the trace of the covariance and the norm of the true gradient to determine the appropriate steplength. This results in a method with a R-linear convergence guarantee to a neighborhood that is not dependent on $S$ (which is called scale-invariant). The algorithm additionally tracks the sum of the gain ratios in order to determine the "effective" number of iterations taken, and cut the steplength appropriately.