Learning Efﬁcient Vision Transformers via Fine-Grained Manifold Distillation
Abstract
In the past few years, transformers have achieved promising performances onvarious computer vision tasks.
Unfortunately, the immense inference overhead ofmost existing vision transformers withholds their from being deployed on edgedevices such as cell phones and smart watches.
Knowledge distillation is a widelyused paradigm for compressing cumbersome architectures via transferring information to a compact student.
However, most of them are designed for convolutionalneural networks (CNNs), which do not fully investigate the character of visiontransformer (ViT).
In this paper, we utilize the patch-level information and proposea ﬁne-grained manifold distillation method.
Speciﬁcally, we train a tiny studentmodel to match a pre-trained teacher model in the patch-level manifold space.
Then, we decouple the manifold matching loss into three terms with careful designto further reduce the computational costs for the patch relationship.
Equippedwith the proposed method, a DeiT-Tiny model containing 5M parameters achieves76.5% top-1 accuracy on ImageNet-1k, which is +2.0% higher than previous distillation approaches.
Transfer learning results on other classiﬁcation benchmarks anddownstream vision tasks also demonstrate the superiority of our method over thestate-of-the-art algorithms.
Introduction
The past decade has witnessed the rise of attention-based models in the ﬁeld of natural languageprocessing (NLP) [1, 2].
Such models belonging to the transformer family [3] can effortlessly buildlong-range dependencies and have achieved remarkable performance.
Inspired by the success in NLP,researchers have made great efforts to introduce the transformer-based architectures to vision domainsand achieved promising results.
In an early attempt, Dosovitskiy et al.
[4] proposed a transformerbased vision model termed ViT, which takes split image patches as the input.
Surprisingly, ViTobtains comparable performance when compared to the state-of-the-art convolutional neural networks(CNNs), demonstrating the immense potential of applying transformers to solve vision tasks.
Inspiredby the design of replacing the whole image with patches as input, various vision transformers havebeen proposed, including Swin [5], T2T [6], Twins [7], and TNT [8].
Although many transformers have shown the excellent capability for various vision tasks, they oftenrequire more parameters with large computational burden.
For example, ViT-B [4], which contains86M parameters and is well pre-trained on JFT-300M, can only achieve comparably performancewhen compared to a ImageNet pre-trained CNN-based architecture [9] with 5M parameters.
The requirement for large computing resource and storage prevents them from being deployed onmemory-bounded edge devices such as cell phones and smart watches.
To facilitate the challengesabove, a series of methods have been proposed to investigate compact deep neural networks, such asnetwork pruning [10, 11], low-bit quantization [12] and knowledge distillation [13].
[Figure]
The patch-level manifold space containing ﬁne-grained information facilitates knowledge transfer.
However, smaller models usually lead to performance degradation.
Knowledge distillation (KD)
[13]is a promising approach for inheriting information from a high-performance teacher to a compactstudent and maintaining the strong performance.
Touvron et al.
[14] ﬁrst proposed a KD-based compression approach for vision transformers.
Theytrained the student transformer to match hard labels provided by a pre-trained CNN teacher.
Althoughthis approach obtains satisfying result, they ignore the intermediate-layer information inherited invision transformers.
There have been works [15, 16] proving the effectiveness of learning fromintermediate layers for transformers in NLP, but these methods require teachers and students to haveexactly the same embedding dimension at corresponding layers, which is a fairly tight constraint andcannot always be satisﬁed.
Manifold learning-based KD [17, 18] can support layers with mismatchingdimensions and make use of inter-sample information concurrently.
However, existing manifolddistillation approaches are designed for CNNs and cannot utilize the patch-level information of visiontransformers.
As shown in Figure 1, patches depict a manifold space in a more ﬁne-grained way.
Such information can facilitate knowledge transfer remarkably.
Based on this consideration, we propose a ﬁne-grained patch-level manifold distillation method.
Inparticular, we regard vision transformers as projectors mapping inputs into multiple manifold spaceslayer by layer.
At each layer, we collect embeddings of patches to build their manifold relation mapand train the student to match the relation map of the teacher.
Since the computational complexity ishigh, we decouple the relation maps into three parts, which reduces the complexity by around twoorders of magnitude.
We evaluate the proposed method on the ImageNet-1k image classiﬁcation task.
The proposed manifold KD outperforms the distillation method in [14] by +2.0% top-1 accuracyon DeiT-Tiny.
We also conduct transfer learning experiments on CIFAR-10/100 and evaluation ourmethod on the downstream vision task such as object detection.
The proposed method outperformsits counterparts on both tasks.
Our contributions are summarized as follows:
• We propose a ﬁne-grained manifold distillation method, which transfers patch-level manifoldinformation between vision transformers.•
We use three decoupled terms to describe the manifold space and simplify the computationalcomplexity signiﬁcantly.•
We conduct extensive experiments to verify the effectiveness of the proposed method.
Theresults also demonstrate the importance of soft-label distillation and ﬁxed-depth students.
Related works
Transformer was originally designed for NLP tasks [19].
Inspired by its remarkable performance,researchers have made efforts to adopt transformer-based models in CV tasks [20, 21].
Among them, Dosovitskiy et al.
[4] proposed to divide an image into patches and used embeddings of the patches asmodel input.
According to their input processing scheme, many variants of vision transformers have been proposed.
For example, Han et al.
[8] proposed TNT to model in-patch attention, Zhou et al.[22] and Touvron et al.
[23] built deeper vision transformers, and some researchers [24, 5] adoptedthe experience in CNN to guide the design of vision transformers.
However, Most well-performedvision transformers are extremely resource-consuming and should be compressed for deployment.
KD is a model compression method proposed by Hinton et al.
[13], which trains a lightweight studentmodel to match soft labels given by a large pre-trained teacher model.
Moreover, there are also workscombining adversarial training with KD [25], or training a student with only a few samples [26].
Dueto its excellent performance, KD has been adopted in various research ﬁelds, such as CV [26, 25], NLP [15], and recommendation systems [27].
To compress vision transformers, Touvron et al.
[14]proposed a KD-based method termed DeiT.
They added a distillation token into the student andtrain the student with hard labels provided by the teacher.
Although DeiT has achieved remarkableperformance, KD for vision transformers has not been well explored yet.
Manifold learning is an approach for non-linear dimensionality reduction.
It learns a smooth manifoldembedded in the original feature space to construct low-dimensional features [28, 29].
Recently,several works introduce manifold learning in KD [17, 18].
These methods train the student to preserverelationships among samples learned by the teacher.
For vision transformers, these primary attemptsare coarse and can be further improved because the basic input elements are patches but not images.
Method3.1 Preliminaries
Vision transformer.
Vision transformers are attention-based models, and each layer of themconsists of a multi-head self-attention (MSA) block and a multi-layer perceptron (MLP) block.
Thesemodels take split images as input, i.e. , the patches.
In particular, assuming the patch size is P × P ,an input image x ∈ R H × W × C of size H × W and channel number C is reshaped into ﬂattenedpatches x p ∈ R N ×
( P 2 · C ) for processing, where N = HW/P 2 .
The patches are then projected intoa D -dimension embedding space and added to positional embeddings.
We denote the summation asx e ∈ R N × D .
Each model layer works as follows:
[Equation]
where LN denotes the layer normalization operation.
Knowledge Distillation.
KD is a widely used model compression method, where we use predictions of a large pre-trained teacher model as the learning target of a tiny student model.
Given asample x corresponding to a label y , representing the prediction of the student and the teacher asf s ( x ) and f t ( x ) , respectively, the loss function of KD can be formulated as:
[Equation]
where H CE is the cross-entropy function, H KL is the Kullback-Leibler divergence function, τ is alabel smoothing hyperparameter termed temperature, and λ is a balancing hyperparameter.
Sometimes the intermediate features of the teacher can also be used for knowledge transfer.
Forexample, Romero et al.
[30] trained the student to output features similar to the teacher at intermediatelayers.
However, such methods require the teacher and the student to have the same embeddingdimension.
Otherwise, additional mapping layers are required for aligning, making the distillationprocess non-transparent.
Manifold learning-based KD methods [18, 17] support mismatched embedding dimensions.
They train the student to learn sample relationships predicted by the teacher butignore patch-level information in vision transformers.
Image
Patches
P a t c h
E m b e dd i ng
B l o c k
B l o c k 𝑛
B l o c k 𝑁
C l a ss i f
i e r ⋯ ⋯
Teacher(fixed)
P a t c h
E m b e dd i ng
B l o c k
B l o c k 𝑚 B l o c k 𝑀
C l a ss i f
i e r ⋯ ⋯
Student
KDloss
Teacher Block 𝑛
MHA &
LN
MLP&
LN
MHA &
LN
MLP&
LN
Student Block 𝑚
Matching in Manifold Space
[Figure]
loss
[Figure]
3.2 Fine-grained manifold distillation
To utilize the patch-level information, we propose a ﬁne-grained manifold distillation method.
Unlikeexisting KD methods [14] for vision transformer only distilling with logits, our method distillspatch-level manifolds at intermediate layers.
Figure 2 provides an overview of the proposed method.
In the ﬁne-grained manifold distillation method, we regard a vision transformer as a feature projectorthat embeds image patches into a series of smooth manifold space layer by layer.
At each pair ofmanually selected teacher-student layers, we aim to teach the student layer to output features havingthe same patch-level manifold structure as the teacher layer.
In particular, given samples of batchsize B , we denote the feature of the student layer and the teacher layer as F T ∈ R B × N × D T and
F S ∈ R B × N × D S , respectively, where D T and D S are embedding dimensions.
We ﬁrst normalizethe feature at the last dimension and then compute the manifold structure, or manifold relation map,as follows:
[Equation]
where ψ : R D 1 × D 2 × D 3 → R D 1 D 2 × D 3 is a tensor reshape operation.
The manifold relation map of
F T can be obtained in a similar way.
After that, we train the student to minimize the gap between
M ( ψ ( F T )) and M ( ψ ( F S )) with the following loss:
[Equation]
However, the computation of manifold relation maps is resource-consuming.
Its computationalcomplexity is O
( B 2 N 2 D ) and a memory space of B 2 N 2 size is required to save the result.
Takingthe settings B = 128 , N = 196 and D = 192 in the DeiT-Tiny model [14] as an example, it requiresmore than 240GFLOPs to compute and 2.5GB of memory space to save a single manifold relationmap.
The remarkable resource consumption limits the ﬁne-grained manifold distillation methodscaling up to multiple layers.
Hence we must simplify the computation.
Inspired by the orthogonal decomposition of matrices, we decouple a manifold relation map into threeparts: an intra-image relation map, an inter-image relation map, and a randomly sampled relation map.
Figure 3 illustrates the decoupling.
We compute the intra-image patch-level manifold distillation lossas follows:
[Equation] [Figure]
Similarly, the inter-image patch-level manifold distillation loss is computed by:
[Equation]
Moreover, to relieve the information loss caused by the decoupling, we relate the intra-image and theinter-image manifolds via a relation map computed across randomly sampled patches.
Speciﬁcally,we sample K rows in the reshaped feature ψ ( F ) to obtain F r ∈ R K × D s , and compute the randomsampled patch-level manifold distillation loss as follows:
[Equation]
Hence, the overall loss function of our proposed method is:
[Equation]
where α , β , and γ are hyperparameters.
The summation over l means that the manifold relation mapmatching is performed on multiple pairs of teacher-student layers.
We will discuss the layer selectingscheme in Section 4.3.
Complexity analysis.
The computational complexity of the decoupled manifold relation map is
O ( BN 2 D + B 2 ND + K 2 D ) , and the memory space requirement is BN 2 + NB 2 + K 2 .
Stilltaking DeiT-Tiny as an example, the ﬂoating-point operations and memory space requirement reduceto 3GFLOPs and 32MB, respectively.
The sampling number K is set to 192, the same as ourexperiments.
With the decoupling scheme, ﬁne-grained manifold distillation across multiple layersbecomes feasible.
3.3 Optimization
Patch merging.
Although the decoupling reduces the computation and memory space remarkably,the computing and storing overhead is still unaffordable when the patch size is too small.
For instance,the patch size at the ﬁrst stage of SwinTransformer [5] is 4 , indicating that the total patch number
N is 3136 under the input size of 224 × 224.
Such a large patch number signiﬁcantly increases thecomputational complexity and memory space requirement of the intra-image patch-level manifoldloss L intra .
To remedy this drawback, we merge adjacent patches and view them as a single patch tofurther simplify the computation.
In particular, given a feature map F ∈ R N × D , we ﬁrst reshape itinto F ∈ R H × W × D , where the height H and the width W are obtained following the original patchsplitting operation.
Then we adopt a merging setting ( H (cid:48) , W (cid:48) ) to merge every ( H/H (cid:48) )
× ( W/W (cid:48) )
adjacent patches in a non-overlapping manner.
Zero-padding is adopted when H/H (cid:48) or W/W (cid:48) isnot an integer.
After merging, the feature map becomes F ∈ R H (cid:48)
× W (cid:48) × ( HWD/H (cid:48) W (cid:48) ) and is ﬁnallyreshaped into F ∈ R ( H (cid:48) W (cid:48) )
× ( HWD/H (cid:48) W (cid:48) ) .
By adjusting the merging setting ( H (cid:48) , W (cid:48) ) , we caneasily strike the trade-off between the complexity and the granularity of manifold relation maps. [Table]
Soft distillation.
Previous work [14] adds an additional distillation token into the student, and usesthis token to learn hard labels provided by a CNN teacher.
However, when teacher and student areboth vision transformers, we propose that it is better to teach the student only with soft labels of theteacher.
This design is based on the assumption that and a larger model can learn more knowledgethan a smaller one, and models of the same family share the same knowledge pattern, i.e. , a studentcan learn most knowledge of a teacher.
In our work, unless the student is larger than the teacher, weset the hyperparameter λ in Equation 3 to 1.
Fixed depth.
Stochastic depth [31] is a regularization method that has becomes an infrastructure intraining vision transformers [14, 5].
We propose that the student should not adopt this regularizationwhen the teacher is trained with stochastic depth, since the soft labels already contain knowledgeabout this regularization.
Otherwise, the repeated regularizations may harm the student performance.
Hence, we adopt a ﬁx-depth student in our method, i.e. , the stochastic depth regularization is not usedfor training the student.
Experiments
We evaluate the proposed ﬁne-grained manifold distillation method on ImageNet-1k [33] imageclassiﬁcation task, CIFAR-10/100 [34] transfer learning task, and COCO [35] object detection task.
In the following, we ﬁrst report the experiment setup and then results of these tasks.
4.1 Setup
Datasets.
We evaluate the proposed method mainly on the ImageNet-1k image classiﬁcation task.
ImageNet-1k is a subset of the ImageNet dataset [33], which consists of more than 1.2M trainingimages and 50K validation images from 1000 classes.
To test the generalization performance ofstudent models trained with the proposed method, we conduct transfer learning experiments on
CIFAR-10 and CIFAR-100 datasets [34].
The two datasets both contain 50K training images and10K testing images, which are categorized into 10 classes and 100 classes, respectively.
Moreover,we conduct experiments on the object detection downstream task with COCO dataset [35].
We usethe COCO 2017 split, which consists of 118K training images and 5K validation images containingobjects from 80 categories.
Baselines & models.
We compare the proposed method with DeiT [14], which adds an additionaldistillation token in the student model to learn hard labels from the teacher.
SwinTransformer studentscontaining no distillation token, so we compare with the original KD method [13] on these model.
Table 1 summarizes the used models in our experiments.
4.2 Distillation results on ImageNet-1k
Implementation details.
On the ImageNet-1k image classiﬁcation task, we train DeiT studentswith CaiT teachers and SwinTransformer students with SwinTransformer teachers.
We slightlymodify the architecture of DeiT students by removing the distillation token and only using the classtoken.
The hyperparameter λ in the KD loss is set to 1, i.e. , the real label is not used to train thestudent.
When the teacher is smaller than the student, to prevent the performance degradation causedby the weak teacher, we set λ to 0.5.
In the ﬁne-grained manifold distillation loss, hyperparameters [Table] [Table] α , β , and γ are set to 4, 0.1, and 0.2, respectively.
The sampling number K in loss term L random isset to 192.
We set the above hyperparameters one by one with the grad search method, indicatingthat their combination may not be optimal.
We select the ﬁrst 4 layers and the last 4 layers of thestudent and the teacher to conduct manifold distillation.
Note that the class token in DeiT is ignoredwhen computing manifold relation maps.
Moreover, to train SwinTransformer students efﬁciently,we adopt a patch merging setting of (14 , 14) .
Other training settings follow those in DeiT [14] and
SwinTransformer [5], except the stochastic depth regularization, which is not used in our experiments.
Each student is trained for 300 epochs with 8 Tesla-V100 GPUs.
Results.
Table 2 presents classiﬁcation results on ImageNet-1k.
Compared with DeiT, the proposedmethod achieves remarkably performance improvements, which are +2 .
0% on DeiT-Tiny and +0 .
9%on DeiT-Small with the CaiT-S24 teacher.
When the teacher is the much weaker CaiT-XXS24, theimprovement becomes +1 .
6% on DeiT-Tiny and +1 .
2% on DeiT-Small.
Note that we report the RegNetY-16GF teacher results in DeiT for comparison, but do not conduct ﬁne-grained manifolddistillation experiments with this CNN model because the proposed method is designed for distillationbetween vision transformers.
Moreover, we conduct experiments on SwinTransformer, one of the state-of-the-art vision transformerarchitecture.
When the teacher is Swin-Small and the student is Swin-Tiny, the proposed methodachieves +0 .
5% performance improvement compared with the original KD method.
4.3 Ablation & parameter comparison
Ablation of main components.
We design experiments to verify the effectiveness of the proposedﬁne-grained manifold distillation method, the soft distillation scheme, and the ﬁxed student depth. [Table] [Table]
[Table] Table 3 presents the results. The results demonstrates that for distillation between vision transformers,soft distillation and ﬁxed depth are both signiﬁcant.
Moreover, combined with the ﬁne-grainedmanifold distillation method, student performance can be further improved.
Ablation of decoupled manifold loss terms.
The decoupled ﬁne-grained manifold distillation lossconsists of three terms.
We study their effectiveness and report the result in Table 6.
The results showthat each component in the decoupled ﬁne-grained manifold distillation loss contributes to improvingthe ﬁnal performance.
Layers for ﬁne-grained manifold distillation.
To study the impact of different layer selections inmanifold relation map matching, we evaluate 5 layer selecting schemes and compare their performance.
In particular, we take CaiT-S24 as the teacher model and DeiT-Tiny as the student model.
Thenumber of selected layers is set to 6.
From the results reported in Table 4, conducting ﬁne-grainedmanifold distillation at the head and the tail of student models are both crucial.
Hence, we speculatethat the “Shallow/Deep” layer selecting scheme is the best because of its outstanding performanceand ease of use.
Parameters in decoupled manifold distillation loss.
There are 4 hyperparameters in decoupledmanifold distillation loss: the weight of inter-image loss α , the weight of intra-image loss β , theweight of randomly sampled loss γ , and the sampling number K .
We compare different settingsof these parameters and report the results in Table 5.
Our default setting outperforms those eitherincreasing or reducing one of the parameters.
However, due to our coarse parameter searchingscheme, we believe that the performance can be further improved by setting these hyperparametersmore carefully. [Table] [Table]
4.4 Transfer learning
To measure the generalization ability of the proposed method [Table] we conduct transfer learning experiments.
We ﬁne-tune a DeiT-Tiny student trained with a CaiT-S24 teacher on CIFAR-10 and
CIFAR-100 datasets for 1000 epochs.
We adopt a batch size of 768 and an AdamW optimizer with alearning rate of 5 × 10 − 6 .
Other settings follow DeiT [14]
[Table]
Table 7 reports the transfer learning results.
Students trained with the ﬁne-grained manifold distillation method generalize better than others (+0.25% on CIFAR-10 and +0.71% on CIFAR-100),demonstrating a favorable generalization ability of the proposed method.
4.5 Downstream task
To further evaluate the effectiveness of our proposed method, we adopt the ﬁne-grained manifolddistillation to train object detection models on COCO 2017 dataset.
Implementation details.
We adopt Mask R-CNN as the detection framework.
The teacher backboneis Swin-Small and the student backbone is Swin-Tiny.
All used backbones are pretrined on ImageNet1k The teacher is trained for 36 epochs and each student is trained for 12 epochs.
When trainingstudents with ﬁne-grained manifold distillation, we adopt the distillation loss on outputs of the lasttwo backbone stages and outputs of the feature pyramid network neck [36].
Results.
Table 8 presents the detection results.
Our manifold distilled student outperforms the studenttrained without distillation (+1.0% box AP), demonstrating that the proposed method beneﬁts theobject detection downstream task.
Conclusion
In this paper, we propose a ﬁne-grained manifold distillation method for vision transformers.
Wematch patch-level intermediate features of the student and the teacher in a manifold space, anddecouple the matching loss into three terms to reduce the computational complexity.
Moreover,we adopt a patch merging scheme to further simplify the computation.
Different from previousworks, we distill the student with soft labels and ﬁxed depth.
We conduct experiments on ImageNet1k image classiﬁcation, CIFAR-10/100 transfer learning, and COCO object detection, and theproposed method outperforms existing methods.
We also conduct substantial ablation experiments todemonstrate the superiority of the ﬁne-grained manifold distillation method.
The large search spaceof hyperparameters is one of the most serious drawbacks of the proposed method.
In the future, wewill study to further simplify the proposed method and help it be parameter insensitive.
Checklist
For all authors...
(a) Do the main claims made in the abstract and introduction accurately reﬂect the paper’s contributions 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?
[No]
(d) Have you read the ethics review guidelines and ensured that your paper conforms to
them?
[Yes]
If you are including theoretical results...
(a) Did you state the full set of assumptions of all theoretical results?
[N/A]
(b) Did you include complete proofs of all theoretical results?
[N/A]
If you ran experiments...
(a) Did you include the code, data, and instructions needed to reproduce the main experi-
mental results (either in the supplemental material or as a URL)?
[No]
(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)?
[No]
(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]
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?
[No]
(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
using/curating?
[No]
(e) Did you discuss whether the data you are using/curating contains personally identiﬁable
information or offensive content?
[No]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) Did you describe any potential participant risks, with links to Institutional Review
Board (IRB) approvals, if applicable?
[N/A]
(c) Did you include the estimated hourly wage paid to participants and the total amount
spent on participant compensation?
[N/A]