Adaptive Second Order Coresets for Data-efficient Machine Learning

First-order coreset methods iteratively select weighted subsets of training data that closely approximate the full gradient at particular values of $w_{t}$, e.g. beginning of every epoch ​(Killamsetty et al., 2021, 2020; Mirzasoleiman et al., 2020):

where $\textbf{g}_{t,j}$ and $\gamma_{t,j}>0$ are the gradient and the weight of element $j$ in the coreset $S$. Such subsets often perform poorly for high-dimensional and non-convex functions, due to the following reasons: (1) the scale of gradient $\mathbf{g}\in\mathbb{R}^{d}$ is often different along different dimensions. Hence, the selected subsets estimate the full gradient closely only along dimensions with a larger gradient scale. This can introduce a significant error in the optimization trajectory for both convex and non-convex loss functions; (2) the loss functions associated with different data points $l_{i}$ may have similar gradients but very different curvature properties at a particular $w_{t}$. Thus, for a small $\delta>0$, the gradients $\nabla l_{i}(w_{t}+\delta)$ at $w_{t}+\delta$ may be totally different than the gradients $\nabla l_{i}(w_{t})$ at $w_{t}$. Consequently, subsets that capture the gradient well at at a particular point during training may not provide a close approximation of the full gradient after a few gradient updates, e.g., mini-batches. This often results in inferior performance, particularly when selecting larger subsets for non-convex loss functions; (3) subsets that only capture the gradient, select one representative example with a large weight from data points with similar gradients at $w_{t}$. Such subsets lack diversity and cannot distinguish different subgroups of the data. Importantly, the large weights introduce a substantial error in estimating the full gradient and result in a poor performance, as we show in Fig. 7 in the Appendix.

To address the above issues, our main idea is to select subsets of training examples that capture the full gradient preconditioned with the curvature of the loss landscape. In doing so, we normalize the gradient by multiplying it by the Hessian inverse, $\mathbf{H}^{-1}\mathbf{g}$, before selecting the subsets. This allows selecting subsets that (1) can capture the full gradient in all dimensions equally well; (2) contain a more diverse set of data points with similar gradients, but different curvature properties; and (3) allow adaptive first and second-order methods trained on the coresets to obtain similar training dynamics to that of training on the full data.

Formally, our goal in AdaCore is to adaptively find the smallest subset $S\subseteq V$ and corresponding per-element weights $\gamma_{j}>0$ that approximates the full gradient preconditioned with the Hessian matrix, with an error of at most $\epsilon>0$ at every iteration $t$, I.e.,:

where $\mathbf{H}^{-1}_{t}\textbf{g}_{t}$ and $\sum_{j\in S}\gamma_{t,j}\mathbf{H}_{t,j}^{-1}\textbf{g}_{t,j}$ are preconditioned gradients of the full data and the subset $S$.

Directly solving the optimization problem (5) requires explicit calculation and storage of the Hessian matrix and its inverse. This is infeasible for large models such as neural networks. In the following, we first address the issue of calculating the inverse Hessian at every iteration. ​Then, we discuss how to efficiently find a near-optimal subset to estimates the full preconditioned gradient by solving Eq. ​(5).

The subset selection problem (5) is NP-hard (Natarajan, 1995). However, it can be considered as a special case of the sparse vector approximation problem that has been studied in the literature, including convex optimization formulations—e.g. basis pursuit (Chen et al., 2001), sparse projections (Pilanci et al., 2012; Kyrillidis et al., 2013), LASSO (Tibshirani, 1996), and compressed sensing (Donoho, 2006). These methods, however, are expensive to solve and often require tuning regularization coefficients and thresholding to ensure cardinality constraints. More recently, the connection between sparse modeling and submodular²²2A set function $F:2^{V}\rightarrow\mathbb{R}^{+}$ is submodular if $F(S\cup\{e\})-F(S)\geq F(T\cup\{e\})-F(T),$ for any $S\subseteq T\subseteq V$ and $e\in V\setminus T$. optimization have been demonstrated (Elenberg et al., 2018; Mirzasoleiman et al., 2020). The advantage of submodular optimization is that a fast and simple greedy algorithm often provides a near-optimal solution. Next, we briefly discuss how submodularity can be used to find a near-optimal solution for Eq. (5). We build on the recent result of (Mirzasoleiman et al., 2020) that showed that the error of estimating an expectation by a weighted sum of a subset of elements is upper-bounded by a submodular facility location function. In particular, via the above result, we get:

