Matching Guided Distillation

respectively at the target distillation positions. Without loss of generality, we assume the feature maps are of the same spatial size $N=HW$ (height $H$ and width $W$) but could consist of different number of channels, e.g. $C_{T}=512$ while $C_{S}=128$. Given a teacher network $f_{T}$ with frozen parameters, we wish to enhance the training of the student networks $f_{S}$ using the hints from $f_{T}$. In a nutshell, the problem of feature distillation seeks for the optimum student network $f_{S}$ that minimizes the loss of the main task together with feature discrepancy penalty:

where $\gamma$ is a trade-off coefficient for distillation loss. The task loss $L_{task}$, for example, could be cross-entropy loss for classification or smooth-$L_{1}$ loss for object localization. The key to feature distillation is the design of the distillation loss, $L_{distill}$, which ensures the similarity between intermediate features $\mathbf{T}$ and $\mathbf{S}$:

where $\sigma_{T},\sigma_{S},d_{p}$ are the teacher and student feature transforms that convert raw feature into an easy-to-transfer form and distance functions respectively. In the past few years, various designs have been proposed to make better use of information contained in teacher networks. In MGD, $\sigma_{T}$ is a marginal ReLU as same as [11]. Based on a recent comprehensive review [11] on these design aspects, we build the pipeline by employing the marginal ReLU for teacher transform $\sigma_{T}$:

where $m<0$ is a margin value, computed as an expectation value over all training images. Following [11], we choose the partial $L_{2}$ loss function to calculate feature distance:

for any pair of matrices $\mathbf{T},\mathbf{S}\in{}^{C\times N}$ of the same dimension.

Our goal is to find a binary matrix $\mathbf{M}\in\{0,1\}^{C_{S}\times C_{T}}$ that encodes the channel-wise correspondence, where $m_{ij}=1$ if $i$-th student and $j$-th teacher channels are pertinent. In the special case when the teacher and student feature maps are of the same dimension (i.e., $C_{T}=C_{S}$), the matching is assumed to be one-to-one and the resulting matrix $\mathbf{M}^{\prime}$ defines a permutation of $C_{T}$ channels:

In general, we resort to a many-to-one matching as the teacher channel number $C_{T}$ is often several times more than the student one $C_{S}$. In order to make the distillation procedure evenly distributed over feature channels, we further constrain that each student channel has to be associated with $\alpha=\lfloor C_{T}/C_{S}\rfloor$ teacher channels. More specifically, the many-to-one balanced matching $\mathbf{M}$ satisfies:

This constraint enforces that $\mathbf{M}$ is a wide-shape matrix, where the sum of each column equals to one because each teacher channel can only be connected to one student channel. On the other hand, the sum of each row needs to be $\alpha$. In another word, each student channel has to be associated with $\alpha$ teacher ones.

Now Eq. 11 is not a standard linear assignment problem, because the Hungarian algorithm works on the square cost matrix. To satisfy this prerequisite, let $\mathbf{D}^{\prime}=[\mathbf{D};\cdots;\mathbf{D}]\in{}^{C_{T}\times C_{T}}$ to be a square matrix by concatenating $\alpha$ matrices $\mathbf{D}$ vertically²²2 For the case when $C_{T}$ is not divisible by $C_{S}$, we simply shave $C_{T}-\alpha C_{S}$ dummy teacher channels off the cost matrix. Although this solution is not optimal, we found the result is still promising on several datasets. . Our solution proceeds by optimizing a standard linear assignment problem to solve out $\mathbf{M}^{\prime}$ first:

using Hungarian algorithm. We then evenly slice the resulting matrix $\mathbf{M}^{\prime}=[\mathbf{M}_{1}^{\prime};\cdots;\mathbf{M}^{\prime}_{\alpha}]$ in row blocks, where each sub-matrix $\mathbf{M}^{\prime}_{i}\in\{0,1\}^{C_{S}\times C_{T}}$ is of the same size. It’s easy to prove that the optimal solution for Eq. 11 is:

Random Drop. The major limitation of the first sparse matching choice is that it only retains a small fraction ($C_{S}/C_{T}$) of information conveyed in the original teacher features. To reduce the information loss, our second choice for teacher reduction is to sample a random teacher channel from the ones associated with each student channel. More specifically, there are $\alpha$ non-zero elements in each row of $\mathbf{M}$ according to the constraint (Eq. 10). The random drop operation modifies the correspondence as $\mathbf{M}^{RD}\in\{0,1\}^{C_{S}\times C_{T}}$ by randomly keeping one non-zero element in each row, i.e., $\sum_{j=1}^{C_{T}}m^{RD}_{ij}=1$ for any $i=1,\cdots,C_{S}$. To have a better understanding, we visualize one case in Fig. 2(b), where the second student channel $C^{\prime}_{2}$ is associated with $C_{1}$ and $C_{2}$ teacher channels after the channel matching step. In random drop reduction, we randomly pick one of them (e.g. $C_{2}$) to match with the student. In order to maximize the randomness, we generate correspondence matrices $\mathbf{M}_{i}^{RD}$ independently for different spatial positions of the feature map. The overall reduction operation can be defined as:

