To Fold or Not to Fold: a Necessary and Sufficient Condition on Batch-Normalization Layers Folding

1 Introduction

Deep Neural Networks (DNNs) are ubiquitous in various subdomains of computer vision, e.g. in Image Classification [He et al., 2016], Object Detection [He et al., 2017] or Semantic Segmentation [Chen and others, 2017]. Such models evolved from narrow architectures [LeCun and others, 1989], to deeper and more complex architectures [Tan and Le, 2019]. The main concern in the early developments of very deep neural networks was the efficiency of the learning process [Krizhevsky et al., 2012] with batch-normalization (BN) layers [Ioffe and Szegedy, 2015] being a prime example of solution to this challenge. BN layers center and reduce the features which solves the internal covariate shift problem arising from the change of distribution of the layer inputs. It was so effective that it became almost mandatory in modern architectures.

However, very deep neural network deployment is still limited due to their heavy computational requirements. To tackle this problem, many solutions for DNN acceleration have been developed including, but not limited to, pruning [Frankle and Carbin, 2018; Lin and others, 2020; Yvinec et al., 2021] and quantization [Zhao and others, 2019; Meller et al., 2019;

∗Contact Author

Finkelstein et al., 2019]. Pruning consists in removing a fraction of the mathematical operations performed within each layers while quantization consists in reducing the temporal cost of scalar operations by using lower bit-wise representations. However most of these methods leverage the dependencies across consecutive convolutional and fully-connected layers. Those dependencies are affected by BN layers. This is the first motivation for BN folding. It consists in removing the BN layers from the graph by updating the remaining parameters to keep the predictive function unchanged. It facilitates the use of most pruning and quantization protocols.

The second motivation is straightforward: the lower the number of layers, the faster the inference, all other factors being equal, which is the case with BN folding [Jacob and others, 2018]. Although BN operations adds very few parameters to the network, the induced computational cost is important because of the sequentiality of the operations. The folding process removes this bottleneck and without changing the predictive function but it is still often overlooked in the literature.

For those two reasons, BN folding is usually presented in a naive way where BN layers are folded if and only if they are immediately connected to a convolutional layer or a fullyconnected layer. Although, this condition is sufficient for a BN layer to be foldable, it is far from necessary meaning that it can't be optimal. In this study, we determine a necessary and sufficient condition for a BN layer to be foldable which translates into an algorithm for optimal batch-normalization layers folding, dubbed BaN-OFF. Our tests are two-folds: first, we validate the systematic improvement over the standard folding protocol. Second, we demonstrate the importance as a compression step in more complex pipeline. Our contributions are:

• From a theoretical standpoint: a novel necessary and sufficient condition on the feasibility of batch normalization layers folding in any neural network architecture.
• From a practical standpoint: an algorithm which implements the proposed optimal condition and outperforms the existing folding methods at virtually no cost in preprocessing time.
• Outstanding inference acceleration results on several tasks and architectures and systematic improvement from the standard folding algorithm, see Figure 1.

In this work, we focused on BN as defined by [Ioffe and Szegedy, 2015], however, all our results extend to any normalization layers that implements static affine transformations such as Group Normalization [Wu and He, 2018] or Switchable Normalization [Luo and others, 2019].

2 RelatedWork

Batch-normalization, introduced by Ioffe et al. in [2015], had and continues to have a tremendous impact in deep learning. Prior to it, architectures such as VGG [Simonyan and Zisserman, 2014] achieved remarkable results on ImageNet [Deng et al., 2009] but were difficult to train. However as discussed in [Martin and Mahoney, 2019] where the authors drew a comparison between the standard VGG and VGG-BN where additional BN layers are used. It appeared that such layers drastically improve the performance while simplifying the training process. Concurrently, these layers have been used in almost all novel architectures ([Ren and others, 2015; He et al., 2016; Szegedy and others, 2016; Lin and others, 2017; Zhao and others, 2017; Sandler et al., 2018; Zhou and others, 2019; Tan and Le, 2019; Tan et al., 2020]) to the point that they are now considered an unavoidable part of the standard convolutional and fully-connected layers.

