Progressive Blockwise Knowledge Distillation for Neural Network Acceleration

1 Introduction

Recent years have witnessed a great development of deep convolutional neural networks (DCNNs) and their various applications [Krizhevsky et al., 2012; Simonyan and Zisserman, 2015; He et al., 2016; Szegedy et al., 2017]. Due to the resource limit of real world devices, DCNN compression and acceleration have emerged as a crucial and challenging problem in practice. Typically, the problem is resolved from the following four perspectives: 1) quantization and binarization [Hubara et al., 2016; Rastegari et al., 2016; Wu et al., 2016; Zhou et al., 2017]; 2) parameter pruning and sharing [Courbariaux et al., 2015; Hu et al., 2016; Li et al., 2017a; Luo et al., 2017; Molchanov et al., 2017; Li et al., 2017c]; 3) matrix factorization [Tai et al., 2016; Lin et al., 2016; Jaderberg et al., 2014; Sainath et al., 2013]; and 4) model distillation [Bucila et al., 2006; Ba and Caruana, 2014; Hinton et al., 2014; Romero et al., 2015; Li et al., 2017b]. In principle, the first three points focus on how to carry out an efficient network inference process using a variety of computational acceleration techniques with low memory usage. In contrast, the last one aims at distilling the original network model into a low-complexity network model in terms of the teacher-student learning strategy. Without sacrificing too much accuracy, the low-complexity network model naturally possesses the properties of high computational efficiency and low memory usage. However, the effectiveness of model distillation is often challenged in the aspects of teacher-student network optimization and student network structure design. Therefore, we mainly focus on constructing an effective network learning strategy with a structurepreserving criterion for model distillation in this paper.

∗ Indicates equal contribution

†Corresponding author

More specifically, the process of model distillation usually involves two components: a teacher network with a complicated network structure as well as a simple student network. In essence, the process seeks for a feasible student network to mimic the output of the teacher network. Usually, conventional approaches (e.g., Knowledge Distillation [Hinton et al., 2014]) rely on a non-convex joint network optimization strategy that converts the teacher network into the desired student network in a one-pass fashion. Implemented in a huge search space of the student network function with a wide variety of network configurations, the aforementioned non-convex joint optimization process is usually intractable and unstable in practice. Following the work [Hinton et al., 2014], Nowak and Corso [Nowak and Corso, 2018] make an attempt to compress a subnetwork block of the teacher network into a student subnetwork block and then design different methods for initialization and training. Moreover, their design criterion for the student subnetwork block is to simply retain a convolution layer and directly remove the other two convolution layers from the teacher subnetwork block. Such a blockwise distillation strategy is simple and easy to optimize, but incapable of effectively modeling the sequential dependency relationships between layer-specific subnetwork blocks. In addition, the student subnetwork block design criterion is also incapable of well preserving the receptive field information on feature extraction.

Motivated by the observations above, we propose a blockwise learning scheme for progressive model distillation. Specifically, the proposed learning scheme converts a sequence of teacher subnetwork blocks into a sequence of student subnetwork blocks after progressive blockwise optimization. Additionally, we propose a structure-preserving criterion for student subnetwork design. This criterion allows us to transform the teacher subnetwork blocks into the student subnetwork blocks without changing the receptive field.

As a result, the proposed progressive learning scheme aims at distilling the knowledge of the entire teacher network by locally extracting the knowledge of each block in a progressive learning manner with the structure-preserving criterion of feature extraction. Therefore, the proposed scheme creates a novel network acceleration strategy in terms of progressive blockwise learning, and has the following advantages: 1) structure-preserving for each subnetwork block; 2) easy to implement for progressive blockwise optimization; 3) fast stagewise convergence; 4) flexible compatability with existing learning modules; and 5) good balance with high accuracy and competitive FLOPs reduction.

As a result, the main contributions of this work are twofold:

• We propose a progressive blockwise learning scheme for model distillation, which is innovative in the area of neural network acceleration. The proposed learning scheme naturally converts the problem of network acceleration into that of progressive blockwise function approximation.
• We propose a structure design criterion for the student subnetwork block, which is able to effectively preserve the original receptive field from the teacher network.

2 Our Approach

2.1 ProblemDefinition

To better understand our representations, we provide the detailed explanations of the main notations and symbols used throughout this paper as shown in Tab. 1.

A neural network is mainly comprised of the convolution layers, the pooling layers, and the fully connected layers. The subnetwork between two adjacent pooling layers is defined as a subnetwork block.

Let a complicated network T be the teacher network, which is composed of N subnetwork blocks:

where ti (i ∈ {1, 2, . . . ,N}) is the mapping function of the i-th block in the sequence and c is the mapping function of the classifier. To simplify the representation of the network, we shorten it as:N∏

