Batch Normalization has Multiple Benefits: An Empirical Study on Residual Networks

1 INTRODUCTION . Skip connections ( He et al. , 2016a ) and batch normalization ( Ioffe & Szegedy , 2015 ) have contributed to the dramatic improvement in the performance of convolutional nets in recent years . However batch normalization also has a number of limitations . It breaks the independence between training examples within a batch , performs poorly with small batch sizes ( Wu & He , 2018 ) , and increases the cost of computing an update ( Wu et al. , 2018 ) . Many authors have tried to understand the role batch normalization plays during training ( Balduzzi et al. , 2017 ; Santurkar et al. , 2018 ; Bjorck et al. , 2018 ; Sankararaman et al. , 2019 ; Yang et al. , 2019 ; Ghorbani et al. , 2019 ) , to replace batch normalization with cheaper alternatives ( Ba et al. , 2016 ; Salimans & Kingma , 2016 ; Ulyanov et al. , 2016 ) , or to remove normalization entirely ( Mishkin & Matas , 2015 ; Krähenbühl et al. , 2015 ; Zhang et al. , 2019 ) . Santurkar et al . ( 2018 ) argued that batch normalization helps primarily because it improves the conditioning of the loss landscape , enabling stable training with larger learning rates . Bjorck et al . ( 2018 ) argued that this not only explains why batch normalization helps during optimization , but also explains why it improves the final test accuracy , since large learning rates are thought to generalize well ( Keskar et al. , 2016 ; Smith & Le , 2017 ; Jastrzębski et al. , 2017 ) . Other authors have remarked that batch normalization introduces noise by computing the batch statistics over a subset of the training data , helping generalization ( Luo et al. , 2019 ) . Meanwhile , both Balduzzi et al . ( 2017 ) and Yang et al . ( 2019 ) argued that the combination of batch normalization and skip connections preserves the correlation between the gradients of different training examples at initialization , and they speculated that this enables the efficient training of deep architectures . Mishkin & Matas ( 2015 ) , Krähenbühl et al . ( 2015 ) and Zhang et al . ( 2019 ) have all demonstrated empirically that some of the benefits of batch normalization can be recovered by improving the initialization scheme . In this work , we argue that one can not build a complete picture of the benefits of batch normalization , unless one studies the training loss and test accuracy achieved by deep networks at a wide range of batch sizes . When the batch size is large , one can train with large learning rates , and the role of batch normalization in enhancing the conditioning of the loss landscape becomes apparent . When the batch size is smaller , one does not need large learning rates to train efficiently , and so we can study the other benefits of batch normalization in isolation . Our results establish the following : 1 . Batch normalization does enable training with significantly larger learning rates , and this enables practictioners to efficiently parallelize training over much larger minibatches . 2 . However this is not the only benefit . Batch normalization also performs well on both the training loss and the test accuracy with small batch sizes and small learning rates . 3 . We propose “ ZeroInit ” , a simple change to the initialization scheme of residual networks , designed to preserve the correlation between gradients at the start of training . We show that this enables us to train very deep residual networks without normalization , and our method recovers most of the benefits of batch normalization when the batch size is not too large . Our results suggest that batch normalization improves conditioning ( Santurkar et al. , 2018 ; Bjorck et al. , 2018 ) , preserves the correlation between gradients ( Balduzzi et al. , 2017 ; Sankararaman et al. , 2019 ; Yang et al. , 2019 ) , and has additional regularization benefits ( Luo et al. , 2019 ) . It is easy to miss one or more of these benefits if one performs experiments at a single batch size . One should also be cautious in claiming to identify a complete replacement for batch normalization without first verifying that it matches the performance of batch normalization for a wide range of batch sizes . This paper is structured as follows . In section 2 we introduce residual networks , batch normalization , and Fixup initialization ( Zhang et al. , 2019 ) . In section 3 we compare the test accuracy and training loss of residual networks trained with and without batch normalization at a range of batch sizes , with the aim of isolating the different benefits batch normalization brings . In section 4 we introduce “ ZeroInit ” , a simple initialization scheme inspired by both Fixup and shattered gradients ( Balduzzi et al. , 2017 ) . This scheme trains very deep residual networks to high accuracies without batch normalization , but it does not replicate the full benefits of batch normalization when training with large batch sizes . In section 5 we evaluate the performance of ZeroInit and Fixup on ImageNet . 2 BACKGROUND . Residual Networks ( ResNets ) : First introduced by He et al . ( 2016b ) , these models are composed of a sequence of residual blocks , which contain both a “ residual branch ” containing a number of convolutions and activation layers , as well as a “ skip connection ” , which is typically just the identity . These skip connections create pathways from the input to the output which have a shorter effective depth than the main path through all the residual branches , and these pathways enable the input signal to pass more easily from the input to the output ( Balduzzi et al. , 2017 ; Sankararaman et al. , 2019 ) . Crucially , the network still retains a high degree of expressivity thanks to the main path through all the residual branches . Most of the experiments in this paper will follow the popular Wide-ResNet architecture , introduced by Zagoruyko & Komodakis ( 2016 ) . However we also consider the original ResNet-V1 ( He et al. , 2016a ) and ResNet-V2 ( He et al. , 2016b ) architectures in section 5 . Batch Normalization : As in most previous work , we apply batch normalization in convolutional layers . The inputs to and outputs from batch normalization layers are therefore 4-dimensional tensors , which we denote by Ib , x , y , c and Ob , x , y , c . Here b denotes the minibatch , c denotes the channels , and x and y denote the two spatial dimensions . Batch normalization applies the same normalization to every input in the same channel ( Ioffe & Szegedy , 2015 ) . This leads to the update : Ob , x , y , c = γc Ib , x , y , c − µc√ σ2c + + βc . Here , µc = 1Z ∑ b , x , y Ib , x , y , c denotes the per-channel mean , and σ 2 c = 1 Z ∑ b , x , y I 2 b , x , y , c−µ2c denotes the per-channel variance of the inputs , and Z denotes the normalization constant , which is summed over the minibatch and spatial dimensions x and y . A small constant is generally added to the variance for numerical stability . The “ scale ” and “ shift ” parameters , γc and βc respectively , are learnt during training and are added to preserve model expressivity after the normalization operation . Running averages of the mean µc and variance σ2c are also maintained during training , and these averages are used at test time to ensure the predictions are independent of other examples in the batch . Batch normalization enables researchers to train significantly deeper residual networks ( Sankararaman et al. , 2019 ) , and these networks have achieved state-of-the-art performance on a number of tasks . However batch normalization also has a number of limitations . It breaks the independence between training samples in a minibatch , which makes it harder to apply in certain models ( Girshick , 2015 ) , and also contradicts the assumptions of most theoretical models of optimization ( Polyak & Juditsky , 1992 ; Mandt et al. , 2017 ; Park et al. , 2019 ) . The normalization operation itself can constitute a significant fraction of the total cost of computing a parameter update in some models ( Wu et al. , 2018 ) . Batch normalization also performs poorly when the batch size is too small ( Ioffe , 2017 ; Wu & He , 2018 ) , which often limits the size of model that can be trained on a single device . Fixup : Zhang et al . ( 2019 ) proposed Fixup initialization which removes batch normalization entirely from residual networks . They demonstrated that this algorithm is able to train ResNet-50 to competitive validation accuracies on ImageNet , while it also displayed promising results when training very deep ( 1000+ layers ) Wide-ResNets on CIFAR-10 . The Fixup algorithm has four components : 1 . The classification layer and final convolution of each residual branch are initialized to zero . 2 . The initial weights of the remaining convolutions are scaled down by L−1/ ( 2m−2 ) , where L is the number of residual branches and m the number of convolutions in each branch . 3 . A scalar multiplier is introduced at the end of each residual branch , intialized to one . 4 . Scalar biases are introduced before every convolution , linear and element-wise activation layer , initialized to zero . It is not clear whether all of these components are required to enable efficient training without batch normalization , or if a subset of these components would be sufficient . The authors claim that component 2 above is crucial , however we will demonstrate below that it is not required in practice . Additionally , the authors only provide results at a single batch size and do not tune the learning rate . We demonstrate below that a significantly simpler initialization scheme ( ZeroInit ) enables the efficient training of very deep residual networks without batch normalization when the batch size is not too large . ZeroInit and Fixup both achieve similar test accuracies to batch normalization when the batch size is small , but they underperform batch normalization when the batch size is very large . 3 THE BENEFITS OF BATCH NORMALIZATION . The primary goal of this paper is to identify the empirical benefits of batch normalization when training residual networks . We perform most of our experiments on CIFAR-10 ( Krizhevsky et al. , 2009 ) , since this enables us to run thorough experiments , averaging across multiple training runs and tuning the learning rate at each batch size . To verify that our conclusions scale to larger datasets , we provide a limited set of experiments on ImageNet in section 5 ( Russakovsky et al. , 2015 ) . Experimental Setup : In figure 1 , we provide the mean performance of a 16-4 Wide-ResNet ( Zagoruyko & Komodakis , 2016 ) , trained on CIFAR-10 for 200 epochs at a range of batch sizes . At each batch size , we train the network 15 times for a range of learning rates on a logarithmic grid , and we independently measure the mean and standard deviation of the best 12 runs on both the training loss and the test accuracy . This ensures that our results are not corrupted by outliers or failed runs . Throughout this paper , we train using SGD with heavy ball momentum , and fix the momentum coefficient m = 0.9 . The optimal test accuracy is the mean performance at the learning rate whose mean test accuracy was highest , while the optimal training loss is the mean training loss at the learning rate whose mean training loss was lowest . We always verify that the optimal learning rates are not at the boundary of our grid search . Although we tune the learning rate on the test set , we emphasize that our goal is not to achieve state of the art results . Our goal is to compare the performance of different training procedures . We apply the same experimental protocol in each case . We adopt the following simple learning rate schedule . We first hold the learning rate constant for 100 epochs , before dropping the learning rate by a factor of 2 every 10 epochs . We find that this simple schedule outperforms the original sequence of 3 drops proposed by He et al . ( 2016a ) . We apply data augmentation including padding , random crops and left-right flips and we also use L2 regularization with a coefficient of 5× 10−4 . We provide results both with and without batch normalization . When using batch normalization , we evaluate the batch statistics over a fixed ghost batch size of 64 ( Hoffer et al. , 2017 ) , before accumulating gradients to form larger batches . Evaluating the batch statistics over a fixed number of training examples independent of the batch size preserves the benefits of noise when the batch size is large ( Luo et al. , 2019 ) , and reduces communication overheads in distributed training ( we provide experiments using the original formalism of batch normalization in figure 3 ) . For completeness , we also provide results here without normalization using a simple modification to the initialization scheme , “ ZeroInit ” . We will describe this modification in full in section 4 . Results : Looking at the optimal test accuracy in figure 1a , we make three observations . First , the optimal test accuracy is substantially higher with batch normalization than without batch normalization , even after we tune both the learning rate and the batch size . Second , both with and without batch normalization , the performance is independent of batch size in the small batch limit , before beginning to fall as the batch size rises . Third , the largest batch size at which one can train efficiently with batch normalization is substantially larger than the corresponding batch size without batch normalization . Since the number of training epochs is fixed , this indicates that one can train Wide-ResNets with significantly fewer parameter updates when using batch normalization . Given access to sufficient hardware , this will enable practitioners to dramatically reduce the wallclock time of training ( Goyal et al. , 2017 ) . ZeroInit significantly improves the optimal test accuracy achieved without normalization , and is competitive with batch normalization ( within one standard deviation ) for batch sizes B . 256 . However it underperforms batch normalization when the batch size is large . Now considering the optimal training accuracy in figure 1b , we find that batch normalization achieves lower optimal training losses at all batch sizes , and continues to perform well when the batch size is large . The optimal training loss without batch normalization is similar both with and without ZeroInit . However when ZeroInit is not used , the optimal training loss degrades rapidly as the batch size is increased . To help interpret these results , in figure 2 we provide the optimal learning rates at which both the training loss is minimized and the test accuracy is maximized . From figure 2a , we conclude that for all three techniques , the optimal learning rate is initially proportional to the batch size , before saturating when the batch size is large ( Krizhevsky , 2014 ; Goyal et al. , 2017 ; Smith et al. , 2017 ; McCandlish et al. , 2018 ; Shallue et al. , 2018 ) . Furthermore , the optimal ratio of the learning rate to the batch size is similar for all three techniques . However , while the optimal learning rate saturates at a maximum value of 2−3 without batch normalization , it reaches a larger value of 21 with ZeroInit , and an even larger value of 24 with batch normalization . This directly explains why batch normalization and ZeroInit perform better when the batch size is large ; they have larger maximum stable learning rates which enables them to scale efficiently to larger batches ( McCandlish et al. , 2018 ) . We observe very similar results when considering the optimal learning rates minimizing the training loss in figure 2b . Discussion : Santurkar et al . ( 2018 ) and Bjorck et al . ( 2018 ) argued that the empirical success of batch normalization is primarily explained by the increase in the maximum stable learning rate , which arises because batch normalization improves the conditioning of the loss landscape . However our results conclusively demonstrate that this is not the case . While a larger maximum stable learning rate is the primary reason why batch normalization performs well with large training batches , when the batch size is small the optimal learning rate both with and without batch normalization is also small . Nevertheless , the optimal test accuracy and the optimal training loss are significantly improved when batch normalization is used in the small batch limit as well as the large batch limit . We conclude that batch normalization increases the largest stable learning rate , and this enables us to reduce the number of parameter updates required to train the model when training with large batches . However batch normalization also improves the final training loss and test accuracy , even when the batch size and learning rate are small . Our results demonstrate that one can not understand batch normalization without studying its performance at a range of batch sizes . ZeroInit recovers much of the performance gap on the test set when the batch size is small , but it does not close the gap in the training loss , and it performs worse with large batches . It is difficult to measure exact wallclock times in our distributed training setup , but we find that 16-4 wide residual networks typically complete 200 epochs roughly 40 % faster at batch size 64 without batch normalization ( Wu et al. , 2018 ) .

Summary: This paper conducts extensive experiments to verify two practical benefits of batch normalization. i) It increases the final test accuracy and the largest stable learning rate; ii) it enables efficient training with larger batches and a larger learning rate. In addition, the authors propose a new initialization scheme, “ZeroInit”, to train a deep ResNet to improve the test accuracy. My detailed comments are as follows.