Inference Acceleration. There are two main approaches to DNN inference acceleration: pruning and quantization. Pruning consists in removing elements of the graph defined by the DNN [Renda et al., 2020]. Quantization consists in mapping the weight values from continuous values to evenly-spread integer values using lower bit-wise representations (e.g. int8 or int4) [Krishnamoorthi, 2018]. Both approaches are often combined or even intrinsically leverage BN folding.

Batch-normalization Folding. BN folding consists in removing such BN layers from the network's graph by updating the appropriate adjacent layers. This operation is performed if and only if the structural changes don't affect the predictive function associated to the network. The naive approach [Jacob and others, 2018] is widely documented and applied in DNN acceleration. However it is not optimal as it doesn't fold every foldable BN layers. To solve this issue, we propose BaN-OFF, a necessary and sufficient condition for such layers to be folded as well as its corresponding algorithm.

Batch-normalization Folding and other Accelerations. In [Yvinec et al., 2021], authors propose to prune similar neurons based on their weight representations. To preserve the predictive function they update the consecutive layers and to do so they perform BN folding. Other pruning mechanisms also leverage BN folding such as [Srinivas and Babu, 2015; Liu and others, 2017; Ye and others, 2018]. [Nagel et al., 2019] aim at balancing the weights, for quantization, across the entire network post BN folding and before performing quantization. Similarly, in [Zhao and others, 2019], the method folds the BN layers before splitting outlier neurons in order to compactify the distribution of the weight values thus reducing the quantization error.

Our novel algorithm allows to fold more BN layers than existing baselines, significantly reducing the inference time of deep neural networks. In addition, BAN-OFF allows to process more layers using existing pruning or quantization methods.

3 Methodology

Let f be a feed forward neural network with L layers (fl)l∈{1,...,L} . In this work, we define any layer as one of the following operations:

1. expressive layers: a layer fl is expressive if and only if it can be expressed as the mathematical operation fl : x ↦→ Wlx + bl.
2. non-affine layers: a layer fl is non-affine if and only if it corresponds to a mathematical operation that doesn't satisfy f (ax + b) = ˜af (x) + ˜b.
3. BN layers [Ioffe and Szegedy, 2015]: a layer fl of BN is defined, in the general case, by four parameters γl, βl, µl and σl such that f : x ↦→ γl x-µl σl+ϵ + βl with ϵ a small constant positive scalar (usually ϵ = 10-3 ).
4. other layers.

For instance, convolutional and fully-connected layers without their activation functions are the most common examples of expressive layers. Non-affine layers correspond and are limited to the activation functions in most deep neural network architectures. Finally, pooling layers, concatenations, additions and flattening are examples of other layers commonly found in said architectures.

3.1 Batch-Normalization Folding Limits

Simplest Case. let's assume f is sequential with the following patterns: if fl is an expressive layer then fl+1 is a BN layer and fl+2 is a non-affine layer, i.e. fl+2(fl+1(fl(x))) = σ (BN(Wlx + bl)) where σ is an activation function and BN is a BN operation, see Figure 2 a). Under those assumptions all the BN layers can be folded using the naive BN-folding algorithm as follows: for each BN layer l + 1 we remove the layer in the graph associated to the network and update the weights of fl with

Then the resulting neural network is mathematically identical to the original network while requiring less operations to be computed.

Sequential Model. Let's assume f is sequential in the most general sense, i.e. f : x ↦→ fL(. . . f1(x)), see Figure 2 b). The difference with the previous case is the possible presence of other layers between the BN layers and the expressive ones as well as the the possibility to have the activation before the BN. The naive BN-folding algorithm consists in folding the BN layers in the previous or next expressive layer as long as their are no intermediate non-affine layers. To fold a BN layer backward (i.e. the expressive layer to update has a lower index) we use equation 1. On the other hand, to fold a BN layer forward (i.e. the expressive layer to update has a higher index l + k) we use

Under the sequential assumption, the algorithm folds all the BN layers that can be folded. Therefore this BN-folding algorithm is optimal in this specific case.

