Organizations: Chandar Research Lab · Mila – Quebec AI Institute · Université de Montréal · Stanford University · Arizona State University · Polytechnique Montréal · Canada CIFAR AI Chair
Large neural networks are commonly trained across many workers, and communication between them is a major cost. Workers therefore take several local optimization steps before sending their net parameter displacements, or pseudo-gradients, to an outer optimizer, which often applies Nesterov momentum. We study how the outer momentum, the number of local steps and the averaging over workers control the training dynamics, and in particular their stability. For random quadratics, we derive an exact mean-square stability criterion in any dimension, which reveals that the inner batch size determines whether increasing the number of local steps eventually stabilizes or destabilizes the outer iteration. Moreover, while averaging the pseudo-gradients of independent workers can make larger outer momentum values admissible, the number of workers required for stability can grow exponentially with the number of local steps. Finally, we show that the mean-square dynamics governing stability can be probed directly at intermediate neural-network checkpoints using shared-minibatch perturbations. We test our theoretical findings across networks, and measure their consequences for training loss.
Figures & tables
Figure 1: Mean-square stable regions in (α,β) . For inner batch size b and inner-batch threshold bc (as defined in Section 5.2 ), more local steps enlarge the mean-square stable region when b>bc ( left ), and shrink it when b<bc ( middle ). Averaging the pseudo-gradient over W workers enlarges the mean-square stable region ( right ), but never past the mean boundary (gray), which averaging leaves unchanged.
Algorithm 1 Two-phase optimization with outer Nesterov momentum
Figure 2: The pseudo-gradient of K inner SGD steps, and the block operator. ( a ) From θt the inner trajectory takes K SGD steps on minibatches ξ0,…,ξK−1 and ends at θK,ω . The pseudo-gradient Pω(θt) is the displacement from θK,ω to θt . The outer step moves θt by −αPω(θt) to θt+1 . ( b ) The block operator DK,ω maps a displacement δ of the starting point to the change DK,ωδ of the pseudo-gradient. Each green arrow is the change DK,ωδ for one draw of the minibatches ω . The draws have mean DˉKδ , and the shaded ellipse marks their covariance. Mean-square stability depends on DK,ω only through DˉK and the second moment SK of Section 3.2 .
Figure 3: Mechanism behind each implication, on the model of Figure 1 . ( i ) The noise budget shrinks as β grows (Proposition 2 ). ( ii ) The block noise grows with K below bc and dies out above it (Proposition 3 ). ( iii ) Averaging divides the loop gain by W (Corollary 4 ). ( iv ) Splitting a fixed total batch raises ρ(LK,W) (Proposition 5 ). Green marks stable settings and red marks unstable ones.
Figure 5
Appendix figures & tables2 assets
Supplementary material from the paper’s appendix.
Appendix
Algorithm 2 The frozen linearization at a checkpoint θ0
Figure 6: Mean stable but mean-square unstable. The mean ∥E[zt]∥2 and the second moment E∥zt∥2 for the linear-regression model above ( η=2.5 , b=1 , K=4 , α=1 , β=0.5 ), starting from a unit error with an empty buffer. Both curves are computed exactly from the moment recurrences ( 9 ). The dotted line marks the initial value.