Don't Use Large Mini-Batches, Use Local SGD

In plain words

Training a neural network on many machines at once is usually done by giving every machine a big chunk of data per step and then combining their results. Making those chunks very large speeds things up, but the resulting model often does worse on new data than one trained with smaller chunks.

The paper's fix is simple. Train normally at first, then let each machine work on its own for a few steps before the machines agree on a shared model. Talking to the other machines less often saves communication, and the extra randomness from working separately seems to help the model generalize and settle in flatter spots. The authors also study how much communication can be saved by working separately for longer, and they sketch a variant for machines that are arranged in groups with fast links inside each group and slow links between groups.

TL;DR

• Large-batch SGD on many workers generalizes worse than small-batch SGD, even when optimization looks fine.
• Post-local SGD runs mini-batch SGD until the first learning-rate decay, then local SGD with H local steps between averages.
• Author-reported: it matches or beats small-batch SGD on CIFAR-10/100 for three CNNs, beats the large-batch baseline on ImageNet ResNet-50 at K·B = 4096 and 8192, and needs fewer communication rounds than mini-batch SGD after the switch.
• Costs: the first phase still communicates every step, and the gains depend on the first-decay switch point and on tuned learning rates for each setting.

Problem

Synchronous mini-batch SGD scales by raising the global batch size with the number of workers, but large batches often lose test accuracy even when the training loss is reached §1 §2. Local SGD cuts communication by taking H local steps between synchronizations, but the earlier comparisons did not show whether it also helps generalization at large batch sizes §1 §2. The paper targets two regimes: a communication-bound setting where the time per step is dominated by synchronization §1, Scenario 1, and a generalization-bound setting where very large global batches are used [§1, Scenario 2] §4.2.

Method

Mini-batch SGD and local SGD

Mini-batch SGD with K workers and local batch size B averages K gradient estimates each step and applies one update §1, eq. 1. In local SGD, each worker keeps its own model, takes H sequential SGD steps on its own local mini-batches, and then all models are averaged; the averaged model becomes the starting point of the next round §1, local update display. With B = H times the local batch, both schemes evaluate the same number of gradients per round, but their updates differ, so their generalization can differ §1. The paper lists local SGD under the names federated averaging and parallel SGD §1.

The appendix gives the same procedure as Algorithm 1 with a per-round difference Δ and a synchronized update alg. 1, app. B.1. The update line in that algorithm differs from the main-text averaging rule; see Pitfalls.

Post-local SGD

Post-local SGD uses H_t = 1 for rounds before a switch time t′ and H for rounds after it §3, post-local display alg. 2, app. C.1. Before t′, the method is identical to mini-batch SGD with local batch size B, so warm-up and learning-rate schedules tuned for large-batch training carry over unchanged §3. The paper sets t′ to the first learning-rate decay, which is 50% of the training samples in the CIFAR runs §3 footnote app. A.4.1. The appendix also starts post-local SGD only after the second decay in a ResNet-20 example. It reports this as more communication and a higher risk of sharper minima app. C.2.

# Post-local SGD  [alg. 2]
# Input:  w0 (initial model), K workers, T rounds, t1 (first LR decay round)
#         H (local steps after t1), B_loc (local batch size per worker)
# Output: consensus model after T rounds

def post_local_sgd(w0, K, T, t1, H, B_loc):
    # Start every worker from the same model  [alg. 2]
    consensus = w0
    local_models = [w0 for k in range(K)]
    for t in range(T):
        # One step per round before t1, H steps after  [§3] [alg. 2]
        H_t = 1 if t < t1 else H
        for k in range(K):  # the workers run in parallel
            for h in range(H_t):
                # Gradient on B_loc samples from this worker's shard  [alg. 2]
                grad = gradient(local_models[k], sample(shard[k], B_loc))
                # Plain update; the code adds weight decay and momentum here
                # [code: distributed_code/pcode/optim/utils.py#L13-L46]
                local_models[k] = local_models[k] - lr(t) * grad
        # Average the local models; restart all workers from the average [§3]
        consensus = sum(local_models) / K
        local_models = [consensus for k in range(K)]
    return consensus