where $\mathbf{t}_{j}\in{}^{C_{T}}$ denotes the $j^{th}$ column of the feature $\mathbf{T}$.
Absolute Max Pooling. Following [11], we place the distillation module before ReLU. Therefore both positive and negative values are transferred from teacher via the partial distance loss $d_{p}$ to student. We hope the reduced teacher features still take the maximum activations in the same spatial position, including positive (usable) and negative (adverse) information. To reach this purpose, we propose a novel pooling mechanism, named Absolute Max Pooling (AMP), as shown in Fig. 2(c). Given a set of feature activations $\mathbf{x}=[x_{1},\dots,x_{C}]^{T}\in{}^{C}$, the AMP is designed to choose the element that yields the largest magnitude:

Similar to the random drop idea, the AMP operation is performed independently for each spatial position $\mathbf{t}_{j}\in{}^{C_{T}}$ of the teacher feature map $\mathbf{T}\in{}^{C_{T}\times N}$. For each student channel $i=1,\cdots,C_{S}$, AMP is used to select the most active teacher channel among the associated $\alpha$ ones. Because this teacher-student association has been encoded as the non-zero elements of $i^{th}$ row $\mathbf{m}^{i}\in\{0,1\}^{\mathbf{C}_{T}}$ of matrix $\mathbf{M}$, we can write the overall reduction operation in matrix form as:

where $\circ$ indicates the element-wise product between two vectors. As shown in Fig. 2(d), AMP pools these $\alpha$ feature nodes into single one along the channel axis.
Theoretical Summary. In the theoretical perspective, we describe the understanding on these three channel reduction methods as following. All these methods belong to variants of the LA problem. With the same matrix notation, their goal can be connected as seeking for the optimal matching matrix $\mathbf{M}$ by optimizing certain linear objective. The first sparse matching (Eq. 14) is a simple modification of the original Hungarian algorithm. The second random drop (Eq. 15) and the last absolute max pooling (Eq. 16) ideas progressively improve sparse matching by different approaches to reduce information loss from $\sigma_{\mathbf{T}}$. To better understand and verify this insight, we have done many experiments and ablation studies in Section. 4.

Distillation Position. In our experiments, we use contemporary models in recent years, including ResNet [9], MobileNet-V1 [14], -V2 [30] and ShuffleNet-V2 [25]. Commonly, these models contains four stages, each of which is composed of repeated unit blocks, as shown in Fig. 3. We apply distillation in the last unit block of each stage. In the cases of distilling MobileNet-V2 and ShuffleNet-V2, the first stage is skipped, because their first stage is only a convolutional layer and also its feature map size ($H\times W$) is distinct from teacher’s.

Coordinate Descent. To optimize the whole system (Eq. 3), we use Coordinate Descent algorithm [35] by alternating between solving the combinatorial matching problem and updating network weights. Postulating the matching is solved, we plug $\mathbf{M}$ in Eq. 3 and employ Stochastic Gradient Descent (SGD) to update student network weights $f_{S}$. After several epochs or iterations of training with SGD, the student is switched into evaluation mode without learning. Then we feed a dataset that is randomly sampled from training data into student and teacher, in order to update matching $\mathbf{M}$ that optimizes Eq. 11 for the next training rounds. Solving Eq. 11 takes the computational complexity of $O(C_{T}^{3})$. This step introduces negligible cost in our implementation.

Classification. Following the competitor [11], student models are randomly initialized in tasks on CIFAR-100. We use WideResNet (WRN) [41] as the teacher, which is in the model setting of depth 28 and wide factor 4, indicated by s:28-4. The student has multiple settings in different compression aspects: depth, wide factor, and architecture. Results comparison is shown in Table 1. WRN with s:16-4 has the same wide factor but smaller depth, so its channels are in same number of teacher’s at the same distillation positions. Since there is no need to reduce teacher channels, we only use MGD-SM to achieve the lowest error rate $20.53\%$. WRN with s:28-2 has the same depth but a larger wide factor, thus teacher channels need to be reduced. The results show that MGD-SM is better than [11] with $+0.76\%$, MGD-RD is a little worse with $+0.23\%$, and MGD-AMP achieves the lowest error of $21.12\%$. In experiments with s:16-2, whose both depth and wide factor are smaller than those of teacher, MGD-SM & RD are worse than the competitor and MGD-AMP is competitive. Another setting for student is using a different architecture, ResNet with s:56-2, which is deeper but with the same wide factor. MGD has the competitive results under this setting as well. The Table 1 shows that MGD can handle various compression types for student.