General DAG. Let's assume the graph of f is a directed acyclic graph (DAG), see Figure 2 c). Compared to a sequential model this corresponds to adding skip connections and possibly multiple leaves to the graph. In this case, the naive BN-folding algorithm consists in only considering sequential portions of the graph and limiting the resolution to the previous case. However a number of BN layers may be involved in non-sequential patterns as illustrated in Figure 2 c). Such BN layers can be folded by the proposed method. As illustrated, this is done by folding the BN layer in the input expressive layers while updating adequately the the consecutive expressive layer (the one below the BN node in Figure 2 c). For this reason the naive BN-folding algorithm is not optimal. Now, we will define a necessary and sufficient condition for BN folding.

3.2 Theoretical Folding Condition

Every feed forward neural network architecture can be expressed as a DAG. Let f be a DAG with n nodes. We recall that, by definition, a layer fl is affine, if and only if fl(ax + b) = ˜afl(x) + ˜b. The intuition is to isolate BN layers from non-affine layers. To do so, we define the connected components C of a node N (corresponding to a BN layer) recursively. We start with the set C = {N} and we add all the adjacent, affine nodes (i.e. nodes not associated with nonaffine layers). For each new, not-expressive, element of C we repeat the process until there are no new elements. The graph defined by C is decomposed in the layers entering node N and the layers leaving node N. We note the first part Cin and the second Cout. For a simple example (Figure 2 c), if we consider the VGG BN case and a BN layer fl, we have Cin = {fl-1, fl} with the input expressive layer fl-1 and Cout = {fl, fl+1} with fl+1 the activation layer. More examples of such connected components are illustrated in Figure 2. Then, for a given BN layer, the following theorem gives us our necessary and sufficient condition on the possibility to fold the layer.

Theorem 1. Let f be a DAG associated to a neural network. A BN layer of node N can be folded if and only if at least one of Cin or Cout satisfies:

1. the set is not limited to {N}
2. all the leaves (aside N) of the set are expressive layers

Proof. Here we will prove the necessity of the condition by assuming that either one of the condition is not satisfied by both Cin and Cout. First, let's assume Cin and Cout are restricted to {N}, then it implies that the BN layer is "surrounded" by non-affine layers and thus can't be folded (second grey node of Figure 2 b). Second, let's assume that both Cin and Cout have at least one leaf that is not associated to an expressive layer aside fromN itself. Under this assumption, in C, there exists a path fρ(l+k) ◦ · · · ◦ fρ(l) of length k and a path fη(l+k′ ) ◦ · · · ◦ fη(l) of length k′ (see Figure 3) such that

1. fη(l+k′ ) corresponds to node N and fη(l+k′ ) ̸= fρ(l+k)
2. fη(l) = fρ(l)

By definition of C all the consecutive layers of fρ(l+k) are non-affine. Then if we fold N, the layer fη(l) won't compute the same operations and by definition, the path fρ(l) ◦ · · · ◦ fρ(l+k) can't be updated to negate this effect. thus the node N can't be folded. Such scenario is illustrated in Figure 4 b). This demonstrates the necessity of the proposed condition.

Following this result, we can see in Figure 2 c) a BN layer that can be folded but not with the naive BN-folding algorithm. To conclude our proof we will demonstrate that the condition is sufficient.

3.3 Proposed Practical Algorithm

Instead of an analytic proof we propose a constructive proof via an algorithm which performs the BN folding as long as the condition from theorem 1 is satisfied. Let's assume Cin satisfies the condition, then we separate the leaves in two categories: the ones with a path to node N, noted I and the others O. The nodes in I are updated using equation 1 and the others negate the previous update with the following update:

If Cin doesn't satisfy the second condition from theorem 1, we use equation 2 to update I and the following for O:

Now we simply need to verify that any path in the graph associated to f computes the same outputs before and after folding. As we only update C, it is enough to do the verification for any path in C. We limit the verification to the case study of Figure 4 a). Before BN folding, the graph corresponds to the mathematical function f . We note Wi and bi the weights and biases of each expressive node of the graph. We note σ the operation performed by the non-affine node and g1 and g2 the operations from the other nodes (green nodes). As such, we have

