LIT: Block-wise Intermediate Representation Training for Model Compression

1 INTRODUCTION . Modern deep networks have achieved increased accuracy by continuing to introduce more layers ( Ioffe & Szegedy , 2015 ; He et al. , 2016 ) at the cost of higher computational overhead . In response , researchers have proposed many techniques to reduce this computational overhead at inference time , which broadly fall under two categories . First , in deep compression ( Han et al. , 2015a ; Zhu et al. , 2016 ; Li et al. , 2016 ; Hubara et al. , 2017 ) , parts of a model are removed or quantized to reduce the number of weights and/or the computational footprint.1 However , deep compression techniques typically require new hardware ( Han et al. , 2016 ) to take advantage of the resulting model sparsity . Second , in student/teacher methods—introduced in knowledge distillation ( KD ) ( Hinton et al. , 2014 ) and further extended ( Romero et al. , 2015 ; Kim & Rush , 2016 ; Furlanello et al. , 2018 ) —a smaller student model learns from a large teacher model through distillation loss , wherein the student model attempts to match the logits of the teacher model . As there are no constraints on the teacher and student models , KD can produce hardware-friendly models : the student can be a standard model architecture ( e.g. , ResNet ) , optimized for a given hardware substrate . Hint training ( i.e. , FitNets ( Romero et al. , 2015 ) ) extends KD by using a teacher ’ s intermediate representation ( IR , i.e. , the output from a hidden layer ) to guide the training of the student model . The authors show that hint training with a single IR outperforms KD in compressing teacher networks ( e.g. , maxout networks ( Goodfellow et al. , 2013 ) ) to thinner and deeper student networks . We ask the natural question : does hint training compress more modern , highly-structured , very deep networks—such as ResNet ( He et al. , 2016 ) , ResNeXt ( Xie et al. , 2017 ) , VDCNN ( Conneau et al. , 2016 ) , and StarGAN ( Choi et al. , 2017 ) ? We find that standard hint training ( i.e. , with a single hint ) and training with multiple hints is not effective for modern deep networks ( Section 4.2 ) . We hypothesize that , for 1In this work , we refer to this class of methods as “ deep compression , ” and methods to reduce model size more generally as “ model compression. ” modern deep networks , hint training causes unstable IRs : the deepest network considered in Romero et al . ( 2015 ) was only 17 layers , achieving 91.61 % on CIFAR10 ; in contrast , a modern 110-layer ResNet achieves 93.68 % on CIFAR10 . In this work , we extend hint training ’ s ability to transfer intermediate knowledge from teacher to student to reduce the depth of modern , highly-structured architectures ( e.g. , compressing a standard ResNeXt-110 to a standard ResNeXt-20 with no loss in accuracy ) . We do this via a novel method called bLock-wise Intermediate representation Training ( LIT ) , a student/teacher compression technique that outperforms training student networks from scratch , hint training , and KD . LIT targets highly structured , modern networks that consist of repetitive blocks ( i.e. , groups of layers ) that can be scaled up/down for accuracy/speed tradeoffs ; for example , ResNets have standard configurations from 20 to hundreds of layers . LIT leverages two key ideas to reduce unstable IRs in deep networks . First , LIT directly trains student networks of the same width as the teacher model ( as opposed to using a single , thinner hint as in hint training ) . Second , LIT avoids unstable student IRs deep in the network by using the IR from the previous block in the teacher model as input to the current student block during training ; each student block is effectively trained in isolation to match the corresponding ( deeper ) block in the teacher . We show that LIT ’ s block-wise training improves accuracy , allows for copying parts of the teacher model directly to the student model , and permits selective compression of networks ( e.g. , compressing one out of three blocks in a network and copying the rest ) . For example , consider compressing a ResNet-56 from a ResNet-110 ( Figure 1 ) , each of which have four stages . The IR loss is applied to the output of each block , and the teacher model ’ s IRs are used as input to the student blocks . We show that LIT outperforms standard KD on a range of models ( ResNet , ResNeXt , VDCNN , StarGAN ) and datasets ( CIFAR10 , CIFAR100 , Amazon Reviews , CelebA ) : empirically , LIT can reduce model sizes from 1.7× to 5.5× with no loss in accuracy . Recent work on Born Again networks ( Furlanello et al. , 2018 ) uses standard KD to train identical student and teacher models to higher accuracies ( i.e. , no compression ) . We show that the benefits of this procedure also apply to LIT student/teacher training , and LIT enables up to 0.64 % higher accuracy than KD-based Born Again networks on the networks we consider . 2 RELATED WORK . Knowledge distillation . Hinton et al . ( 2014 ) ; Bucilu et al . ( 2006 ) introduced knowledge distillation ( KD ) in which a teacher ensemble or models outputs are used to train a smaller student model , which inspired a variety of related methods , e.g. , for cross-modal distillation or faster training ( Gupta et al. , 2016 ; Chen et al. , 2015 ; Frosst & Hinton , 2017 ; Romero et al. , 2015 ; Furlanello et al. , 2018 ) . FitNets extends KD by regressing a student model ’ s IR to a teacher model ’ s IR , as the student models they consider are thinner and deeper . Wang et al . ( 2018 ) extends FitNets by training networks iteratively using hints . In contrast , LIT uses the teacher IRs as input to the student model in training and directly penalizes deviations of the student model ’ s IRs from teacher model ’ s IRs , which helps guide training for higher accuracy and improved inference performance . In Born Again networks ( Furlanello et al. , 2018 ) , the same network architecture is used as both the teacher and student in standard KD , resulting in higher accuracy . We show that LIT outperforms the Born Again procedure on ResNet , ResNeXt , and VDCNN . Deep compression . In deep compression , parts of a network ( weights , groups of weights , kernels , or filters ( Mao et al. , 2017 ) ) are removed for efficient inference ( Han et al. , 2015a ) , and the weights of the network are quantized , hashed , or compressed ( Hubara et al. , 2016 ; Rastegari et al. , 2016 ; Zhu et al. , 2016 ; Hubara et al. , 2017 ) . These methods largely do not take advantage of a teacher model and typically require new hardware for efficiency gains ( Han et al. , 2016 ) . Methods that prune filters ( Li et al. , 2016 ) can result in speedups on existing hardware , but largely degrade accuracy . We show that LIT models can be pruned , and thus these methods are complementary to LIT . Additionally , deep compression does not perform as well on modern networks : Han et al . ( 2015a ) compressed VGG by ~10× , but ResNet-110 can only be compressed ~1.6× ( Li et al. , 2016 ) , compared to LIT ’ s 4× compression . Network architectures for fast inference . Researchers have proposed network architectures ( e.g. , MobileNet ( Howard et al. , 2017 ) ) and new operations for fast inference ( e.g. , ShuffleNet ( Zhang et al. , 2017 ) ) on specific hardware . However , these architectures and operations are largely designed for power/resourceconstrained mobile devices and sacrifice accuracy for low power . We focus on highly accurate , very deep networks in this work . 3 METHODS . LIT uses an augmented loss function and training procedure to distill a teacher model into a student model . In its training procedure , LIT both 1 ) penalizes deviations of the student model ’ s IRs from the teacher model ’ s IRs ( IR loss ) and 2 ) uses the KD loss ( for the entire student network ) . As LIT directly penalizes deviations in IRs , LIT requires that the teacher model and student model have outputs of the same size at some intermediate layer . A key challenge in the LIT procedure is that the student network will not have meaningful IRs for a large part of the training ( e.g. , at the start of training when the weights are initialized randomly ) . To address this issue , LIT uses the teacher model ’ s IRs as inputs to the student model ( described below ) . We describe the overall LIT procedure , describe the KD loss , describe how IRs are used in LIT , and discuss hyperparameter optimization . LIT . In LIT , we combine the KD and IR loss . We show that combining the losses results in smaller models for a fixed accuracy in Section 4 . Specifically , for teacher T and student S the full LIT loss is : β·LKD , α ( T , S ) + ( 1−β ) ·LI ( T , S ) ( 1 ) with α , β∈ [ 0,1 ] ( α is described below , β is an interpolation parameter ) . In some cases , we use β=0 , i.e. , we only use the IR loss ( e.g. , for GANs ) . As the IRs have matching dimensions , LIT also allows parts of the teacher model to be copied directly into the student model . For example , for ResNets , we copy the teacher ’ s first convolution ( before the skip connections ) and fully connected layer to the student model . LIT can also be used to compress specific parts of a model , as we do with StarGAN ’ s generator ( Choi et al. , 2017 ) . Finally , after we train the student model with combined loss , we fine-tune the student model with the original loss ( KD loss for classification , generator loss for GANs ) . Knowledge distillation loss . In KD , a ( typically larger ) teacher model or ensemble is used to train a ( typically smaller ) student model . Specifically , the KL-divergence between the probabilities of the student and teacher model is minimized , in addition to the standard cross-entropy loss . Formally , denote ( for the teacher model ) qτi = exp ( zi/τ ) ∑ jzj/τ where zi are the inputs to the softmax and τ is a hy- perparmeter that “ softens ” the distribution . Denote pτi to be the corresponding quantity for the student model . Then , the full KD loss is : LKD ( pτ , qτ , y ) =α·H ( y , pτ ) + ( 1−α ) ·H ( pτ , qτ ) ( 2 ) for y to be the true labels , H to be the cross-entropy loss , and α to be the interpolation parameter . Hinton et al . ( 2014 ) sets α=0.5 , but we show that the choice of α can affect performance ( Section 4.4 ) . Training via intermediate representations . In LIT , we logically divide the student and teacher networks into k sub-networks such that the input and output dimensions match for the corresponding sub-networks ( an example is shown in Figure 1 ) . Denote the full teacher network and student network to be T and S respectively . Denote the teacher sub-networks to be Ti and the student sub-networks to be Si such that Ti , Si : Rni−i→Rni and that the composition of the sub-networks is the full network , namely that Tk ( Tk−1 ( ···T1 ( x ) ) ) =T ( x ) . We will omit the argument when convenient . Denote the loss on the IR loss l ( e.g. , L2 loss ) . The full intermediate loss ( given the set of splits ) is : LI ( T , S ) : =l ( S1 , T1 ) + ∑k i=2l ( Si ( Ti−1 ) , Ti ) ( 3 ) Concretely , consider a ResNet-110 as the teacher and a ResNet-56 as the student , each with three “ stages ” , i.e. , layers in the network with downsampling , and an L2 intermediate loss . Here , the first teacher ResNet “ stage ” is T1 , etc . and the L2 deviation from the feature maps , across all the downsampling feature maps , is the full intermediate loss . A schematic is shown in Figure 1 . This procedure has two key decisions : 1 ) where to logically split the teacher and student models , and 2 ) the choice of IR loss . We discuss these settings in the hyperparameter optimization below . Hyperparameter optimization . LIT inherits two hyperparameters from KD and introduces one more : τ ( the temperature in KD ) , α ( the interpolation parameter in KD ) , and β ( the interpolation parameter in LIT ) , along with an intermediate representation loss and split points . In this work , we only consider adding the IR loss between natural split points , e.g. , when a downsampling occurs in a convolutional network . We have additionally found that L2 loss works well in practice , so we use the L2 loss for all experiments unless otherwise noted ( Section 4.4 ) . We have found that iteratively setting τ , then α , then β to work well in practice . We have found that the same hyperparameters work well for a given student and teacher structure ( e.g. , ResNet teacher and ResNet student ) . Thus , we use the same set of hyperparameters for a given student and teacher structure ( e.g. , we use the same hyperparameters for a teacher/student of ResNet-110/ResNet-20 and ResNet-110/ResNet-32 ) . To set the hyperparameters for a given structure , we first set τ using a small student model , then α for the fixed τ , then β for the fixed α and τ ( all on the validation set ) .

Summary: This paper introduces LIT, a network compression framework, which uses multiple intermediate representations from a teacher network to guide the training of a student network. Experiments are designed such that student networks are shallower than teacher networks, while maintaining their width. The method is validated on CIFAR-10 and 100 as well as on Amazon Reviews.