Large-batch training

Overview

The problem: keep test accuracy while raising the global batch size K·B, the number of workers times the per-worker batch, so that more workers can be used 1808.07217 §1. The paper's observation is that large global batches often lose test accuracy even when the training loss reaches the same level as with small batches 1808.07217 §1.

The shared structure of the variants is one question: how often do workers exchange information, and what is exchanged? Plain mini-batch SGD averages gradients every step. Local SGD lets each worker take H steps on its own and then averages the models 1808.07217 §1, §3. The invariant all variants keep is the large-batch recipe from the paper's baselines: the learning rate is scaled linearly with the number of workers and warmed up linearly over the first 5 epochs 1808.07217 app. A.3. The paper credits this recipe to Goyal et al., who are not covered by a card here.

Variants

Mini-batch SGD with linear scaling and warm-up

The baseline. Each step averages K local gradients and applies one update 1808.07217 eq. 1. The peak learning rate is K times the single-worker rate, reached by a linear warm-up 1808.07217 app. A.3. For ResNet-20 on CIFAR-10 the single-worker rate is 0.2 at local batch 128 [arxiv:1808.07217 app. A.4.1]. On the paper's CIFAR-10 setting this variant falls off at large batches: with K=4 it scores 92.6% at KB=512 and 89.5% at KB=8192, and with K=16 it scores 92.5% at KB=2048 and 76.3% at KB=16384 1808.07217 §4.2. Card: Don't Use Large Mini-Batches, Use Local SGD.

Batch-size ramp

Start with a small global batch and grow it during training. The paper's table reports 92.5% for K=4 ramped from KB=512 to 8192, and 92.0% for K=16 ramped from KB=2048 to 16384 1808.07217 §4.2. The ramp end is the same batch that gives 76.3% with a constant batch of 16384 in the same table, so the ramp recovers most of the gap without changing the final batch size.

Local SGD

Each worker takes H steps on its own mini-batches, then all models are averaged, and the average starts the next round 1808.07217 §1. If the mini-batch baseline uses a batch H times the local batch, both schemes evaluate the same number of gradients per round, but their updates differ 1808.07217 §1. In the paper's CIFAR-10 table, local SGD with K=4 and H=16 reaches 92.5% at an effective batch of 8192, against 89.5% for mini-batch SGD at the same batch 1808.07217 §4.2. With K=16 and H=8 it reaches 92.0% at 16384, against 76.3% for mini-batch SGD 1808.07217 §4.2.

Post-local SGD

Run mini-batch SGD until the first learning-rate decay, then switch to local SGD with H steps between averages 1808.07217 §3 [arxiv:1808.07217 alg. 2]. The first phase equals the baseline, so the warm-up and learning-rate schedule tuned for large batches carry over unchanged 1808.07217 §3. The switch is set at the first decay, which is 50% of the samples in the CIFAR runs 1808.07217 §3 1808.07217 app. A.4.1. The paper's headline CIFAR-10 result uses ResNet-20 with K·B=2048 and H=32: 93.02 ± 0.24, against 92.48 ± 0.17 for large-batch SGD and 92.63 ± 0.26 for small-batch SGD 1808.07217 §4.2. At K·B=4096 on CIFAR-100, H=16 gives 68.30 ± 0.30, against 67.34 for large-batch SGD and 68.84 for small-batch SGD 1808.07217 app. C.5.1. On ImageNet with ResNet-50 at K·B=4096 the paper reports 76.18 against 75.87 for mini-batch SGD, and at K·B=8192 the two are tied, 75.65 against 75.64 1808.07217 §4.2 [arxiv:1808.07217 app. C.5.3].

The cost: the first phase still synchronizes every step. The paper reports a whole-run speed-up of about 1.3×, with larger speed-ups for the second phase alone 1808.07217 app. C.3. The speed-up measurements come from 16 GPUs on an 8×2 cluster with 10 Gbps Ethernet 1808.07217 §4 [arxiv:1808.07217 app. C.3].

Hierarchical local SGD

