AdaScale SGD: A Scale-Invariant Algorithm for Distributed 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 ; 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 . 2 PROBLEM FORMULATION . 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 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 ·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 can not 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 .

Summary: 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.