Dynamic Sparse Training via Balancing the Exploration-Exploitation Trade-off

We formalize the sparse training process of the proposed DST-EE as follows. We define a $L$-layer deep neural network with dense weight $\bf W={[\bf W_{1},\bf W_{2},...,\bf W_{L}]}$. During the training process, the weight of $i$-th layer at $t$-th iteration is denoted by ${\textbf{W}}_{i}^{t}$. We randomly initialize sparse weight tensor as ${\textbf{W}^{\prime}}$ = ${[\bf{W}^{\prime}_{1},\bf{W}^{\prime}_{2},...,\bf{W}^{\prime}_{L}]}$ with sparsity distribution of $P$ using ERK [12] initialization. Each sparse weight tensor within a layer has a corresponding mask tensor (zero elements masked by 0 and other elements masked by 1) with the same size. We define zero elements in weight tensor as non-active weights and others as active weights. For each iteration, we only update the active weights. In addition, every $\Delta T$ iteration, we update the mask tensor, i.e., for $i$-th layer, we drop the $k_{i}$ weights that are closest to zero (i.e., smallest positive weights and the largest negative weights), the dropped weights are denoted by ArgTopK$(\textbf{W}^{\prime}_{i},k_{i})$. We denote $\textbf{N}_{i}^{t}$ as the counter tensor that collects the occurrence frequency for each 1 mask. We initialize $\textbf{N}_{i}^{t}$ as a zero tensor with the same size as the corresponding weight tensor. Every $\Delta T$ iteration, the counter tensor is updated by adding the counter tensor with the existing mask tensor. We use ${\textbf{S}}_{i}^{t}$ to denote the importance score tensor in $q$-th mask update. We design the following acquisition function to compute the importance score tensor

where the first term $|\frac{\partial l(\textbf{W}_{i}^{t},\mathcal{X})}{\partial\textbf{W}_{i}^{t}}|$ is the absolute gradient tensor of $i$-th layer at $t$-th iteration. $\partial l(\textbf{W}_{i}^{t},\mathcal{X})$ is the loss of $i$-th layer. $\mathcal{X}$ is the input training data. In the second term $c\frac{\ln t}{\textbf{N}_{i}^{t}+\epsilon}$, $c$ is the coefficient to balance between the two terms and $\epsilon$ is a positive constant to make the remainder as nonzero. For each importance score tensor, we identify the $k$ highest absolute values and select the indices. These corresponding mask values with the same indices will be set to 1s. In the next iteration, we update the weights using the new mask tensor. In the whole process, we maintain that the newly activated weights are the same amount as the previously deactivated weights. We repeat the aforementioned iterations till the end of training. The details of our method are illustrated in Algorithm 1, where $\cdot$ means tensor matrix multiplication.

Figure 2 shows the training data flow of one layer using the proposed method. We use ${\textbf{W}}^{t}$ and ${\textbf{G}}^{t}$ to denote the weight and gradient tensor, respectively. $n$ is the total number of rounds of mask updates. $l^{t}$ is the loss to compute the gradient tensor. In the first iteration of each $\Delta T$, the weight tensor has a corresponding binary mask tensor, where zero elements are masked by 0 in the mask tensor and other elements are masked by 1. ${\textbf{N}}_{t}$ is the counting tensor, indicating the number of non-zero occurrences in previous mask updates.

In proposed sparse training, we exploit current knowledge (weights and gradients) and define the exploitation score to help decide the mask with the highest accuracy. More specifically, we define the exploitation score $\textbf{S}_{\text{exploi}}$ in $q$-th mask update as the first item of Eq. (1), i.e., $\textbf{S}_{\text{exploi}}=|\frac{\partial l(\textbf{W}_{i}^{t},\mathcal{X})}{\partial\textbf{W}_{i}^{t}}|,\quad t=q\Delta T,\quad i=1,2,...,L$.

We further propose an evaluation metric to quantify the degree of exploitation for weight growth. With high degree of exploitation, the policy will find a model with local minimal with large loss reduction in a short time. Therefore, a growth policy is designed to have a high exploitation degree if it leads to a fast reduction in losses in the next iteration.

To formulate the evaluation metric, we denote $\bm{W}=[w_{1}^{(1,1)},w_{1}^{(1,2)},...,w_{1}^{(m_{1},n_{1})},...,w_{j}^{(p,q)},...,w_{L}^{(m_{L},n_{L})}]$ as weight of a model, where $w_{j}^{(p,q)}$ denotes the weight element in the $p$-th row and $q$-th column of $j$-th layer in the model. $j$-th layer has $m_{j}$ rows and $n_{j}$ columns. We further define $\bm{W}_{jpq,-jpq}=[0,...,0,w_{j}^{(p,q)},0,...,0]$ with same size of $\bm{W}$. The degree of exploitation is denoted as $\Delta\mathcal{L}_{g}^{jpq}$ when the weight element in the $p$-th row and $q$-th column of $j$-th layer is grown in sparse mask update iteration, then

To generalize, we use $\Delta\mathcal{L}_{g}$ to denote the degree of exploitation of the model if k weights with indices of ${I_{1},I_{2},...,I_{k}}$ are grown, then

