PaperlayerView exact revision on arXiv

Source-linked AI summary

High-Performance Large-Scale Image Recognition Without Normalization

Andrew Brock, Soham De, Samuel L. Smith, Karen Simonyan

arXiv:2102.06171v1cs.CVcs.LGstat.ML

TL;DR

Batch normalization has useful training properties but depends on minibatch interactions and can be costly, unstable, or inconvenient. This paper introduces Adaptive Gradient Clipping and improved Normalizer-Free ResNets, which achieve competitive and state-of-the-art ImageNet performance, including stronger fine-tuning results after large-scale pre-training.

  • Problem

    Batch normalization introduces computational costs, training–inference discrepancies, and interactions between minibatch examples, while prior normalization-free models lacked competitive accuracy or stability.

  • Method

    The paper introduces Adaptive Gradient Clipping, which uses unit-wise gradient-to-parameter norm ratios, and develops improved Normalizer-Free ResNets.

  • Results

    NFNets match EfficientNet-B7 accuracy while training 8.7× faster, reach 86.5% top-1 ImageNet accuracy, and attain 89.2% after fine-tuning from 300 million labelled images.

  • Takeaways & Limitations

    Normalizer-Free models can match or exceed batch-normalized models on the reported ImageNet training and large-scale fine-tuning evaluations while remaining faster to train.

  • Takeaways & Limitations

    AGC requires fine-grained clipping-threshold tuning across model depth, batch size, and learning rate.

Abstract

from arXiv · show

Batch normalization is a key component of most image classification models, but it has many undesirable properties stemming from its dependence on the batch size and interactions between examples. Although recent work has succeeded in training deep ResNets without normalization layers, these models do not match the test accuracies of the best batch-normalized networks, and are often unstable for large learning rates or strong data augmentations. In this work, we develop an adaptive gradient clipping technique which overcomes these instabilities, and design a significantly improved class of Normalizer-Free ResNets. Our smaller models match the test accuracy of an EfficientNet-B7 on ImageNet while being up to 8.7x faster to train, and our largest models attain a new state-of-the-art top-1 accuracy of 86.5%. In addition, Normalizer-Free models attain significantly better performance than their batch-normalized counterparts when finetuning on ImageNet after large-scale pre-training on a dataset of 300 million labeled images, with our best models obtaining an accuracy of 89.2%. Our code is available at https://github.com/deepmind/ deepmind-research/tree/master/nfnets

1. Introduction

Batch normalization brings substantial training benefits but introduces computational, statistical, and reproducibility drawbacks. The paper addresses these limitations with AGC and NFNets, achieving competitive or state-of-the-art ImageNet results and improved fine-tuning performance.

  • Motivation: Batch normalization is computationally expensive, creates training–inference discrepancies, and breaks independence between minibatch examples.These properties add memory and gradient-evaluation costs, hidden hyperparameters, and implementation or task-specific complications.
  • Motivation: Batch-statistic interactions can hinder reproducibility, enable information leakage, and make performance sensitive to batch size.The paper notes particular concerns for distributed training, contrastive learning, sequence modeling, and small batches.
  • Contributions: Adaptive Gradient Clipping clips gradients using the unit-wise ratio of gradient norms to parameter norms.The technique is intended to stabilize Normalizer-Free Network training with larger batches and stronger augmentations.
  • Contributions: 8.7× faster training gives NFNet-F1 accuracy comparable to EfficientNet-B7, while the largest NFNet reaches 86.5% top-1 ImageNet accuracy.NFNets are presented as a family of Normalizer-Free ResNets spanning multiple training latencies.
  • Contributions: NFNets achieve substantially higher ImageNet validation accuracy than batch-normalized networks after pre-training on 300 million labelled images, reaching 89.2% top-1 after fine-tuning.This comparison concerns fine-tuning on ImageNet after large-scale private-dataset pre-training.

2. Understanding Batch Normalization

