PaperlayerView exact revision on arXiv

Source-linked AI summary

Large-Scale Distributed Second-Order Optimization Using Kronecker-Factored Approximate Curvature for Deep Convolutional Neural Networks

Kazuki Osawa, Yohei Tsuji, Yuichiro Ueno, Akira Naruse, Rio Yokota, Satoshi Matsuoka

arXiv:1811.12019v5cs.LGcs.CVstat.ML

TL;DR

Large distributed training suffers a generalization gap as effective mini-batch sizes grow. The paper uses distributed K-FAC, a second-order optimizer with approximate curvature, and reports 75% top-1 validation accuracy even at mini-batch size 131,072.

  • Problem

    Increasing effective mini-batch sizes in distributed training can degrade validation accuracy, creating a large-batch generalization gap.

    Source: §1 · p. 1

  • Method

    The paper develops a distributed K-FAC optimizer using approximate Fisher curvature, hybrid data/model parallelism, and reduced overhead from stale or diagonal curvature approximations.

    Source: §1 · p. 2, §3.2 · p. 3, §3.3 · p. 4

  • Results

    75% top-1 validation accuracy was achieved with mini-batch size 131,072 in 978 iterations, while sizes under 16,384 converged in 35 epochs.

    Source: §1 · p. 2

  • Takeaways & Limitations

    Second-order optimization can match the generalization capability of highly optimized SGD while supporting extremely large mini-batches.

    Source: §1 · p. 2, §6 · p. 8

Abstract

from arXiv · show

Large-scale distributed training of deep neural networks suffer from the generalization gap caused by the increase in the effective mini-batch size. Previous approaches try to solve this problem by varying the learning rate and batch size over epochs and layers, or some ad hoc modification of the batch normalization. We propose an alternative approach using a second-order optimization method that shows similar generalization capability to first-order methods, but converges faster and can handle larger mini-batches. To test our method on a benchmark where highly optimized first-order methods are available as references, we train ResNet-50 on ImageNet. We converged to 75% Top-1 validation accuracy in 35 epochs for mini-batch sizes under 16,384, and achieved 75% even with a mini-batch size of 131,072, which took only 978 iterations.

1. Introduction

Large mini-batch training can degrade validation accuracy, motivating a second-order K-FAC approach that aims to preserve generalization while converging faster and scaling distributed computation. The paper combines distributed implementation choices with curvature approximations to train ImageNet ResNet-50 at extreme batch sizes.

Motivation

  • Large distributed data-parallel systems increase effective mini-batch size with process count, and validation accuracy can degrade beyond a certain batch-size threshold.Prior approaches include learning-rate scaling, warmup, layer-wise learning rates, adaptive batch sizes, and batch-normalization modifications.Source: §1 · p. 1

Motivation and approach

  • K-FAC addresses the large-mini-batch generalization gap with a second-order method designed to converge faster and tolerate large batches without ad hoc modifications.The method uses a more mathematically rigorous curvature approximation than crude Hessian approximations.Source: §1 · pp. 1–2

Results

  • 75% top-1 validation accuracy was achieved with mini-batch size 131,072 in 978 iterations, while sizes under 16,384 converged in 35 epochs.The reported results indicate similar generalization to highly optimized SGD at substantially larger mini-batch sizes.Source: §1 · p. 2

Distributed implementation

  • The distributed optimizer reduces K-FAC overhead by reducing Fisher-matrix update frequency after a few hundred iterations, training ResNet-50 on ImageNet in 10 minutes to 74.9% top-1 accuracy using 1,024 Tesla V100 GPUs (Table 2).Source: §1 · p. 2

Curvature approximation

  • Batch Normalization Fisher matrices can be approximated as diagonal, reducing K-FAC computation and memory consumption.The approximation is presented as an additional overhead-reduction scheme.Source: §1 · p. 2

2. Related work

