Weighted Mutual Learning with Diversity-Driven Model Compression
Abstract
Online distillation attracts attention from the community as it simplifies the traditional two-stage knowledge distillation process into a single stage.
Online distillation collaboratively trains a group of peer models, which are treated as students, and all students gain extra knowledge from each other.
However, memory consumption and diversity among students are two key challenges to the scalability and quality of online distillation.
To address the two challenges, this paper presents a framework called Weighted Mutual Learning with Diversity-Driven Model Compression ( WML ) for online distillation.
First, at the base of a hierarchical structure where students share different parts, we leverage the structured network pruning to generate diversified students with different models sizes, thus also helping reduce the memory requirements.
Second, rather than taking the average of students, this paper, for the first time, leverages a bi-level formulation to estimate the relative importance of students with a close-form, to further boost the effectiveness of the distillation from each other.
Extensive experiments show the generalization of the proposed framework, which outperforms existing online distillation methods on a variety of deep neural networks.
More interesting, as a byproduct, WML produces a series of students with different model sizes in a single run, which also achieves competitive results compared with existing channel pruning methods.
Introduction
Deep neural networks (DNNs) have achieved impressive success for a variety of application fields, including computer vision [27, 45], natural language processing [2], time-series analytics [51] and so on [6, 47, 52].
However, the performance of DNNs is heavily dependent on model parameters and computations, which hinders its application on intelligent edge systems [20].
To relieve this issue, various techniques have been proposed to design lightweight models, including network pruning [18], network quantization [7], and neural architecture search [42].
Knowledge distillation [25], as a complementary method to further improve the performance of the lightweight models, can transfer the discriminative ability from a high-capacity yet cumbersome teacher model to a small and compact student model by encouraging the student to mimic the behavior of the teacher.
The vanilla knowledge distillation [25], as shown in Figure 1 (a), includes two stages: it first trains the cumbersome teacher model, whose soft logits that contain informative dark knowledge [25] are then transferred to a compact student model.
As the teacher is fixed during the knowledge transferring, the student can only learn from the teacher once it is trained, and the vanilla knowledge distillation is also called offline distillation [15].
To further simplify the distillation process and relieve the restriction when the larger-capacity teacher model is not available, the online distillation framework has been
ResBlock
ResBlock ResBlock
Logit Logit
ResNet ResNet
ResBlock ResBlock
ResBlock
Logit4 ResBlock
Logit
Logit Logit
Logit Logit
ResNet
ResNet (a): Vanilla Knowledge Distillation (b): Online Distillation (c): The Proposed
[Figure]
(WML)
[Figure]
proposed [4, 69, 70], which simultaneously trains a group of student models (a.k.a. peers) with similar structures.
As each student can gain extra knowledge from each other, it is also called deep mutual learning [69], as shown in Figure 1 (b).
Although online distillation conducts the distillation with a one-phase training scheme without pretraining a teacher, it has high memory consumption when jointly training several independent models.
In addition, it is supposed to yield more informative knowledge when diversity presents among students, while the similar structures in online distillation may hinder the students from learning more informative knowledge from peers [9, 16].
In order to simultaneously overcome the two aforementioned difficulties, memory consumption and diversity among peers, this paper proposes a novel framework, called Weighted Mutual Learning with
Diversity-Driven Model Compression ( WML ) for online distillation, as present in Figure 1 (c).
Our
WML considers a hierarchical structure where each peer share different parts, so as reduce the memory consumption compared with online distillation.
In addition, we leverage the structured network pruning to generate diversified peer models under different pruning ratios, which also further reduces the memory consumption.
On the other hand, simply treating the peers with different capacities to be equally important might neglect the diversity of all peer models, and limit the effectiveness of online distillation when learning from a group of peers with different capacities [9, 16].
In this paper, we formulate the weighted mutual learning based online distillation as a bilevel optimization problem, where the inner loop is to learn compact peers with the online distillation and the outer loop is to find suitable weights to determine a better optimization direction for knowledge distillation from a group of peers.
The weights are optimized by the mirror descent [3, 5], and the hypergradient is analyzed in a theoretical way with a close-form.
Our formulation can be seen as a dynamic weighting strategy for online distillation, which adaptively distills knowledge from the peer models to a student.
A summary of our main contributions follows.
• First, this paper proposes a Weighted Mutual Learning with Diversity-Driven Model Compression ( WML ) for online distillation, where the peer models share different parts in a hierarchical structure and the structured network pruning is leveraged to further generate diversified peers and reduce the memory consumption.
Second, we formulate the weighted mutual learning based online distillation as a bilevel optimization problem, and the hypergradient of optimizing the weights is derived with a close-form.
The dynamic weighting strategy enables the online distillation to adaptively determine the importance of each peer in distilling knowledge from a group of peers.
• Third, extensive experiments on several neural networks and three datasets verify the effectiveness of the proposed framework, which significantly outperforms existing online distillation and self-distillation methods.
More interesting, as a byproduct, WML can get a series of pruned student models with different model sizes in a single run, which also achieves competitive results when compared with existing channel pruning methods.
Related work 2.1 Knowledge Distillation
Knowledge Distillation (KD) is the concept [25] that distills knowledge from a larger teacher model to help the training of a smaller student model, by steering the student’s logits towards teacher’s logits, as shown in Figure 1 (a).
FitNet [43] further extends the feature maps as the learning target instead of logits, as shown in Figure 4 (a).
Rather than only learning from a single teacher, many recent studies [9, 16] investigate learning from a group of teachers, where the outputs of teachers are usually averaged to guide the student training, while how to effectively integrate different types of knowledge from multiple teachers is less investigated [9, 16].
The traditional knowledge distillation is also called offline distillation since the teacher models are first trained before distillation.
Differently, online distillation simultaneously trains a group of student models that can gain extra knowledge from each other without a pretrained teacher model, so it is also called the teacher-absent knowledge distillation.
Deep mutual learning (DML) [69], in Figure 1 (b), is the first proposed online setting for knowledge distillation, which trains multiple peer models in a collaborative way that each model can be seen as a teacher during others’ training process.
KDCL [16] further ensembles the output of all students as the teacher knowledge in the online distillation setting since the ensemble is supposed to yield better generalization ability than a single model [1].
To further reduce the training requirements, several recent works [70] enforce all students to share the same early blocks while equipped with individual auxiliary head branches, shown in Figure 4 (b).
Although this paradigm can relieve the heavy memory consumption in online distillation, it suffers from the drawback of diversity as all peers share the same backbone and with a similar structure.
Differently, Song and Chai [46] proposed a hierarchical structure so that the peers can share different parts in the early blocks, as show in Figure 4 (d).
Self-distillation is another recently-proposed distillation paradigm, which can be regarded as a special case of online distillation [15, 54, 59], shown in Figure 4 (c).
Specifically, [59] leverages the knowledge from the deeper layers to guide the training for the shallow sections in the backbone.
In this paper, we consider a hierarchical structure [46] that each peer shares different parts and gains weighted knowledge from other peers instead of the average ensemble.
Different from self-distillation that reduces the “depth” to get peers, this paper leverage the channel pruning to “slim” models to get diversified peers.
More descriptions of the difference between the proposed framework and the existing distillation frameworks are also described in Appendix B.
2.2 Network Pruning
Network pruning is an effective way to compress overparameterized neural networks by removing parameters with minimal performance degradation [55].
Based on the granularity, network pruning can be divided into two types: unstructured and structured pruning, where the former individually removes parameters in a network while the latter removes entire filters, channels, or even layers [39].
Although the unstructured pruning can mask most parameters in a network without performance drop [12], it requires specific hardware or software to accelerate the sparse network’s inference time [10, 19].
Differently, structured pruning preserves the structural regularity and can accelerate network on commodity hardware platforms [22, 32].
In this paper, we consider structured pruning for generic applications.
Next, the network pruning can happen after training vs. before training (a.k.a.
at initialization) [48].
Most conventional pruning methods identify the redundant weights that are removed with least degrading the performance of the trained model.
In contrast, recent works show that the randomly initialized neural network can be pruned without incurring any training [30], where
SNIP [29], GraSP [49], and SynFlow [48] are three well-known metrics for the unstructured pruning at initialization.
[39] further extends these metrics for structured pruning through summing scores of all parameters θ in a filter to determine its importance.
In this paper, we extend SNIP to determine the pruning ratio for each layer at initialization as [39] demonstrates it is the most effective among the three to determine the filter importance for channel pruning at initialization.
Methodology
This section describes our Weighted Mutual Learning with Diversity-Driven Model Compression ( WML ) framework for online distillation, which consists of three main parts: the overall structure and the optimization for learning a group of peer models, the pruning approach to get a series of peers with different model sizes, and a bi-level formulation based weighted mutual learning.
3.1 Overview Structure and Loss Function
Figure 1 (c) depicts the overall structure of the proposed WML framework, where we take ResNet-32, which is divided into three sections in terms of depth, as an example for easy understanding.
Different from the most commonly-used auxiliary heads [70] that all peer models share the same part, we consider a hierarchical structure [46] where each peer model shares different parts.
The framework is similar to the binary tree, where the branches at the same levels are copies of each other, and each path is considered as a peer model.
In addition, we further prune each block in WML with different ratios to further enhance the diversity among peer models and also reduce the memory consumption.
Different from the vanilla knowledge distillation, which requires a much larger pretrained model as the teacher, the online distillation directly learns from each other or the ensembles of students [1].
In this paper, we consider the linear combination of knowledge from all peer models, and the weight of each peer model is determined by a bi-level formulation which will be discussed in Sec. 3.
This paper focuses on the multi-class classification problems, and X = { x n } Nn =1 is the training images with N samples, Y = { y n } Nn =1 is the corresponding labels.
The output logit for i-th peer model is represented as z i , and its corresponding cross entropy loss is L CE which is computed with the z
i and labels Y .
The overall loss to train the whole framework containing M peers is defined as:
[Equation]
where ω i indicates the importance of the i-th peer model and α is the hyperparameter to balance the supervision from labels and teachers.
(cid:80)
Mj
=1 ω j KL
( z i , z j ) is the weighted knowledge from peers, which is one of the main differences between our WML and DML [69] in loss function, where DML considers each peer equally important in distillation.
Another alternative distillation strategy is to calculate the KL divergence between each student with the ensemble of all students [16], while [69] found it is not as competitive as directly learning from multiple teachers.
3.2 Diversifying Peer Models with Channel Pruning
Although the hierarchical structure can proportionally reduce the memory consumption as peers share different blocks to reduce complexity, it still has higher computational cost than traditional knowledge distillation.
To further reduce the memory requirements, we introduce the channel pruning, the most commonly-used paradigm in structured pruning, to our framework, which also encourage the diversity among homogenized peers.
The training-free network pruning has shown its superiority in finding competitive subnetwork with extreme efficiency [29, 48, 49], and [39] has demonstrates SNIP [29] is informative to determine the filter importance for channel pruning.
To be more specific, the relative importance of each filter f based on SNIP can be calculated as:
I
[Equation]
, where (cid:12)(cid:12)(cid:12) ∂
L (
[Equation]
i )
∂θ i ⊙ θ i (cid:12)(cid:12)(cid:12) is the SNIP score of parameter θ i .
This value gives an expectation of the impact on loss when the filter is pruned, and filters with smaller
SNIP scores tend to produce weak impact to the loss as compared to the other filters in that layer [39].
In this paper, we further propose a relative importance I l to determine the importance of each layer at initialization, which is defined as:
[Equation]
where N l is the number of parameters in layer l .
Given a pruning ratio P of a model, the number of pruned parameters of each layer is then obtained by p l = P × N × ˜ I l , where ˜ I l is a normalization of 1 I l and N is the total number of parameters in the model, and the number of preserved filter is calculated accordingly which aims to preserve more filters in a more important layer.
In the practical implementation, we set a upper bound of the pruned ratio for each layer to avoid layer collapse [49].
In this way, given different P to peers, we can get a series of peer models with different structures.
3.3 Weighted Distillation with Bi-level Formulation
The proposed framework contains M peer models whose outputs can be also combined to generate a powerful ensemble.
However, the generalization capacity of the peer model is different from the model size.
Intuitively, we would expect to raise the importance of model
i with the highest capacity in distillation, i.e., enlarging ω i when peer model i contains most parameters [4, 9].
In this paper, we consider a dynamic weighting method to determine the importance of each peer, and the weighted mutual learning for online distillation can be formulated as a bilevel optimization problem:
[Equation]
where z ∗ i is obtained by training the whole framework θ with loss function in Eq.(1).
The inner loop is to optimize the whole network θ with the loss defined in Eq.(1), and the outer loop is to adjust the importance ω to make the weighted ensemble match the ground truth so that it reflects the importance of each peer more precisely.
We define the loss function of inner loop optimization as L 1 and outer loop as L 2 for simplicity.
Since z ∗ i is the local optimizer in the inner loop, Eq.(3) is a complicated nested optimization problem [36, 41, 52, 62, 67, 68].
In this paper, we consider a common one-step unroll learning paradigm [36, 41] that θ ∗ =
θ − γ ∇ θ L 1 with learning rate γ , and the gradient for the outer loop (which is also called as hypergradient) can be formulated as:
[Equation]
However, it is still impractical to calculate the second-order information for a large neural network, e.g., ∂ 2 L 1 ∂ω∂θ .
The finite-difference approximation is usually leveraged to roughly approximate this term [36], while it only considers the first two terms with the Taylor expansion.
In the following theorem, we show that we can exactly calculate ∂ 2 L 1 ∂ω∂θ without approximation in our framework.
Theorem 1 With one-step unroll learning paradigm, the gradeint for ω i in Eq.
(3) is formulated as:
[Equation]
where L a =
(1 − α )
L CE ( z i , Y ) + α (cid:80) Mj
=1 KL ( z j , z i ) .
Proof: Based on the definition of L 1 in Eq.(1), we have
[Equation]
Accordingly, ∂ 2 L 1 ∂ω∂θ can be calculate with autograd of L a on θ , and Theorem 1 is then proved.
□
Theorem 1 shows that the second-order term
[Table] L [Table] can be analyzed in a closed-form without Taylor expansion approximation when L 1 is with a specific structure.
Now that we can simply get the gradient to update model parameters θ and the weights for peers ω with autograd .
Since ω is a probability simplex that (cid:80) Mi =1 ω i = 1 , we use the mirror descent to update ω [3, 5].
Algorithmoutlines our weighted mutual learning for online distillation.
To be more specific, we first run several
[Table] [Table]
[Table] [Table]
[Table] [Table]
[Table] [Table]
[Table] [Table]
[Table] [Table]
[Table] [Table]
[Table]
SCAN
[Table] [Table]
[Table] [Table]
[Table] [Table]
WML 77.
[Table]
[Table] [Table]
ResNet
DSN [28] 77.
[Table] [Table]
SD [59] 77.
[Table] [Table]
[Table]
SCAN [60] 77.
[Table] [Table]
WML 77.
[Table]
steps of SGD based on the loss function in Eq.(1) to update model parameters θ with a fixed ω .
Then we calculate the gradient of ω i based on Eq.5, and run one step of mirror descent to update ω i that:
[Equation]
where η is the step size, and ω ki is the model importance for i-th peer in k-th step.
We repeat this loop until the stopping criterion meets.
Experiments
We perform a series of experiments to evaluate the effectiveness of our framework on seven convolutional neural networks with three image classification datasets.
We first compare our framework with two self-distillation baselines, SD [59]and SCAN [60], and the deeply supervised net (DSN) [28], since they share similar structures as our WML.
Specifically, SD and SCAN add early-exit blocks to the shallow blocks to get peer models, where SCAN contains additional attention modules.
DSN shares a similar structure as SD while without knowledge distillation.
After that, we compared
WML with existing knowledge distillation and online distillation baselines.
More interesting, as a byproduct, we also compared the pruned models produced by our framework with channel pruning baselines.
A series of ablation studies on dynamic weighting strategy and pruning are also discussed.
4.1 Comparison with Self-Distillation
For fair comparisons in this experiment, we consider 4 peer models in our WML, and also prune peer models into similar model size as classifiers in self-distillation (SD) [59], where Model1-3 are under
[Table]
different pruning ratios to match Classifier1-3 in self-distillation [59], and Model4 and Classifier are the same as the baseline network.
More experiments with other pruning ratios for WML can be found in Sec.4.4 and Appendix D. We adapt our framework with several ResNets [21] and also
the
MobileNetV2 [44], where results are present in Table 1 and Table 2, respectively.
There are two interesting pieces of evidence showing the advantage of the proposed WML on
CIFAR100.
Firstly , we can observe that WML leads to significant accuracy boosts, where our peer
Model4 outperforms baseline networks with an increment of 2.29% to 3.56% on CIFAR100 for the three networks.
In addition, our peer models beat the most classifiers in the three self-distillation baselines under similar model sizes, verifying the superiority of the proposed framework over selfdistillation.
Although SCAN achieves better results in Model1 and Model2 for ResNet18, we should notice that SCAN adds additional attention modules over SD, which has more parameters than
SD and our WML.
When the peer model is small, the additional attention module that contains relatively numerous parameters plays a more important role to improve the performance, while whose advantage decreases with the peer models become larger.
Secondly , compared with self-distillation baselines, which usually encounter significant performance drops in the shallower classifiers, our peer models can maintain similar performance with moderate network pruning, especially on the large networks.
Taking the Model2 on ResNet50 as an example, whose model size is only 50% as the baseline, our WML obtains 79.69% accuracy, while SD only achieved 74.21% accuracy.
SCAN added additional attention modules to further improve self-distillation performance, achieving 78.34% accuracy while still worse than our WML.
The above results demonstrate that reducing the width is a better way to get more competitive peers than the self-distillation paradigm using the shallow models.
The comparison results on ImageNet in Table 2 again verify the advantage of the proposed framework, where our WML also outperforms the self-distillation baseline under different model sizes on the ImageNet dataset.
As shown in the last column of Table 1, which is the ensemble of the four peer models, our WML outperforms the three baselines by large margins.
This superiority comes from two aspects.
One is that our weighted mutual learning can boost all peers’ performance, thus enhancing the ensemble.
The other is that our optimized weights can capture the peers’ importance more precisely, where more ablation studies can be found in the Appendix D.
4.2 Comparison with Online Distillation
Table 3 compares the proposed WML with several online distillation methods, including RKD[40],
CTSL-MKT[40], DML [69], ONE [70], and self-distillation (SD) [59] on CIFAR10, CIFAR100, and ImageNet.
In these online distillation methods, we consider four peer models as our WML for fair comparisons, and all reported results for baselines can also be found in [61].
We can see that knowledge distillation and online distillation both outperform the directly trained baseline model with a large margin, showing the effectiveness of knowledge distillation.
More interesting, different from traditional knowledge distillation which requires a powerful and large teacher model, online distillation methods can learn more meaningful information from peer models and gain performance improvements without a teacher.
These comparison results again verified the effectiveness of the proposed framework, which outperforms most online distillation methods on different datasets and neural networks.
More comparison with traditional knowledge distillation methods and selfdistillation methods can be found in the Appendix E.
4.3 Comparison with Channel Pruning
In our WML, we leverage channel pruning to generate peer models with different pruning ratios at initialization.
As a byproduct, we can get a series of pruned models with different FLOPs as peer models.
In this subsection, we compared the performance of the pruned models generated by our WML with existing channel pruning methods.
For fair comparisons, we conduct experiments with ResNet32 and ResNet56 on CIFAR10, and MobileNetV2 on ImageNet, which are the common
[Table] [Table]
setting to compare with existing works [61].
One clear advantage of WML over existing channel pruning methods is that we can get a series of pruned models with different FLOPs in a single run.
Table 7 and 5 compared our pruned peer models with existing baselines on CIFAR10 and ImageNet.
We make the peer models toward { 30% , 40% , 50% } pruning ratios on FLOPs for ResNet32 and
ResNet56 on CIFAR10, and { 30% , 50% , 70% } for MobileNetV2 on ImageNet.
As shown in Table
[Table]
there are two interesting findings.
First
[Table]
WML obtains higher accuracy than all channel pruning competitors on CIFAR10 after pruning similar FLOPs.
For example, the model with 40% pruned FLOPs by WML only has a 0.40% accuracy drop for ResNet32 and even 0.20% performance increase for ResNet56, which both outperform the competitors on CIFAR10.
Second , our pruned models with moderate pruning achieve better performance than the baseline model while with much fewer FLOPs.
This benefit comes from the online distillation as we find the SNIP baseline in [39], which adopts the same pruning method as us while without knowledge distillation, obtains worse results than those elaborately designed pruning methods.
However, combined with online distillation, our pruned model can achieve better performance than elaborately pruned models.
A straightforward direction for future work is to apply those competitive pruning methods into our framework for further performance improvements.
Table 5 presents the comparison results of our pruned peer models on ImageNet, which also achieves competitive results compared with other SOTA channel pruning methods under different FLOPs constraints.
For example, when the FLOPs are reduced by about 30%, WML30 only has 0.08% Top accuracy drop, which is better than MetaPruning’s 0.60%, AMC’s 1.00%, and DMCP’s 0.10% at the similar pruning ratio.
A similar advantage can also be found when pruning 50% and 70% FLOPs.
[Table]
[Figure]
4.4 Ablation Studies
In this section
[Figure]
we first conduct ablation studies to investigate whether the dynamic weighting strategy and peer pruning help improve the performance of online distillation.
Figure 2 only reports the results on ResNet50, and more experiments on other networks can be found in Appendix D.
Then, we show increasing the number of peers gradually enhance the performance of the proposed framework.
Figure 2 (a) presents how the performance of Model4 and ensemble change with different pruning levels.
As shown, when slightly pruning peers (only prune less than 10% parameters for Model1-3),
WML can get better results than the unpruned framework in terms of both Model4 and ensemble performance.
An underline reason is that slightly pruning rarely decreases the performance of peer models, while the diversity it brings can further enhance the distillation.
However, when the pruning ratio is too high (the last data point in the Figure 2 (a) prunes more than 70% parameters for Model1-3),we can see a clear performance drop since the pruning damages the performance of
Model1-3 seriously, which also further affects the distillation for Model4.
Figure 2 (b) presents how the dynamic weighting strategy works in our WML, where the pruning ratio is set as Sec.4.1.
As shown, the first data point in Figure 2 (b) indicates that we follow deep mutual learning [69] that remove the ω in Eq.(1), and the remaining points are with different η to update ω .
We can see a clear performance boost from unweighted distillation to weighted distillation, which verifies the effectiveness of our dynamic weighting strategy, which is also robust to the step size η .
In Figure(c), we prune Model 1-3 with 75%, 55%, and 20% parameters to match with Classifier 1-3 in [59], and the four lines in Figure 2 (c) track the change of weights ω for each peer model.
As shown, the weight of Model4 and Model4 consistently increase since they contains much more parameters and with better generalization ability.
In contrast, the weights for the other two models both decrease, and
Model1 with the least parameters descents the fastest.
This phenomenon also verifies our weighting strategy can adaptively adjust the importance of peers in distillation.
The main experiments show good results of the proposed framework with four peers.
We also verify the superiority of WML with different number of peers in Figure 3.
For M = 2 , we leave one peer unpruned and the other is pruned with 50%, and for M = 4 , 8 , we leave one peer unpruned and the remained peer models are pruned gradually from 30% to 70%.
As show in Figure 3, WML scale well with more peers although gain decrease with the number of peers increase.
The phenomenon verifies that more knowledge from peers can further enhance the performance of the students.
[Figure]
Conclusion and Future Work
This paper proposes a new framework for online distillation, called Weighted Mutual Learning with
Diversity-Driven Model Compression ( WML ), which introduces the structured network pruning to obtain diversified peer models, and a bi-level formulation is further leveraged to estimate the relative importance of peers in forming a powerful teacher during the online distillation.
The WML innovatively combines network pruning and online distillation to enhance performance and compress models simultaneously.
Extensive experimental results have verified the proposed framework, which leads to consistent accuracy improvements on various neural networks and datasets under different model sizes.
This paper considers a simple and efficient training-free channel pruning method to get the peer models while which is not as competitive as recent training-based channel pruning methods.
Combining more advanced channel pruning methods with online distillation for further performance and efficiency boost and extending WML to feature distillation or other applications, e.g., object detection, natural language processing, and architecture search, are among directions for future work.
Acknowledgments and Disclosure of Funding
This work was partially supported by Independent Research Fund Denmark under agreements 802200246B and 8048-00038B, the VILLUM FONDEN under agreement 34328, and the Innovation Fund
Denmark centre, DIREC.
This work was also in collaboration with Digital Research Centre Denmark – DIREC, supported by Innovation Fund Denmark.
Checklist
For all authors... (a) Do the main claims made in the abstract and introduction accurately reflect the paper’s contribu- tions and scope?
[Yes] (b) Did you describe the limitations of your work?
[Yes] (c) Did you discuss any potential negative societal impacts of your work?
[Yes] (d) Have you read the ethics review guidelines and ensured that your paper conforms to them?
[Yes] 2.
If you are including theoretical results... (a) Did you state the full set of assumptions of all theoretical results?
[No] (b) Did you include complete proofs of all theoretical results?
[Yes] 3.
If you ran experiments... (a) Did you include the code, data, and instructions needed to reproduce the main experimental results (either in the supplemental material or as a URL)?
[Yes] (b) Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)?
[Yes] (c) Did you report error bars (e.g., with respect to the random seed after running experiments multiple times)? [Yes] (d) Did you include the total amount of compute and the type of resources used (e.g., type of GPUs, internal cluster, or cloud provider)?
[Yes] 4.
If you are using existing assets (e.g., code, data, models) or curating/releasing new assets... (a) If your work uses existing assets, did you cite the creators?
[Yes] (b) Did you mention the license of the assets?
[Yes] (c) Did you include any new assets either in the supplemental material or as a URL?
[No] (d) Did you discuss whether
and how consent was obtained from people whose data you’re us- ing/curating?
[N/A] (e) Did you discuss whether the data you are using/curating contains personally identifiable informa- tion or offensive content?
[N/A] 5.
If you used crowdsourcing or conducted research with human subjects... (a) Did you include the full text of instructions given to participants and screenshots, if applicable?
[N/A] (b)
[Table]
(c): Self-Distillation
ResBlock ResBlock
Logit
Logit3 ResBlock
Bottleneck Bottleneck
Logit
ResBlock ResBlock
Logit
ResBlock ResBlock
ResBlock
Logit
ResBlock transform transform transform (a): Feature Distillation
Shared Backbone
Head
Head Head
A v er ag e E n s e m b l e Logit4 Head
Logit Logit
Logit (b): On-the-fly Native Ensemble (ONE) (d): Collaborative Learning
ResBlock
ResBlock ResBlock
Logit Logit
Logit Logit
ResBlock ResBlock
ResBlock ResBlock
A v er ag e E n s e m b l e
A Implementation Details
In this paper, we conducted a series of experiments on several image classification tasks to evaluate the effectiveness of our framework.
There are three types of networks, including ResNet18, ResNet50, and
ResNet101 on CIFAR10, CIFAR100, and ImageNet to compare with self-distillation methods and knowledge distillation baselines; ResNet32 and ResNet56 on CIFAR10 to compare with network pruning baselines; and also the MobileNetV2 on ImageNet dataset to compare with self-distillation baseline and also network pruning baselines.
CIFAR10 contains 50K training images and 10K test images with 10 classes and CIFAR100 is similar with CIFAR100 with only has 100 classes.
ImageNet contains over 1.2 million training images and 50K test images with 1000 classes.
We follow the typical data augmentation of these three datasets, where only random resized cropping and horizontal flipping are applied on CIFAR10 and CIFAR100 to get 32 × 32 images.
In contrast, on ImageNet, we also use the typical random resized crop and horizontal flipping, and the brightness, contrast, saturation data augmentation, to get 224 × 224 images.
In the CIFAR experiments, the number of initial filters is set as 64 for ResNet18, ResNet50, and ResNet101, while 16 for ResNet32 and ResNet56.
The SGD optimizer is selected with a momentum of 0.9 and a weight decay of 0.0005.
The batch size is set at 128 and the model is trained with 200 epochs.
For the MobileNet on ImageNet, the model is also trained by SGD with 90 epochs with 256 batch size.
Following [59], we set α 0.
during the knowledge distillation in all experiments.
The default step size η to update ω is also set as 0.1.
In our WML, we specify different pruning ratios to increase the diversity among students, as it is supposed to yield more informative knowledge in online distillation[9, 16], where similar phenomenon also exists in the neural architecture search [63–66]. We default pruned peer models to be of similar size as Classifiers in self-distillation [59] for fair comparison in the main experiments.
For example, the Classifier1-4 contain { 32 . 2% , 53 . 5% , 76 . 9% , 100% } parameters as the baseline ResNet50, so we set the pruning ratio as P = { 0 .
5 , 0 . 3 , 0 } in our WML in the main experiments to get four peer models.
In our channel pruning, we leverage SNIP at initialization to get the relative importance of each layer in Eq.(2), and the number of to be pruned parameters in each layer is calculated as p l = P × N × ˜ I l , where ˜ I l is a normalization of 1 I l and N is the total number of parameters in the model.
However, when the pruning ratio is too high, e.g., when P is larger than 0.8, the supposed pruned parameter p l may be larger than N l (the total number of parameters in this layer), which will result in layer collapse [49].
In our practical implementation, we set an upper-bounded ratio P up =0.9 that only maximal 90% parameters in a layer can be pruned.
B Ablation Study on the Structure
The Figure 1 and Figure 4 summarized existing popular paradigm of knowledge distillation.
Sec. 2 has detailed described the difference between the proposed framework and the existing works.
Apart from the self-distillation, the peers in On-the-fly Native Ensemble (ONE) [70] and collaborative learning (CL) [46] also partially share blocks.
In this paper, each head in ONE contains the last block and fully-connected layer in the ResNet, and CL
[Table]
[Figure]
share the same hierarchy structure as our WML. From the structure perspective, the main difference between ONE and our WML is that the backbone of ONE is the same for all peers, which means that it is impossible to prune peers under different ratios.
Although collaborative learning also considers the same hierarchy structure as us, its peers are with the same structure, while WML introduces channel pruning to diversify the structures of peers.
From the optimization perspective, the primary difference between our WML with all existing frameworks is that we propose a novel weighted mutual learning strategy in which each peer dynamically learns from the remaining peers, while the average ensemble is the most popular paradigm to get the knowledge source during the distillation for existing works.
Figure 4 (a) presents another paradigm in distillation, feature distillation, which extends the intermediate feature map as the learning target instead of logits, so as steer the student’s intermediate feature map towards the teacher rather than only the output logits.
Recent work on feature distillation [61] verified the superiority of this learning paradigm.
In this paper, we focus on knowledge distillation and leave feature distillation to our future work.
This section conducts experiments to show the superiority of our hierarchy structure with pruning to diversify the peers.
In this section, we consider four peers and replace the weighted mutual learning with the average ensemble for distillation, as denoted by WML-P, where all peers are under { 0 . 3 , 0 . 2 , 0 . 1 , 0 } pruning ratio.
Different from CL, which is additionally equipped with a backpropagation rescaling, we empirically found that it could only bring ignorable improvements to our hierarchy structure with pruning.
In this way, we dismiss this operator in WML-P. Table 6 presents the reproduced comparison between WML-P and two online distillation competitors, ONE and CL, under the same experimental settings.
As shown, we can find that WML-P and
CL both outperform ONE.
One potential reason is that the frameworks of WML-P and CL both contain more parameters than ONE, which further enhance the discriminative ability of the online distillation framework.
More interesting, we found that our WML-P, which is also the pruned version of CL with fewer parameters, can achieve slightly better results, especially in larger networks, e.g., ResNet101.
One potential reason is that those larger networks can maintain their discriminative ability under slight pruning.
In contrast, the diversity that the pruning brings can further enhance the ensemble performance so as to improve the distillation.
C Discussion on the Generalization of Weighted Ensemble
One of the key contributions of this paper is that we proposed weighted mutual learning, which adaptively distills knowledge from the peer models to a student by assigning different weights to different peers under
[Figure]
[Figure]
different capacities. To further verify the effectiveness and generalization of the proposed weighted mutual learning, we conduct a series of experiments to compare WML-P and WML-S, where WML-P only leverages the pruning to diversify the peers without the weighted mutual learning.
For fair comparison, in this section, we only slightly prune the neural networks, e.g. with { 0 . 3 , 0 . 2 , 0 . 1 , 0 } pruning ratio for four peers in WML-P and WML-S (the pruning ratio here is smaller than the pruning ration in the maintext) [Figure] and Table 6 presents the comparison results.
Compared with results in Table 1, which pruned much more parameter, we can find that WML-S can obtain higher accuracy in Table 6, especially for small networks, e.g. ResNet18.
A similar finding can be found in the next section on the ablation study on pruning ratio.
As shown in Table 6, WML-S gains performance improvements over WML-P for all peers in ResNet18 from 0.35% to 1.49%, and increase the ensemble performance by 2.96%.
Similar results can also be found in larger neural networks, e.g., ResNet50 and
ResNet101.
However, we can see a decrease in the performance gains when network becomes larger, which shares a similar trend as the number of heads.
A similar finding can be found in the analysis of η .
In the following, we further investigate whether the learned weights capture the importance of each peer by comparing the optimization trajectory of average-ensemble and weighted-ensemble.
In this experiment, we consider a similar experimental setting to our main experiments, so making the peers with the identical model sizes as the self-distillation.
As shown in Figure 5, we can find that leveraging the learned weight can further enhance the performance of the ensemble by tuning the importance of each peer.
This advantage is more obvious in the small ResNet18.
The above phenomenons imply that tuning the weights is another option to improve the discriminative ability of the ensemble, especially when each peer model is not powerful enough.
D More Ablation Study on Peer Pruning and Dynamic Weighting Strategy
This section gives the complementary results of the ablation study on the dynamic weighting strategy and peer pruning.
Figure 2 in the Sec.4.4 only reports the results on ResNet50, and we release more results of
ResNet18 and ResNet101 in Figure 6 and 7.
Interesting, we can find a controversial phenomenon in ResNet18 and ResNet101 on the pruning ratio, as shown in Figure 6 (a) and Figure 7 (a), where we can see a sharp drop on performance with the higher pruning ratio in ResNet18 while ResNet101 can maintain or even improve the performance with the pruning.
The potential reason is that the ResNet101 is much larger than Resnet18, and it also contains much more redundant parameters that pruning 30% parameters will not deteriorate its discriminative ability.
On the contrary, pruning 30% parameters clearly affect the discriminative ability of
ResNet18 so as deteriorate the performance.
In addition, for a larger neural network, our WML is more robust to the η , as shown in Figure 6 (b) and Figure 7 (b).
In Figure 6 (c) and Figure 7 (c), we also prune Model 1-4 with different pruning ratio, e.g., { 0 . 7 , [Table]
[Table]
consistently increase since they contain much more parameters and with better generalization ability in ResNet18 and ResNet101.
In contrast, the weights for the other two models both decrease.
Results in ResNet18 are very similar to ResNet50, while ResNet101 presents a bit different results.
We can find that Model3 gets the highest weight in ResNet101, while Model4 gets the highest weight for both ResNet18 and ResNet50, which contains the most parameters.
This interesting phenomenon also suggests that moderate pruning of an overparameterized model could hardly affect its performance, but further improve its discriminative ability.
In addition, we also investigate the effects of channel pruning methods.
In Table 7, we consider another two training-free network pruning methods, SynFlow [48] and GraSP [49], for the channel pruning in our WML, called as WML_SynFlow and WML_GraSP, respectively.
The two pruning approaches are similar with SNIP, with only replace the parameter score for θ i in to ∂
R SF ∂θ i ⊙ θ
i and − ( H ∂ L ∂θ i ) ⊙ θ
i , where H is Hessian matrix and R SF is obtained by setting the input as the all-ones vector 1 and applying an element-wise absolute on all network parameters θ .
Similarly, we also set a upper bound of the pruned ratio for each layer to avoid layer collapse as WML_SNIP.
Table 7 compares the three different training-free pruning approaches in channel pruning for our WML, where we set the pruning ratio as about 40%.
As shown, we find that WML_SynFlow achieves similar results as our WML_SNIP for both ResNet32 and ResNet56.
In contrast, GraSP obtains much worse results, showing that an appropriate channel pruning method is also important to enhance the performance of our WML.
E More Comparison with Knowledge Distillation
In this subsection, we compare our WML with several recent knowledge distillation methods, including FitNet[43], AT[58], and LS-GDFD[11], and more recent self-distillation approaches, including FRSKD[26],
CS-KD[57], TESKD [33],DDGSD[53], BAKE[13], and MUSE[14].
We conduct comparison experiments on CIFAR100 with ResNet18 and ResNet34, respectively.
We report the results from the cited papers and [33].
We can see that our proposed WML outperforms these knowledge distillation methods, especially those recent self-distillation approaches with large margins, again verifying the effectiveness of our weighted mutual learning paradigm.