The pseudocode omits weight decay, Nesterov momentum and the learning-rate warm-up; these are described in Hyperparameters.

Learning-rate scaling and warm-up

The large-batch recipe used throughout is two rules: scale the learning rate linearly with the global batch size, and warm up from a small value §4 setup app. A.3. Concretely, with K workers the peak rate is K times the single-worker rate, and the rate rises linearly from the single-worker rate to the peak over the first 5 epochs app. A.3. The single-worker rate is tuned per model on one GPU app. A.4.1. For large-batch baselines the paper additionally grid-searches the multiplier around the linear factor, e.g. the values 15, 15.5, 16, 16.5 and 17 for K = 16, extending the grid if the best value sits on an edge app. A.4.1. Post-local SGD reuses the same hyperparameters as the mini-batch setting §3 app. A.4.1.

Hierarchical local SGD

Hierarchical local SGD nests two levels. Inside each GPU block, workers synchronize after H local steps; across blocks, the blocks synchronize after H_b block steps §3 app. D.1. The appendix gives the procedure as an equation app. D.1 and as Algorithm 5 alg. 5, app. D.2. The paper's experiments use ResNet-20 on CIFAR-10 with a learning rate that starts at 0.1, decays by 10 at 50% and 75% of samples, and Nesterov momentum 0.9 applied per local model app. D.3. The repository does not contain a hierarchical optimizer, so this variant is described from the paper only.

Noise interpretation

Writing the mini-batch update as a true gradient plus noise, the noise has covariance roughly (1/B − 1/N) times the gradient second-moment matrix §5. The authors argue that the local steps inject extra, structured noise. During local phases the noise covariance is about K times that of mini-batch SGD with the same local batch §5. They use this to explain why post-local SGD generalizes like small-batch SGD §5.

Hyperparameters

Values are as stated in the paper; "not stated" marks anything missing.

Results

All numbers are author-reported. Setups are given in each row.

Local SGD vs. mini-batch SGD at equal effective batch (CIFAR-10, ResNet-20, K=4 and K=16, fine-tuned mini-batch baselines). §4.2

Post-local SGD on CIFAR, global batch K·B = 2048, K=16 (or small/large baselines), 3 runs, ±std. §4.2

Post-local SGD with ResNet-50 on ImageNet with LARS, K=32 and H=4. §4.2

ImageNet, ResNet-50, post-local SGD vs. mini-batch SGD (top-1). §4.2 app. C.5.3

• K·B = 4096: post-local 76.18 vs. mini-batch 75.87.
• K·B = 8192: post-local 75.65 vs. mini-batch 75.64.
• Post-local training phase is reported at 1.35× speedup §4.2.

Larger global batch on CIFAR-100 (K·B = 4096, K=32, 3 runs). app. C.5.1

Training length for ResNet-20, CIFAR-100, K·B = 4096. app. C.5.1

Language modelling, WikiText-2 LSTM, validation perplexity (lower is better), K=16, K·B = 1024, 3 runs. app. C.5.2

Speedup of training on 16 GPUs (8×2 cluster, 10 Gbps), CIFAR-10 and CIFAR-100. app. C.3

Speedup is T_1 / T_16, the wall-clock time on one GPU divided by the time on 16 GPUs. The phase-two-only speedups are larger, e.g. ResNet-20 CIFAR-10: 17.33 (H=16) and 20.80 (H=32) app. C.3.

Sign-based compression plus post-local SGD (ResNet-20, K·B = 2048, K=16, 3 seeds). §4.2 app. C.5.5

Isotropic gradient noise baseline (Neelakantan et al.) vs. post-local SGD. app. C.5.4