Prior large-scale distributed training studies using second-order methods were limited, and earlier K-FAC work on ResNet-50 did not establish a clear advantage over strong SGD baselines. This work compares K-FAC with state-of-the-art large-mini-batch SGD, scales to much larger batches, and uses decentralized collective communication.

  • Few studies had applied second-order methods to large-scale distributed deep-network training, while Table 1 reports ResNet-50 ImageNet K-FAC training epochs, iterations, and Top-1 validation accuracy.Source: §2 · p. 2, Table 1 · p. 2
  • Earlier K-FAC results on ResNet-50 used an SGD reference with only around 70% Top-1 validation accuracy, making K-FAC’s claimed advantage unclear.Source: §2 · p. 2
  • This work compares Top-1 validation accuracy against state-of-the-art SGD methods for large mini-batches using the ResNet-50 ImageNet benchmark.Table 2 is identified as the comparison source in the text.Source: §2 · p. 2
  • K-FAC studies trained with mini-batch size 512 on 8 GPUs, whereas this work evaluates sizes up to 131,072 and achieves 75% Top-1 validation accuracy.The largest batch corresponds to 32 samples per GPU on 4096 GPUs; gradient accumulation can mimic multi-GPU execution without using that many GPUs.Source: §2 · p. 2
  • Earlier K-FAC implementations used TensorFlow and parameter servers, whose centralized gradient and parameter exchanges create communication bottlenecks at large scale.Source: §2 · p. 2
  • The implementation instead uses decentralized MPI/NCCL collective communication, an approach described as scalable to thousands of GPUs without modification.Source: §2 · p. 2

3. Distributed K-FAC

The paper defines K-FAC as a second-order optimizer that approximates Fisher curvature with Kronecker factors, then distributes its computation through a hybrid data/model-parallel design. Communication, mixed precision, and symmetry-aware methods reduce K-FAC’s overhead for large-scale training.

3.2. K-FAC

  • K-FAC approximates the Fisher information matrix with layer-wise Kronecker factors derived from activations and output gradients.The inverse is approximated through the inverses of the individual Kronecker factors.Source: §3.2 · p. 3

3.3. Our design

  • K-FAC’s extra inverse-FIM computation creates substantial overhead, motivating distributed parallelization and a relaxation technique that reduce it to almost negligible levels.The design targets K-FAC overhead specifically because it is considerable compared with SGD.Source: §3.3 · p. 3
  • The distributed design averages factors with ReduceScatterV, assigns inverse and matrix-matrix computations across layers, and restores data parallelism with AllGatherV.This hybrid scheme distributes the model during curvature processing while allowing all processes to update parameters afterward.Source: §3.3 · p. 4

3.4. Further acceleration

  • The hybrid approach addresses two hotspots: factor construction remains data-parallel, while factor communication is reduced through collective operations.Factor construction does not become faster simply by adding processes, so acceleration targets both computation and communication.Source: §3.4 · p. 4, §3.3 · p. 4
  • Tensor Cores more than double Kronecker-factor calculation speed, and communicating only upper triangles reduces traffic for symmetric factors.A symmetric N × N matrix requires sending N(N + 1)/2 upper-triangular entries.Source: §3.4 · p. 5

4. Training schemes

The training schemes combine stronger augmentation with damping, learning-rate, momentum, and weight-scaling schedules to stabilize and regularize large-scale K-FAC training. These choices address limited prior guidance for tuning K-FAC on large models and mini-batches.

4. Training schemes

  • The paper studies training schemes because K-FAC’s behavior and hyperparameter tuning on large models, datasets, and mini-batches remain insufficiently explored.The schemes are evaluated as ways to achieve higher accuracy in large-mini-batch K-FAC experiments.Source: §4 · p. 5

4.1. Data augmentation

  • Running mixup combines raw samples with virtual samples from the previous step to increase mixup’s regularization effect.The mixing coefficient λ is sampled from a Beta distribution, with α = β = αmixup.Source: §4.1 · p. 5
  • Random erasing replaces the erased region with zero values, using probability p = 0.5 and specified area and aspect-ratio ranges.The erased area ratio is Se ∈ [0.02, 0.25] and the aspect ratio is re ∈ [0.3, 1].Source: §4.1 · p. 5

