Training Skinny Deep Neural Networks with Iterative Hard Thresholding Methods

Introduction

Deep neural networks (DNNs) have achieved remarkable success in various applications. This has been driven by the rapid growth in the size of datasets, increasingly deeper network architectures and the development of various techniques in training deep models.

Despite their strong capability of learning rich and discriminative representations, DNNs usually suffer from following two problems caused by their inherent huge parameter space when applied in practice. First, most of state-of-the-art DNN architectures are prone to over-fitting even trained on large datasets . Secondly, as they usually consume large storage memory and computational resource, it is difficult to embed modern DNN models into devices with limited power and memory, e.g., mobile phones.

A lot of regularization techniques have thus been proposed to decrease the risk of over-fitting for DNNs, e.g., dropout . However, those techniques are unable to reduce the storage cost. Recently, several works have been devoted to compressing and speeding-up DNNs by pruning internal layer connections. However, most of those methods achieve efficiency gain at the cost of performance deterioration.

In this work, we propose a novel approach for training skinny DNNs (SDNNs) with explicit size constraints to address the above problems simultaneously, i.e., reducing the risk of overfitting to improve the generalization capability of deep models and meanwhile reducing the model size to decrease storage memory cost. Compared with other sparsity pursuit methods for training DNNs , SDNN is superior as it is able to boost the performance even when the model compression rate is high. Briefly, our proposed approach for training SDNNs contains two alternative phases:

Phase I Hard thresholding over connections and sub-network fine-tuning. We apply hard thresholding over connections (weight parameters) at each layer to select the most prominent ones for the DNN model. The hard thresholding preserves the top $k$ weight parameters with the largest magnitude and disables the others by zeroing their values. Then, we fine-tune the non-disabled parameters for compensating the performance loss caused by reducing the number of filters.

Phase II Connection restoration and training the entire network. The frozen connections are re-activated and all the parameters are learned through training the entire network. The goal of this phase is to restore the truncated parameters and involve them in learning better representations.

Alternating the above two phases in training DNNs is able to produce SDNNs which have a stronger generalization capability with fewer parameters compared with the counterpart trained in the conventional way. We term such a method compositing of the above two phases as iteratively hard thresholding (ITH). Note that these two alternative phases are only needed in training SDNNs, while in the testing stage SDNN only takes one feedforward pass for inputs to make prediction.

To verify the effectiveness of SDNNs and the ITH training approach, we conduct extensive experiments on four public datasets with various scales, i.e., MNIST , CIFAR10 , CIFAR100 and ImageNet for two DNN architectures with different complexities including Network in Network , and AlexNet . The experimental results clearly demonstrate that SDNN with ITH does not only improve the generalization capability of deep models and provide state-of-the-art performance, but also reduce the size of parameters at the same time. Therefore, ITH is a quite appealing approach for training SDNNs in real-world scenarios concerning limited computational and storage cost.

Preliminaries: Gradient Hard Thresholding Revisit

The gradient hard thresholding (GHT) algorithm was proposed by to solve the following sparsity-constrained convex optimization problem:

It is proved that under mild conditions, the GHT algorithm converges geometrically to the point with bounded deviation from global optimum, with a high probability .

Skinny Deep Neural Networks

where $h_{y^{*}}(\mathbf{x},\mathbf{W})$ denotes the probability score predicted for $\mathbf{x}$ on the ground truth category of $y^{*}$, and $\lambda$ is the weight decay factor. In traditional methods, a deep model is optimized through minimizing the loss function without any constraint over $\mathbf{W}$:

In this work, we propose to impose explicit cardinality constraints to layer-wise parameters during training in order to reduce the parameter size. Thus the optimization problem for training an SDNN is formulated as

2 From GHT to SDNN

Inspired by the success of GHT on solving sparsity-constrained convex problems , we apply it to train DNNs to alleviate the over-parameterization issue. However, straightforwardly applying GHT to train SDNNs cannot provide desirable results since the loss function is non-convex and highly complex. Therefore applying hard thresholding operation to the parameters of deep model at each iteration in the same way as GHT will cause the model to diverge and be unable to learn meaningful parameters.

In this paper, we propose a novel approach for training SDNNs with cardinality constraints. A summary of details on training SDNNs is presented in Algorithm 1. Compared with GHT, there are two main differences in the approach used for training SDNNs:

Multi-Step Update. GHT performs hard thresholding operation at each iteration. However in training SDNNs, we update the parameters for $s_{1}$ iterations before performing hard thresholding. Such a strategy yields following two advantages. First, the training is accelerated by largely reducing the time of performing hard thresholding. Secondly, by updating all parameters (including those of connections truncated by the last hard thresholding operation) sufficiently, the SDNNs are likely to learn more discriminative representation with more parameters. Otherwise with a single-step update, it is highly possible that connections which are truncated by the first hard thresholding operation would be always truncated in all subsequent hard thresholding operations due to the little change in magnitudes.

3 Training SDNNs