Setting the upper bound in the right-hand side of Eq. (11) to be less than $\epsilon$ results in the smallest weighted subset $S^{*}$ that approximates full preconditioned gradient by an error of at most $\epsilon$, at iteration $t$. Formally, we wish to solve the following optimization problem:

By introducing a phantom example $e$, we can turn the minimization problem (12) into the following submodular cover problem, with a facility location objective $F(S)$:

where $C_{1}=L(\{e\})$ is a constant upper-bounding the value of $L(S)$. The subset $S^{*}$ obtained by solving the maximization problem (13) is the medoid of the preconditioned gradients, and the weights $\gamma_{j}$ are the number of elements that are closest to the medoid $j\in S^{*}$, i.e. $\gamma_{j}=\!\sum_{i\in V}\mathbb{I}[j=\min_{s\in S}\|\mkern 1.5mu\overline{\mkern-1.5mu\mathbf{H}\mkern-1.5mu}\mkern 1.5mu_{t,i}^{-1}\overline{\mathbf{g}}_{t,i}\!\!-\!\mkern 1.5mu\overline{\mkern-1.5mu\mathbf{H}\mkern-1.5mu}\mkern 1.5mu_{t,s}^{-1}\overline{\mathbf{g}}_{t,s}\|]$. For the above submodular cover problem, the classical greedy algorithm provides a logarithmic approximation guarantee $|S|\leq\big{(}1+\ln(\max_{e}F(e|\emptyset))\big{)}|S^{*}|$ (Wolsey, 1982). The greedy algorithm starts with the empty set $S_{0}=\emptyset$, and at each iteration $t$, it chooses an element $e\in V$ that maximizes the marginal utility $F(e|S_{t})=F(S_{t}\cup\{e\})-F(S_{t})$. Formally, $S_{t}=S_{t-1}\cup\{{\arg\max}_{e\in V}F(e|S_{t-1})\}$. The computational complexity of the greedy algorithm is $\mathcal{O}(nk)$. However, its complexity can be reduced to $\mathcal{O}(|V|)$ using stochastic methods (Mirzasoleiman et al., 2015), and can be further improved using lazy evaluation (Minoux, 1978) and distributed implementations (Mirzasoleiman et al., 2013). The pseudocode can be found in Alg. 1 in Appendix A.3.

Here, we analyze the convergence rate of first and second order methods applied to the weighted subsets $S$ found by AdaCore. By minimizing Eq. (13) at every iteration $t$, AdaCore finds subsets that approximate full preconditioned gradient by an error of at most $\epsilon$, i.e. $\|\mathbf{H}^{-1}_{t}\textbf{g}_{t}-\sum_{j\in S}\gamma_{t,j}\mathbf{H}_{t,j}^{-1}\textbf{g}_{t,j}\|\leq\epsilon$. This allows us to effectively analyze the reduction in the value of the loss function ${\cal{L}}$ at every iteration $t$. Below, we discuss the convergence of a first and second-order gradient methodapplied to subsets extracted by AdaCore.

Convergence for Newton’s Methods and AdaHessian We first provide the convergence analysis for the case where the function ${\cal{L}}$ in Problem (1) is strongly convex, i.e. there exist a constant $\alpha>0$ such that $\forall w,w^{\prime}\in\mathbb{R}^{d}$ we have ${\cal{L}}(w)\geq{\cal{L}}(w^{\prime})+\langle\nabla{\cal{L}}(w^{\prime}),w-w^{\prime}\rangle+\frac{\alpha}{2}\|w^{\prime}-w\|^{2}$, and each component function has a Lipschitz gradient, i.e. $\forall w\in\mathcal{W}$ we have $\|\nabla{\cal{L}}(w)-\nabla{\cal{L}}(w^{\prime})\|\leq\!\beta\|w-w^{\prime}\|$. We get the following results by applying Newton’s method and AdaHessian to the weighted subsets $S$ extracted by AdaCore.

The proofs can be found in Appendix A.1.

Convergence for (S)GD in Over-parameterized Case Next, we discuss the convergence behavior of gradient descent applied to the subsets found by AdaCore. In particular, we build upon the recent results of (Liu et al., 2020) that guarantees convergence for first-order methods on a broad class of general over-parameterized non-linear systems, including neural networks for which the tangent kernel, defined as $\mathbf{J}^{T}\mathbf{J}$ are not close to constant, but satisfy the Polyak-Lojasiewicz (PL) condition. Where $\mathbf{J}=\partial f/\partial w$ is the Jacobian of the function $f$ with respect to the parameters $w$. A loss function ${\cal{L}}$ is $\mu$-PL^∗ on a set $\mathcal{W}$, if $\frac{1}{2}\|\nabla{\cal{L}}(w)\|^{2}\geq\mu{\cal{L}}(w),\forall w\in\mathcal{W}$.