• CIFAR-10, K·B = 2048: noise injection 92.46 ± 0.27; mini-batch SGD 92.48 ± 0.17; post-local 93.02 ± 0.24.
• CIFAR-100, K·B = 4096: noise injection 67.15 ± 0.26; mini-batch SGD 67.38 ± 0.34; post-local 68.30 ± 0.48.

ResNet-50 with local SGD, ImageNet, time-to-accuracy. §4.1 app. B.3.2

• Local SGD reaches 75% top-1 at least 1.5× faster than mini-batch SGD on a 16×2-GPU cluster §4.1.

Local SGD speed-up on CIFAR-10 (ResNet-20, B_loc=128). §4.1 fig. 2

• Reported in the text: 2.59× time-to-accuracy for H=8 at 91.2% test accuracy, and 1.76× for H=2 at 92.57% §4.1.
• Local SGD on 8 GPUs reaches about 2× lower time-to-accuracy than mini-batch SGD on 16 GPUs §4.1.

Convex logistic regression on w8a (appendix illustration). app. B.2

• With communication 25× more expensive than a gradient step, (B_loc=16, H=16) converges more than 2× faster than (B_loc=64, H=1) for K=16 app. B.2.

Hierarchical local SGD (ResNet-20, CIFAR-10, 16 GPUs, H·H_b = 16, 3 runs). app. D.3.1

Training time for flat local SGD, CIFAR-10, ResNet-20, B_loc=128, 8×2 GPUs, H from 1 to 1024: 20.07 minutes at H=1, lowest at 8.32 minutes for H=32, and 10.65 minutes at H=1024 app. D.3.1. The paper attributes the hierarchical gains to synchronizing inside each node app. D.3.1.

Flatness. Post-local SGD reaches low Hessian-spectrum minima where large-batch SGD does not, and linear interpolation between the two solutions shows the difference on both training and test loss §5.1 fig. 4 app. C.4.

Paper vs code

The repository is epfml/LocalSGD-Code at commit 3d4811d01673af205a00176f5389ed008a1ddb37, licensed Apache-2.0. The paper's footnote links it §4, implementation.

• Averaging rule. The paper averages the local models §3 app. B.1. The code computes the consensus update as consensus minus the averaged difference (consensus − local), which equals the average of the local models. No extra learning-rate factor is applied on top of the differences local_sgd.py. The appendix pseudocode applies γ again (see Pitfalls); the code does not.
• Switch point. The paper switches at the first LR decay §3. The code syncs every step while the float epoch is below turn_on_local_step_from and every local_step iterations afterwards local_sgd.py. The README demo uses 150 for a 300-epoch run, which equals the first decay README.md.
• Learning rate. The code computes learning rate = lr / base_batch_size × local batch × lr_scaleup_factor, with the scaleup factor defaulting to the number of workers create_scheduler.py. The README sets lr 0.1 with base batch 64 and local batch 128 README.md. That gives 0.2 at B=128, the same as the paper's ResNet-20 value.
• Warm-up. The code warms up linearly from the base rate to the scaled rate over lr_warmup_epochs; the README uses 5 create_scheduler.py README.md. If the flag is unset, the code uses min(scale factor, 150), not 5 parameters.py.
• Momentum. Nesterov momentum is per worker. Momentum buffers are kept in each worker's optimizer state and are not averaged at synchronization utils.py. The paper says momentum is applied independently to each local model app. A.4.1.
• Weight decay on BatchNorm. The code skips weight decay for any parameter whose name contains "bn" create_optimizer.py. In the ResNet code, the projection-shortcut BatchNorm sits inside a Sequential and is named like downsample.1, which does not contain "bn" resnet.py. So those BatchNorm parameters do receive weight decay, which is not what the paper's text says app. A.4.1. This is inferred from naming and was not run.
• Hierarchical local SGD. No hierarchical optimizer is present in the repository (checked by searching the Python sources for "hierarch"). The hierarchical results are not verified against code.
• Compression. The repository contains sign-SGD and EF-sign-SGD variants, including local versions local_sign_sgd.py local_ef_sign_sgd.py. I did not check them line by line against the appendix pseudocode.