The training of SDNN mainly consists of two iteratively alternating phases, each of which contains an operation for the parameters of deep models followed by an optimization process which is mutually different. In this section, we explain the details of Algorithm 1.

Network Initialization: This corresponds to Step 1. At the beginning of the training stage, all the parameters of the SDNN are trained for $s_{1}$ epochs without considering cardinality constraints. Therefore the optimization formulation in this step is exactly the same as the one in Eqn. (3). The aim of this step is to provide a good initialization for the following steps, which prevents the SDNN from diverging or getting stuck in a bad local minimum.

Phase II: This phase corresponds to Step 5 and Step 6, which conduct a weight restoration operation and trains the entire network, respectively. Step 5 removes the cardinality constraints over parameters which are set to be 0 in Step 3 so that all parameters are updated in Step 6 freely. This step is critical in SDNN for the following reason: by updating all parameters including those set to be 0 in Step 3, the SDNN is able to restore some connections that are beneficial for learning feature representation with strong discriminative power. As can be seen in Figure 2, at the initial training stage, a large proportion of connections which are truncated in the last round of hard thresholding operation become significant (with parameters in large magnitudes) at this phase. Thus SDNN has a strong capability to search among a large parameter space for seeking a better local optimum. Besides, we observe that the ratio decrease as the training progresses due to the deep model converges to a good optimum.

Related Works

Early approaches for deep model compression including optimal brain damage and optimal brain surgeon that prune the connections in networks based on the second order information. However, those methods are not feasible for deep networks due to high computational complexity. Recent works aiming at network pruning include , which prune connections in a progressively greedy way or using sparsity related regularizer . Although those works can reduce model size significantly, they suffer from the dramatic performance loss. In contrast, SDNN does not only offer significant compression ratios but also improves the performance simultaneously.

Our work is also in line with model compression. For example, proposes to quantize the deep model by minimizing L2 error and seeks an low-rank approximation of the model. Recently, combined pruning, quantization and Huffman coding techniques and provided rather high compression ratios. However, those methods also introduce performance drop. There are also works trying to compress a model by using binning methods , but they can only be applied over fully connected layers. In contrast, our method can be applied for compressing both convolution layers and fully connected layers.

Experiments and Analysis

We justify our method on four scale-various object classification benchmarks, i.e. three small-scale ones including CIFAR10 , CIFAR100 and MNIST and the large-scale ImageNet . Two evaluation metrics, the classification performance and the number of parameters in a model, are used in comparison with other methods.

Deep Models In our experiments, we train and test two deep models which are with different complexities, i.e. Network in Network (NIN) , and AlexNet . Briefly, NIN has only convolutional weight layers by replacing the single linear convolution layers in the conventional CNNs by multilayer perceptrons and using the global average pooling layer to generate feature maps for each category. Compared with NIN, AlexNet is wider and deeper containing 60M parameters with five convolution layers and three fully connected layers.

Implementation All of our experiments are conducted on a NVIDIA TITAN GPU using Caffe. The numbers of training epochs for Step 1 and Step 4 in Algorithm 1 are set to $s_{1}=15$ and $s_{2}=150$ for small-scale datasets and $s_{1}=15$ and $s_{2}=40$ for the ImageNet dataset to reduce training time. The hyperparameters of NIN and AlexNet including learning rate, momentum and weight decay follow and , respectively. All the results of our methods in the paper are based on the models reaching convergence.

Data augmentation is used in many models for object classification to prevent overfitting. The horizontal flipping is used for CIFAR10 and CIFAR100. For ImageNet, we use random crop and horizontal flipping as in . No data augmentation method is applied for MNIST.

Memory Usage To efficiently utilize the sparsity property of SDNN to reduce the memory storage size, we refer to the “Bitmask” storage format proposed in , which stores the nonzero parameters as well as a mask whose number of bits is equal to the number of total parameters. The bit value will be set to 1 if the corresponding parameter is nonzero, otherwise to 0. As indicated by , such memory usage methods can be directly used in deep models at runtime and reduce the storage size efficiently at the same time.

2 Model Analysis

Hard Thresholding vs Random Thresholding

We conduct experiments to justify the hard thresholding by replacing it with random thresholding in our method while keeping the other configurations unchanged. Specifically, we set $r=0.5$ in our experiments and replace Step 4 in Algorithm 1 with random thresholding in which half of parameters in each layer are randomly set to zero. However, we observe that deep models diverge quickly in Step 5 after random thresholding. The reason is that the random thresholding places deep models in bad conditions by truncating those important parameters.

3 Results