Therefore, T is rewritten as:

The parameters of the teacher network are denoted as:

where Wc and Wti (i ∈ {1, 2, . . . , N}) are the parameters of c and ti.

Our objective is to design a student network with high computational efficiency and low memory usage, and to learn the corresponding optimal parameters. The student network composed of N student subnetwork blocks can be written as:

where si denotes a student subnetwork block. The corresponding optimal parameters of the student network are denoted as:

In essence, the problem is to design N student subnetwork blocks sequence S and optimize the corresponding parameters WS using the prior knowledge of N teacher subnetwork blocks sequence T :

The main challenges to solve this problem are: 1) The joint optimization of the student network function with a wide variety of parameters is usually intractable and unstable in practice; 2) Designing a feasible student network structure from scratch is difficult. In Sectiom 2.2 and Section 2.3, we propose our solutions to these two challenges.

2.2 Progressive Blockwise Learning

To reduce the optimization difficulty described in Eq. (7), we propose a progressive blockwise learning scheme. As shown in Fig. 1, our blockwise learning scheme learns the sequence of student subnetwork blocks by N block learning stages, and only optimizes one block at each learning stage while keeping the other blocks fixed.

To better introduce our blockwise scheme, we use the auxiliary function Ak (k ∈ {0, 1, . . . , N}) to represent our intermediate network at the k-th block learning stage:

where sj is the optimized student network block and ti is the teacher network block. The parameters of Ak are denoted as below:

As can be noted from the description of Ak , A0 is the teacher network T and AN is the optimized network S. Hence, the problem defined in Eq. (7) can be solved as below:

1 st block learning stage

2 nd block learning stage

N th block learning stage

More specifically, we use the teacher-student learning strategy at each block learning stage. At the block learning stage 1, we consider A0 /A1 as the teacher/student network, and then learn A1 from A0 . Similarly, we learn A2 from A1 at the block learning stage 2. By analogy, using such a progressive way, we will eventually solve the problem in Eq. (10).

In order to solve the teacher-student network optimization problem at each block learning stage, we use two terms to compose the objective loss function. Taking the k-th block learning stage for example, the first term of the loss compares the output of the block sk with its corresponding block in the teacher network tk, and is formulated as:

where I is the input of the network (e.g., image) and ‖ · ‖F denotes the Frobenius norm.

The second classification loss term Lkcls is to make the output of the student network approximate the ground truth, which is defined as:

where y is the ground truth for I and softmax(.) means the softmax loss between the network's output and y as described in [Simonyan and Zisserman, 2015].

Then the objective loss function for one training sample

(I, y) at the k-th block learning stage is as below:

where λlocal is a hyper-parameter to balance these two terms of the loss function. Therefore, our objective loss function for the training data {(I (1) , y(1) ), . . . , (I (M) , y(M)) } is:

By optimizing this loss function, the student subnetwork block can be trained under the ground truths and knowledge of the teacher subnetwork block. The details of our progressive blockwise learning scheme are shown in Alg. 1.

Moreover, the optimization order of our method can be flexible. Typically, we try three optimization orders: 1) Bottom-Up (1, 2, . . . ,N); 2) Top-Down (N,N - 1, . . . , 1); 3) Skipping-First (2, 3, . . . ,N, 1). Our experimental results show that optimizing the network in a Bottom-Up fashion gives better results than other orders. Therefore, we use the Bottom-Up optimization order as the default optimization order. The details are discussed in Section 3.

2.3 Student Subnetwork Block Design

We propose a blockwise design scheme. Based on our blockwise design scheme, the student subnetwork block we design should not change the input/output size of the next/previous block and maintain the receptive field at the same time. In addtion, this student subnetwork block is based on the teacher subnetwork block but contains fewer parameters and FLOPs (floating point operations).

Algorithm 1: Progressive Blockwise Learning

Input: A teacher network T = tN ◦ tN-1 ◦ · · · ◦ t1 and the corresponding parameters WT = {Wc,WtN ,WtN-1 , . . . ,Wt1 }.

Output: The objective student network S and the

corresponding optimal parameters WS.

• 1 Design the structure of si by the structure of ti as described in Section 2.3 with the corresponding randomly initialized weights Wsi (i = 1, 2, . . . , N);
• 2 Build the auxiliary function A0 and its parameters WA0 by the teacher network T ;
• 3 for k=1, 2, . . . ,N do

• 4 Build the auxiliary function Ak and its parameters WAk using Eq. (8) and Eq. (9);
• 5 Obtain the optimal parameters Wsk by minimizing Eq. (14);
• 6 Update WAk using Eq. (9);

7 end

• 8 return S = AN ;