4.2. Warmup damping

  • Warmup damping starts with a large damping rate and gradually decreases it toward a target because the Fisher matrix changes rapidly early in training.The damping schedule uses γ(0), γtarget, and a warmup duration twarmup; batch-normalization damping is scaled by ρBN > 1.Source: §4.2 · p. 5

4.3. Learning rate and momentum

  • Polynomial learning-rate decay is paired with momentum whose rate is adjusted to keep its ratio to the learning rate fixed during training.The schedule defines decay start and end epochs and a decay-rate parameter.Source: §4.3 · pp. 5–6

4.4. Weights rescaling

  • Weights are rescaled after updates to have norm √2 · dout, with ϵ = 1·10−9 used to stabilize computation.dout denotes the layer’s output dimension or number of channels.Source: §4.4 · p. 6

5. Results

The experiments evaluate distributed K-FAC for ResNet-50 on ImageNet across scalability, extremely large mini-batches, and Fisher-information approximations. K-FAC maintains competitive accuracy at large batch sizes while using stale and diagonal Fisher approximations to reduce overhead.

5.1. Experiment environment

  • The experiments used ResNet-50 on ImageNet with 32 images per GPU, while accumulation mimicked larger global mini-batches that could not be run directly.Hyperparameters were tuned on fewer nodes because full-machine access was limited to a 24-hour window.Source: §5.1 · p. 6, §5 · p. 6

5.2. Scalability

  • Scaling from 1 to 64 GPUs was superlinear, with 64 GPUs 131.1% faster, but performance degraded beyond 256 GPUs because of communication overhead.The superlinear behavior is attributed to the hybrid data/model-parallel design.Source: §5.2 · p. 6, Figure 2 · p. 6

5.3. Large mini-batch training with K-FAC

  • K-FAC achieved at least 75% top-1 validation accuracy with mini-batches from 4K through 131K, converging in much less than 90 epochs for sizes through 65K.At 32K, K-FAC converged in half as many epochs as SGD; at 131K, SGD-based ImageNet training did not sustain top-1 validation accuracy.Source: §5.3 · p. 6

5.4. Fisher information and large mini-batch training

  • Larger mini-batches made the Fisher information matrix fluctuate less, supporting more aggressive refresh reduction; the input-correlation factor A fluctuated less than the gradient factor G.The refresh schedule reduced updates after the first roughly 500 iterations, when the Fisher information changed rapidly.Source: §5.4 · p. 7
  • Using 1,024 Tesla V100 GPUs, K-FAC reached 74.9% top-1 accuracy on ImageNet in 10 minutes, with training time and accuracy competitive with reported SGD results.The run used 45 epochs, including validation after each epoch.Source: §5.4 · p. 8
  • Approximating batch-normalization Fisher matrices diagonally reduced memory from 1017MiB to 587MiB without affecting training curves, even with stale Fisher information.This result suggests that only diagonal Fisher values are essential for batch-normalization training in the tested setting.Source: §5.4 · p. 8, Figure 4 · p. 8

6. Conclusion

The paper presents K-FAC as a distributed second-order optimizer for extremely large mini-batches, achieving over 75% Top-1 accuracy with fewer epochs or iterations than prior large-batch SGD approaches.

6. Conclusion

  • Figure 5 tracks percentile changes in Kronecker factors for convolutional layers and the FIM for batch-normalization layers during large-mini-batch training.The plotted percentiles are 5th, 25th, 50th, 75th, and 95th across ResNet-50 layers, with the red line marking iteration 500.Source: Figure 5 · p. 8
  • K-FAC matched the generalization capability of highly optimized SGD even with extremely large mini-batches, contrary to claims that second-order methods generalize worse.The benchmark was ResNet-50 on ImageNet.Source: §6 · p. 9
  • The study reports a relationship between the Fisher information matrix and large-mini-batch training across mini-batch sizes from 4K to 131K.The authors state this is the first work to observe this relationship for ResNet-50.Source: §6 · p. 9
  • The distributed design still has a computation/communication bottleneck, although the Kronecker factors may be approximated more aggressively without accuracy loss.The authors also propose further eigenvalue analysis to understand FIM-based preconditioning.Source: §6 · p. 9
Loading 1811.12019v5…