Continuation KD : Improved Knowledge Distillation through the Lens of Continuation Optimization (2022)

Abstract: Knowledge Distillation (KD) has been extensively used for natural language understanding (NLU) tasks to improve a small model’s (a student) generalization by transferring the knowledge from a larger model (a teacher). Although KD methods achieve state-of-the-art performance in numerous settings, they suffer from several problems limiting their performance. It is shown in the literature that the capacity gap between the teacher and the student networks can make KD ineffective. Additionally, existing KD techniques do not mitigate the noise in the teacher’s output: modeling the noisy behaviour of the teacher can distract the student from learning more useful features. We propose a new KD method that addresses these problems and facilitates the training compared to previous techniques. Inspired by continuation optimization, we design a training procedure that optimizes the highly non-convex KD objective by starting with the smoothed version of this objective and making it more complex as the training proceeds. Our method (Continuation-KD) achieves state-of-the-art performance across various compact architectures on NLU (GLUE benchmark) and computer vision tasks (CIFAR-10 and CIFAR-100).

## introduction deep neural networks have achieved great success in many challenging tasks including the ones in natural language processing @xcite and computer vision @xcite . most successful neural networks are usually large and overparametrised @xcite . the big size of these models limits them from deploying on computers with low computational power such as edge devices. compressing the models can empower us to provide a variety of machine learning-based services offline on low-resource machines. this issue is even more severe in models designed for nlp. although, after the introduction of the transformers @xcite the performance of models in this field improved dramatically, the size of the models increased exponentially. nowadays, some of these models have more than 100 billion parameters @xcite , and they are still increasing. one way to address the expensive computational complexity of deep networks and their overparameterization is neural model compression @xcite . among all of the compression methods, knowledge distillation (kd) @xcite is one of the prominent techniques which have been used to compress a variety of models in different deep learning applications such as natural language processing @xcite @xcite , speech processing @xcite , and computer vision @xcite . in kd, we have a large accurate network, which is called a teacher, and a small network, which is called a student, that we desire to train. the innate capacity gap between the student and the teacher was speculated @xcite to impede the training and was addressed by multiple works @xcite . however, the existing solutions have several complications: for example, @xcite requires an extra intermediate network to be trained, and @xcite has a rigid twostage structure that requires the careful design for successful training. moreover, none of these techniques is robust against the noisy data and noisy teacher's outputs, which can distract the student with the limited capacity from learning more useful features and make it overfit to the noise. we propose a new solution, continuation kd, inspired by the continuation method from optimization and nonlinear equations. during the training, we gradually move from the smoothed objective function, which is robust to overfitting to the noise, to the original highly non-convex function. we conduct extensive experiments on the glue benchmark for distilroberta (6-layer) @xcite and bert-small (4-layers) @xcite core models and show a significant improvement over the previous baselines. besides that, we demonstrate that in the computer vision setting continuation kd also outperforms its competitors. overall, our contributions are the following: ## methodology in this section, we describe our continuation-kd technique, which addresses both the capacity gap problem and the lack of robustness against the noise in the teacher's output. continuation-kd uses a loss function with two objectives (eq. 3). the first objective l ce is a cross-entropy loss term that trains the student based on the given hard labels. the second objective l cn t kd is our proposed annealed hinge loss function that gradually trains the student to mimic the behaviour of the teacher. inspired by the continuation method @xcite , continuation-kd starts with an easy objective to train at the beginning of training. as the training proceeds, the whole objective function becomes more and more complex. formally, the loss function of the continuation-kd is defined as follows: note that in eq. 4, ∥z sϕ(t i )z t ∥ 2 2 is a mean square loss between the student's logits and the annealed version of the teacher's logits. also, max{0, ∥z sϕ(t i )z t ∥ 2 2ϕ(t i )m} is a hinge loss with an annealed margin ϕ(t i )m. this loss function avoids penalizing negligible differences (the ones less than m ϕ(t i )) between the outputs of the student and teacher. this feature helps the student to learn a meaningful behaviour of the teacher rather than focusing on higher frequency fluctuations. at the beginning of training, we set t 1 = t max which leads to the most softened version of the teacher's output (ϕ(t 1 ) = 1 tmax ). since we have ψ(1) = 0, the student only learns the behaviour of the teacher's smoothest version, which is an easy target to learn. then, during training, we decrease the temperature, at this phase, functions ψ(i) and ϕ(t i ) are both increasing which in turn leads to increasing the sharpness of the teacher and smoothly shifting from the hinge loss to the cross-entropy loss. both of these operations increase the complexity of the whole loss function. note that at the function ϕ(t i ) also anneals margin m. its reason is that smoothing the teacher with ϕ(t i ) damps its noise as well. therefore we damp the margin m with ϕ(t i ) to apply a margin proportional to the amount of noise in the smoothed version of the teacher. figure 1 visualize different components of continuation-kd. also note that, if we set @xmath0, then continuation-kd becomes identical to annealing-kd, where k is the number of epochs in the first stage and nk is the number of epochs in the second stage. this fact shows that the annealing-kd is actually a special case of our continuation-kd. algorithm 1 demonstrates the details of continuation-kd. it requires a student s, a teacher t , a dataset d, max temperature t max , number of epochs n, an increasing function ψ, and margin m as inputs and returns the trained student at the output. at the beginning, it sets variables @xmath1. also, variable k indicates the number of epochs before updating t during training. φ and ψ are the output values of ϕ(i) and ψ(i). function get-mini-batch(d) retrieves a minibatch (x, y ) from the dataset d. then these data 5263 samples feed into loss functions in the next lines to get the outputs of l ce and l cn t kd . then the linear combination of line 14 combines these two losses to get the continuation loss l. finally, l is fed into optimization-back-propagation(.) function to optimize the student network based on the back propagation of the gradient of this loss function and update the weights of the student. this part of the training is identical to regular training of the neural networks. save-best-checkpoint(.) function checks the performance of the current student model on a validation dataset. if it is better than the previous checkpoints, it saves the checkpoint. in the end, we load the best checkpoint and return it. in the next section, we report the experimental results of continuation-kd method. ## experiments this section demonstrates our evaluation results comparing continuation-kd with other baselines for natural language processing and computer vision tasks. we compare our method with stateof-the-art techniques such as annealing-kd (jafari et al., 2021), takd @xcite and other baselines like vanilla-kd @xcite and training the student only with hard labels. in the following sub-sections, we will discuss each in more detail. ## analysis to investigate how continuation-kd works, we did two ablation studies explaining different aspects of this technique: effects of the dynamically changing factors and noise mitigation. effects of dynamic factors ϕ and ψ. in the first ablation, we scrutinize the effect of each dynamically changing component of the continuation-kd loss function. it basically considers the effect of functions: ψ(i) from eq. 3, ϕ(t i ) when it anneals the outputs of the teacher in eq. 4, and ϕ(t i ) when it anneals the margin of the hinge loss in eq. 4. to investigate the effects of these components, we repeat the nlu experiments with distilroberta model described in section 5.3 on mrpc and rte datasets three times. in each trial, we fixed two of the components, and we only let one of them dynamically change during training. as shown in ## conclusion in this work, we present continuation-kd, a novel kd method inspired by continuation optimization. our continuation-kd technique starts optimizing a smoothed version of the objective function and gradually increases the complexity of the loss towards the original, highly non-convex one. we demonstrated that our method alleviates the capacity gap problem: an innate problem of kd resulting from the different capacities of a student's and a teacher's networks which detriments the performance. besides that, we show that the method can lessen the student's overfitting to noise in the teacher's output. our technique is stable because it doesn't require two stages in the training and is efficient. it outperforms its competitor kd methods for different backbone models' architectures in both computer vision and nlp. for this investigation, we proposed to implement continuation method by smoothing the objective with the hinge loss. however, it can potentially be done differently. investigating other realizations of continuation optimization for improving small models' performance is an interesting next step.