More specifically, we prune half of the channels of all convolution layers which belong to the teacher subnetwork blocks to construct a simple student subnetwork block. Then we apply one 1 × 1 convolution layer at the end of the block to expand the number of output channels to the original number. Fig. 2 shows our student subnetwork block design based on the teacher subnetwork block of two K ×K convolution layers. The figure also illustrates that our design criterion preserves the receptive field from the teacher subnetwork block.

Blockwise Complexity Analysis

The main complexity of the model's computation costs lies in the convolution layers. For network acceleration, our method mainly focuses on the convolution layers. Taking the VGG-

16 model [Simonyan and Zisserman, 2015] for example, we will analyze the FLOPs reduction of the general student subnetwork block in this section.

Let the FLOPs of a K ×K (K = 3) convolution layer be Convfl. If the teacher subnetwork block has L (L ∈ {2, 3}) layers described above, then its total FLOPs TBfl is L × Convf l .

As shown in Fig. 2, we have three types of convolution layers based on our design. The first type of convolution layer is the same input and half the number of filters as the original convolution layer (Layer 1 in Fig. 2). The second type of convolution layer is half the input and half the number of filters as the original convolution layer (Layer 2 in Fig. 2). The third type of convolution layer is half the input and the same number of filters as the original convolution layer (Layer 3 in Fig. 2), but its kernel size is 1 × 1. So the FLOPs SBfl of the student subnetwork block which we design based on the previous teacher subnetwork block is:

In summary, the FLOPs reduction Cfl of the student subnetwork block we design is:

3 Experiments

3.1 Datasets

CIFAR10 [Krizhevsky and Hinton, 2009] is a labeled subset of the 80 million tiny images dataset for object recognition. This dataset contains 60000 32×32 RGB images in 10 classes, with 5000 images per class for training and 1000 images per class for testing.

CIFAR100 [Krizhevsky and Hinton, 2009] is also a labeled subset of the 80 million tiny images dataset for object recognition. It contains 100 classes including 600 32 × 32 images each, with 500 images for training and 100 images for testing.

ImageNet [Krizhevsky et al., 2012] is a dataset for ImageNet Large Scale Visual Recognition Challenge 2012. It contains 1.28 million training images and 50k validation images in 1000 classes.

3.2 Implementation Details

Network Architecture

In this paper, we utilize the widely used network VGG-16 [Simonyan and Zisserman, 2015] as the teacher network sample. We divide the VGG-16 network into five blocks using the pooling layers as the boundaries. For each block, we use our block design criterion described in Section 2.3 to obtain the student subnetwork block.

Data Preprocessing

On CIFAR10 and CIFAR100, the only preprocessing we do is subtracting the mean RGB value computed on the training data from each pixel. On ImageNet, we employ the data augmentation strategy as the practice in [Krizhevsky et al., 2012; Simonyan and Zisserman, 2015], including the 224 × 224 random cropping, the horizontal flip, and the per-pixel mean subtracted.

Training Details

We implement our architecture using Caffe [Jia et al., 2014] and use an NVIDIA TITAN X GPU to train the network. To validate our proposed learning scheme, we first conduct four ablation experiments on CIFAR10 which will be fully described in Section 3.3 in the following aspects: 1) Optimization Orders Comparison; 2) Receptive Fields Comparison; 3) Block Structure Comparison; 4) Block Training Strategy Comparison. Then we verify our method on CIFAR100 and ImageNet which will thoroughly be described in Section 3.4, and compare our proposed method to recent state-of-the-art approaches.

On CIFAR10 and CIFAR100, we use SGD with a minibatch size of 100 at each block learning stage. The initial learning rate is set to 0.01 and is divided by 10 after 3 epochs. We train the network using a weight decay of 0.005 and a momentum of 0.9. From our experiments, we notice that each learning stage converges in less than 6 epochs. And so we terminate each learning stage after 6 epochs.

On ImageNet, we use SGD with a mini-batch size of 32 at each block learning stage. The momentum parameter is chosen as 0.9, the initial learning rate is set to 0.01 and the weight decay is 0.0005. It takes 50000 training iterations for our method to converge at every learning stage.

3.3 Ablation Experiments

In this section, we utilize four ablation experiments to validate our method on CIFAR10. And we obtain the original VGG-16 model by fine-tuning the pre-trained model [Simonyan and Zisserman, 2015] on CIFAR10.

Optimization Orders Comparison. We evaluate the performance of different optimization orders. We try three optimization orders: Bottom-Up, Top-Down, Skipping-First as described in Section 2.2. The results in Fig. 3 (a) show that the Bottom-Up order performs better than the other two optimization orders. For Top-Down optimization order, the accuracy of the model drops rapidly at the stage five of optimizing block one which indicates that the first-block features have a significant impact on the higher block features.