The section identifies four benefits of batch normalization that alternatives must recover: residual-branch scaling, mean-shift elimination, regularization, and efficient large-batch training.

  • Benefits of Batch Normalization: Batch normalization downscales residual-branch activations at initialization, biasing signals toward skip connections and supporting deep networks.This produces well-behaved gradients early in training.
  • Benefits of Batch Normalization: Batch normalization eliminates the mean-shift caused by non-antisymmetric activations, which otherwise compounds with network depth.The supplied passages connect this shift to correlated activations across independent training examples.
  • Benefits of Batch Normalization: Batch normalization acts as a regularizer through noise in batch statistics, and test accuracy can improve with batch-size tuning or ghost batch normalization.The passage presents this regularizing effect as widely believed and empirically consistent with those practices.
  • Benefits of Batch Normalization: Batch normalization smooths the loss landscape and increases the largest stable learning rate, enabling efficient large-batch training.Larger learning rates are described as essential for efficient training with large batches.

3. Towards Removing Batch Normalization

Prior Normalizer-Free ResNets recover several batch-normalization benefits through variance control, scaled weight standardization, and residual-branch initialization. They match batch-normalized ResNets in some settings but lag behind state-of-the-art networks and large-batch performance.

  • Normalizer-Free ResNets: NF-ResNets use variance-preserving residual branches, a small residual scale α, and β_i normalization based on activation variance.Their residual update is h_i+1 = h_i + αf_i(h_i/β_i), with variance tracked across blocks and reset at transitions.
  • Normalizer-Free ResNets: SkipInit adds a learnable scalar initialized to zero at each residual branch’s end, while Scaled Weight Standardization prevents hidden-activation mean-shift.Activation-specific gains are used so the scaled nonlinearity and standardized layer preserve variance.
  • Prior Results: At batch size 1024, regularized Normalizer-Free ResNets match batch-normalized pre-activation ResNets on ImageNet.The regularization uses Dropout and Stochastic Depth.
  • Prior Results: Normalizer-Free ResNets outperform batch-normalized counterparts at very small batch sizes but perform worse at batch sizes of 4096 or higher.The passage identifies large-batch performance as a remaining limitation.
  • Prior Results: Prior Normalizer-Free ResNets do not match state-of-the-art networks such as EfficientNets.This gap motivates the improved architectures and training techniques developed in the paper.

4. Adaptive Gradient Clipping for Efficient Large-Batch Training

AGC stabilizes Normalizer-Free ResNets by clipping unit-wise updates relative to parameter norms, enabling larger batches and stronger augmentation. Its benefits are strongest at large batch sizes, where smaller clipping thresholds are needed.

  • AGC mechanism: AGC clips each unit’s gradient according to the ratio of gradient norm to parameter norm.The method uses unit-wise rather than layer-wise norm ratios, with a small epsilon preventing zero-initialized parameters from being permanently clipped.
  • Scaling stability: AGC enables stable NF-ResNet training with batch sizes up to 4096 and strong RandAugment augmentation.NF-ResNets without AGC fail to train with such strong augmentation.
  • Accuracy across batch sizes: AGC maintains performance comparable or better than batch-normalized networks on ResNet50 and ResNet200 at large batch sizes.Its benefits are smaller when the batch size is small.
  • Clipping sensitivity: Smaller clipping thresholds are necessary for stability at higher batch sizes.The optimal λ depends on the optimizer, learning rate, and batch size, and should generally decrease as batches grow.
  • Layer-wise ablations: AGC should be omitted from the final linear layer, while clipping all four residual stages is required for stability at batch size 4096 with the default learning rate.The initial convolution can sometimes be left unclipped.

5. Normalizer-Free Architectures with Improved Accuracy and Training Speed