The proofs can be found in Appendix A.2.

We show an exponential convergence for GD (Theorem 4.3) and SGD (Theorem 4.4) under the $\mu$-PL^∗ condition, as well as for second order methods (Theorems 4.1, 4.2), under $\alpha$-strongly convex and $\beta$-smooth assumptions on the loss.

In our convex experiments, we apply AdaCore to select a coreset to classify the Ijcnn1 dataset using L2-regularized logistic regression: $f_{i}(x)=ln(1-\text{exp}(-w^{T}x_{y}y_{i}))+0.5\mu w^{T}w$. Ijcnn1 includes 49,990 training and 91,701 test data points of 22 dimensions, from 2 classes with 9-to-1 class imbalance ratio. In the convex setting, we only need to calculate the curvature once to find one AdaCore subset for the entire training. Hence, we utilize the complete Hessian information, computed analytically, as discussed in Appendix B.3. We apply an exponential decay learning schedule $\alpha_{k}=\alpha_{0}b^{k}$ with learning rate parameters $\alpha_{0}$ and $b$. For each model and method (including the random baseline) we tuned the parameters via a search and reported the best results.

AdaCore achieves smaller loss residual with a speedup Figure 1 compares the loss residual for SGD and Newton’s method applied to coresets of size 10% extracted by AdaCore (blue), Craig (orange), and random (green) with that of full dataset (red). We see that AdaCore effectively minimizes the training loss, achieving a better loss residual than Craig and random sampling. In particular, AdaCore matches the loss achieved on the full dataset with more than a 2.5x speedup for SGD and Newton’s methods. We note that training on random 10% subsets of the data cannot effectively minimize the training loss. We show the superior performance of training with SGD on subsets of size 10% to 90% found with AdaCore vs Craig in Appendix Fig. 6.

AdaCore better estimates the full gradient Fig. 2 shows the normalized gradient difference between gradient of the full data vs. weighted gradient of subsets of different sizes obtained by AdaCore vs Craig and Random, at the end of training by each method. We see that by considering curvature information, AdaCore obtains a better gradient estimate than Craig and Random subsets.

Datasets We use CIFAR10 (60k points from 10 classes) , class imbalanced version of CIFAR10 (32.5k points from 10 classes) and CIFAR100 (32.5k points from 100 classes) (Krizhevsky et al., 2009), BDD100k (100k points from 7 classes) (Yu et al., 2020). The results on MNIST (70k points from 10 classes) (Deng, 2012) can be found in Appendix C.6. Images are normalized to [0,1] by division with 255.

Models and Optimizers We train ResNet-20 and ResNet-18 (He et al., 2016), with convolution, average pooling and dense layers with softmax outputs and weight decay of $10^{-4}$. We use a batch size of 256 in all experiments (except Table 3, Fig. 4a), and train using SGD with momentum of 0.9 (default), or AdaHessian. For training, we use standard learning rate scheduler for ResNet starting with 0.1 and exponentially decaying by factor 0.1 at epochs 100 and 150. We used linear learning rate warm-up for the first 20 epochs to prevent weights from diverging when training with subsets. All experiments were ran on a 2.4GHz CPU and RTX 2080 Ti GPU.

Calculating the Curvature To calculate the Hessian diagonal using Eq. (8), we use a batch size of $b_{H}\!=\!64$ to calculate the expected Hessian diagonal over the training data. We observed that a smaller batch size provides a higher quality Hessian compared to larger batch sizes, ​as shown in Table 1.

Baseline Comparison and Ablation Study Table 1 shows the accuracy of training ResNet-20, using SGD with momentum of 0.9 and AdaHessian, for 200 epochs on $S$=1% subsets of CIFAR-10 chosen every $R$=1 epoch by different methods. For SGD+momentum, AdaCore outperforms Craig by 12%, Random by 10%, GradMatch by 6%, and Glister by 16.8%. Note that in total, AdaCore selects 74% of the dataset during the entire training process, whereas Random visits 87%. Thus, AdaCore effectively selects subsets contributing the most to generalization. We see that the accuracy gap between the baselines and AdaCore shrinks when applying more powerful optimizers such as AdaHessian. Table 1 also shows the effect of exponential averaging of gradients and Hessian diagonal, and larger batch sizes for calculating the Hessian diagonal $b_{H}$. We see that exponential averaging help AdaCore achieving better performance, and smaller $b_{H}$ provides better results.

