Block-wise Intermediate Representation Training for Model Compression
Abstract
Knowledge distillation (KD) is a popular method for reducing the computational overhead of deep network inference, in which the output of a teacher model is used to train a smaller, faster student model.
Hint training (i.e., FitNets) extends KD by regressing a student model’s intermediate representation (IR) to a teacher model’s IR.
In this work, we introduce b L ock-wise
I ntermediate representation T raining (LIT), a novel model compression technique that extends the use of IRs in deep network compression, outperforming KD and hint training.
LIT has two key ideas: 1) LIT trains a student of the same width (but shallower depth) as the teacher by directly comparing the IRs, and 2) LIT uses the IR from the previous block in the teacher model as an input to the current student block during training, avoiding unstable IRs in the student network.
We show that LIT provides substantial reductions in network depth without loss in accuracy — for example, LIT can compress a ResNeXt-110 to a ResNeXt- (5.5 × ) on CIFAR10 and a VDCNN-29 to a VDCNN-9 (3.2 × ) on Amazon Reviews, outperforming KD and hint training in network size for a given accuracy.
Finally, we show that LIT can effectively compress GAN generators, which are not supported in the KD framework because GANs output pixels as opposed to probabilities.
Introduction
Modern deep networks have achieved increased accuracy by continuing to introduce more layers [1, 2] 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 [3, 4, 5, 6], parts of a model are removed or quantized to reduce the number of weights and/or the computational footprint.
However, deep compression techniques typically require new hardware [7] to take advantage of the resulting model sparsity.
Second, in student/teacher methods—introduced in knowledge distillation (KD) [8] and further extended [9, 10, 11]—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 [9]) 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 [12]) 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 [2], VDCNN [13], and StarGAN [14]?
We ﬁnd that standard hint training (i.e., with a single hint) and training with multiple hints is not effective for modern deep networks.
[Figure]
We hypothesize that, for modern deep networks, hint training causes unstable IRs: the deepest network considered in [9] was only 17 layers, achieving 91.61% on CIFAR10; in contrast, a modern 110-layer ResNet achieves 93.68% on CIFAR10.
Further experiments validating this hypothesis are given in [15]
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 b L ock-wise I ntermediate representation T raining (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 trade-offs; for example, ResNets have standard conﬁgurations fromto 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 sections.
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.
Because it is possible to transfer IRs directly, LIT is, to our knowledge, the ﬁrst student/teacher compression method that works for GAN generators.
LIT can compress GAN generators by only compressing the repetitive blocks present in certain GANs [14].
In contrast, KD does not apply directly as the KL divergence in KD loss operates on probabilities but not the pixels output by GAN generators.
LIT can compress GANs by leveraging LIT’s key property that, by matching the teacher IR dimensions, parts of the teacher network can be directly copied to the student network.
Experiments
We evaluate LIT’s efﬁcacy at compressing models on a range of tasks and models, including image classiﬁcation, sentiment analysis, and image-to-image translation (GAN).
Throughout, we use student and teacher networks with the same broad architecture (e.g., ResNet to ResNet).
We consider
ResNet [2], ResNeXt [16], VDCNN [13], and StarGAN [14].
We use standard architecture depths, widths, and learning rate schedules; details are given in [15].
LIT is effective at compressing a range of datasets and models.
We ran LIT on a variety of models and datasets for image classiﬁcation and sentiment analysis (CIFAR10, CIFAR100, Amazon Reviews).
We additionally performed KD and hint training on these datasets and models.
Figure 2 shows the results for ResNet and ResNeXt for CIFAR10 and CIFAR100, and VDCNN on
Amazon Reviews (full, polarity).
LIT can compress models by up to 5.5 × (CIFAR10, ResNeXtto 20) on image classiﬁcation and up to 3.2 × on sentiment analysis (Amazon Reviews, VDCNNto 9), with no loss in accuracy.
LIT outperforms KD and hint training on all settings.
We have found that, in some cases, training sequences of models using LIT results in higher performance.
Thus, for VDCNN, we additionally compressed using LIT a VDCNN-29 to a VDCNN-17, and using this [Figure]
[Figure] [Figure] [Figure]
[Figure] cc u
[Figure]
a c y
[Figure]
[Figure] trained from scratch (two residual blocks) StarGANs. As shown (column two, four), LIT can appear to improve GAN performance while signiﬁcantly compressing models. Best viewed in color.
[Table]
VDCNN-17, we trained a VDCNN-9 and VDCNN-17.
We also found that in some cases, KD degrades the accuracy of student models when the teacher model is the same architecture (ResNeXt-110 on
CIFAR100, VDCNN-29 on Amazon Reviews polarity).
This corroborates prior observations in [17].
LIT can compress GANs.
We compressed StarGAN’s generator [14] using the LIT procedure with β =0 (i.e., only using the intermediate representation loss).
The original StarGAN has 18 total convolutional layers (including transposed convolutional layers), with 12 of the layers in the residual blocks (for a total of six residual blocks).
We compressed the six residual blocks to two residual blocks (i.e., 12 to four layers) while keeping the rest of the layers ﬁxed.
The remaining layers for the teacher model were copied to the student model and ﬁne-tuned.
The discriminator remained ﬁxed.
As shown in Table 1, LIT outperforms all baselines in inception and FID score.
Additionally, as shown in Figure 3, the student model appears to perceptually outperform both the teacher model and equivalent model trained from scratch, suggesting LIT can serve as a form of regularization. Acknowledgments
This research was supported in part by afﬁliate members and other supporters of the Stanford DAWN project—Ant Financial, Facebook, Google, Intel, Microsoft, NEC, Teradata, SAP, and VMware—as well as DARPA under No. FA8750-17-2-0095 (D3M), industrial gifts and support from Toyota
Research Institute, Keysight Technologies, Hitachi, Northrop Grumman, NetApp, and the NSF under grants DGE-1656518 and CNS-1651570.