Transfer Learning. We use fine-grained categorization on CUB-200 to investigate distillation for transfer learning. We implement MGD and our competitor Overhaul for using ResNet-50 to distill light students. The teacher ResNet-50 has been pre-trained on IN-1K and then trained on CUB-200. We use two prevailed lightweight models, MobileNet-V2 and ShuffleNet-V2. Conspicuously, they have fewer parameters than ResNet-50. The main results have been shown in Table 3. We have two baselines for each student: one is trained from scratch (Baseline-FS), and the other is fine-tuned (Baseline-FT) from IN-1K. We summarize the experimental results in two folds.

Object Detection & Instance Segmentation. To verify the generalization of MGD, we extend it to object detection and instance segmentation. RetinaNet [23] is a modern one-stage detector, which has excellent performance in both precision and speed. EmbedMask [40] is a novel framework for instance segmentation, which utilizes embedding strategy to generate instance masks on a unified one-stage structure. In this section, we experiment with RetinaNet and EmbedMask respectively, using three different backbones: ResNet-50 as teacher, ResNet-18 and MobileNet-V2 as students. All these backbones are pre-trained on IN-1K. We train the models on COCO train2017 set and test them on val2017. Baselines are trained without distillation. As comparison, we also train with [11, 42, 12] under same configurations. The main results are presented in Table 4. For object detection, in both cases of distilling ResNet-18 and MobileNet-V2, MGD has stable improvements by $0.47-0.91$ point. In segmentation, in the case of ResNet-18, we freeze the first two backbone stages to avoid OOM. MGD can bring about $0.8$ mAP point for ResNet-18 and $0.6$ point for MobileNet-v2, it outperforms the competitor. These results prove that MGD works more stable than previous works in object detection and instance segmentation.

Frequency of Updating $\mathbf{M}$. To find the best practices, we experiment on how intense the frequency of updating $\mathbf{M}$ could affect the MGD performance. Here all experiments adopt MGD-AMP. The baseline is updating $\mathbf{M}$ in the end of every training epoch (frequency=1) as same as validation does. In Fig. 5, updating in every 2 training epochs achieves the best results both in MobileNet-V2 and ShuffleNet-V2 on CUB-200. If the frequency becomes larger than 2 in epochs (frequency=4), it will produce higher error rates. This study suggests that MGD should not be updated either fast or lazily. On IN-1K, we update $\mathbf{M}$ after every epoch due to its large dataset scale.

Absolute Max Pooling. Absolute max pooling (AMP) leads to largest improvements compared to proposed SM and RD. Now we compare it with basic pooling operations, max pooling (MP) and average pooling (AvgP), to perform reduction along the channel axis. Table 6 shows that AMP behaves stably better than MP and AvgP. AvgP performs worst, because average pooling operation is easily to counteract the sharpness feature for a group of channel tensors. Albeit MP works closely with AMP, it’s still not perfect because it would shave negative feature values by positive ones. This result shows the effect of AMP for preventing feature information loss from reduction. For a better illustration, we have a fundamental and intuitive experiment on these three pooling operations in Section. A1.1.
With vs. Without Matching. This ablation checks the importance of channel mathing mechanism. We remove matching process and simply use AMP as a feature reducer along channel axis. The right table shows model top-1 err. with matching? ✓ ResNet-50 20.02 MobileNet-V2 22.10 20.64 ShuffleNet-V2 27.62 25.95 the results of MGD with and without channels matching. It proves the effectiveness of channels matching in MGD. Both of distilling MobileNet-V2 and ShuffleNet-V2 without matching are worse than that with matching about 1.5% in top-1 error.
Capacity Analysis. Next, we illustrate MGD is more efficient for training than other methods. We investigate the capacity of joint training with MGD and [11]. In experiments, we use four GeForce 1080Ti GPU cards to run training. Under the same experimental settings, MobileNet-V2 has less parameters and memory consumption than teacher without distillation. Table 6 shows [11] brings too many parameters to cause OOM. As well known, additional distillation module not only bring the learnable weights for training, but also additional gradients and optimizer states into GPU memory [27]. In contrast, training with MGD has little bit of additional parameters due to its basic nature of parameter-free. Moreover it has better results.
Optimization Analysis.