Fig. 3a compares the performance of ResNet-18 on 1% subsets selected from CIFAR-10 with different methods. We compare the performance of training on AdaCore, Craig, GradMatch, Glister, and Random subsets for 200 epochs, with training on full data for 15 epochs. This is the number of iterations required for training on the full data to achieve a comparable performance to that of AdaCore subsets. We see that training on AdaCore coresets achieves a better accuracy 2.5x faster than training on the full dataset, and more than 4.5x faster than the next best subset selection algorithm for this setting (c.f. Fig. 8b in Appendix for complete results).

Frequency and size of subsets selection Table 2, 4 shows the performance of different methods for selecting subsets of size $S\%$ of the data every $R$ epochs, from CIFAR-10 and imbalanced CIFAR-10. Table 2 shows that selecting subsets of size 1% every $R=5,10,20$ epochs with AdaCore achieves a superior performance compared to the baselines. Table 4 shows that AdaCore can successfully select larger subsets of size $S=10\%,30\%$ and outperform Craig (Std is reported in Appendix, Table 5).

AdaCore speeds up training Fig 4 compares the speedup of various methods during training ResNet18 on 10% subsets selected every $R=20$ epochs from BDD 100k and CIFAR-100. All the methods are trained to achieve a test accuracy between 72% and 74% on BDD 100k, and between 57% and 50% on CIFAR-100. On BDD 100k, AdaCore achieves 74% accuracy in 100 epochs and training on full data achieves a similar performance in 45 epochs. For CIFAR-100, AdaCore achieves 59% accuracy in 200 epochs and training on full data achieves a similar performance in 40 epochs. Complete results on speedup and test accuracy of each method can be found in Appendix C.3, C.4. We see that AdaCore achieves 2.5x speedup over training on full data and 1.7x over that of training on random subsets on BDD 100k. For CIFAR-100, AdaCore achieves 4.2x speedup over training on random subsets and 2.9x over training on full data. Compared to the baselines, AdaCore can achieve achieve the desired accuracy much faster.

Effect of batch size Table 3 compares the performance of training with different batch sizes on subsets found by various methods. We see that training with larger batch size on subsets selected by AdaCore can achieve a superior accuracy. As AdaCore selects more diverse subsets with smaller weights, one can train with larger mini-batches on the subsets without increasing the gradient estimate error. In contrast, Craig subsets have elements with larger weights and hence training with fewer larger mini-batches has larger gradient error and does not improve the performance.

In summary, see that AdaCore consistently outperforms the baselines over various architectures, optimizers, subset sizes, selection frequency, and batch sizes.

AdaCore selects more diverse subsets Fig. 3b shows the number of times different methods selected a particular elements during the entire training. We see that AdaCore successfully selects a more diverse set of examples compared to Craig. We note that GradMatch may not be able to select subsets with the desired size, and instead augments the selected subset with randomly selected examples. Hence, it has a normal-shaped distribution. Fig. 3c shows mean forgetting score for all examples within a class ranked by AdaCore at the end of training, over sliding window of size 100. We see that AdaCore prioritizes selecting less forgettable examples. This shows that indeed AdaCore is able to distinguish different groups of easier examples better, and hence can prevent catastrophic forgetting by including their representatives in the coresets.

AdaCore vs Forgettability and Uncertainty Fig. 5a, 5b show mean forgettability and uncertainty in sliding windows of size 100, 200 over examples sorted by AdaCore at the end of training. We see that AdaCore heavily biases its selections towards forgettable and uncertain points, as training proceeds. Interestingly, 5a reveals that AdaCore avoids the most forgettable samples in favor of slightly more memorable ones, suggesting that AdaCore can better distinguish easier groups of examples. Figure 5b shows similar bias towards uncertain samples. Fig. 5c, 5d show the most and least selected images by AdaCore, respectively. We see the redundancies in the never selected images, whereas images frequented by AdaCore are quite diverse in color, angles, occluded subjects, and airplane models. This confirms the effectiveness of AdaCore in extracting the most crucial subsets for learning and eliminating redundancies.