Group Whitening: Balancing Learning Efficiency and Representational Capacity
Lei Huang, Yi Zhou, Li Liu, Fan Zhu, Ling Shao
Introduction
Batch normalization (BN) represents a milestone technique in deep learning , and has been extensively used in various network architectures . BN standardizes the activations within a mini-batch of data, which improves the conditioning of optimization and accelerates training . Further, the stochasticity of normalization introduced along the batch dimension is believed to benefit generalization . However, this stochasticity also results in differences between the training distribution (using mini-batch statistics) and the test distribution (using estimated population statistics) , which is believed to be the main cause of BN’s small-batch-size problem — BN’s error increases rapidly as the batch size becomes smaller . To address this issue, a number of approaches have been proposed . One representative method is group normalization (GN), which divides the neurons into groups and then applies the standardization operation over the neurons of each group, for each sample, independently. GN provides a flexible solution to avoid normalization along the batch dimension, and benefits visual tasks limited to small-batch-size training .
As a widely used operation in data pre-processing, whitening not only standardizes but also decorrelates the data , which further improves the conditioning of the optimization problem . A whitened input has also been shown to make the gradient descent updates similar to the Newton updates for linear models . Motivated by this, Huang et al. proposed batch whitening (BW) for deep models, which performs whitening on the activations of each layer within a mini-batch. BW has been shown to achieve better optimization efficiency and generalization than BN . However, BW further amplifies the disadvantage of BN in estimating the population statistics, where the number of parameters to be estimated with BW is quadratic to the number of neurons/channels. Thus, BW requires a sufficiently large batch size to work well.
To exploit whitening’s advantage in optimization, while avoiding its disadvantage in normalization along the batch dimension, this paper proposes group whitening (GW). GW divides the neurons of a sample into groups for standardization over the neurons in each group, and then decorrelates the groups. Unlike BW, GW has stable performance for a wide range of batch sizes, like GN, and thus can be applied to a variety of tasks. GW further improves the conditioning of optimization of GN with its whitening operation.
One important hyperparameter of GW is the group number. We observe that GW/GN has a significantly degenerated training performance when the group number is large, which is similar to the small-batch-size problem of BW/BN. We attribute this to the constraints on the output imposed by the normalization operation, which affect the model’s representational capacity. As such, this paper defines the constraint number of normalization (as will be discussed in Section 4) to quantitatively measure the magnitude of the constraints provided by normalization methods. With the support of the constraint number, we analyze how the batch size (group number) affects the model’s representational capacity for batch (group) normalized networks. Our analysis also presents a new viewpoint for understanding the small-batch-size problem of BN.
We apply the proposed GW to two representative deep network architectures (ResNet and ResNeXt ) for ImageNet classification and COCO object detection and instance segmentation . GW consistently improves the performance for both architectures, with absolute gains of ~ in top-1 accuracy for ImageNet and ~ in bounding box AP for COCO.
Preliminaries
During training, batch normalization (BN) standardizes the layer input within a mini-batch, for each neuron, asBN and other normalization methods discussed in this paper all use extra learnable scale and shift parameters . We omit this for simplicity.:
Here, and , where is the variance over mini-batches for the i-th neuron, is a column vector of all ones, and is a small number to prevent numerical instability. During inference, the population statistics are required for deterministic inference, and they are usually calculated by running average over the training iterations, as follows:
Such an estimation process can limit the usage of BN in recurrent neural networks , or harm the performance for small-batch-size training .
To avoid the estimation of population statistics shown in Eqn. 2, Ba et al. proposed layer normalization (LN) to standardize the layer input within the neurons for each training sample, as:
Here, and , where is the variance over the neurons for the i-th sample. LN has the same formulation during training and inference, and is extensively used in natural language processing tasks .
To exploit the advantage of whitening over standardization in improving the conditioning of optimization, Huang et al. proposed decorrelated BN , which performs zero-phase component analysis (ZCA) whitening to normalize the layer input within a mini-batch, as:
where indicates the trace of and is calculated iteratively as:
Here, . Other BW methods also exist for calculating the whitening matrix ; please refer to for more details.
It is necessary for BW to estimate the population statistics of the whitening matrix during inference, like BN. However, the number of independent parameters in of BW is , while of BN is . This amplifies the difficulty in estimation and requires a sufficiently large batch size for BW to work well (Figure 1). Although group-based BW — where neurons are divided into groups and BW is performed within each one — can relieve this issue, it is still sensitive to the batch size (Figure 1) due to its inherent drawback of normalizing along the batch dimension.
Group Whitening
GW avoids normalization along the batch dimension, and it works stably across a wide range of batch sizes (Figure 1). GW also ensures that the normalized activation for each sample has the properties: and , which should improve the conditioning, like BW, and benefit training. We conduct several experiments to validate this, and the results in Figure 2 show that the group whitened output (by GW) has significantly better conditioning than the group standardized one (by GN), which is similar to normalization along the batch dimension . Note that the condition number of BW is . We also find that GN/GW has better conditioning with increasing group number. Besides, we find that BN has better conditioning than GN/GW, which suggests that normalizing along the batch dimension is better for decorrelating the data than normalizing along the channel dimension.
Our method is significantly different from the group-based BW , in which the whitening operation is also applied within mini-batch data. Specifically, group-based BW has difficulty in estimating the population statistics, as discussed in Section 2. Note that group-based BW is reduced to BN if the channel number in each group , while GW is reduced to GN if the group number .
Revisiting the Constraint of Normalization
Here, we define the constraint number of normalization to quantitatively measure the magnitude of the constraints provided by the normalization method.
BN’s benefits in accelerating the training of DNNs are mainly attributed to two reasons: 1) The distribution is more stable when fixing the first and second momentum of the activations, which reduces the internal covariant shifts ; 2) The landscape of the optimization objective is better conditioned , by improving the conditioning of the activation matrix with normalization. Based on these arguments, GW/GN should have better performance when increasing the group number, due to the stronger constraints and better conditioning. However, we experimentally observe that GN/GW has significantly degenerated performance when the group number is too large (Figure 1 (b)), which is similar to the small-batch-size problem of BN/BW. We investigate the reason behind this phenomenon.
We first show that the batch size/group number has a value range, which can be mathematically derived. The normalization operation can be regarded as a way to find a solution satisfying the constraints . To ensure the solution is feasible, it must satisfy the following condition:
where is the number of variables in . Based on Eqn. 12, we have for BN to ensure a feasible solution. We also provide the ranges of batch size/group number for other normalization methods in Table 1. Note that the batch size should be larger than or equal to to achieve a numerically stable solution for BW when using ZCA whitening in practice . This also applies to GW, where should be less than or equal to .
We then demonstrate that normalization eventually affects the feature representation in a certain layer. Figure 3 shows the histogram of normalized output , by varying of GNNote that the channel number in each group . We vary , rather than , for simplifying the discussion. and of BN. We observe that: 1) the values of are heavily constrained if or is too small, e.g., the value of is constrained to be if ; 2) is not Gaussian if or is too small, while BN/GN aims to produce a normalized output with a Gaussian distribution. We believe that the constrained feature representation caused by GN/GW with a large group number is the main factor leading to the degenerated performance of a network. Besides, we also observe that the normalized output of GN is more correlated than that of BN, which supports the claim that BN is more capable of improving the conditioning of activations than GN, as shown in Section 3.
In summary, our qualitative and quantitative analyses show that group/batch based normalizations have low diversity of feature representations when is small. We believe these constrained feature representations affect the performance of a network, and can lead to significantly deteriorated results when the representation is over-constrained.
2 Effect on Representational Capacity of Model
The constraints introduced by normalization are believed to affect the representational capacity of neural networks , and thus the learnable scale and shift parameters are used to recover the representations . However, such an argument is seldom validated by either theoretical or empirical analysis. Theoretically analyzing the complexity measure (e.g., VC dimensions or the number of linear regions ) of the representational capacity of neural networks with normalization is a challenging task, because normalized networks do not follow the assumptions for calculating linear regions or VC dimensions. Here, we conduct preliminary experiments, seeking to empirically show how normalization affects the representational capacity of a network, by varying the constraints imposed on the feature.
We follow the non-parametric randomization tests fitting random labels to empirically compare the representational capacity of neural networks. To rule out the optimization benefits introduced by normalization, we first conduct experiments using a linear classifier, where normalization is also inserted after the linear module. We train over epochs using stochastic gradient descent (SGD) with a batch size of 16, and report the best training accuracy among the learning rates in in Figure 5 (a). We observe that GN and GW have lower training accuracy than when normalization is not used, which suggests that normalization does indeed reduce the model’s representational capacity in this case. Besides, the accuracy of GN/GW decreases as the group number increases. This suggests that the model may have weaker representational ability when increasing the constraints on the feature. Note that we have the same observations regardless of whether or not the learnable scale and shift parameters of GN/GW are used.
3 Discussion of Previous Work
Previous analyses on BN are mainly derived from the perspective of optimization . One argument is that BN can improve the conditioning of the optimization problem , either by avoiding the rank collapse of pre-activation matrices or alleviating the pathological sharpness of the landscape . This argument has been further investigated by computing the spectrum of the Hessian for a large-scale dataset . The improved conditioning enables large learning rates, thus improving the generalization . Another argument is that BN is scale invariant , enabling it to adaptively adjust the learning rate , which stabilizes and further accelerates training . Other analyses focus on investigating the signal and gradient propagation, either by exploiting mean-field theory , or a neural tangent kernel (NTK) .
Different from these works, we are the first to investigate how BN/GN affects a model’s representational capacity, by analyzing the constraint on the representation of internal features. This opens new doors in analyzing and understanding normalization methods. We also investigate how batch size affects the training performance of batch normalized networks (Figure 1 (a)), from the perspective of a model’s representational capacity. Several works have shown that batch size is related to the magnitude of stochasticity introduced by BN, which also affects the model’s training performance. However, the stochasticity analysis is specific to normalization along the batch dimension, and cannot explain why GN with a large group number has significantly worse performance (Figure 1 (b)). Our work provides a unified analysis for batch and group normalized networks.
Large-Scale Visual Recognition Tasks
We investigate the effectiveness of our proposed GW on large-scale ImageNet classification , as well as COCO object detection and segmentation . We use the more efficient and numerically stable ‘ItN’ (with ) to calculate the whitening matrix for both GW and BW, in all experiments. Our implementation is based on PyTorch .
We experiment on the ImageNet dataset with 1,000 classes . We use the official 1.28M training images as a training set, and evaluate the top-1 accuracy on a single-crop of 224x224 pixels in the validation set with 50k images. We investigate the ResNet and ResNeXt models.
We follow the same experimental setup as described in , except that we use two GPUs and train over 100 epochs. We apply SGD with a mini-batch size of 256, momentum of 0.9 and weight decay of 0.0001. The initial learning rate is set to 0.1 and divided by 10 at 30, 60 and 90 epochs. Our baseline is the 50-layer ResNet (ResNet-50) trained with BN .
We investigate the effects of group number for GW/GN, which we use to replace the BN of ResNet-50. We vary the group number ranging in (we use the channel number if it is less than the group number in a given layer), and report the training and validation accuracies in Figure 6. We can see that GW has consistent improvement over GN in training accuracy, across all values of , which indicates the advantage of the whitening operation over standardization in terms of optimization. Besides, GW also has better validation accuracy than GN. We believe this may be because the stronger constraints of GW contribute to generalization. We also observe that both GN and GW have significantly reduced training accuracy when the group number is too large (e.g., g=128), which is consistent with the previous results in Figure 5.
Although GW (=64) provides slight improvement over the BN baseline ( vs. ), it has a additional time costNote that our implementations are based on the APIs provided by PyTorch and are not finely optimized. For more discussion on time costs, please refer to the Appendix F.2. on ResNet-50. Based on the analysis in Section 4, it is reasonable to only partially replace BN with GW in networks, because 1) normalization within a batch or a group of channels both have their advantages in improving the optimization and generalization; 2) whitening can achieve better optimization efficiency and generalization than standardization , but at a higher computational cost .
Here, we investigate the position at which to apply GW (=64) in ResNet-50. ResNet and ResNeXt are both composed primarily of a stem layer and multiple bottleneck blocks . We consider: 1) replacing the BN in the stem layer with GW (referred to as ‘S1’); and 2) replacing the , , , and BNs in all the bottleneck blocks, which are referred to as ‘B1’, ‘B2’, ‘B12’ and ‘B3’, respectively. We investigate five architectures, S1, S1-B1, S1-B2, S1-B3 and S1-B12, which have 1, 17, 17, 17 and 33 GW modules, respectively. We also perform experiments using BW and BWΣ (employing a covariance matrix to estimate the population statistics of BW) for contrast.
We report the results in Table 2. BW/BWΣ improve their BN counterparts on all architectures by a clear margin, which demonstrates the advantage of the whitening operation over standardization . GW provides significant improvements over BW/BWΣ on S1-B1, S1-B2, S1-B3 and S1-B12 (an absolute improvement of on average). We attribute this to the advantage of GW in avoiding the estimation of population statistics. We also observe that GW has a slightly worse performance on S1-B12 than on S1-B1/S1-B2. We believe there is a trade-off between GW and BN, in terms of affecting the model’s representational capacity, optimization efficiency and generalization.
We also investigate the effect of inserting a GW/BW layer after the last average pooling (before the last linear layer) to learn the decorrelated feature representations, as proposed in . This can slightly improve the performance ( on average) when using GW, though the net gain is smaller than using BW () or BWΣ (). Please refer to the Appendix F.1 for details.
1.2 Validation on Larger Models
In this section, we further validate the effectiveness of GW on ResNet-101 , ResNeXt-50 and ResNeXt-101 . We apply GW (=64) in these models following the S1-B2 architecture, which achieves the best performance (Table 2) without significantly increasing the computational cost (it is only increased by roughly ). For comparison, we also apply BWΣ following the ‘S1-B2’ architecture, combining the learning of decorrelated features (BWΣ has a slightly improved performance compared to BW ). Our baselines are the original networks trained with BN, and we also provide the results trained with GN.
The results are shown in Table 3. We can see that 1) our method improves the baseline (BN) by a significant margin (between and ); and 2) BWΣ has consistently better performance than BN, but the net gain is reduced on wider networks (RexNeXt-50 and ResNeXt-101), which is probably caused by the difficulty in estimating the population statistics. We also conduct experiments using more advanced training strategies (e.g., cosine learning rate decay , label smoothing and mixup ) and GW again improves the baseline consistently. Please refer to the Appendix F.3 for details.
2 Object Detection and Segmentation on COCO
We fine-tune the models trained on ImageNet for object detection and segmentation on the COCO benchmark . We experiment on the Faster R-CNN and Mask R-CNN frameworks using the publicly available codebase ‘maskrcnn-benchmark’ . We train the models on the COCO set and evaluate on the COCO set. We report the standard COCO metrics of average precision (AP), AP50, and AP75 for bounding box detection (APbbox) and instance segmentation (APm) . For BN, we use its frozen version (indicated by ) when fine-tuning for object detection .
For the Faster R-CNN framework, we use the ResNet-50 models pre-trained on ImageNet (Table 3) as the backbones, combined with the feature pyramid network (FPN) . We consider two setups: 1) we use the box head consisting of two fully connected layers (‘2fc’) without a normalization layer, as proposed in ; 2) following , we replace the ‘2fc’ box head with ‘4conv1fc’, which can better leverage GN, and apply GN/GW to the FPN and box head. We use the default hypeparameter configurations from the training scripts provided by the codebase for Faster R-CNN. The results are reported in Table 4. The GW pre-trained model improves and GN by and AP, respectively. By adding GW/GN to the FPN and ‘4conv1fc’ head box, GW improves and GN by and AP, respectively.
For the Mask R-CNN framework, we use the ResNeXt-101 models pre-trained on ImageNet (Table 3) as the backbones, combined with FPN. We use the ‘4conv1fc’ box head, and apply GN/GW to the FPN, box head and mask head. We again use the default hypeparameter configurations from the training scripts provided by the codebase for Mask R-CNN . The results are shown in Table 5. GW achieves in box AP and in mask AP, an improvement over of and , respectively.
Here, we further show that the network mixed with BNs and GWs can still work well under small-batch-size scenarios. As illustrated in , one main cause of BN’s small-batch-size problem is the inaccurate estimation between training and inference distributions, which is amplified for a network with increased BN layers (these inaccuracies are compounded with depth). We believe inserting GW (which ensures the same distribution between training and inference) between consecutive BN layers will ‘break’ these compounding inaccuracies, thus relieving the small-batch-size problem of BNs in a network. We train Faster R-CNN from scratch and use normal BN that is not frozen. We follow the same setup as in the previous experiment (e.g., two images/GPUs). We find that using all BNs only obtains 25.10 AP, while 28.37 AP is achieved using our mixture of BNs and GWs (the S1-B2 architecture). Note that using all GNs (GWs) obtains 28.19 (28.79) AP. This experiment further validates that our mixture of BNs and GWs may also help mitigate the small-batch-size problem of the BNs in a network.
Conclusion and Future Work
In this paper, we proposed group whitening (GW), which combines the advantages of normalization within a group of channels and the whitening operation. The effectiveness of GW was validated on large-scale visual recognition tasks. Furthermore, we also analyzed the feature constraints imposed by normalization methods, enabling further understanding of how the batch size (group number) affects the performance of batch (group) normalized networks from the perspective of representational capacity. This analysis can provide theoretical guidance for applying GW and other normalization methods in practice. It would be interesting to build a unified framework to further investigate the effects of normalization in representation, optimization and generalization, by combining the proposed constraint analysis with the previous conditioning analysis and stochasticity analysis . Our GW has also the potentialities to be used as a basic module in the switchable normalization methods to improve their performance.
References
Appendix A Algorithms
The forward pass of the proposed group whitening (GW) method is shown in Algorithm 1, and its corresponding backward pass is shown in Algorithm 2.For GW, we also use the extra learnable dimension-wise scale and shift parameters, like BN . We omit this in the algorithms for simplicity. Note that we need to specify the method for calculating the whitening matrix in Line 8 of Algorithm 1, as well as its backward operation shown in Line 5 of Algorithm 2. As stated in the paper, we use zero-phase component analysis (ZCA) whitening and its efficient approximation by Newton’s iteration (‘ItN’) . Here, we provide the details.
ZCA whitening calculates the whitening matrix by eigen decomposition as: , where and are the eigenvalues and associated eigenvectors of , i.e..
The corresponding backward operation is as follows:
where sets the off-diagonal elements of as zero.
‘ItN’ whitening calculates the whitening matrix by Newton’s iteration as: , where indicates the trace of and is calculated iteratively as:
Here, .
The corresponding backward operation is as follows:
Here, can be calculated by the following iterations:
We also provide the python code of GW using ItN whitening, based on PyTorch , in Figure A1.
Appendix B More Results on Effects of Batch Size and Group Number
In Figure 1 of the paper, we show the effects of batch size (group number) for batch (group) normalized networks, where the results are obtained with a learning rate of 0.1. Here, we provide more results using different learning rates, shown in Figure A2. We obtain similar observations.
Appendix C More Results on Conditioning Analysis
In Figure 2 of the paper, we perform a conditioning analysis on the normalized output, where we report and use a one-layer and two-layer multilayer perceptron (MLP) as to obtain the activations. Here, we provide more results, shown in Figure A3. We obtain similar observations.
Appendix D Derivation of Constraint Number of Normalization Methods
BW ensures that the normalized output is centered and whitened, which has the constraints as:
where is a -dimensional column vector of all zeros, and is a matrix of all zeros. Note that there are independent equations in the system of equations . Let’s denote . We have , and thus is a symmetric matrix. Therefore, there are independent equations in the system of equations . We thus have independent equations in , and the constraint number of BW is .
where, and . This ensures that the normalized output for each sample has the constraints:
In the system of equations 23, the number of independent equations is . Therefore, the constraint number of GN is , when given samples.
Following the analysis for BW, the number of independent equations is from Eqns. 27 and 28. Therefore, the constraint number of GW is , when given samples.
Appendix E Investigating Representational Capacity on CNNs
Figure A4 shows the results of GN with varying group number (we use a batch size of 256), where we report the difference in training accuracy between GN and the model without normalization (‘Base’). We observe that: 1) GN has significantly degenerated performance when the group number is too large (relative to the channel number), e.g., GN has worse performance than ‘Base’ when ; 2) The net gain of GN over ‘Base’ is amplified as the depth increases. These observations are consistent with the experiments on the MLPs shown in Section 4.2 of the paper.
Figure A5 gives the results of BN with varying batch size, wher we report the difference in accuracy between BN and ‘Base’. We observe that: 1) BN has significantly degenerated performance when the batch size is too small, e.g., BN has worse performance than ‘Base’ when or ; 2) The net gain of BN over ‘Base’ is amplified as the depth increases. These observations suggest that there is also a trade-off for BN between the benefits of normalization on optimization and its constraints on representation.
Appendix F More Experimental Results on ImageNet
As described in Section 5.1.1 of the paper, we investigate the effect of inserting a GW/BW layer after the last average pooling (before the last linear layer) to learn the decorrelated feature representations, as proposed in . We provide the results in Table A1. This can slightly improve the performance ( on average) when using GW (comparing Table A1 to Table 2 of the paper). We note that benefits the most from this kind of architecture.
F.2 Running Time Comparison
In this section, we compare the wall-clock time of the models described in Section 5.1 of the paper. We run the experiments on GPUs (NVIDIA Tesla V100). All implementations are based on the API provided by PyTorch, with CUDA (version number: 9.0). We use the same experimental setup as described in Section 5.1 of the paper. We evaluate the training time for each iteration, averaged over 100 iterations. The ResNets-50 baseline (BN) costs . Replacing the BNs of ResNet-50 with our GWs (=64) costs , a additional time cost on ResNet-50. This is one factor that drives us to investigate the position at which to apply GW.
Table A2 shows the time costs of five architectures, S1, S1-B1, S1-B2, S1-B3 and S1-B12, which have 1, 17, 17, 17 and 33 GW modules, respectively. Note that applying GW in the S1-B3 architecture results in a clearly increased computational cost, compared to S1-B1/S1-B2. This is because the channel number of the third normalization layer is larger than that of the first/second normalization layer, in the bottleneck blocks .
Table A3 shows the time costs of ResNets and ResNeXts (the corresponding models in Table 3 of the paper) for ImageNet classification.
F.3 Results on Advanced Training Strategies
In Section 5.1.2 of the paper, we show the effectiveness of our GW on ResNets and ResNeXts , under the standard training strategy (e.g., using learning rate step decay). Here, we also conduct experiments using more advanced training strategies: 1) We train over 100 epochs with cosine learning rate decay ; 2) We add the label smoothing tricks with a smoothing constant ; 3) We use mixup training with in the Beta distribution . The results are shown in Table A4, where GW improves the baselines consistently.
Appendix G Additional Experiments on Neural Machine Translation
Our GW does indeed generalize layer normalization (LN), which is a widely used technique in NLP tasks. We thus believe our GW has the potential to improve the performance of LN in NLP tasks. We conduct additional experiments to apply our GW on Transformer (where LN is the default normalization) for machine translation tasks using fairseq-py . We evaluate on the public IWSLT14 German-to-English (De-EN) dataset using BLEU (higher is better). We use the hyper-parameters recommended in fairseq-py for Transformer and train over 50 epochs with five random seeds. The baseline LN has a BLEU score of . GW (replacing all the LNs with GWs) has a BLEU score of . Note that the hyperparameters were designed for LN, and may not be optimal for GW.