Besides exploitation, we simultaneously choose masks that have never been explored so the model will not be stuck in a bad local optimum. We define our exploration score $\textbf{S}_{\text{explor}}$ as the second item in Eq. (1), i.e., $\textbf{S}_{\text{explor}}=\frac{\ln t}{\textbf{N}_{i}^{t}+\epsilon},\quad t=q\Delta T,\quad i=1,2,...,L$, where $\textbf{N}_{i}^{t}$ is a counter tensor that collects the active (nonzero) occurrence frequency of each element. If an element with an active (nonzero) occurrence frequency of zero, it will have a corresponding higher exploration score than explored elements, thus being grown.

Inspired by RigL-ITOP [1], we use an evaluation metric to quantify the degree of exploration for weight growth. Assume $\bm{B}=[b_{1}^{(1,1)},b_{1}^{(1,2)},...,b_{1}^{(m_{1},n_{1})},...,b_{j}^{(p,q)},...,b_{L}^{(m_{L},n_{L})}]$ is a binary vector to denote if the corresponding parameter in $\bm{W}$ is explored (1) or not (0) throughout the process of sparse training. For exploration rate [1], we use the same formulation as RigL-ITOP [1], i.e., $R=\frac{\sum_{{j=1}}^{L}\sum_{{p=1}}^{m_{j}}\sum_{{q=1}}^{n_{j}}b_{j}^{(p,q)}}{\sum_{{j=1}}^{L}{m_{j}}\times{n_{j}}}$.

The mask tensor search task is challenging in sparse training. Firstly, the mask search task is a high-dimensional problem due to a large number of weights in DNNs. Secondly, the search space has many local minima and saddle points [25, 26] because of the non-convex loss function of DNNs [25, 26]. Therefore, the mask tensor search process is easily trapped in a bad local optimal because of its low global exploration efficiency [18] or needs a longer time to fully explore the loss landscape.

A better balance between exploration and exploitation can encourage search algorithms to better understand the loss landscape and help the sparse model escape from the bad local optima. The importance and challenges of balancing the exploration and exploitation tradeoff have been emphasized in many studies [27, 28]. However, they have not gained enough attention in sparse training. Therefore, there is a strong need to better control the balance and we propose to consider both the exploration and exploitation scores when choosing the mask. And our importance score in Eq. (1) combines the two scores and overcome the limitations of previous work.

We evaluate VGG-19 and ResNet-50 on CIFAR-10/CIFAR-100 and evaluate ResNet-50 on ImageNet. The model training and evaluation are performed with CUDA 11.1 on 8 Quadro RTX6000 GPUs and Intel(R) Xeon(R) Gold 6244 @ 3.60GHz CPU. We use a cosine annealing learning rate scheduler with an SGD optimizer. For CIFAR-10/100, we use a batch size of 128 and set the initial learning rate to 0.1. For ImageNet, we use a batch size of 128. We use the same sparsity initialization method ERK in the state-of-the-art sparse training method such as RigL [14] and ITOP [1]. To further validate the generalizability of the proposed method, we conduct experiments on graph neural network for link prediction tasks on ia-email [33] and wiki-talk [34] datasets.

CIFAR-10/CIFAR-100. The results of CIFAR-10/100 are shown in Table I. We compare our method with SOTA on VGG-19 and ResNet-50 models at sparsity of 90%, 95%, and 98%. To demonstrate the effectiveness of the proposed method, we compare it with three types of methods (i.e., pruning-at-initialization (SNIP, GraSP, SynFlow), dense-to-sparse training (STR, SIS), and dynamic sparse training (DeepR, SET, RigL)) from top to bottom. The results of baselines are obtained from the GraNet [35] paper. Overall, both pruning-at-initialization and dense to sparse methods have higher accuracy than dynamic sparse training (except for RigL (using ITOP [1] setting)). Among the various sparsity ratios, the proposed method achieves the highest accuracy for both VGG-19 and ResNet-50. Using typical training time (total training epochs is 160), there is almost no accuracy loss compared to the dense model at sparsity of 90% on both CIFAR-10 and CIFAR-100. On both VGG-19 and ResNet-50, the proposed method has the highest accuracy compared with SOTA sparse training methods at different sparsity on both CIFAR-10 and CIFAR-10 datasets. For VGG-19, our method has up to 3.3%, 4.6% and 6.7% increase in accuracy on CIFAR-10 and up to 11.1%, 15.3% and 18.8% higher performance in accuracy on CIFAR-100, at sparsity ratios 90%, 95% and 98%, respectively. For ResNet-50, our proposed method has accuracy improvement than RigL with the same training epochs. More specifically, on CIFAR-10, our method has 0.51, 0.86, 0.94 higher accuracy score at sparsity ratio 90%, 95%, 98%, respectively. On CIFAR-100, the accuracy improvements of the proposed method compared to the SOTA sparse training method are 2.2%, 2.0%, 0.83% at sparsity ratios of 90%, 95%, and 98%, respectively.

We investigate the effect of coefficients on exploration degree and test accuracy on VGG-19, CIFAR-10 / CIFAR-100 datasets as shown in Figure 3. The left subfigure in Figure 3a shows the different exploration degree curves generated using different tradeoff coefficients on CIFAR-100 with a sparsity of 0.95. We could see the larger $c$, the higher degree of exploration of the sparse model. The right subfigure in Figure 3a illustrates the test accuracy curves for different coefficients. Within the coefficient range, the larger $c$, the higher test accuracy. The combination of these two subfigures unveils the observation that the higher the exploration degree or higher weights coverage, the higher the test accuracy score. Similar observations are shown in Figure 3b, which validate our methods.