NFNets are manually designed Normalizer-Free ResNets optimized for the Pareto frontier of ImageNet accuracy and on-device training latency. Their design combines architectural changes, scaling rules, and stronger explicit regularization.

  • Design objective: NFNets target improved ImageNet top-1 accuracy versus actual training latency on accelerators, rather than theoretical FLOPS alone.The models are manually designed to improve the holdout-accuracy and training-latency Pareto front.
  • Backbone architecture: The NFNet backbone uses modified width and depth patterns, a second spatial convolution, and an SE-ResNeXt-D starting point with GELU activations.The second 3×3 grouped convolution adds little FLOPS and has almost no training-time impact on target accelerators.
  • Depth scaling: The smallest NFNet uses stage depths [1, 2, 6, 3], with deeper variants multiplying every stage depth by a scalar.This replaces the default nonuniform ResNet depth pattern.
  • Width scaling: The preferred width pattern is [256, 512, 1536, 1536], increasing capacity in the third stage relative to the default pattern.This was the only tested alternative reported as better than the default width pattern.
  • Regularization: NFNets increase Dropout regularization with model capacity because the absence of batch normalization’s implicit regularization otherwise causes dramatic overfitting.Changing weight decay or stochastic depth was not effective for this purpose.

6. Experiments

NFNets improve ImageNet accuracy and training efficiency, including strong gains under aggressive augmentation and transfer learning after large-scale pre-training. They also outperform batch-normalized counterparts in reported comparisons, while batch-normalized large variants can be unstable.

  • 6.1. Evaluating NFNets on ImageNet: 20% to 40% slower training is observed for batch-normalized versions, which also achieve slightly lower test accuracies than NFNets.The larger F4 and F5 batch-normalized variants were not stable with or without AGC.
  • 6.1. Evaluating NFNets on ImageNet: 84.7% top-1 accuracy is matched by NFNet-F1 against EfficientNet-B7 with RA, while NFNet-F1 is 8.7 times faster to train.NFNet-F5 reaches 86.0% top-1, slightly exceeding EfficientNet-B8 with MaxUp.
  • 6.1. Evaluating NFNets on ImageNet: 86.5% top-1 accuracy is attained by NFNet-F6 with SAM, improving over the existing ImageNet state of the art without extra data.NFNet-F5 reaches 86.3% top-1 with SAM.
  • 6.1. Evaluating NFNets on ImageNet: Strong augmentation with MixUp, CutMix, and RandAugment progressively benefits NFNets but does not similarly benefit EfficientNets beyond well-tuned RandAugment.The augmentation combination produces an intense augmentation level for NFNets.
  • 6.2. Evaluating NFNets under Transfer: Around 1% absolute top-1 improvement is typical for Normalizer-Free networks over Batch-Normalized counterparts across transfer-learning cases.The experiment pre-trains models for 10 epochs on 300 million labeled images, then fine-tunes all layers on ImageNet.
  • 6.2. Evaluating NFNets under Transfer: 89.2% ImageNet top-1 accuracy is achieved by NFNet-F4+ after 20 epochs of pre-training, the highest accuracy reported using transfer learning.This result is second highest among results using extra training data in the cited comparison.

Conclusion

The paper presents NFNets as normalization-free image-recognition models that match or exceed batch-normalized accuracy while training faster. Adaptive Gradient Clipping stabilizes large-batch training and supports strong augmentation in the proposed models.

  • Conclusion: NFNets match or exceed the classification accuracies of the best batch-normalized models on large-scale datasets while remaining faster to train.The models attain state-of-the-art ImageNet performance without extra data and are better suited to fine-tuning after very large-scale pre-training.
  • Conclusion: Adaptive Gradient Clipping stabilizes large-batch training and enables optimization of unnormalized networks with strong data augmentations.The technique is described as a simple clipping algorithm.

A.1. ImageNet Experiment Settings

