SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Jie Hao
jhao6@gmu.edu
Kaiyi Ji
Department of CSE
kaiyiji@buffalo.edu
Mingrui Liu*
mingruil@gmu.edu
Coreset is a small set that provides a data summary for a large dataset, such that training solely on the small set achieves competitive performance compared with a large dataset. In rehearsal-based continual learning, the coreset is typically used in the memory replay buffer to stand for representative samples in previous tasks, and the coreset selection procedure is typically formulated as a bilevel problem. However, the typical bilevel formulation for coreset selection explicitly performs optimization over discrete decision variables with greedy search, which is computationally expensive. Several works consider other formulations to address this issue, but they ignore the nested nature of bilevel optimization problems and may not solve the bilevel coreset selection problem accurately. To address these issues, we propose a new bilevel formulation, where the inner problem tries to find a model which minimizes the expected training error sampled from a given probability distribution, and the outer problem aims to learn the probability distribution with approximately  $K$  (coreset size) nonzero entries such that learned model in the inner problem minimizes the training error over the whole data. To ensure the learned probability has approximately  $K$  nonzero entries, we introduce a novel regularizer based on the smoothed top-  $K$  loss in the upper problem. We design a new optimization algorithm that provably converges to the  $\epsilon$ -stationary point with  $O(1 / \epsilon^4)$  computational complexity. We conduct extensive experiments in various settings in continual learning, including balanced data, imbalanced data, and label noise, to show that our proposed formulation and new algorithm significantly outperform competitive baselines. From bilevel optimization point of view, our algorithm significantly improves the vanilla greedy coreset selection method in terms of running time on continual learning benchmark datasets. The code is available at https://github.com/MingruiLiu-ML-Lab/Bilevel-Coreset-Selection via-Regularization.
Deep Neural Networks (DNNs) have achieved tremendous successes in various domains, including computer vision [41, 30], natural language processing [72, 15], generative modeling [26] and
Figure 1: Illustration of our algorithm. There are two neural network models, one is  $M_{tr}$  for model training, and the other is  $M_{cs}$  for coreset selection. A coreset is selected from the current data stream  $B_{t}$  by conducting six steps. ① Feed a stream mini-batch  $B_{t}$  and sampled buffer data to  $M_{tr}$ . ② Copy model parameters from  $M_{tr}$  to  $M_{cs}$ :  $\theta_{cs} \gets \theta_{tr}$ . ③ Feed a mini-batch  $B_{t}$  into model  $M_{cs}$ . ④ Conduct bilevel optimization to update the model parameter  $\theta_{cs}$  in  $M_{cs}$  and output a probability distribution  $w$ . ⑤ Sample a coreset from  $B_{t}$  based on the distribution of  $w$  and add the sampled data into buffer. ⑥ Calculate stochastic gradient based on  $B_{t}$  and sampled data in the buffer in Step 1, and update  $\theta_{tr}$  based on gradient information. Repeat the above steps for each stream mini-batch.
games [68]. However, in continual learning, where DNNs are trained on a sequence of tasks with possibly non-i.i.d. data, the performance will be degraded on the previously trained tasks. This is referred to as catastrophic forgetting [53, 52, 60]. To alleviate catastrophic forgetting, one of the effective ways is rehearsal-based continual learning, where a small replay buffer is maintained and revisited during the continuous learning process. There is a line of works studying how to efficiently maintain the replay buffer using the coreset selection approach [6, 78, 83], in which a small set of data is selected as representative samples to be used in continual learning.
The coreset selection in continual learning is formulated as a cardinality-constrained bilevel optimization problem which is solved by incremental subset selection [6]. This greedy approach is computationally expensive and hence is not scalable when the coreset size is large. To address this issue, Zhou et al. [83] propose a relaxation of the bilevel formulation in [6], which drops the nested nature of bilevel formulation and actually becomes two sequential optimization problems. Tiwari et al. proposed a gradient approximation method in [71], which selects a coreset that approximates the gradient of model parameters over the entirety of the data seen so far. Yoon et al. [78] proposes an online coreset selection method by maximizing several similarity metrics based on data pairs within each minibatch, and sample pairs between each minibatch and coreset. These approaches do not directly address the algorithmic challenges caused by the nested nature of the bilevel optimization problem, and may not solve the original bilevel coreset selection problem efficiently.
The key challenges in the bilevel coreset selection problems are two folds. First, the bilevel formulation in [6] needs to directly perform optimization over cardinality constraint, which is a nonconvex set and greedy approaches are expensive when the coreset size is large. Second, the bilevel formulation in [6] has a nested structure: one problem is embedded within another, and the outer and inner functions both have dependencies over the same set of decision variables. It remains unclear how to design efficient algorithms to solve constrained bilevel optimization algorithms for the coreset selection problem with provable theoretical guarantees.
Our proposed solution addresses these challenges with a novel bilevel formulation and provably efficient optimization algorithms. The proposed new bilevel formulation is referred to as Bilevel Coreset Selection via Regularization (BCSR). The main differences between our approach and the standard bilevel approach in [6] are: (i) unlike the standard bilevel formulation which requires performing optimization based on a cardinality constraint, we propose to solve a bilevel optimization on a probability simplex over training examples; (ii) to make sure the probability distribution lies in a low dimensional manifold, we propose to add a smoothed top- $K$  loss as a regularizer to the upper-level problem; (iii) due to our new formulation, we are able to design a simple and effective first-order method to solve this new bilevel problem with provable non-asymptotic convergence guarantees. The first-order method is easy to implement and much faster than the greedy approach as in [6]. Our main contribution is listed as follows.
Algorithm 1 PyTorch-style pseudocode for BCSR
1 # tasks: the task sequences; T: the number of tasks  
2 # L(): loss function; x: image batch, y: label batch  
3 # model, model-proxy: for training and coreset selection  
4 # update_w: a function updating outer variable w  
5 M = [[] * T] # memory buffer  
6 for ind, task in tasks:  
7 for (x, y) in dataloader:  
8 model-proxy.load_state_dict(model.params())  
9 (x_m, y_m) = next_batch(M)  
10 out1 = model(x)  
11 out2 = model(x_m)  
12 loss = L(out1, y) + L(out2, y_m)  
13 loss_backward()  
14 (x_core, y_core) = find CORESSET(x, y, K, model-proxy)  
15 M[ind].append((x_core, y_core))  
16 update(model.params)  
17 def find CORESSET(x, y, K, model-proxy):  
18 coreset_w = 1.0/y.size()  
19 for j in range(outer_loops):  
20 for i in range(inner_loops):  
21 lower_loss = L(coreset_w * model(x, y))  
22 lower_loss_Backward()  
23 update(model-proxy.params)  
24 coreset_w = update_w(coreset_w)  
25 return torch multinomial(coreset_w, K)
- We propose a new bilevel formulation, namely BCSR, for the coreset selection in rehearsal-based continual learning. Instead of directly learning the binary masks for each sample, the new formulation tries to learn a probability distribution in a low-dimensional manifold by adding a smoothed top- $K$  loss as a regularizer in the upper problem. This formulation is designed to satisfy two important features in continual learning with DNNs: (i) keeping the nested structure in the coreset selection; (ii) being amenable to first-order algorithms, which makes it easy to implement in modern deep learning frameworks such as PyTorch and TensorFlow. Based on the new formulation, we propose an efficient first-order algorithm for solving it. The main workflow of our algorithm is illustrated in Figure 1, and the corresponding Pytorch-style pseudocode is presented in Algorithm 1.
- We have conducted extensive experiments among various scenarios to verify the effectiveness of our proposed algorithm, including balanced, imbalanced, and label-noise data. Our algorithm outperforms all baselines for all settings in terms of average accuracy, and it is much better than all other coreset selection algorithms. For example, on imbalanced data of Multiple Datasets, BCSR is better than the best coreset selection algorithm by  $4.65\%$  in average accuracy. From bilevel optimization point of view, our algorithm significantly improves the vanilla greedy coreset selection method [6] in terms of running time on continual learning benchmark datasets.
- Under the standard smoothness assumptions of the loss function, we show that our algorithm requires at most  $O(1 / \epsilon^4)$  complexity for finding an  $\epsilon$ -stationary point in the constrained case $^2$ . Notably, the  $O(1 / \epsilon^4)$  complexity consists of  $O(1 / \epsilon^2)$  backpropagations and  $O(1 / \epsilon^4)$  samplings from Gaussian distribution, where the latter cost is computationally cheap.
Continual Learning There are different classes of continual learning methods, including regularization-based approaches [39, 81, 10, 1, 59, 64, 17], dynamic architecture methods [61,
In the constrained setting, the definition of  $\epsilon$ -stationary point is defined with gradient mapping, i.e.,  $w$  is a  $\epsilon$ -stationary point of the function  $\phi$  if  $\frac{1}{\beta} \| w - \mathcal{P}_{\Delta}(w - \beta \nabla \phi(w)) \| \leq \epsilon$ , where  $\beta$  is the stepsize,  $\mathcal{P}$  is the projection operator,  $\Delta$  is the probability simplex. This matches the best iteration complexity as in the single level optimization problem [24].
79, 62, 51, 76, 44, 77], and rehearsal-based methods [48, 57, 11, 58, 31, 3, 18, 28, 80, 6, 83, 78, 84]. In the rehearsal-based continual learning, the memory is either reproduced experience replay [48] or generative replay [67]. Our work focuses on the aspect of the coreset selection on the replay memory and can be flexibly integrated into rehearsal-based methods in continual learning.
**Coreset Selection** The coreset selection methods were used frequently in supervised and unsupervised learning, such as  $k$ -means [19], Gaussian mixture model [49], logistic regression [33] and bayesian inference [9]. They were also used frequently in the active learning literature [74, 63]. The coreset selection in continual learning is related to the sample selection [34, 2, 3]. Nguyen et al. [55] introduce variational continual learning which is combined with coreset summarization [5]. Borsos et al. [6] proposed the first bilevel formulation for the coreset selection in continual learning, which is later improved by [83, 78]. Compared with these works, our work focuses on improved bilevel coreset selection: we provide a better bilevel formulation than [6] and design a provably efficient optimization algorithm.
Bilevel optimization Bilevel optimization is used to model nested structure in the decision-making process [73]. Recently, gradient-based bilevel optimization methods have broad applications in machine learning, including meta-learning [20], hyperparameter optimization [56, 22], neural architecture search [46], and reinforcement learning [40, 32]. These methods can be generally categorized into implicit differentiation [16, 56, 45, 4] and iterative differentiation [50, 21, 20, 65, 27] based approaches. Recently, various stochastic bilevel algorithms have been also proposed and analyzed by [12, 37, 25, 32, 29, 4, 14]. A comprehensive introduction can be found in the survey [47]. In this work, we propose a novel stochastic bilevel optimizer with very flexible parameter selection, which shows great promise in the coreset selection for continual learning.
3 New Bilevel Formulation for Coreset Selection in Continual Learning
In this section, we first introduce our new bilevel formulation, namely Bilevel Coreset Selection via Regularization (BCSR). The key idea of this approach is to learn a probability distribution over the whole dataset such that the best model parameter obtained by minimizing the loss on the sampled dataset (i.e., the minimizer for the lower-level problem) is also the best for the whole dataset (i.e., the minimizer for the upper-level problem), and then a coreset can be sampled based on the learned probability distribution. In addition, the learned probability distribution is expected to lie in a low dimensional manifold (i.e., with  $K$  nonzero entries where  $K$  is the coreset size). To achieve this, we added a smooth top-  $K$  loss as a regularizer to promote the probability distribution to have  $K$  nonzero entries. Specifically, the objective function of BCSR is:
$$
\min_{\substack{0\leq w_{(i)}\leq 1\\ ||w||_{1} = 1}}\Bigl[\phi (w) = \sum_{i = 1}^{n}\ell_{i}(\theta^{*}(w)) - \lambda \sum_{i = 1}^{K}\mathbb{E}_{z}(w + \delta z)_{[i]} \Bigr]
$$
$$
s. t., \theta^ {*} (w) = \arg \min  _ {\theta} \left[ L (\theta , w) = \sum_ {i = 1} ^ {n} w _ {(i)} \ell_ {i} (\theta) \right] \tag {1}
$$
where  $n$  is the sample size,  $\theta$  is the model parameter,  $w$  is the sample weights,  $\ell_i(\theta)$  denote the loss function calculated based on  $i$ -th sample with model parameter  $\theta$ ,  $w_{(i)}$  is the  $i$ -th coordinate of  $w$ ,  $w_{[i]}$  is the  $i$ -th largest component of  $w$ ,  $\lambda > 0$  is the regularization parameter,  $w + \delta z$  denote to adding  $\delta z$  on each coordinate of  $w$  where  $z \sim \mathcal{N}(0,1)$ . Note that  $R(w,\delta) := -\lambda \sum_{i=1}^{K} \mathbb{E}_z(w + \delta z)_{[i]}$  denote the smoothed top-  $K$  regularization. We add this regularization to make sure the summation of the top-  $K$  entries of the learned probability vector is large, such that we can confidently choose a coreset with size  $K$ . The goal of employing Gaussian noise to the regularizer is for the ease of algorithm design: this Gaussian smoothing technique can make the regularizer to be smooth such that it is easier to design efficient first-order bilevel optimization solvers. Otherwise, the upper-level problem would become nonconvex and nonsmooth, and it would be difficult for algorithm design under this case.
Discussion and Comparison with Prior Works In this part, we illustrate how this new formulation addresses the drawbacks of the previous approaches. The work of [6] does not use any regularizer and regards the weight of each sample as a binary mask. This formulation needs to solve a combinatorial
Algorithm 2 Bilevel Coreset Selection via Regularization (BCSR)
Input: Dataset  $\mathcal{D}$    
Initialize: model parameter  $\theta_0$  , memory  $\mathcal{M} = \{\}$  1: for batch  $\mathcal{B}_t\sim \mathcal{D}$  do 2: Compute coreset  $S_{t} =$  Find-coreset  $(\theta_{t - 1},\mathcal{B}_t)$  3:  $\mathcal{M} = \mathcal{M}\cup S_t$  4: end for
Algorithm 3 Find-coreset  $(\theta, \mathcal{B})$
Input: the current model parameter  $\theta$ , the current batch  $\mathcal{B}$ , the iteration parameters  $J, N, Q$   
Initialize: coreset size  $K$ ,  $n = |\mathcal{B}|$ , and  $v_0$   
1:  $w_0 = \left[\frac{1}{n}, \ldots, \frac{1}{n}\right]$  (uniform probability initialization)  
2:  $\theta_1^0 = \theta$   
3: for  $j = 0, 2, \ldots, J - 1$  do  
4: Set initialization  $\theta_j^0 = \theta_{j-1}^N$  if  $j > 0$  and  $\theta_1^0$  otherwise  
5: for  $k = 1, \ldots, N$  do  
6: update  $\theta_j^k$  according to eq. (2).  
7: end for  
8: Set initialization  $v_j^0 = v_{j-1}^Q$  if  $j > 0$  and  $v_0$  otherwise  
9: Compute estimate  $v_j^Q$  by solving QP in (4) by GD with stepsize  $\eta$  and initial  $v_j^0 = v_{j-1}^Q$ .  
10: Compute hypergradient estimate in (3)  
11: Update  $w_{j+1}$  and project it onto simplex by (5)  
12: end for  
13:  $S \gets$  a set of  $K$  data points sampled from  $\mathcal{B}$  according to the probability distribution  $w_J$   
14: Return  $S$
optimization problem and their approach of incremental subset selection is computationally expensive. The work of [83] relaxes the bilevel formulation in [6] to minimize the expected loss function over the Bernoulli distribution  $s$ , i.e.,  $\min_{s\in \mathcal{C}}\Phi (s)$ , and develops a policy gradient solver to optimize the Bernoulli variable. Their gradient  $\nabla_s\Phi (s) = \mathbb{E}_{p(m|s)}L(\theta^* (m))\nabla_s\ln p(m|s)$  does not include the implicit gradient of  $L(\theta^{*}(m))$  in terms of  $s$ . However,  $\theta^{*}(m)$  actually depends on the mask  $m$ , and  $m$  depends on the Bernoulli variable  $s$ . In contrast, our bilevel optimization computes the hypergradients for the coreset weights  $w$  ( $0\leq w\leq 1$  and  $\| w\| _1 = 1$ ), which considers the dependence between  $\theta (w)$  and  $w^3$ . In addition, Zhou et al. [83] assume that the inner loop can obtain the exact minimizer  $\theta^{*}(m)$ , which may not hold in practice. In contrast, we carefully analyze the gap between the estimated  $\theta^{*}(w)$  and itself by our algorithm and analysis.
4 Algorithm Design
Equipped with the new formulation, the entire algorithm is presented in Algorithm 2, which calls Algorithm 3 as a subroutine. Each time the algorithm encounters a minibatch  $\mathcal{B}$ , a coreset is selected within this minibatch by invoking Algorithm 3. Algorithm 3 is a first-order algorithm for solving the bilevel formulation (1). In Algorithm 3, the model parameter  $\theta$  and the weight distribution  $w$  are updated alternatively. We first perform  $N$  steps of gradient descent steps to find a sufficiently good  $\theta$  for the lower-level problem (lines 5-7) by the update rule:
$$
\theta_ {j} ^ {k} = \theta_ {j} ^ {k - 1} - \alpha \nabla_ {\theta} L \left(\theta_ {j} ^ {k - 1}, w _ {j}\right), \tag {2}
$$
where  $\theta_j^k$  denotes the model parameters at the  $j$ -th outer loop and the  $k$ -th inner loop. To update the outer variable  $w$ , BCSR approximates the true gradient  $\nabla \phi(w)$  of the outer function w.r.t  $w$ , which is called hypergradient [56]. BCSR constructs a hypergradient estimator:
$$
\varphi_ {j} = \frac {1}{| \mathcal {B} |} \sum_ {\tilde {z} \in \mathcal {B}} \nabla_ {w} R (w, \delta ; \tilde {z}) - \nabla_ {w} \nabla_ {\theta} L \left(\theta_ {j} ^ {N}, w _ {j}\right) \left[ \left(\nabla_ {\theta} ^ {2} L \left(\theta_ {j} ^ {N}, w _ {j}\right)\right) ^ {- 1} \left(\sum_ {i = 1} ^ {n} \nabla_ {\theta} \ell_ {i} \left(\theta_ {j} ^ {N}\right)\right) \right], \tag {3}
$$
3The coreset weight  $w$  in our formulation is equivalent to sample mask  $s$  in [83]
Table 1: Experiment results on Split CIFAR-100
Methods	Balanced	Imbalanced	Label Noise
AT	FGTT	AT	FGTT	AT	FGTT
K-means Features	57.82±0.69	0.070±0.003	45.44±0.76	0.037±0.002	57.38±1.26	0.098±0.003
K-means Embedding	59.77±0.24	0.061±0.001	43.91±0.15	0.044±0.001	57.92±1.25	0.091±0.016
Uniform	58.99±0.54	0.074±0.004	44.73±0.11	0.033±0.007	58.76±1.07	0.087±0.006
iCaRL	60.74±0.09	0.044±0.026	44.25±2.04	0.042±0.019	59.70±0.70	0.071±0.010
Grad Matching	59.17±0.38	0.067±0.003	45.44±0.64	0.038±0.001	59.58±0.28	0.073±0.008
SPR	59.56±0.73	0.143±0.064	44.45±0.55	0.086±0.023	58.74±0.63	0.073±0.010
MetaSP	60.14±0.25	0.056±0.230	43.74±0.36	0.079±0.014	57.43±0.54	0.086±0.007
Greedy Coreset	59.39±0.16	0.066±0.017	43.80±0.01	0.039±0.007	58.22±0.16	0.066±0.001
GCR	58.73±0.43	0.073±0.013	44.48±0.05	0.035±0.005	58.72±0.63	0.081±0.005
PBCS	55.64±2.26	0.062±0.001	39.87±1.12	0.076±0.011	56.93±0.14	0.100±0.003
OCS	52.57±0.37	0.088±0.001	46.54±0.34	0.022±0.003	51.77±0.81	0.103±0.007
BCSR	61.60±0.14	0.051±0.015	47.30±0.57	0.022±0.005	60.70±0.08	0.059±0.013
where  $R(w, \delta; \widetilde{z}) \coloneqq -\lambda \sum_{i=1}^{K} (w + \delta \widetilde{z})_{[i]}$  and  $\widetilde{z} \sim \mathcal{N}(0, 1)$ . Solving the Hessian-inverse-vector product in eq. (3) is computationally intractable. We denote  $v^* \coloneqq (\nabla_\theta^2 L(\theta_j^N, w_j))^{-1} (\sum_{i=1}^n \nabla_\theta \ell_i(\theta_j^N))$  in eq. (3), where  $v^*$  can be approximated by solving the following quadratic programming problem efficiently by  $Q$  steps of gradient descent (line 9):
$$
\min  _ {v} \frac {1}{2} v ^ {T} \nabla_ {\theta} ^ {2} L \left(\theta_ {j} ^ {N}, w _ {j}\right) v - v ^ {T} \sum_ {i = 1} ^ {n} \nabla_ {\theta} \ell_ {i} \left(\theta_ {j} ^ {N}\right). \tag {4}
$$
Next, the hypergradient estimate (line 10) is computed based on the output of approximated quadratic programming. Note both model parameters and sample weights need to use warm start initialization (line 4 and line 8). Then the weight is updated and projected onto simplex (line 11):
$$
\hat {w} _ {j + 1} = w _ {j} - \beta \varphi_ {j}, w _ {j + 1} = \mathcal {P} _ {\Delta^ {n}} (\hat {w} _ {j + 1}), \tag {5}
$$
where  $\Delta^n \coloneqq \{w \in \mathbb{R}^n : 0 \leq w_{(i)} \leq 1, ||w||_1 = 1\}$ . In terms of other experimental hyperparameters, we allow very flexible choices of hyperparameters (e.g.,  $N$ ,  $Q$ ) as shown in our theory to achieve polynomial time complexity for finding a  $\epsilon$ -stationary point.
The selected coresets for each task are stored in a memory buffer with a fixed size  $m$ . There is a separated memory slot for each task with size  $[m / i]$  when the task  $i$  comes. After each task  $i$ , the memory slots before  $i$  will randomly remove some samples to adjust all the memory slots to  $[m / i]$ . That means the memory size for each task decreases as the task ID increase to maintain the total buffer size  $m$  unchanged. The same memory strategy is also used in the greedy coreset approach [6].
We conduct extensive experiments under various settings, including balanced data, imbalanced data, and label-noise data. The empirical results demonstrate the effectiveness of our method in rehearsal-based continual learning.
Datasets We use commonly-used datasets in the field of continual learning, including Split CIFAR-100, Permuted MNIST, Multiple Datasets, Tiny-ImageNet, and Split Food-101. We follow the experimental settings as that in prior work [78] and [83]. Each dataset is processed with three approaches: balanced, imbalanced, and label-noise. Please refer to Appendix L for more details about data processing and settings.
Baselines We compare our algorithm BCSR with other continual learning methods based on coreset strategy, including  $k$ -means features [55],  $k$ -means embedding [63], Uniform Sampling, iCaRL [57], Grad Matching [9], Greedy Coreset [6], PBCS [83], GCR [71], and OCS [78]. We also compare with non-coreset reply method, SPR [38], MetaSP [70]. All algorithms are built upon episodic memory, which stores coreset selected from stream data. Then a model, such as ResNet-18 [30], is trained over the data from the current stream and the episodic memory.
Metrics Average accuracy and forgetting measure [10] are two primary evaluation metrics that are used in continual learning literature. AVG ACC  $(A_{T})$  is the average accuracy tested on all tasks after finishing the task  $T$ :  $A_{T} = \frac{1}{T}\sum_{i=1}^{T}a_{T,i}$ , where  $a_{T,i}$  is the test accuracy of task  $i$  after training