where G1 and G2 correspond to the transformation of the inputs by the first green node and second green node, respectively. In this example, the BN layer satisfies the conditions of theorem 1. We fold the BN layer in the layers of set I, i.e. the second and third expressive layers, using equation 1. We note their weights and biases ˜Wi and ˜bi. Finally, we adapt the layer of O, i.e. the fourth expressive layer, to its new processed inputs using equation 3 and note ¯W4 and ¯b4. The folded network will compute:

(6)

where the folded operator ˜G2 is equal to BN(G2) before folding. Because we negated the folding in W4 and b4 the predictive function remains unchanged.

Consequently, the algorithm based on the proposed necessary and sufficient condition from theorem 1 is optimal, i.e. a BN layer can be folded if and only if

1. the set is not limited to {N}
2. all the leaves (aside N) of the set are expressive layers We now validate the empirical, cost-free, benefits resulting from BaN-OFF.

4 Experiments

4.1 Toy

We implemented the toy architectures described in Figure 2. Table 1 lists our results on an Intel CPU m3. The inference speed-up is measured as the following ratio Told-Tnew Told ∈ [0; 100%] where Tnew is the inference time of the network after BN folding and Told is the original inference time, i.e. the higher the better.

As expected the two algorithms, naive and BaN-OFF, share the same performance on the two first architectures as they perform the exact same transformations of the network. On the other hand, for the last configuration only the proposed method performed BN folding and resulted in inference speed increase. We also observe that the reduction of inference time is not correlated with the proportion of removed parameters as BN layers don't require many parameters but may put sequential computations of the architecture on hold during the forward pass. Furthermore, the predictions were only changed by up to 10-6 in L1 norm on vector outputs of 1000 elements. This is due to numerical approximations. Overall, this shows the importance of BN folding in neural networks acceleration as it doesn't change the predictive function and provide a significant boost in terms of speed.

4.2 Main Results

Cifar10. Table 2 summarizes our results on ResNets [He et al., 2016] trained on Cifar10 [Krizhevsky et al., 2009]. For both algorithms, the folded networks share identical predictions with the original models. We conducted our first experiments on a small CPU as it corresponds to the most challenging use-case. On a small device, sequential operations, e.g. BN layers, are not bottleneck compared to large expressive layers. Furthermore, using a small device induces instability on the results for both algorithms. Standard deviation values (over the percentage of relative reduced inference time per frame) range from 3.97 to 13.65. However the proposed method systematically outperforms naive BN-folding.

Furthermore, we confirm the previous observation on the decorrelation between the proportion of removed parameters and the improvement in inference speed. We also observe that the inference boost can't be naively deduced from simple architecture properties such as depth and wideness.

ImageNet. Table 3 contains our results on ResNets [He et al., 2016], MobileNet v2 [Sandler et al., 2018] and EfficientNets [Tan and Le, 2019] trained on ImageNet [Deng et al., 2009]. Folded networks, with the naive and the proposed method, the resulting networks share identical predictions with the original model. Similarly to Cifar10, the results were obtained on an Intel m3 CPU. Standard deviations values range between 1.14 for the least accelerated models and

12.68 for the most accelerated models. As expected, the proposed method systematically outperforms the naive approach.

In terms of proportion of removed parameters, the results are comparable to our results on Cifar10 which demonstrates that, even for more challenging tasks it matters more to remove a few meaningful parameters (e.g. BN layers) than a lot of parameters. This is based on the results obtained with pruning methods such as [Frankle and Carbin, 2018; Lin and others, 2020]. We also note that the correlation between depth and acceleration due to BN-folding on ResNets may be a coincidence.

4.3 Stability over Hardware

Central Processing Unit (CPU) comparison. Our results, in Table 4, show the stability of the inference boost across different CPUs. We also note that the maximum standard deviation on the results drops from 13.65 to 2.06 on the larger CPU. These results show the difficulties to work on very small devices and the importance of efficient acceleration.