The CIFAR-10 dataset consists of 60,000 color images of 32 $\times$ 32 pixels in 10 classes. The total dataset is split into 50,000 training images and 10,000 testing images. Table 1 compares the performance and # parameters of SDNN and other state-of-the-art methods either when data augmentation is used or not. SDNN with different sparsity ratios are denoted with SDNN-#$\times$ where # is the reciprocal of the sparsity ratio. It is observed that when data augmentation is not applied, SDNN with sparsity ratio $r=0.5$ achieves the best result among all methods, reducing the error rate (ER) of the baseline model NIN by 1.71%. Compared with RCNN-96 and RCNN-128 which have model sizes of 0.67M and 1.19M, respectively, SDNN-2$\times$ with a much smaller model size (0.49M) outperforms them by 0.61% and 0.28%, respectively, demonstrating that SDNN is able to significantly improve the generalization capability of deep models. To further test SDNN’s performance in deep models with a larger size, we evenly increase the parameter of each layer in the original NIN by two times, resulting in a model which is four times as large as the original NIN. We denote the enlarged NIN as $\mbox{NIN}_{2}$ and test SDNN with sparsity ratio $r=0.5$ (denoted as $\mbox{SDNN}_{2}$-2$\times$) on it. Compared with $\mbox{NIN}_{2}$, $\mbox{SDNN}_{2}$-2$\times$ is able to reduce the ER by 0.77% and 1.72% when data augmentation is applied or not, respectively, again verifies SDNN’s capability to reduce overfitting. Note compared with ResNet which gets the lowest ER against all other methods, our method is only with a 0.02% higher ER, but our model is much faster in testing since ResNet with 1.7M parameters has 110 layers while ours only has 9 layers.

We also test SDNN with larger sparsity ratios when data augmentation is used. When $r=0.9$, SDNN-10x is slightly worse than the baseline NIN model (0.28% higher on ER) with a significantly reduced model size (0.1M). We further train the SDNN-20x which has a sparsity ratio $r=0.95$ with model size $0.05M$ to achieve the ER of 10.53%, which is 1.64% higher than the baseline NIN. Considering NIN is a network consisting of all convolution layers, this is still a satisfiable result.

CIFAR100

CIFAR100 has 50,000/10,000 training/testing color images with resolution of $32\times 32$ pixels. Since the number of training images of each class in CIFAR100 is only one tenth of that in CIFAR10, deep models trained on this dataset are prone to overfitting. Table 2 summaries the state-of-the-art methods on this dataset. It is obviously observed that either when the data augmentation is used or not, SDNN-2$\times$ achieves the lowest ER among all methods with the smallest model size. Specifically, compared with NIN model, SDNN-2$\times$ is able to reduce the ER by 5.18% and 3.19% either when data augmentation is used or not, respectively, which again demonstrates the advantages of SDNN. Moreover, compared with RCNN-96 which has a model size of 0.67M and achieves 34.18% ER without using data augmentation, our method outperforms it significantly by reducing the ER by 3.68% while simultaneously attaining a smaller model size (0.49M versus 0.67M). Like on CIFAR10, we also test SDNN2-2$\times$ by adopting half of the model size of NIN2. As can be seen from Table 2, our method is able to further enhance the generalization capability of NIN2 by reducing the ER significantly by 2.16% /1.98% when trained with / without data augmentation, respectively.

MNIST

MNIST is one of the standard datasets in the machine learning community. It consists of 70,000 handwritten digits of 0 to 9 with 28$\times$28 resolution of pixels in gray scale format, which are split into 60,000/10,000 training/testing set. Table 4 shows the comparison results, from which we can see SDNN-2$\times$ is able to achieve the best performance using the least parameters. Compared with the baseline model NIN, SDNN-2$\times$ greatly reduces the error rates from 0.47% to 0.19% on this heavily benchmarked dataset with only half of the NIN model size. To our best knowledge, this is the best performance ever achieved by methods without using ensembles and other preprocessing methods, which could further boost the performance of ours.

ImageNet

To test the scalability of SDNN to large-scale datasets and model with bigger sizes, we conduct experiments on a much more challenging image classification task on 1000-class ImageNet dataset. As a renowned dataset in the vision research community, ImageNet contains about 1.2M training images, 50,000 validation images and 10,0000 testing images. To compare with other network pruning methods , we also use AlexNet as our baseline model according to the publicly available configurations in Caffe1. Table 4 lists the performance of SDNN and other methods on ImageNet using AlexNet. Compared with the baseline model, our model with a two and four times compression ratio can further reduce the top-5 error rates by 1.66% and 0.81%, respectively, demonstrating the efficacy of SDNN on large-scale datasets. Compared with , both report higher error rates when pruning the deep model. SDNN is superior by boosting the performance of deep models while largely reducing the model size at the same time. For example, compared with , SDNN-4$\times$ reduces its error rates significantly by 0.85% while obtaining models with the same size.

Conclusion

In this paper, we proposed an iterative hard thresholding method (IHT) to improve the performance of the deep neural network and reduce the size of parameters simultaneously. The training of SDNN using IHT consists of two alternative phases, i.e. firstly performing hard thresholding to set connections with small magnitudes to zero and fine-tune the significant filters, and secondly, re-activating the freezing. Experiments conducted on four scale-various datasets, i.e. CIFAR10, CIFAR100, MNIST and ImageNet using deep networks with different complexities, i.e. NIN and AlexNet demonstrated that our method is able to significantly boost the discriminative capability of deep models while largely reducing their sizes simultaneously.

References