SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Yilin Lyu $^{1*}$  Liyuan Wang $^{2*}$  Xingxing Zhang $^{2\dagger}$  Zicheng Sun $^{1}$  Hang Su $^{2}$  Jun Zhu $^{2}$  Liping Jing $^{1\dagger}$
$^{1}$ Beijing Key Lab of Traffic Data Analysis and Mining, Beijing Jiaotong University  $^{2}$ Dept. of Comp. Sci. & Tech., Institute for AI, BNRist Center, THBI Lab, Tsinghua-Bosch Joint Center for ML, Tsinghua University {yilinlyu, zichengsun, lpjing}@bjtu.edu.cn wly19@tsinghua.org.cn, xxzhang1993@gmail.com {suhangss, dcszj}@tsinghua.edu.cn
Continual learning entails learning a sequence of tasks and balancing their knowledge appropriately. With limited access to old training samples, much of the current work in deep neural networks has focused on overcoming catastrophic forgetting of old tasks in gradient-based optimization. However, the normalization layers provide an exception, as they are updated interdependently by the gradient and statistics of currently observed training samples, which require specialized strategies to mitigate recency bias. In this work, we focus on the most popular Batch Normalization (BN) and provide an in-depth theoretical analysis of its sub-optimality in continual learning. Our analysis demonstrates the dilemma between balance and adaptation of BN statistics for incremental tasks, which potentially affects training stability and generalization. Targeting on these particular challenges, we propose Adaptive Balance of BN  $(\mathrm{AdaB}^2\mathrm{N})$ , which incorporates appropriately a Bayesian-based strategy to adapt task-wise contributions and a modified momentum to balance BN statistics, corresponding to the training and testing stages. By implementing BN in a continual learning fashion, our approach achieves significant performance gains across a wide range of benchmarks, particularly for the challenging yet realistic online scenarios (e.g., up to  $7.68\%$ ,  $6.86\%$  and  $4.26\%$  on Split CIFAR-10, Split CIFAR-100 and Split Mini-ImageNet, respectively). Our code is available at https://github.com/lvyilin/AdaB2N.
Continual learning aims to acquire, accumulate and exploit knowledge from sequentially arrived tasks. As the access to old training samples is typically limited, the acquired knowledge tends to be increasingly relevant to the more recent tasks. Such recency bias, also known as catastrophic forgetting [17], is generally attributed to the nature of gradient-based optimization in deep neural networks (DNNs) [9, 30], where the learning of network parameters depends on a tug-of-war over currently observed training samples. Numerous efforts have been devoted to addressing catastrophic forgetting of these learnable parameters through rectifying their gradients. However, as an important component of DNNs, the normalization layers employ particular updating rules that are only partially gradient-based, and their recency bias in continual learning remains to be addressed.
Table 1: Comparison of normalization methods for continual learning
Method	Statistics	Training Adaptability	Online	Offline
BN [11]	Unbalanced	X	X	X
CN [19]	Unbalanced	✓	✓	X
BNT [34]	Balanced	X	X	✓
TBBN [8]	Balanced	✓	X	✓
AdaB2N (Ours)	Adaptively balanced	✓	✓	✓
The normalization layers typically include affine transformation parameters and statistics parameters to normalize the internal representations, in order to accelerate convergence and stabilize training [11, 33, 2, 24]. Since these two kinds of parameters are updated respectively from the gradient and statistics (i.e., the empirical mean and variance) of currently observed training samples, the normalization layers require specialized strategies to overcome recency bias in continual learning. In particular, the most popular Batch Normalization (BN) [11] performs an exponential moving average (EMA) of statistics parameters for testing, which preserves historical information to some extent, and thus outperforms many other alternatives in continual learning [19, 8]. As the BN statistics of old tasks still decay exponentially, some recent work [19, 8, 34] enforced a balance in the contribution of individual tasks to the BN statistics, but with limited effectiveness and generality (see Table 1 for a conceptual comparison of their targets).
So, how to update normalization statistics for better continual learning: balance or adaptation? In this work, we present an in-depth theoretical analysis of the most popular BN, demonstrating its potential challenges posed by continual learning. First, a regular EMA with constant momentum cannot guarantee the balance of BN statistics for testing, which yields more recency bias as more training batches of each task are observed. Although a modified EMA with normalized momentum of batch number can enforce balanced statistics, they tend to be outdated for learning each task and thus limit model adaptability. Besides, the sharp changes in task distributions can interfere with the benefit of BN in stabilization of training, which further affects continual learning.
Based on these theoretical insights, we propose Adaptive Balance of BN (AdaB $^2$ N) to improve continual learning. Specifically, we leverage a Bayesian-based strategy to decide an appropriate way of accommodating task-wise contributions in normalized representations to stabilize training. We further incorporate an adaptive parameter to reconcile the constant and normalized momentum of EMA, in order to strike a balance of BN statistics for testing while training each task with adequate adaptability. Extensive continual learning experiments demonstrate the superiority of our approach, which provides substantial improvements especially for the challenging yet realistic online scenarios.
Continual Learning: The major focus of this area is to address recency bias aka catastrophic forgetting in DNNs [17, 30, 28, 29]. Representative strategies are generally designed for the learnable parameters with gradient-based optimization [30], such as weight regularization [13, 1, 27], experience replay [32, 6], gradient projection [15, 20], etc. In particular, experience replay of a few old training samples has proven to be usually the most effective strategy and is applicable to large-scale applications [30, 32, 31]. However, its reliance on the old training samples may further deteriorate the imbalance and overfitting of normalization statistics. Corresponding to the real-world dynamics, continual learning gives rise to many meaningful settings and particular challenges [30]. For example, a desirable continual learning model should perform well in both task-incremental and class-incremental settings, which depend on whether the task identity is provided for testing. In contrast to offline scenarios of learning each task with multiple epochs, the model should also accommodate online scenarios of learning all tasks in a one-pass data stream, which is realistic but extremely more challenging.
Normalization Layer: The normalization layers have become a critical component for many representative DNN architectures, where Batch Normalization (BN) [11] is the most popular choice along with many other competitive alternatives designed for single-task learning [33, 24, 2]. However, the normalization layers require specialized strategies to perform appropriate parameter updates in continual learning. As the EMA of statistics parameters retains a certain amount of past information, BN is naturally the dominant choice for continual learning and indeed a strong baseline for overcoming
recency bias in normalization layers [19, 8]. To further improve BN, Continual Normalization (CN) [19] removed differences in task-wise distributions and then performed BN to acquire task-balanced statistics, providing benefits for online continual learning. BNT [34] and TBBN [8] constructed task-balanced batches to update BN, but only effective for offline continual learning. However, none of the current efforts can improve both online and offline scenarios [8], with the effectiveness of CN varying further with batch size in different online benchmarks. This motivates our work to provide a more in-depth investigation of BN with respect to the particular challenges of continual learning.
3 Analysis of Normalization Statistics in Continual Learning
In this section, we introduce the formulation of continual learning and representative normalization strategies, and then provide an in-depth theoretical analysis to diagnose the most popular BN. Additionally, a notation table is included in Appendix A, and proofs of theorems can be found in Appendix B.
Let's consider a general setting for continual learning, where a neural network with parameters  $\theta$  needs to learn a sequence of tasks from their training sets  $\mathcal{D}_1,\dots ,\mathcal{D}_T$  in order to perform well on their test sets. The training set of task  $t$  is defined as  $\mathcal{D}_t = \{(x_{t,n},y_{t,n})\}_{n = 1}^{|D_t|}$ , where  $(x_{t,n},y_{t,n})$  is a data-label pair of task  $t\in [1,T]$  and  $|\cdot |$  denotes its amount. Since the old training sets  $\mathcal{D}_{1:t - 1} = \bigcup_{i = 1}^{t - 1}\mathcal{D}_i$  are unavailable when learning task  $t$ ,  $\theta$  tends to catastrophically forget the previously learned knowledge. An effective strategy is to save a few old training samples in a small memory buffer, denoted as  $\mathcal{M} = \{\mathcal{M}_k\}_{k = 1}^{t - 1}$ , in order to recover the old data distributions when learning the current task  $t$ . We denote the  $m$ -th training batch as  $B_{m}$ , which includes  $N_{t}$  training samples from  $\mathcal{D}_t$  and  $N_{k}$  training samples from  $\mathcal{M}_k$  for  $k = 1,\dots ,t - 1$ . Here  $N = \sum_{k = 1}^{t}N_{k} = |B|$  denotes the batch size. The batch index  $m\in (m_{t - 1},m_t]$ , where  $m_t = \lfloor |\mathcal{D}_{1:t}| / N\rfloor E$  is the index of the last batch of task  $t$ ,  $m_0 = 0$ , and  $E$  is the number of training epochs. We denote  $r_m = N_t / N$  as the proportion of the current task  $t$  in  $B_{m}$ , where  $r_m < 1$  indicates the use of replay while  $r_m\equiv 1$  indicates that no replay is used. Without loss of generality, we assume that the old training samples in each batch are task-balanced, i.e.,  $N_{1} = N_{2} = \dots = N_{t - 1}$ . The test set follows the same distribution as the training set, with the task identity provided in task-incremental learning (Task-IL) while not provided in class-incremental learning (Class-IL) [25].
Inside the network, the internal representation of the  $\ell$ -th layer of the  $m$ -th batch is denoted as  $a_{m}^{(\ell)} \in \mathbb{R}^{N \times C^{(\ell)} \times D^{(\ell)}}$ , where  $C^{(\ell)}$  and  $D^{(\ell)}$  represent the number of channels and feature dimensions, respectively. For the sake of notation clarity, we omit the layer and batch index hereafter when no confusion arises. Then, the operation of a normalization layer can be generally defined as:
$$
\hat {a} = \gamma a ^ {\prime} + \beta , \quad a ^ {\prime} = \frac {a - \mu (a)}{\sqrt {\sigma^ {2} (a , \mu (a)) + \varepsilon}}, \tag {1}
$$
where the statistics parameters (i.e., the empirical mean and variance) are calculated from the corresponding functions  $\mu (\cdot)$  and  $\sigma^2 (\cdot ,\cdot)$ , respectively,  $\varepsilon$  is a small positive constant to avoid the denominator being zero,  $\gamma$  and  $\beta$  are affine transformation parameters that are learned by gradient-based optimization. Representative normalization strategies differ mainly in the ways of computing statistics parameters of internal representations, i.e., the exact forms of  $\mu$  and  $\sigma^2$ , as detailed below.
Batch Normalization (BN) [11] is the most popular choice for a wide range of DNN architectures. During the training phase, BN calculates the statistics parameters of internal representations for each training batch through  $\mu : \mathbb{R}^{N \times C \times D} \to \mathbb{R}^C$  and  $\sigma^2 : \mathbb{R}^{N \times C \times D} \times \mathbb{R}^C \to \mathbb{R}^C$ , where
$$
\mu (a) = \frac {1}{N D} \sum_ {n = 1} ^ {N} \sum_ {d = 1} ^ {D} a _ {[ n,: d ]}, \quad \sigma^ {2} (a, \mu (a)) = \frac {1}{N D} \sum_ {n = 1} ^ {N} \sum_ {d = 1} ^ {D} \left(a _ {[ n,: d ]} - \mu (a)\right) ^ {2}. \tag {2}
$$
During the testing phase, to make the normalization independent of other samples in the same batch, BN replaces the batch statistics of Eq. (1) with the population statistics estimated by exponential moving average (EMA). For ease of notation, here we define a convenient function  $S:\mathbb{R}^{N\times C\times D}\to \mathbb{R}^{2\times C}$  to obtain the two kinds of statistics, where
$$
S (\cdot) = \left[ \mu (\cdot), \sigma^ {2} (\cdot , \mu (\cdot)) \right] ^ {\top}. \tag {3}
$$
Then we denote the statistics of internal representation  $a_{m}^{(\ell)}$  as  $S_{m} \coloneqq S(a_{m})$ , with the layer index omitted. Now the population statistics of BN can be written as
$$
\hat {\mathbb {E}} [ S _ {m} ] := \hat {\mathbb {E}} [ S | a _ {1}, \dots , a _ {m} ] = (1 - \eta) \hat {\mathbb {E}} [ S | a _ {1}, \dots , a _ {m - 1} ] + \eta S _ {m}, \tag {4}
$$
where  $\eta \in (0,1)$  is the momentum of EMA and usually set to 0.1 in practice.
There are many alternative improvements of BN designed for single-task learning without the use of EMA. Group Normalization (GN) [33] can be seen as a general version, which divides the channels into  $G$  groups and then normalizes the internal representations of each group. Specifically, let  $\xi : \mathbb{R}^{N \times C \times D} \to \mathbb{R}^{N \times G \times K \times D}$  be a reshaping function that groups the channels, where  $G$  is the group number and  $K = \lfloor \frac{C}{G} \rfloor$ . Then the mean and variance of GN are obtained through  $\mu : \mathbb{R}^{N \times C \times D} \to \mathbb{R}^{N \times G}$  and  $\sigma^2 : \mathbb{R}^{N \times C \times D} \times \mathbb{R}^{N \times G} \to \mathbb{R}^{N \times G}$ , where
$$
\mu (a) = \frac {1}{K D} \sum_ {k = 1} ^ {K} \sum_ {d = 1} ^ {D} \xi (a) _ {[ :,:, k, d ]}, \quad \sigma^ {2} (a, \mu (a)) = \frac {1}{K D} \sum_ {k = 1} ^ {K} \sum_ {d = 1} ^ {D} (\xi (a) _ {[ :,:, k, d ]} - \mu (a)) ^ {2}. \tag {5}
$$
In particular, GN is equivalent to Layer Normalization (LN) [2] by putting all channels into one group  $(G = 1)$ , and Instance Normalization (IN) [24] by separating each channel as a group  $(G = C)$ . In contrast to these alternatives, Continual Normalization (CN) [19] is designed particularly for (online) continual learning, which performs GN with  $\gamma = 1$  and  $\beta = 0$  to remove task-wise differences and then performs BN to acquire task-balanced statistics.
3.2 Conflict of Balance and Adaptation
Although BN has become a fundamental component for many DNN architectures, its particular recency bias in continual learning remains under addressed. To quantitatively analyze the task-wise contributions in BN statistics, we define a statistical weight (of a certain batch) of each task. Then we diagnose the recency bias of BN statistics with a specified dilemma between balance and adaptation.
Definition 1 (Statistical weight of a task). The statistical weight of task  $t$  w.r.t. the population statistics  $\hat{\mathbb{E}}[S_m]$  is the weight  $w_t \in \mathbb{R}^+$  such that  $w_t = \sum_i^m w_i^t / Z$ .  $w_i^t \in \mathbb{R}^+$  is the statistical weight of task  $t$  of the  $i$ -th batch such that  $\sum_{t=1}^{T} \sum_{i=1}^{m} w_i^t S_i^t = \hat{\mathbb{E}}[S_m]$ , where  $S_i^t \coloneqq S(a_{i,1[(\cdot,y_n) \in \mathcal{D}_t]})$ .<sup>3</sup>  $Z$  is the normalizing constant  $Z = \sum_t^T \sum_i^m w_i^t$ .
We now give rigorously the statistical weight for EMA in order to demonstrate the recency bias:
Theorem 1 (Statistical weight for EMA). Let  $T, Z$  be defined as above, and  $\{m_t\}_{t=1}^T$  be an increasing sequence (i.e.,  $m_t < m_{t+1}$ ), where  $m_t \in \mathbb{Z}^+$ ,  $\eta : [1, m_T] \to (0,1)$ , and  $r : [1, m_T] \to (0,1]$ . The statistical weight of task  $t \in [1, T]$  w.r.t. the population statistics  $\hat{\mathbb{E}}[S_{m_T}]$  estimated by Eq. (4) is
$$
w _ {t} = \left[ \sum_ {i = m _ {t - 1} + 1} ^ {m _ {t}} \eta_ {i} r _ {i} \prod_ {j = i + 1} ^ {m _ {T}} \left(1 - \eta_ {j}\right) + \sum_ {i = m _ {t} + 1} ^ {m _ {T}} \frac {\eta_ {i} \left(1 - r _ {i}\right)}{T - 1} \prod_ {j = i + 1} ^ {m _ {T}} \left(1 - \eta_ {j}\right) \right] / Z, \tag {6}
$$
where we define  $\eta_{i} \coloneqq \eta(i)$  and  $r_{i} \coloneqq r(i)$  for brevity.
Eq. (6) specifies the impact of each task on the EMA population statistics, which exposes a dilemma between balance and adaptation, as we demonstrate below.
Corollary 2 (Adaptation of BN statistics). Let  $\{m_t\}_{t=1}^T$  be defined as above. If  $r(\cdot) \equiv 1$  and  $\eta(\cdot) \equiv \eta = 1 - \bar{\eta}$ , the statistical weight of task  $t \in [1, T]$  w.r.t. the population statistics  $\hat{\mathbb{E}}[S_{m_T}]$  is  $w_t = \frac{\bar{\eta}^{m_T - t} - \bar{\eta}^{m_T - t + 1}}{1 - \bar{\eta}^{m_T}} / Z$ . When  $m_T - m_{T-1} = m_{T-1} - m_{T-2} = \dots = m_1$ , we have  $w_t \approx \bar{\eta}^{m_1(T-t)} / Z$ , with approximation error  $|\epsilon| \leq O\left(\frac{\bar{\eta}^{m_1}}{1 - \bar{\eta}^{m_1T}}\right)$ .
Corollary 2 shows that without the use of replay, a model dealing with evolving task distributions has a strong preference for the distribution of the current task, as the statistical weights of past tasks decrease exponentially. Replay of old training samples can balance statistical weights of all seen tasks to some extent. However, it cannot eliminate the recency bias except in an unrealistic case:
3For simplicity, we assume zero covariance among the internal representations of different tasks.
Corollary 3 (Balance of BN statistics). Let  $\{m_t\}_{t=1}^T$  be defined as above. If  $r(\cdot) \equiv r$ ,  $\eta(\cdot) \equiv \eta = 1 - \bar{\eta}$ , and  $m_T - m_{T-1} = m_{T-1} - m_{T-2} = \dots = m_1$ , then the statistical weight of each task  $t \in [1, T]$  w.r.t. the population statistics  $\hat{\mathbb{E}}[S_{m_T}]$  achieves (approximately) equal if and only if  $r \approx 1/T$ , with approximation error  $|\epsilon| \leq \mathcal{O}(\bar{\eta}^{m_1})$ .
Corollary 3 shows that it is possible to balance the statistical weights by constructing a task-balanced batch with the memory buffer. However, this will severely slow down the convergence. Besides, as the memory buffer is usually limited in size, the model can easily overfit to the preserved old training samples [4], which is detrimental to generalization. Corollary 3 also demonstrates the importance of keeping flexibility of  $\eta$ . For instance, if we substitute  $\eta(i) = 1/(1 + i)$  into Eq. (4), this leads to a normalized momentum of batch number. In this case the population statistics can be considered as the cumulative moving average (CMA) [19] of the batch statistics, which can balance the statistical weights as below:
Corollary 4 (CMA). Let  $\{m_t\}_{t=1}^T$  be defined as above. If  $r(\cdot) \equiv r$ , then the statistical weight of each task  $t \in [1, T]$  w.r.t. the population statistics  $\hat{\mathbb{E}}[S_{m_T}]$  achieves equal if  $\eta(i) = 1/(1 + i)$ .
However, the balanced statistical weights do not necessarily lead to better performance, as the statistics of past tasks may be outdated relative to the current model. This is empirically validated in [19], where CMA does not outperform EMA in continual learning. Therefore, an improved strategy of updating population statistics for testing is urgently needed to reconcile balance and adaptation.
In addition to the update of population statistics during the testing phase, the current strategy of calculating batch statistics during the training phase is also sub-optimal for continual learning, as detailed in the next subsection.
3.3 Disturbed Training Stability
BN has been shown to stabilize training and improve generalization under traditional i.i.d. data setting [3, 21, 16]. However, it is in fact not clear whether BN can achieve the same benefit in continual learning, especially when the model has just switched to a new task with drastic changes in loss and parameters. We perform a preliminary analysis by the following theorem:
Theorem 5 (Theorem 4.1 of [21], restated). Let  $\mu, \sigma^2, \gamma, \ell, a, a'$  be defined as above. Let  $\mathcal{L}$  be the loss function of a BN model (i.e., using the normalized representations via Eq. (1)) and  $\mathcal{L}'$  be the identical loss function of an identical non-BN model (i.e.,  $a' \equiv a$ ). The gradient w.r.t. the internal representation  $a$  of the  $\ell$ -th layer follows that
$$
\left\| \nabla_ {a _ {[ :, j ]}} \mathcal {L} \right\| ^ {2} \leq \frac {\gamma^ {2}}{\sigma_ {[ : , j ]} ^ {2} C} \left[ C \| \nabla_ {a _ {[ :, j ]}} \mathcal {L} ^ {\prime} \| ^ {2} - \left(\mathbf {1} ^ {\top} \nabla_ {a _ {[ :, j ]}} \mathcal {L} ^ {\prime}\right) ^ {2} - \left(\nabla_ {a _ {[ :, j ]}} ^ {\top} \mathcal {L} ^ {\prime} \cdot a _ {[ :, j ]} ^ {\prime}\right) ^ {2} \right]. \tag {7}
$$
Theorem 5 demonstrates that BN can contribute to reducing the Lipschitz constant of the loss and thus stabilizing the training of a single task. However, as the model underfits to the distribution of each continual learning task, the upper bound of the gradient magnitude will be loosened by the increase of  $\{\nabla_{a}^{\top}\mathcal{L}^{\prime}\cdot a^{\prime}\}$  within the third term, i.e., the similarity of the normalized representation and its corresponding gradient, which potentially affects the benefit of BN. Here we present empirical results to validate the above claim. As shown in Fig. 1, a dramatic change in loss  $\mathcal{L}$  occurs whenever the task changes. We observe an increase of cosine similarity between the gradient and the normalized representation at this time, accompanied by a dramatic change in the gradient magnitude. This clearly demonstrates the importance of properly adapting BN statistics to each continual learning task, so as to better leverage the benefit of BN in stabilization of training.
Following the theoretical analysis, we propose Adaptive Balance of BN (AdaB $^2$ N) from two aspects: (1) training aspect: we apply a Bayesian-based strategy to adjust inter-task weights and construct the normalized representation, allowing for learning an appropriate way of weighting task distributions to stabilize training; and (2) testing aspect: instead of using a fixed momentum  $\eta$  in EMA, we turn to a generalized version that balances the statistical weights of past tasks while reducing the impacts of past batches. We now present both ideas as follows with a pseudo-code in Algorithm 1.
(b) Gradient Similarity
Figure 1: Dynamics of the loss of training batches, gradient similarity (i.e., cosine similarity between gradients and normalized representations) and gradient magnitude.
(c) Gradient Magnitude
Algorithm 1 Training Algorithm of Adaptive Balance of BN (AdaB $^2$ N)
Input: training dataset  $\mathcal{D}_{1:T}$ , memory buffer  $\mathcal{M}$ , batch size  $N$ , neural network  $f_{\vartheta, \omega} = f_{\vartheta(L+1)} \circ f_{\omega^{(L)}} \circ f_{\vartheta^{(L)}} \cdots \circ f_{\omega^{(\ell)}} \circ f_{\vartheta^{(\ell)}} \cdots \circ f_{\omega^{(1)}} \circ f_{\vartheta^{(1)}}$ , where  $f_{\vartheta^{(\ell)}}$  and  $f_{\omega^{(\ell)}}$  are the  $\ell$ -th network layer and normalization layer, respectively,  $\omega = \{\gamma^{(\ell)}, \beta^{(\ell)}, \psi^{(\ell)}\}_{\ell=1}^{L}, \vartheta = \{\vartheta^{(\ell)}\}_{\ell=1}^{L+1} = \theta \setminus \omega$ , initial parameters  $\vartheta_0, \omega_0$ , the objective function  $\mathcal{L}$  in Eq. (15), and hyperparameters  $\tilde{\eta}, \kappa, \lambda$ .
Output: parameters  $\vartheta, \omega$