The ImageNet experiments combine standard preprocessing with strong augmentation, AGC, large-batch optimization, and mixed-precision hardware evaluation. Training latency is measured as observed wallclock time, while NFNets remain competitive with EfficientNet variants at comparable test compute.

  • Compute and evaluation: NFNet models are competitive with large EfficientNet variants for a given test GFLOPs budget, using single-model, single-crop ImageNet validation.The models are optimized for training latency rather than test FLOPs.
  • Optimization: AGC is applied with λ = 0.01 and ϵ = 10^-3, alongside learning-rate warmup and cosine decay.The maximal learning rate is chosen as 0.1 × B/256 after warmup.
  • Optimization: 4096 examples are used per batch for 360 epochs, corresponding to roughly 112,000 total training steps.The schedule is fixed for reported results, without early stopping.
  • Latency measurement: Median wallclock time over 5000 training steps is used to measure training latency, excluding dataloading from the measurement.The setup is intended to reflect speed during a full training run.
  • Hardware evaluation: TPUv3 latency uses 32 devices with batch size 32 per device and synchronized gradients, while V100 measurements use float16 training.Larger models may require rematerialization on TPU or gradient accumulation on GPU.
  • ImageNet preprocessing and augmentation: Training uses distorted bounding-box crops and horizontal flips, with RandAugment, MixUp, and CutMix added in the full NFNet recipe.RandAugment uses four layers, while MixUp and CutMix are applied to separate halves of each batch.
  • Augmentation results: NFNets progressively benefit from the combination of MixUp, CutMix, and RandAugment, unlike EfficientNets beyond a well-tuned RandAugment baseline.The paper hypothesizes that reduced implicit regularization makes NFNets more amenable to stronger augmentation.

A.5. Large Scale Pre-Training Details

Large-scale transfer experiments pre-train ResNet and NFNet variants on JFT-300m before ImageNet fine-tuning. NFNets outperform batch-normalized counterparts in every reported transfer case, and wider models can reach the same performance with shorter pre-training.

  • Pre-training setup: JFT-300m contains 300 million labeled images spanning roughly 18,000 classes and is used for large-scale pre-training.Models are pre-trained at resolution 224.
  • Pre-training setup: Very small weight decays are important during pre-training because large weight decay constrains models when capturing information from the large-scale dataset.The authors also report SGD as more performant than Adam in this setting.
  • NFNet pre-training: NFNet-F4 and NFNet-F4+ use batch size 4096, with 40 and 20 pre-training epochs respectively.Both have around 830ms per-step training latency at per-core batch size 32.
  • NFNet pre-training: The wider NFNet-F4+ reaches the same final performance as NFNet-F4 with less pre-training time, given sufficient pre-training data.The authors interpret this as more efficient than training a smaller model for longer.
  • Fine-tuning: 15,000-step ImageNet fine-tuning uses batch size 2048 and learning rate 0.1 with warmup and cosine decay, plus SAM and mild regularization for NFNets.The NFNet fine-tuning setup uses weight decay 10^-5, DropOut 0.25, and stochastic depth 0.1.

B. Downsides of Batch Normalization

Batch normalization introduces implementation and reproducibility challenges, including train–test discrepancies, hardware-sensitive behavior, and interactions between training examples that can affect learning objectives.

  • Implementation challenges: Batch normalization can update statistics during testing despite testing mode being enabled in certain Keras versions.This contradicts user expectations when freezing layers for downstream fine-tuning.
  • Train–test behavior: Batch normalization’s train–test discrepancy has caused published systems to use training-mode normalization during testing, making reported results depend on sampling batch size.The examples cited are DCGAN and SAGAN.
  • Reproducibility: Subtle implementation differences can hamper exact reproduction across hardware, including EfficientNet’s cross-replica statistics and nested moving averages.The number of devices affects statistic computation nonlinearly, while moving averages compound the averaging horizon.
  • Interactions between examples: Batch normalization breaks example independence in contrastive learning, creating potential intra-batch information leakage that affects objectives such as SimCLR and MoCo.MoCo addresses this by shuffling examples between devices, increasing implementation complexity and reducing exact reproducibility.

C. Model Details