Pitfalls

• Extra learning-rate factor in the appendix pseudocode. Algorithms 1 and 2 (and the hierarchical one) apply the learning rate γ to the averaged difference Δ, which is already a sum of γ-scaled steps. Read literally, this double-scales the update. The main-text rule averages the local models §3, and the code does the same local_sgd.py. Implement model averaging, not the algorithm box as printed.
• Index typos in the algorithms. The local update in Algorithms 1 and 2 writes w_(t)+h−1 without the worker index k; the local update should use the worker's own model alg. 1 alg. 2. Treat as a typo.
• Two meanings of t. The main text uses t for SGD iterations in H_(t) §3. The algorithms use t for synchronization rounds alg. 2. The code uses the float epoch index against an epoch threshold local_sgd.py. Make the switch in one consistent unit.
• Switch point depends on the schedule. The paper defines t′ as the first decay §3 footnote. Other schedules need the switch set explicitly; the paper gives no rule for that case.
• Phase one is not communication-saving. Before t′, the method is mini-batch SGD with one synchronization per step. Communication savings come only in the second phase §3 §4.2.
• Whole-run speedups are modest. The reported whole-training speedup is about 1.3× §4.2 app. C.3, because the first phase still synchronizes every step. The larger speedups are for phase two only app. C.3.
• Mislabelled model in the speedup table. The speedup table labels the model WideResNet-28-12, while the accuracy experiments use WideResNet-28-10 app. C.3 §4.2 app. A.2. Treat the speedup row as the same model family and check the setup before reusing it.
• Repeated baseline values. The CIFAR-100 K·B=4096 table reports a small-batch value of 68.84 ± 0.06 for ResNet-20 app. C.5.1, which is also the value in the K=2 column of the K·B=2048 table §4.2. The large-batch baseline is 67.34 in the K·B=4096 table but 67.38 in the isotropic-noise table app. C.5.4. The paper does not explain these differences. Compare runs by setting, not by value.
• Small LM gain. Post-local H=16 gives 86.85 perplexity vs. 86.90 for large batch, with overlapping standard deviations app. C.5.2. This is not a clear gain.
• Learning-rate tuning costs. The large-batch baselines are fine-tuned, and the grid for the multiplier is extended when the optimum is at an edge app. A.4.1. Matching the paper requires the same tuning budget.
• ImageNet learning-rate rule is garbled. The appendix writes the scaling as "Number of GPUs × 0.1/256 × B_g", where B_g is not defined app. A.4.2. The exact global batch used for the scaling is not stated.
• Hardware-bound claims. Speedup and time-to-accuracy results depend on an 8×2 GPU cluster with 10 Gbps Ethernet and Kubernetes §4 app. C.3. They are not transferable to other interconnects without measurement.
• Plot-only numbers. Some results appear only in figures, and the text quotes them (2.59×, 1.76×, 1.35×). I did not check these against the plots.
• Version. The source is arXiv v6 (ICLR 2020 style). Numbers may differ from earlier versions [raw/arxiv metadata].
• Noise-injection baseline search. The isotropic-noise baseline failed at the paper's suggested values and was searched only in the grid given app. C.5.4. A stronger search could change the comparison.
• Hierarchical variant. Its learning-rate and block settings differ from the flat runs (starts at 0.1, 300 epochs, H·H_b = 16) app. D.3. Compare only within that table.

Lineage

• Builds on: local SGD and federated averaging, as cited in §2 §2. Large-batch recipe: linear learning-rate scaling and warm-up from Goyal et al. §4 setup app. A.3.
• Compared against: large-batch and small-batch mini-batch SGD, isotropic gradient noise (Neelakantan et al.) app. C.5.4, adaptive local-step schedules (Wang et al.) app. C.5.4.
• Combined with: signSGD and EF-signSGD compression §4.2 app. C.5.5, LARS §4.2.
• Superseded by: not stated.