Graphical Processing Unit (GPU) comparison. Table 4 presents the comparison between CPU and GPU devices response to BN folding. In contrast with CPU devices, GPU are very sensitive to sequential computations and very efficient on parallel operations. This property translates as significantly larger inference speed-up on GPU when compared to CPU. As expected from a larger device, the standard deviations on the results represent less than 5.04% of the average value which shows the stability of the inference speed-up across different runs.

Furthermore, such results motivate the use of acceleration frameworks comprising BN folding for efficient inference on cloud servers using large GPUs.

4.4 Other AccelerationMethods

BN usually constitutes a complementary step in pruning and quantization. However it is not systematically applied. We put the emphasis on the efficiency of BN-folding for DNN acceleration by comparing the cost in terms of accuracy to reach similar acceleration with pruning and quantization methods.

Pruning. Figure 5 shows the comparison of the influence of pruning (merge step from [Yvinec et al., 2021]) alone and BN folding on a ResNet 56 trained on Cifar10. For the sake of comparison, we implemented three classical pruning methods: Merge (which consists in merging similar neurons [Yvinec et al., 2021]), K-means (which consists in clustering neurons by their weight values) and magnitude-based unstructured pruning (which consists in removing weight values based on their absolute value). These methods allows inference speed-ups up to 50.23% but often significantly degrade the model accuracy. In comparison, the proposed BAN-OFF method allows similar runtime reduction without altering the predictive function. This suggests that, while the community focused on removing a lot of parameters, removing a few parameters with BN folding saves a lot of computations. Furthermore, BN folding and such pruning techniques can be combined for improved runtime reduction (up to 59.2%).

Quantization. We considered the naive quantization operator introduced in [Krishnamoorthi, 2018]. Such methods are usually more effective than pruning in terms of inference speed-up, which is confirmed in Figure 5. Similarly to the previous results, BN folding and quantization can be used in conjunction for improved runtime reduction (50.5 - 65.8%).

4.5 Other Data Applications

MS COCO. Table 5 summarizes our results on the detection task of MS COCO [Lin and others, 2014] on an RTX 3090. We considered widely used architectures such as Faster RCNN [Ren and others, 2015], Mask RCNN [He et al., 2017], EfficientDet [Tan et al., 2020] and CenterNet [Zhou and others, 2019]. Similarly to ImageNet, the folded models and the original models share the same predictions and the same accuracy. We see results very comparable to the results from table 4 (column RTX 3090) which is understandable as the detection algorithm are very similar, in the sense of the presence of BN layers, to their corresponding backbones.

CityScapes. From our results, listed in Table 6, on the semantic segmentation task of CityScapes [Cordts and others, 2016], we observe similar results to object detection. On both FPN [Lin and others, 2017] and PSPN [Zhao and others, 2017] models, the proposed method achieves remarkable results with 55% reduction of the inference time on an RTX 3090. Similar results are obtained on U-Net [Ronneberger et al., 2015] with an EfficientNet B0 backbone with great stability across runs (standard deviation below 2 points). Finally, on U-Net, we reach our most remarkable result of almost 63% reduction of latency with a standard deviation below 1 point.

5 Conclusion

In this work, we highlighted and addressed the limitations of existing batch normalization (BN) layer folding approaches. We argued that BN folding is often overlooked in the deep neural network acceleration literature, as it allows drastic runtime reduction at inference time, especially compared to other acceleration methods, e.g. pruning or quantization.

From a theoretical point of view, we proposed a novel necessary and sufficient condition for BN folding working for any architecture. We also derived a practical algorithm, dubbed BAN-OFF, for optimal BN folding. BAN-OFF allows to significantly reduce the inference time of any deep neural network, regardless of its task (classification, object detection or semantic segmentation), architecture or the hardware of which it shall run on, at virtually no cost. The source code is available in gitlab and as a pip package (pip install tensorflow-batchnorm-folding).

Acknowledgments

This work has been supported by the french National Association for Research and Technology (ANRT) and the company Datakalab (CIFRE convention C20/1396).