Receptive Fields Comparison. We evaluate the performance of the student networks with different receptive fields. The method of maintaining each block's receptive field is fully described in Section 2.3, and repeating/reducing one 3 × 3 convolutional layer on the basis of the former to increase/reduce the receptive field of each block. As shown in Fig. 3 (b), the model's accuracy is the highest by maintaining the receptive field of each block at every block learning stage. After all block learning stages, the final model's Top-1 accuracy is improved by 1% compared to the other two methods. This indicates that keeping the student subnetwork block's receptive field the same as that of the teacher network block is feasible. Block Structure Comparison. We compare our student subnetwork block structure design's performance with Deep Net's [Nowak and Corso, 2018]. The Deep Net's method approximates the functions learned by the two or three layers with a single layer in each block. Fig. 3 (c) shows that the structure we design can improve the Top-1 accuracy by 17% over Deep Net's after all block learning stages.

Block Training Strategy Comparison. We evaluate the performance of the loss in Eq. (14). We employ two different training strategies at each block learning stage: 1) cls loss, we train the block only with the classification loss; 2) cls loss+local loss, we train the block with the weighted sum of the classification loss and the local loss described in Eq. (14). As shown in Tab. 2, the final model's accuracy is roughly 97% of the original model using our strategy and our strategy makes faster convergence than another strategy.

3.4 State-of-the-Art Performance Comparison

We evaluate the performance of our method on CIFAR100 and ImageNet, against the state-of-the-art methods, including APoZ-1, APoZ-2 [Hu et al., 2016], Taylor-1, Taylor2 [Molchanov et al., 2017], ThiNet-Conv and ThinNetGAP [Luo et al., 2017]. In principle, the original VGG16 model on CIFAR100 is obtained by fine-tuning the pre-trained model [Simonyan and Zisserman, 2015]; APoZ prunes the zero activation neurons to compress network; Taylor uses Taylor expansion to approximate the loss, and then only utilizes the first order gradient information; ThiNet reduces the number of convolution-layer channels by evaluating channel importance.

Subsequently, we design three variants based on our method: 1) Ours-Conv: it doesn't optimize the last three layers (conv5-1, conv5-2, conv5-3) of the original VGG-16 model similarly to ThiNet-Conv; 2) Ours-Conv-FC: it compresses the fully connected layers by truncated singular value decomposition [Girshick, 2015] on Ours-Conv; 3) Ours-GAP: it replaces the fully connected layers of Ours-Conv with a global average pooling layer [Lin et al., 2014] similarly to ThiNetGAP. We mainly focus on the VGG-16 model's acceleration, so we only compare our proposed method with the VGG-16 based frameworks as shown in Tab. 3.

In principle, our method is based on the progressive blockwise learning scheme with the structure-preserving design criterion. This criterion focuses on improving the accuracy of the compressed model with a competitive FLOPs reduction. From the comparison of these state-of-the-art methods in Tab. 3, we can observe that the accuracy of our proposed method is the highest and our method's FLOPs reduction is comparable to those of others. In addition, Fig. 4 shows that our method makes fast convergence for different learning stages. Taking ThiNet-Conv [Luo et al., 2017] for example, it demonstrates that ThiNet-Conv requires one or two epochs to converge for each layer, but we only need about 1.5 epochs to converge for each block containing two or three layers. Moreover, as shown in Tab. 3, the results of three variant experiments (Ours-Conv, Ours-Conv-FC, Ours-GAP) demonstrate that our method can be flexibly combined with other methods (e.g., fully connected layers decomposition [Girshick, 2015] and GAP [Lin et al., 2014]) to further reduce the parameters and achieve feasible results.

In summary, our scheme sets up an effective and practical network acceleration pipeline from a new viewpoint of progressive blockwise learning. In contrast to other types of pipelines, ours is a powerful tool in the aspects of structurepreserving for knowledge distillation, easy implementation for progressive blockwise optimization, fast stagewise convergence, flexible compatability with existing learning modules, and good balance with high accuracy and competitive FLOPs reduction.

4 Conclusion

In this paper, we have proposed a novel progressive blockwise knowledge distillation scheme for neural network acceleration, which distills the teacher network model by progressively extracting the knowledge of each block in local optimization. We have also proposed a structure-preserving criterion for the student subnetwork block design. The proposed criterion is able to keep the original receptive field unchanged from the teacher network. Therefore, our proposed progressive blockwise learning scheme provides a novel theoretical perspective as well as a promising practical alternative for neural network acceleration. Moreover, the experimental results on several datasets also demonstrate the effectiveness of our scheme.

Acknowledgments

This work was supported in part by the National Natural Science Foundation of China under Grants (U1509206, 61472353, and 61751209), in part by the National Basic Research Program of China under Grant Grant 2015CB352302.