Two nested levels. Workers inside a block synchronize after H local steps, and the blocks synchronize with each other after H_b block steps [arxiv:1808.07217 §3] 1808.07217 app. D.1. In the paper's ResNet-20 CIFAR-10 runs on 16 GPUs with H·H_b=16, the gain appears in one of three layouts. On 8×2 nodes, H=8 with H_b=2 gives 91.41 ± 0.23, against 90.02 ± 0.28 for H=1 with H_b=16. On 4×4 and 2×8 nodes, the best row is H=1 with H_b=16, at 91.65 and 92.14 1808.07217 app. D.3.1. The paper attributes the 8×2 gain to synchronizing inside each node 1808.07217 app. D.3.1.

Combined with sign-based compression

Local steps can be combined with sign-based gradient compression. For ResNet-20 on CIFAR-10 with K·B=2048 and K=16, EF-signSGD goes from 92.45 at H=1 to 92.76 at H=64. Plain signSGD goes from 91.61 at H=1 to 92.22 at H=32 and 92.00 at H=64 1808.07217 app. C.5.5. Use this only when bytes per round are the limit, and see the warnings below.

When to use which

The paper's large-batch baselines were tuned, including a learning-rate grid around the linear factor that was extended when the best value sat at an edge 1808.07217 app. A.4.1. Matching the reported gains needs a comparable tuning budget. The paper's speed-ups depend on its interconnect and are reported as measured, not as transferable 1808.07217 app. C.3.

Implementation sketch

Post-local SGD as the paper states it. Average models, not gradients, at each synchronization. Other schedules need the switch point set by hand.

# Post-local SGD [arxiv:1808.07217 alg. 2]
# K workers, T rounds, t1 = first learning-rate decay round,
# H local steps after t1, B = local batch size

def post_local_sgd(w0, K, T, t1, H, B):
    # every worker starts from the same model
    local = [w0 for k in range(K)]
    for t in range(T):
        # one step before the switch, H steps after it [arxiv:1808.07217 §3]
        H_t = 1 if t < t1 else H
        for k in range(K):
            for h in range(H_t):
                # gradient on this worker's own batch, learning rate
                # scaled by K with a linear warm-up [arxiv:1808.07217 app. A.3]
                g = grad(local[k], sample_batch(k, B))
                local[k] = local[k] - lr(t) * g
        # average the models and restart every worker from the average
        avg = sum(local) / K
        local = [avg for k in range(K)]
    return avg

Warm-up over 5 epochs and a decay ×0.1 at 50% and 75% of samples are the paper's settings for the CIFAR runs 1808.07217 app. A.3 [arxiv:1808.07217 app. A.4.1]. Momentum is Nesterov at 0.9 per local model, and whether momentum buffers are averaged is not stated 1808.07217 app. A.4.1.

Contradictions

• Switch point. The main text switches at the first learning-rate decay 1808.07217 §3. The appendix starts post-local SGD after the second decay in a ResNet-20 example, and notes more communication and a higher risk of sharper minima for that choice 1808.07217 app. C.2.
• Averaging rule. The main text and the appendix's Algorithm 1 differ. The main text averages the local models 1808.07217 §3. Algorithms 1 and 2 apply the learning rate γ to an averaged difference Δ that already contains γ-scaled steps, which as printed scales the update twice 1808.07217 alg. 1 1808.07217 alg. 2 [arxiv:1808.07217 app. B.1]. The paper does not say which one its experiments use.
• Repeated baseline value. Large-batch ResNet-20 on CIFAR-100 at K·B=4096 is 67.34 in one table 1808.07217 app. C.5.1 and 67.38 in another 1808.07217 app. C.5.4. The paper does not explain the difference.
• Model label. The speed-up table names WideResNet-28-12, while the accuracy runs use WideResNet-28-10 1808.07217 app. C.3 [arxiv:1808.07217 §4.2].
• Small language-model gain. Post-local H=16 gives validation perplexity 86.85 against 86.90 for large-batch SGD, with overlapping standard deviations 1808.07217 app. C.5.2. The paper does not present this as a gain.
• Isotropic noise as a substitute. Injecting isotropic gradient noise scores 92.46 on CIFAR-10 at K·B=2048, close to mini-batch SGD at 92.48 and below post-local SGD at 93.02 1808.07217 app. C.5.4. The noise baseline was searched only over the paper's grid, so a stronger search could change this comparison 1808.07217 app. C.5.4.

Papers

• Don't Use Large Mini-Batches, Use Local SGD: introduces post-local SGD and the hierarchical variant, and reports CIFAR and ImageNet results against mini-batch SGD at large global batch sizes.