Continual Normalization: Rethinking Batch Normalization for Online Continual Learning

1 INTRODUCTION . Continual learning ( Ring , 1997 ) is a promising approach towards building learning systems with human-like capabilities . Unlike traditional learning paradigms , continual learning methods observe a stream of tasks and simultaneously perform well on all tasks with limited access to previous data . Therefore , they have to achieve a good trade-off between retaining old knowledge ( French , 1999 ) and acquiring new skills , which is referred to as the stability-plasticity dilemma ( Abraham & Robins , 2005 ) . Continual learning has been a challenging research problem , especially for deep neural networks because of their ubiquity and promising results on many applications ( LeCun et al. , 2015 ; Parisi et al. , 2019 ) . While most previous works focus on developing strategies to alleviate catastrophic forgetting and facilitating knowledge transfer ( Parisi et al. , 2019 ) , scant attention has been paid to the backbone they used . In standard backbone networks such as ResNets ( He et al. , 2016 ) , it is natural to use Batch Normalization ( BN ) ( Ioffe & Szegedy , 2015 ) , which has enabled the deep learning community to make substantial progress in many applications ( Huang et al. , 2020 ) . Although recent efforts have shown promising results in training deep networks without BN in a single task learning ( Brock et al. , 2021 ) , we argue that BN has a huge impact on continual learning . Particularly , when using an episodic memory , BN improves knowledge sharing across tasks by allowing data of previous tasks to contribute to the normalization of current samples and vice versa . Unfortunately , in this work , we have explored a negative effect of BN that hinders its performance on older tasks and thus increases catastrophic forgetting . Explicitly , unlike the standard classification problems , data in continual learning arrived sequentially , which are not independent and identically distributed ( noni.i.d ) . Moreover , especially in the online setting ( Lopez-Paz & Ranzato , 2017 ) , the data distribution changes over time , which is also highly non-stationary . Together , such properties make the BN ’ s running statistics heavily biased towards the current task . Consequently , during inference , the model normalizes previous tasks ’ data using the moments of the current task , which we refer to as the “ cross-task normalization effect ” . Although using an episodic memory can alleviate cross-task normalization in BN , it is not possible to fully mitigate this effect without storing all past data , which is against the continual learning purposes . On the other hand , spatial normalization layers such as GN ( Wu & He , 2018 ) do not perform cross-task normalization because they normalize each feature individually along the spatial dimensions . As we will empirically verify in Section 6.1 , although such layers suffer less forgetting than BN , they do not learn individual tasks as well as BN because they lack the knowledge transfer mechanism via normalizing along the mini-batch dimension . This result suggests that a continual learning normalization layer should balance normalizing along the mini-batch and spatial dimensions to achieve a good trade-off between knowledge transfer and alleviating forgetting . Such properties are crucial for continual learning , especially in the online setting ( Lopez-Paz et al. , 2017 ; Riemer et al. , 2019 ) , but not satisfied by existing normalization layers . Consequently , we propose Continual Normalization ( CN ) , a novel normalization layer that takes into account the spatial information when normalizing along the mini-batch dimension . Therefore , our CN enjoys the benefits of BN while alleviating the cross-task normalization effect . Lastly , our study is orthogonal to the continual learning literature , and the proposed CN can be easily integrated into any existing methods to improve their performances across different online protocols . In summary , we make the following contributions : • We study the benefits of BN and its cross-task normalization effect in continual learning . • We identify the desiderata of a normalization layer for continual learning and propose CN , a novel normalization layer that improves the performance of existing continual learning methods . • Through extensive experiments on different online continual learning benchmarks , we demonstrate that CN is a plug-in replacement for BN and offers promising improvements across different metrics with limited computational overhead . 2 NOTATIONS AND PRELIMINARIES . This section provides the necessary background of continual learning and normalization layers , including the batch normalization ( BN ) layer and its variants . 3 CONTINUAL LEARNING PROTOCOL . We focus on the online continual learning protocol ( Lopez-Paz & Ranzato , 2017 ) over a data stream of T tasks . At each step , the model only observes training samples drawn from the current training task data and has to perform well on the testing data of all observed tasks so far . Moreover , data comes strictly in a streaming manner , resulting in a single epoch training setting . Throughout this work , we conduct experiments on different continual learning protocols , ranging from the taskincremental ( Lopez-Paz & Ranzato , 2017 ) , class-incremental ( Rebuffi et al. , 2017 ) , to the task-free scenarios ( Aljundi et al. , 2019b ) . We will describe each scenario in more details in Section 6 . 3.1 NOTATIONS . We focus on the image recognition problem and denote an image as x ∈ RW×H×C , where W , H , C are the image width , height , and the number of channels respectively . A convolutional neural network ( CNN ) with parameter θ is denoted as fθ ( · ) . A feature map of a mini-batch B is defined as a ∈ RB×C×W×H . Finally . each image x is associated with a label y ∈ { 1 , 2 . . . , Y } and a task identifier t ∈ { 1 , 2 , . . . , T } . 3.2 NORMALIZATION LAYERS . Normalization layers are essential in training deep neural networks ( Huang et al. , 2020 ) . In general , a normalization layer takes a mini-batch of feature maps a = ( a1 , . . . , aB ) as input and perform the Z-normalization as : a′ = γ ( a− µ√ σ2 + ) + β , ( 1 ) where µ and σ2 are the mean and variance calculated from the features in B , and is a small constant added to avoid division by zero . The affine transformation ’ s parameters γ , β ∈ RC are |C|-dimensional vectors learned by backpropagation to retain the layer ’ s representation capacity . For brevity , we will use “ moments ” to refer to both the mean µ and the variance σ2 . Batch Normalization BN is one of the first normalization layers that found success in a wide range of deep learning applications ( Ioffe & Szegedy , 2015 ; Santurkar et al. , 2018 ; Bjorck et al. , 2018 ) . During training , BN calculates the moments across the mini-batch dimension as : µBN = 1 BHW B∑ b=1 W∑ w=1 H∑ H=1 abcwh , σ 2 BN = 1 BHW B∑ b=1 W∑ w=1 H∑ H=1 ( abcwh − µBN ) 2 . ( 2 ) At test time , it is important for BN to be able to make predictions with only one data sample to make the prediction deterministic . As a result , BN replaces the mini-batch mean and variance in Eq . ( 2 ) by an estimate of the global values obtained during training as : µ← µ+ η ( µB − µ ) , σ ← σ + η ( σB − σ ) , ( 3 ) where η is the exponential running average ’ s momentum , which were set to the standard value of 0.1 in our experiments . Spatial Normalization Layers The discrepancy between training and testing in BN can be problematic when the training mini-batch size is small ( Ioffe , 2017 ) , or when the testing data distribution differs from the training distributions . In such scenarios , the running estimate of the mean and variance obtained during training are a poor estimate of the moments to normalize the testing data . Therefore , there have been tremendous efforts in developing alternatives to BN to address these challenges . One notable approach is normalizing along the spatial dimensions of each sample independently , which completely avoid the negative effect from having a biased global moments . In the following , we discuss popular examples of such layers . Group Normalization ( GN ) ( Wu & He , 2018 ) When the mini-batch size is small , BN ’ s running estimates of the global moments become inaccurate . Therefore , GN proposes to normalize each feature individually by dividing each channel into groups and then normalize each group separately . Moreover , GN does not normalize along the batch dimension , thus performs the same computation during training and testing . Computationally , GN first divides the channels C into G groups as : a′bgkhw ← abchw , where k = b∗c C G , ( 4 ) Then , for each group g , the features are normalized with the moments calculated as : µ ( g ) GN = 1 m K∑ k=1 W∑ w=1 H∑ h=1 a′bgkhw , σ ( g ) GN 2 = 1 m K∑ k=1 W∑ w=1 H∑ h=1 ( a′bgkhw − µ ( g ) GN ) 2 ( 5 ) GN has shown comparable performance to BN with large mini-batch sizes ( e.g . 32 or more ) , while significantly outperformed BN with small mini-batch sizes ( e.g . one or two ) . Notably , when putting all channels into a single group ( setting G = 1 ) , GN is equivalent to Layer Normalization ( LN ) ( Ba et al. , 2016 ) , which normalizes the whole layer of each sample . On the other extreme , when separating each channel as a group ( setting G = C ) , GN becomes Instance Normalization ( IN ) ( Ulyanov et al. , 2016 ) , which normalizes the spatial dimension in each channel of each feature . Figure 1 provides an illustration of BN , GN , IN , and the proposed CN , which we will discuss in Section 5 . 4 BATCH NORMALIZATION IN CONTINUAL LEARNING . 4.1 THE BENEFITS OF NORMALIZATION LAYERS IN CONTINUAL LEARNING . We argue that BN is helpful for forward transfer in two aspects . First , BN makes the optimization landscape smoother ( Santurkar et al. , 2018 ) , which allows the optimization of deep neural networks to converge faster and better ( Bjorck et al. , 2018 ) . In continual learning , BN enables the model to learn individual tasks better than the no-normalization method . Second , with the episodic memory , BN uses data of the current task and previous tasks ( in the memory ) to update its running moments . Therefore , BN further facilitates forward knowledge transfer during experience replay : current task data is normalized using moments calculated from both current and previous samples . Compared to BN , spatial normalization layers ( such as GN ) lack the ability to facilitate forward transfer via normalizing using moments calculated from data of both old and current tasks . We will empirically verify the benefits of different normalization layers to continual learning in Section 6.1 . 4.2 THE CROSS-TASK NORMALIZATION EFFECT . Recall that BN maintains a running estimate of the global moments to normalize the testing data . In the standard learning with a single task where data are i.i.d , one can expect such estimates can well-characterize the true , global moments , and can be used to normalize the testing data . However , online continual learning data is non-i.i.d and highly non stationary . Therefore , BN ’ s estimation of the global moments is heavily biased towards the current task because the recent mini-batches only contain that task ’ s data . As a result , during inference , when evaluating the older tasks , BN normalizes previous tasks ’ data using the current task ’ s moments , which we refer to as the crosstask normalization effect . We consider a toy experiment on the permuted MNIST ( pMNIST ) benchmark ( Lopez-Paz & Ranzato , 2017 ) to explore the cross-task normalization effect in BN . We construct a sequence of five task , each has 2,000 training and 10,000 testing samples , and a multilayer perceptron ( MLP ) backbone configured as Input ( 784 ) -FC1 ( 100 ) -BN1 ( 100 ) -FC2 ( 100 ) -BN ( 100 ) -Softmax ( 10 ) , where the number inside the parentheses indicates the output dimension of that layer . We consider the Single and ER strategies for this experiment , where the Single strategy is the naive method that trains continuously without any memory or regularization . For each method , we implement an optimal BN variant ( denoted by the suffix -BN∗ ) that calculates the global moments using all training data of all tasks before testing . Compared to BN , BN∗ has the same network parameters and only differs in the BN ’ s moments used to normalize the testing data . We emphasize that although BN∗ is unrealistic , it sheds light on how cross-task normalization affects the performance of CL algorithms . We report the averaged accuracy at the end of training ACC ( ↑ ) ( Lopez-Paz et al. , 2017 ) , forgetting measure FM ( ↓ ) ( Chaudhry et al. , 2018 ) , and learning accuracy LA ( ↑ ) ( Riemer et al. , 2019 ) in this experiment . Table 1 reports the evaluation metrics of this experiment . Besides the standard metrics , we also report the moments ’ difference magnitudes between the standard BN and the global BN variant . We observe that the gap between two BN variants is significant without the episodic memory . When using the episodic memory , ER can reduce the gap because the pMNIST is quite simple and even the ER strategy can achieve close performances to the offline model ( Chaudhry et al. , 2019b ) . Moreover , it is important to note that training with BN on imbalanced data also affects the model ’ s learning dynamics , resulting in a suboptimal parameter . Nevertheless , having an unbiased estimate of the global moments can greatly improve overall accuracy given the same model parameters . Moreover , this improvement is attributed to reducing forgetting rather than facilitating transfer : BN∗ has lower FM but almost similar LA compared to the traditional BN . In addition , we observe that the hidden features ’ variance becomes inaccurate compared to the global variance as we go to deeper layers ( ∆ ( 1 ) σ2 < ∆ ( 2 ) σ2 ) . This result agrees with recent finding ( Ramasesh et al. , 2021 ) that deeper layers are more responsible for causing forgetting because their hidden representation deviates from the model trained on all tasks . Overall , these results show that normalizing older tasks using the current task ’ s moments causes higher forgetting , which we refer to as the “ cross-task normalization effect ” .

Summary: The paper studies the role of normalization layers in continual learning. The main argument is that vanilla BatchNorm is not very suitable since the statistics of data change across tasks. To alleviate this, the authors propose the CntinualNorm layer which is essentially the combination of GroupNorm followed by BatchNorm. The experiments compare different normalization schemes across various CL benchmarks.