NFNets are modified SE-ResNeXt-D models that replace activation normalization with variance-control mechanisms, structured residual stages, and carefully specified convolutional blocks.

  • Overall architecture: NFNets use a modified SE-ResNeXt-D backbone with a four-convolution stem and variance-preserving nonlinearities.The stem uses stride-2 convolutions at its beginning and end, with GELU as the default nonlinearity.
  • Residual stages: The residual stages use [1, 2, 6, 3] baseline blocks, with transition blocks handling downsampling before standard residual blocks.Subsequent variants multiply the baseline block counts by N.
  • Residual blocks: Each bottleneck contains 1 × 1 convolutions around two 3 × 3 grouped convolutions, using a 0.5 bottleneck ratio and group width 128.The first grouped convolution is strided in transition blocks.
  • Residual blocks: Squeeze-and-Excite rescales activations channel-wise using 2σ(FC(GELU(FC(pool(h)))))×h to maintain signal variance.The scalar multiplier of 2 is non-standard and follows Brock et al. (2021).
  • Classifier: The architecture ends with a 1 × 1 channel-expansion convolution, global average pooling, and a 1000-way classifier whose weights use standard deviation 0.01.Zero initialization of the classifier was sometimes unstable with very large numbers of output classes.
  • Variance control: NFNets use no activation normalization layers in residual blocks; instead, variance downscaling divides the main-path input by β and scales block outputs by α, typically 0.2.Transition shortcuts operate on x_l/β_l to reset variance and ensure unit signal variance at each stage’s start.
  • Initialization and parameterization: NFNet convolutions use Scaled Weight Standardization, while SkipInit provides a learnable zero-initialized gain that improves stability in very deep networks.Affine gains, biases, and SkipInit gains are excluded from weight decay.
  • Ablation reporting: Table 6 reports ImageNet Top-1 accuracy and TPUv3 training latency for each model-ablation entry.The left value is accuracy and the right value is training latency.

D. Additional AGC Ablations

AGC stability depends on the clipping threshold and batch size: larger batch sizes require smaller thresholds, motivating λ = 0.01 for batch size 4096 experiments.

  • Clipping-threshold ablation: Smaller AGC clipping thresholds are necessary for stable ResNet200 training at higher batch sizes.Performance was evaluated across clipping thresholds and batch sizes using five independent runs, reporting the average of the best four.
  • Clipping-threshold ablation: For experiments using batch size 4096, the AGC clipping threshold is λ = 0.01.This setting follows the observed stability requirement at larger batch sizes.

E. Negative Results

The authors report which architecture choices failed to improve the accuracy–training-speed trade-off, favoring largely homogeneous designs and dense 1 × 1 convolutions. Aggressive downsampling improved speed at an accuracy cost, while attention variants were not explored.

  • Architecture heterogeneity: Heterogeneous group widths, group counts, and bottleneck ratios were broadly unnecessary and complicated interpretable manual design.The final NFNet designs vary mainly width and stage depth between stages.
  • Aggressive downsampling: Aggressive downsampling using 8 × 8 DCT coefficients improved model speed but reduced accuracy, even with invertible transforms.The authors suggest model internal resolution may matter more than information loss alone.
  • Grouped 1 × 1 convolutions: Squeeze & Excite layers restore cross-channel connectivity more cheaply than fully connected layers but do not fully recover dense 1 × 1 convolution expressivity.Both standard and straddling forms provide scalar-per-channel cross-channel connectivity; straddling yielded slightly better accuracy in the grouped-convolution experiments.
  • Grouped 1 × 1 convolutions: Grouped 1 × 1 convolutions reduced FLOPS and parameter counts but failed to preserve comparable accuracy while reducing training latency.Small group counts slightly reduced accuracy, whereas high group counts substantially reduced model performance; straddled Squeeze & Excite was slightly better than standard Squeeze & Excite.
  • Unexplored directions: The authors did not experiment with attention variants, which they expect could improve their models.This leaves attention-based architecture changes outside the reported experiments.
Loading 2102.06171v1…