Learning Efficient Vision Transformers via Fine-Grained Manifold Distillation
Abstract
In the past few years, transformers have achieved promising performance on various computer vision tasks.
Unfortunately, the immense inference overhead of most existing vision transformers withholds them from being deployed on edge devices such as cell phones and smart watches.
Knowledge distillation is a widely used paradigm for compressing cumbersome architectures into compact students via transferring information.
However, most of them are designed for convolutional neural networks (CNNs), which do not fully investigate the character of vision transformers.
In this paper, we fully utilize the patch-level information and propose a fine-grained manifold distillation method for transformer-based networks.
Specifically, we train a tiny student model 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 design to further reduce the computational costs for the patch relationship.
Equipped with the proposed method, a DeiT-Tiny model containing 5M parameters achieves 76.5% top-1 accuracy on ImageNet-1k, which is +2.0% higher than previous distillation approaches.
Transfer learning results on other classification benchmarks and downstream vision tasks also demonstrate the superiority of our method over the state-of-the-art algorithms.
Introduction
The past decade has witnessed the rise of attention-based models in the field of natural language processing (NLP) [1, 2].
Such models belonging to the transformer family [3] can effortlessly build long-range dependencies and have achieved remarkable performance.
Inspired by the success in
NLP, researchers have made great efforts introducing the transformer-based architectures to vision domain and achieved promising results.
In an early attempt, Dosovitskiy et al.
[4] proposed a transformer-based vision model termed ViT, which takes split image patches as the input.
ViT obtains comparable performance when compared to the state-of-the-art convolutional neural networks (CNNs), demonstrating the immense potential of applying transformers to vision tasks.
Inspired by the design of splitting the whole image into patches as input, various vision transformers have been proposed, including Swin [5], T2T [6], Twins [7], and TNT [8].
Although transformers have shown the excellent capability for various vision tasks, they require large amounts of parameters, resulting in heavy computational burden.
For example, ViT-B [4] with 86M parameters pretrained on JFT-300M can only achieve comparable performance with CNN-based
EfficientNet [9], while the latter is trained on ImageNet and contains only 5M parameters.
The requirement for inordinate computing resource and storage prevents them from being deployed on memory-bounded edge devices such as cell phones and smart watches.
To facilitate the challenges above, a series of methods have been proposed to investigate compact deep neural networks, such as network pruning [10, 11], low-bit quantization [12, 13] and knowledge distillation [14, 15].
[Figure]
The patch-level manifold space containing fine-grained information facilitates knowledge transfer.
However, smaller models usually lead to performance degradation.
Knowledge distillation (KD)
[14] is a promising approach for inheriting information from a high-performance teacher to a compact student and maintaining the strong performance.
Touvron et al.
[16] first proposed a KD-based compression approach for vision transformers.
They trained the student transformer to match hard labels provided by a pre-trained CNN teacher.
Although this approach obtains satisfying results, they ignore the intermediate-layer information inherited in vision transformers.
There have been works [17, 18] proving the effectiveness of learning from intermediate layers for transformers in NLP, but these methods require teachers and students to have exactly the same embedding dimension at corresponding layers, which is a fairly tight constraint and cannot always be satisfied.
Manifold learning-based KD [19, 20] can support layers with mismatching dimensions and make use of inter-sample information concurrently.
However, existing manifold distillation approaches are designed for CNNs and cannot utilize the patch-level information of vision transformers.
As shown in Figure 1, patches depict a manifold space in a more fine-grained way.
Such information can facilitate knowledge transfer remarkably.
Based on this consideration, we propose a fine-grained manifold distillation method in both patchlevel and batch-level.
In particular, we regard vision transformers as projectors mapping inputs into multiple manifold spaces layer by layer.
At each layer, we collect embeddings of patches to build their manifold relation map and train the student to match the relation map of the teacher.
Since the computational complexity is high, we decouple the relation maps into three parts, which reduces the complexity by two orders of magnitude approximately.
We evaluate the proposed method on the
ImageNet-1k classification task.
The proposed manifold KD outperforms the distillation method in [16] by +2.0% top-1 accuracy on DeiT-Tiny.
We also conduct transfer learning experiments on
CIFAR-10/100 and evaluate our method on downstream tasks such as object detection and semantic segmentation.
The proposed method outperforms its counterparts on both tasks.
Our contributions are summarized as follows:
• We propose a fine-grained manifold distillation method, which transfers patch-level and batch-level manifold information between vision transformers.
• We use three decoupled terms to describe the manifold space and simplify the computational complexity significantly.
• We conduct extensive experiments to verify the effectiveness of the proposed method.
The results also demonstrate the importance of soft-label distillation and fixed-depth students.
Related works
Vision transformer.
Transformer is originally designed for NLP tasks [21].
Inspired by its remarkable performance, researchers have made efforts to adopt transformer-based models in CV tasks [22, 23].
Among them, Dosovitskiy et al.
[4] proposed to divide an image into patches and used embeddings of the patches as model input.
Based on their proposed input processing scheme for images, many variants of vision transformers have been proposed.
For example, Han et al.
[8] proposed TNT to model in-patch attention, Zhou et al.
[24] and Touvron et al.
[25] built deeper vision transformers, and some researchers [26, 5, 27] adopted the experience in CNN to guide the design of vision transformers.
However, Most well-performed vision transformers are extremely resource-consuming and should be compressed for deployment.
Knowledge distillation.
KD is a model compression method proposed by Hinton et al.
[14], which trains a lightweight student model to match soft labels given by a large pre-trained teacher model
Moreover, there are also works using structure of a information flow in the teacher model as knowledge [29, 30, 31], combining self-supervised training with KD [32], or training a student with only a few samples [11].
Due to its excellent performance, KD has been adopted in various research fields, such as CV [11, 32], NLP [17], and recommendation systems [33].
To compress vision transformers, Touvron et al.
[16] proposed a KD-based method termed DeiT.
They added a distillation token into the student and trained the student with hard labels provided by the teacher.
Although DeiT has achieved remarkable performance, KD for vision transformers has not been well explored yet.
Manifold learning.
Manifold learning is an approach for non-linear dimensionality reduction.
It learns a smooth manifold embedded in the original feature space to construct low-dimensional features [34, 35].
Recently, several works introduce manifold learning to KD [19, 20].
These methods train the student to preserve the relationships among samples learned by the teacher.
For vision transformers, these primary attempts are coarse and can be further improved because the basic input elements are patches not images.
Method 3.1 Preliminaries
Vision transformer.
Vision transformers are attention-based models inherited from NLP, and each layer of transformer consists of a multi-head self-attention (MSA) block and a multi-layer perceptron (MLP) block.
These models take split images as input, i.e. , the patches.
In particular, assuming the patch is of size P × P , an input image x ∈ R H × W × C of size H × W and channel number C is reshaped into flattened patches x p ∈ R N ×
( P 2 · C ) for processing, where N = HW/P 2 .
The patches are then projected into a D -dimension embedding space and added to positional embeddings.
We denote the summation result as x 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 a sample x corresponding to a label y , representing the prediction of the student and the teacher as f 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 a label smoothing hyper-parameter termed temperature, and λ is a balancing hyper-parameter.
Sometimes the intermediate features of the teacher can also be used for knowledge transfer.
For example, Romero et al.
[36] trained the student to output features similar to the teacher at intermediate layers.
However, such methods require the teacher and the student to have the same embedding dimension.
Otherwise, additional mapping layers are required for aligning, making the distillation process non-transparent.
Manifold learning-based KD methods [20, 19] support mismatched embedding dimensions.
They train the student to learn sample relationships predicted by the teacher but ignore patch-level information in vision transformers.
3.2 Fine-grained manifold distillation
To utilize the batch-level and patch-level information, we propose a fine-grained manifold distillation method.
Unlike existing KD methods [16] for vision transformer only distilling with logits, our method distills batch and patch level manifolds at intermediate layers.
Figure 2 provides an overview of the proposed method.
In the fine-grained manifold distillation method, we regard a vision transformer as a feature projector that embeds image patches into a series of smooth manifold space layer by layer.
At each pair of manually selected teacher-student layers, we aim to teach the student layer to output features having the same patch-level manifold structure as the teacher layer.
In particular, given samples of batch
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]
size B , we denote the feature of the student layer and the teacher layer as F S ∈ R B × N × D S and
F T ∈ R B × N × D T , respectively, where D S and D T are embedding dimensions.
We first normalize the 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 computational complexity is O
( B 2 N 2 D ) and a memory space of B 2 N 2 size is required to save the result.
Taking the settings B = 128 , N = 196 and D = 192 in the DeiT-Tiny model [16] as an example, it requires more than 240GFLOPs to compute and 2.5GB of memory space to save a single manifold relation map.
The remarkable resource consumption limits the fine-grained manifold distillation method scaling up to multiple layers.
Hence we must simplify the computation.
Inspired by the orthogonal decomposition of matrices, we decouple a manifold relation map into three parts: 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 loss as follows:
[Equation]
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 the inter-image manifolds via a relation map computed across randomly sampled patches.
Specifically, we sample K rows in the reshaped feature ψ ( F ) to obtain F r ∈ R K × D , and compute the randomly sampled patch-level manifold distillation loss as follows: [Equation]
[Equation]
Hence, the overall loss function of our proposed method is:
[Equation]
[Figure]
where α , β , and γ are hyper-parameters.
The summation over l means that the manifold relation map matching is performed on multiple pairs of teacher-student layers.
We will discuss the layer selecting scheme 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 , where each one is reduced by nearly BN/ ( B + N ) times when ignoring the lower order terms.
Still taking DeiTTiny as an example, the floating-point operations and memory space requirement reduce to 3GFLOPs and 32MB, respectively.
The sampling number K is set to 192, the same as our experiments.
With the decoupling scheme, fine-grained manifold distillation across multiple layers becomes 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 first 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 significantly increases the computational complexity and memory space requirement of the intra-image patch-level manifold loss L intra .
To remedy this drawback, we merge adjacent patches and view them as a single patch to further simplify the computation.
In particular, given a feature map F ∈ R N × D , we first reshape it into F r ∈ R
H × W × D , where the height H and the width W are obtained following the original patch splitting operation.
Then we adopt a merging setting ( H ′ , W ′ ) to merge every ( H/H ′ )
× ( W/W ′ )
adjacent patches in a non-overlapping manner.
Zero-padding is adopted when H/H ′ or W/W ′ is not an integer.
After merging, the feature map becomes F rm ∈ R H ′ × W ′ ×
( HWD/H ′ W ′ ) and is finally reshaped into F m ∈ R ( H ′ W ′ ) ×
( HWD/H ′ W ′ ) .
By adjusting the merging setting ( H ′ , W ′ ) , we can easily strike the trade-off between the complexity and the granularity of manifold relation maps.
Soft distillation.
Previous work [16] adds an additional distillation token into the student, and uses this token to learn hard labels provided by a CNN teacher.
However, when teacher and student are both vision transformers, we propose that it is better to teach the student only with soft labels of the teacher.
This design is based on the assumption that a larger model can learn more knowledge than a smaller one, and models of the same family share the same knowledge pattern, i.e. , a student can learn most knowledge of a teacher.
In our work, unless the student is larger than the teacher, we set the hyper-parameter λ in Equation 3 to 1.
Fixed depth.
Stochastic depth [37] is a regularization method that has becomes an infrastructure in training vision transformers [16, 5].
We propose that the student should not adopt this regularization when the teacher is trained with stochastic depth, since the soft labels already contain knowledge about this regularization.
Otherwise, the repeated regularizations may harm the student performance.
Hence, we adopt a fix-depth student in our method, i.e. , the stochastic depth regularization is not used for training the student.
Experiments
We evaluate our fine-grained manifold distillation method on ImageNet-1k [39] classification task,
CIFAR-10/100 [40] transfer learning task, COCO [41] object detection task, and ADE20K [42] [Table]
[Table]
semantic segmentation task.
Our implementation is based on Pytorch framework [43] and the
MindSpore 2 Lite tool [44].
4.1 Setup
Datasets.
We evaluate the proposed method mainly on the ImageNet-1k image classification task.
ImageNet-1k is a subset of the ImageNet dataset [39], which consists of more than 1.2M training images and 50K validation images from 1000 classes.
To test the generalization performance of student models trained with the proposed method, we conduct transfer learning experiments on
CIFAR-10 and CIFAR-100 datasets [40].
The two datasets both contain 50K training images and 10K 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 [41].
We use the COCO 2017 split, which consists of 118K training images and 5K validation images containing objects from 80 categories.
Baselines and models.
We compare the proposed method with DeiT [16], which adds an additional distillation token in the student model to learn hard labels from the teacher.
SwinTransformer students containing no distillation token, so we compare with the original KD method [14] on these models.
Table 1 summarizes the used models in our experiments.
4.2 Distillation results on ImageNet-1k
Implementation details.
On the ImageNet-1k image classification task, we train DeiT students with
CaiT teachers and SwinTransformer students with SwinTransformer teachers.
We slightly modify the architecture of DeiT students by removing the distillation token and only using the class token.
The hyper-parameter λ in the KD loss is set to 1, i.e. , the real label is not used to train the student.
When the teacher is smaller than the student, to prevent the performance degradation caused by the weak teacher, we set λ to 0.5.
In the fine-grained manifold distillation loss, hyper-parameters α , β , and γ are set to 4, 0.1, and 0.2, respectively.
The sampling number K in loss term L random is set to 192.
We set the above hyper-parameters one by one with the grid search method, indicating that their combination may not be optimal.
We select the first 4 layers and the last 4 layers of the student and the teacher to conduct manifold distillation.
Note that the class token in DeiT is ignored when computing manifold relation maps.
Moreover, to train SwinTransformer students efficiently, we adopt a patch merging setting of (14 , 14) .
Other training settings follow those in DeiT [16] 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 classification results on ImageNet-1k.
When compared with the hard logits distillation proposed in DeiT [16], our manifold method achieves remarkable performance improvements.
For example, Deit-Tiny distilled via manifold method outperforms Deit by +2 .
0% top-1 accuracy, and Deit-Small outperforms Deit by +0 .
9% with the CaiT-S24 teacher.
Although the “soft logits distillation" can boost the result by 0.4% compared to the “hard logits distillation", our proposed manifold still obtains better performance by +0 . 6% .
When the teacher is CaiT-XXS24, a much weaker architecture, the corresponding improvements are +1 .
6% on DeiT-Tiny and +1 .
2% on DeiT-Small, respectively.
Note that we report the RegNetY-16GF teacher results in DeiT for comparison, but do not conduct fine-grained manifold distillation experiments with this CNN model because the proposed method is designed for distillation between vision transformers. [Table]
[Table]
Moreover
[Table]
we conduct experiments based on Swin Transformer, one of the state-of-the-art vision transformer architecture.
When the teacher is Swin-Small and the student is Swin-Tiny, our proposed method surpasses the original student by +1 .
and achieves +0 .
5% performance improvement compared with the original soft logits label based KD method, demonstrating the effectiveness of our fine-grained manifold distillation.
4.3 Ablation and parameter comparison
Ablation of main components.
We design experiments to verify the effectiveness of our proposed fine-grained manifold distillation method, the fixed student depth, together with the “hard" and “soft" logits label in previous distillation methods.
As shown in Table 3, soft logits label can bring +0.5% top-1 accuracy compared to hard logits label [Table] and fixed depth can improve the baseline by +0.8-1.0%, which is a practical strategy to distill transformer.
Moreover, when combined with the fine-grained manifold distillation, student performance can be further improved by +0.7%.
[Table]
Layers for fine-grained manifold distillation.
To study the impact of different layer selections in manifold 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.
The number of selected layers is set to 6.
From the results reported in Table 5, conducting fine-grained manifold distillation at the head and the tail of student models are both crucial.
Hence, we speculate that the “Shallow/Deep” layer selecting scheme is the best because of its outstand-
[Table] [Table]
[Table]
ing performance and ease of use.
Following this selecting strategy, we further ablate the influence of the number of selected layers.
We denote the indices of layers selected from a L -layer model as { 1 , 2 , ..., k, L − k + 1 , ..., L − 1 , L } , and conduct experiments with the CaiT-small teacher ( L = 24 ) and the DeiT-Tiny ( L = 12 ) student to compare different k setting.
The corresponding are shown in Table 4.
We can find that our method is robust with various layer numbers, selecting 4 layers ( k = 2 ) obtains 76.4 top-1 accuracy, and selecting 8 ( k = 4 ) and 12 ( k = 6 ) layers achieve 76.5 topaccuracy, respectively.
Ablation of decoupled manifold loss terms.
The decoupled fine-grained manifold distillation loss consists of three terms.
We study their effectiveness and report the results in Table 6.
The results show that each component in the decoupled fine-grained manifold distillation loss contributes to improving the final performance.
In particular, only transferring inter-image relation maps leads to an unstable training process (mixed precision training will corrupt the student model, leading to a NAN loss), we speculate that these relations are highly related to the randomly sampled image batches.
When equipped with intra-image and random-image maps, the inter-image relation maps can further improve the accuracy by +0.1%.
Parameters in decoupled manifold distillation loss.
There are 4 hyper-parameters in decoupled manifold distillation loss: the weight of inter-image loss α , the weight of intra-image loss β , the weight of randomly sampled loss γ , and the sampling number K .
We compare different settings of these parameters and report the results in Table 7.
Our default setting outperforms those either increasing or reducing one of the parameters.
However, due to our coarse parameter searching scheme, we believe that the performance can be further improved by setting these hyper-parameters more carefully.
4.4 Transfer learning
To measure the generalization ability of the proposed method, we conduct transfer learning experiments.
We fine-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 a learning rate of 5 × 10 − 6 .
Other settings follow DeiT [16].
[Table] [Table]
Table 8 reports the transfer learning results.
Students trained with the fine-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
[Table]
we adopt the fine-grained manifold distillation to train object detection models on COCO 2017 dataset and semantic segmentation models on ADE20K dataset.
Implementation details on COCO.
We adopt Mask R-CNN [45] as the detection framework.
The teacher backbone is Swin-Small and the student backbone is Swin-Tiny.
All used backbones are pretrined on ImageNet-1k.
The teacher is trained for 36 epochs and each student is trained forepochs.
When training students with fine-grained manifold distillation, we adopt the distillation loss on outputs of the last two backbone stages and outputs of the feature pyramid network neck [46].
Results on COCO.
Table 9 presents the detection results.
Our manifold distilled student outperforms the student trained without distillation (+1.0% box AP), demonstrating that the proposed method benefits the object detection downstream task.
[Table]
Implementation details on ADE20K.
We adopt UPerNet [47] as the segmentation framework.
The student is trained for 160K iterations following [5].
Other settings follow experiments on COCO.
Results on ADE20K. Table 10 presents the segmentation results.
The proposed method outperforms the ditect training method and FitNet [48], an intermediate feature distilling approach, indicating that the manifold distillation approach can also help train a semantic segmentation model.
Conclusion
This paper proposes a fine-grained manifold distillation method for vision transformers.
We match patch and batch level features of student and teacher in a manifold space, and decouple 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 previous works, the student is distilled with soft labels and fixed depth.
Experiments are conducted on ImageNet-1k, CIFAR-10/100,
MS COCO, and ADE20K.
Corresponding results demonstrate that our manifold KD consistently outperforms existing methods.
The large search space of hyper-parameters is one of the most serious drawbacks of the proposed method.
In the future, we will study to further simplify the proposed method and help it be more parameter insensitive.
Acknowledgment.
This work
Checklist
For all authors... (a) Do the main claims made in the abstract and introduction accurately reflect 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] 2.
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] 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)?
[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] 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?
[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 identifiable 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]