P ROGRESSIVE C OMPRESSED A UTOENCODER FOR S ELF - SUPERVISED R EPRESENTATION L EARNING
Anonymous authors
Paper under double-blind review
A BSTRACT
Masked Image Modeling (MIM) methods are driven by recovering all masked patches from visible ones.
However, patches from the same image are highly correlated and it is redundant to reconstruct all the masked patches in MIM.
This redundancy is neglected by existing methods and causes non-negligible overheads in computation and storage that do not necessarily benefit self-supervised learning.
In this paper, we present a novel approach named Progressively Compressed
AutoEncoder (PCAE) to address this problem by progressively compacting tokens and retaining the least necessary information for representation.
In particular, we propose to mitigate the performance degradation caused by token reduction through exploiting the vision transformer to leak information from discarded tokens to the retained ones.
Besides, we also propose the progressive discarding strategy to achieve a better trade-off between performance and efficiency.
Identifying redundant tokens plays a key role in redundancy reduction.
We resolve this issue using a simple yet effective criterion, i.e. , we identify redundant tokens according to their similarity to the mean of token sequence.
Thanks to the flexible strategy, PCAE can be employed for both pre-training and downstream finetuning and, consequently, reduces the computing overhead non-trivially throughout the training pipeline.
Experiments show that PCAE achieves comparable performance while at most accelerates 1.9 times throughput compared with MAE for self-supervised learning, and accelerates 15%-57% throughput while the performance drop is within 0.6% for downstream classification.
I NTRODUCTION
Self-supervised learning has witnessed great progress in recent years and even outperforms its supervised counterpart in downstream tasks (Chen et al., 2020b; Caron et al., 2021;
Grill et al., 2020).
However, self-supervised methods usually take more epochs and computational overheads for considerable performance, which limits their applications.
Recently, Masked Image Modeling (MIM) becomes a popular topic for its scalability as well as promising performance (Bao et al., 2022;
He et al., 2022; Dong et al., 2021; Chen et al., 2022; Xie et al., 2022), especially for MAE (He et al., 2022) which significantly accelerates training via only operating on 25% visible patches in the encoder.
MIM methods are driven by recovering all masked patches from visible patches, but this reconstruction task could be a huge burden for them.
For example, MAE (He et al., 2022) masks 75% input patches and recovers them in the decoder.
As the encoder only handles 25% tokens, the overhead of the decoder (5.28G Flops) exceeds that of the encoder (4.3G Flops) during training, even though the decoder is relatively lightweight.
SimMIM (Xie et al., 2022) keep masked tokens in the encoder and thus suffer more heavy reconstruction tasks.
Therefore, the efficiency bottleneck of MIM methods de facto lies on the reconstruction task.
Note that MIM methods assume that the input image is redundant, and this redundancy suppose to be consistent for the input and the reconstruction targets, yet neglected by previous methods.
This observation motivates us to relax the efficiency bottleneck of MIM methods by reducing the redundancy in reconstruction targets.
In this paper, we present a novel approach named Progressive Compressed Auto-Encoder (PCAE) to reduce the redundancy in reconstruction targets.
The core idea is that: instead of recovering all masked patches, we neglect redundant ones and only reconstruct representative targets.
However, naively removing part of the reconstruction targets suffers serious performance degradation (please refer to Table 9 for more details) as part of supervisory signals are discarded along with redundancy.
Therefore, we propose to mitigate this issue by exploiting the information propagation capability of the vision transformer to leak information from the discarded patches to the retained ones.
Specifically, we bring in the momentum encoder in pre-training and discard tokens produced by the momentum encoder, rather than the original reconstruction targets.
To further alleviate information loss, we propose the progressive discarding strategy where redundant tokens are progressively discarded at different layers of the momentum encoder (please refer to Table 4 for more details).
Identifying redundant tokens lays the foundation of this discarding strategy.
Matrix decomposition precisely identifies the redundant component of token sequence but exhibits prohibitive complexity in practice.
As an alternative, we propose a simple yet effective criterion where we identify the redundant token according to their similarity to the mean of the token sequence.
Experimental results and visualizations demonstrate that this criterion is very effective and introduces negligible overhead.
Benefiting from the proposed strategy, PCAE reduces the reconstruction targets from 147 tokens to 18 tokens, which significantly improves the efficiency.
Experiments over multiple benchmarks demonstrate the effectiveness of the proposed method.
Specifically, PCAE accelerates 1.9 times throughput (624 img/s vs 328 img/s) or saves nearly 2/3 memory usage (10.82G vs 28.12G) under the same batch size compared with MAE
(He et al., 2022) while enjoying much faster converge speed (PCAE 300 ep 83.6 vs MAE 1600 ep 83.6) or higher performance (PCAE 800 ep 83.8 vs
MAE 1600 ep 83.6).
Besides, we extend the proposed strategy to inference as well as downstream fine-tuning for classification tasks, and the results show this strategy still works well.
In practical deployments, it is important to provide a continuous performance-overhead curve.
Thanks to the flexibility, we could provide models with different inference speeds by adjusting the number of discarded tokens.
Specifically, compared with the vanilla model, we accelerate 15% to 57% throughput during inference while the drop of accuracy is within 0.6%.
In summary, we list our contributions as follows:
• We revisit the reconstruction in MIM and point out that the computational overhead is redundant, which is ignored by previous methods.
• We propose a novel self-supervised learning approach, termed PCAE to accelerate the selfsupervised pretraining with negligible performance loss.
• Our proposed PCAE can be adapted to downstream fine-tuning and achieves unified acceleration among pre-training and downstream tasks.
R ELATED W ORK
Self-supervised learning.
Self-supervised learning builds the learning target from the input signal itself and expects to generalize the obtained knowledge to downstream tasks.
One taxonomy for self-supervised learning is to divide it into contrastive and generative (Liu et al., 2021).
Contrastive learning methods (Chen et al., 2020a; Wu et al., 2018;
He et al., 2020; Chen et al., 2020b; Caron et al., 2020; 2021; Zbontar et al., 2021;
Li et al., 2021; Grill et al., 2020) have been widely studied and achieve great progress over multiple downstream tasks including classification, object detection, and semantic segmentation (Chen et al., 2020b).
Generative self-supervised learning methods include auto-regressive model (Van Oord et al., 2016;
Van den Oord et al., 2016), flow-based model (Dinh et al., 2014) and the auto-encoding model (Ballard, 1987).
Among these methods, the denoising auto-encoding (DAE) model (Vincent et al., 2008) inspires both MLM and MIM methods.
The core idea of DAE is corrupting the input signal and recovering the input signal via auto-encoder which is similar to current MIM methods.
Recent studies (Chen et al., 2022; He et al., 2022; Wei et al., 2022) show that the MIM methods outperform previous contrastive learning methods in downstream fine-tuning tasks based on the modern vision transformer.
In this paper, we mainly focus on
MIM methods.
Masked image modeling.
Predicting the masked part of the image is a classic self-supervised task, and has been researched by previous works (Doersch et al., 2015; Pathak et al., 2016)
based on the convolutional neural network.
However, due to their limited performance, these methods are located in a non-mainstream position for a long time, especially after the advent of contrastive learning.
Recently, the vision transformer (Dosovitskiy et al., 2021) bridges the architecture gap between vision and language, and more MIM methods (Bao et al., 2022; Dong et al., 2021; Xie et al., 2022;
He et al., 2022; Wei et al., 2022; Chen et al., 2022) are proposed inspired by Masked Language
T r a n s f o r m er
E n c o d er
T r a n s f o r m er
D ec o d er
T r a n s f o r m er
E n c o d er
EMA a. Self-supervised Pretraining b. Downstream Fine-tuning
T r a n s f o r m er
E n c o d er
FC visible token masked token discarded token
Student
Teacher
Cross Entropy GT
Pred
Stagerecovery token positionembedding
PTR PTR
Stage
Stagec.
PTR
Less Important
Color Bar
Output 𝒔𝒄𝒐𝒓𝒆 = 𝟏−𝒔𝒊𝒎(𝑿,ഥ𝑿)
Modeling.
BEiT
(Bao et al., 2022), PeCo (Dong et al., 2021) and CAE (Chen et al., 2022) rely on the pre-trained dVAE to produce the prediction target.
MAE (He et al., 2022) and SimMIM (Xie et al., 2022) directly predict the raw pixel value, and MaskFeat (Wei et al., 2022) predicts the HOG (Dalal & Triggs, 2005) feature of masked patches.
Despite predicting different kinds of targets, these methods reconstruct all masked patches at the output end, which still be caught in the low information density problem.
Dynamic Vision Transformer.
Several methods (Rao et al., 2021; Meng et al., 2022; Yin et al., 2022; Liang et al., 2022) have been proposed to dynamically discard inattentive tokens for efficiency in vision transformers.
However, these methods are only accessible for the supervised setting due to their complicated design and requirement on supervisory signals.
Specifically, Dynamic ViT (Rao et al., 2021), A-ViT (Yin et al., 2022) and Ada-ViT (Meng et al., 2022) require to initialize from pre-trained models and only accelerates the inference.
E-ViT (Liang et al., 2022) measures token importance via the attention to the class token, but it requires supervisory signals on the class token.
A recent study (Zeng et al., 2022) proposes to merge patches that shared similar semantic concepts while keeping fine resolution for regions containing critical details, but it still relies on specific architectures and supervisory signals.
In summary, although dynamic vision transformers have been explored in supervised learning, their application to self-supervised learning is still uncovered.
M ETHOD
This section is organized as follows: in section 3.1, we present the overall framework and objective of PCAE.
In section 3.2, we elaborate on how we build the compact reconstruction targets including identifying redundant tokens and the progressive discarding strategy.
In section 3.3, we demonstrate the effectiveness of PCAE through illustrating the diversity of token sequence at different layers and visualizing the remaining tokens at each stage.
In section 3.4, we extend PCAE to downstream classification task so that the proposed strategy could improve efficiency across pre-training, finetuning and final inference.
3. O VERVIEW
The framework of PCAE is illustrated in Figure 1 (a), and the overall objective of PCAE is:
[Equation]
where ⊙ is the element-wise product; M is the random mask and M ⊙ x denotes unmasked patches while (1 − M ) ⊙ x denotes masked patches; f θ is the encoder to be pre-trained while ¯ f is the exponential moving average model of f θ ; k is the number of tokens we expect to keep.
p is the positional embedding of tokens requested to reconstruct and ω is the expanded learnable mask token which has the same shape as p ; h ( · , · ) is the concatenation operation; g ϕ is the decoder network;
M ( · , · ) is the similarity measurement, and we instantiate it with cosine similarity here.
The EMA model ¯ f operates on masked patches and progressively discards redundant tokens from the token sequence, thereby ultimately output much fewer tokens.
On the other end, the encoder process unmasked patches and output f θ ( M ⊙ x ) .
Then, the decoder g ϕ concatenates the output of the encoder and the position embedded mask tokens to reconstruct ¯ f ((1 − M ) ⊙ x ) .
Note that, as a reminiscent method, MAE (He et al., 2022) expects to reconstruct all patches, so the number of mask tokens is 147.
In contrast, PCAE only reconstructs 18 tokens, and the computational overhead on the decoder could be largely reduced.
3. P ROGRESSIVELY TOKEN REDUCTION
In this section, we elaborate on the ¯ f where we progressively discard redundant tokens.
We discuss the progressive discarding strategy and redundant token identification in the following.
Progressive discarding strategy.
As vision transformers are stacked by self-attention blocks, it is important to decide which layer to discard tokens.
Intuitively, for a given keep ratio, discarding all tokens at once suffers serious information loss, and progressive discarding them at different layers could largely alleviate it through exploiting the information leakage brought by the self-attention mechanism.
Formally, we decompose ¯ f into ¯ f 1 ◦ ¯ f 2 ◦ · · · ◦ ¯ f L where ¯ f
i indicates the i th layer in ¯ f .
Then, for the output of i th layer x i ∈ R N i × D where N i indicates the number of tokens in the i t h layer
and D indicates the number of dimensions, we have:
[Equation]
where I is the set of layer index we expect to discard tokens; T ( · ; · ) is the redundant token identification function and we will discuss it in later; k i is the number of tokens we expect to keep in i
th layer.
In practice, we apply the consistent drop ratio τ for all stages, so k i =
⌊ (1 − τ ) ∗ N i ⌋ .
Note that as we perform the token reduction in the momentum encoder, the gradient backpropagation is not taken into consideration.
Thus, once a token is discarded, it will never be used in later layers, which accelerates the inference of the momentum encoder.
Redundant token identification.
Matrix decomposition precisely identifies the redundant component in token sequence but exhibits prohibitive complexity in practice (slow down training five times in our implementation).
As an alternative, we turn to identify redundant tokens according to their similarity to the average token.
We are motivated by two reasons.
First, removing tokens most similar to the average token has the smallest influence on the average of the new token sequence.
Specifically, considering the average of the original token sequence ¯ a , the average of remained tokens a r and the average of discarded tokens a d , it is obvious that ∥ ¯ a − a r ∥
∝ ∥ ¯ a − a d ∥ .
Therefore, this strategy has minimal impact on the overall distribution of tokens.
Second, according to the
Discrete Fourier Transform (DFT), the average token corresponds to the low-frequency component of the image which usually presents as the background, and it is reasonable to remove them as they contribute little to visual perception.
Formally, for a token sequence x ∈ R N i × D , the redundant token identification can be given as:
[Equation]
[Figure] where M is the similarity measurement and instantiated with cosine similarity; rank ( · ) denotes the rank function in ascending order; S is the rank of tokens; k is the number of tokens we expect to keep.
Eq 3 shows that we determine whether a token is redundant according to its relative similarity to the average token.
Tokens that are more similar to the average token than other k tokens are regarded as redundant tokens, and vice versa.
The output of ¯ f could be given by combining Eq 3 and Eq 2.
3. A NALYZE 3.3. D IVERSITY ENHANCEMENT BY PCAE
We run randomly initialized ViT-B networks on ImageNet validation set under different settings and measure the average diversity at each layer (0 indicates the input token sequences) with the nuclear norm (more details please refer to the supplementary) over the whole ImageNet validation set.
As illustrated in Figure 3a, the proposed strategy effectively enhances the diversity of token sequence compared with cases where we do not discard tokens or randomly discard tokens.
In addition, we also apply this setting to pre-trained checkpoints to reveal the differences between pre-training methods.
As illustrated in Figure 3b, we find that all pre-training methods no matter supervised or unsupervised enhance the diversity of token sequence at intermediate layers.
We hypothesize that meaningful pre-training makes the network learn to focus on important things and neglect redundant components in input signals.
Another interesting finding is that MAE suffers low diversity at both the input end (index 0) and the output end (index 12).
We hypothesize that it is caused by the symmetry of its input and output.
In contrast, PCAE overcome the low diversity problem at its output end.
3.3. V ISUALIZATION OF PCAE
We visualize patches corresponding to the remaining tokens after each stage, and the results are presented in Figure 2.
We find that this visualization is exactly in accord with the frequency perspective we discussed above.
Taking the fifth image in the third row as an example, PCAE filters most of the sky background and eventually keep the detail-rich patches such as the beak and claws.
It also holds for other images, such as the three chicks in the first image, the snake’s head in the second image, and the text in the third image.
All these details are well preserved.
We also find that PCAE not only focuses on the salient object, but retains a variety of objects within the image.
As illustrated in Figure 2, tokens kept by PCAE are diverse and distributed throughout the image, which preserves the diversity within the image.
We attribute it to the stage-wise discard strategy.
Take the sixth image as an example, the first two stages discard most patches belonging to the sky.
Therefore, in the third stage, patches belonging to the sky are no longer redundant, thereby being preserved eventually.
The stage-wise strategy helps preserve the diversity within images which is important for the MIM-based method.
3. D OWNSTREAM F INE - TUNING
As the proposed strategy is easy to implement and does not require any extra parameters, it can be seamlessly incorporated into inference as well as downstream fine-tuning.
Specifically, we transfer the implementation of ¯ f that we presented in section 3.2 to downstream tasks and discard tokens both in fine-tuning and inference stages.
We observe two important differences between pre-training acceleration and downstream acceleration.
First, downstream tasks are more sensitive to token reduction.
The number of remaining tokens can be as less as 18 in the pre-training stage, but we observe serious performance degradation in downstream tasks in the such extreme case.
Therefore, we at most discard 40% tokens each stage and discard tokens after 0 th , 7 th , 10 th layer.
Second, the discarded tokens could be used to improve performance.
The discarded tokens are not used in the pre-training stage, but we find they are useful for downstream tasks.
Therefore, we append the average of discarded tokens to the retained tokens in downstream tasks.
However, we remove these appended tokens before final global pooling.
Besides, we apply a warmup strategy on the drop ratio parameter in fine-tuning following
Liang et al.
Specifically, we do not discard any token in the previous 10 epochs and then gradually increase the drop ratio to the target value in the following 25 epochs.
Then, we keep the target drop ratio until fine-tuning ends.
E XPERIMENTS
In this section, we first compare PCAE with other MLM methods on multiple benchmarks including classification, object detection, and segmentation.
Then, we extend PCAE to downstream tasks for acceleration.
After that, we ablate on the hyper-parameters of PCAE.
We follow MAE He et al.
(2022) for the settings in data augmentation and optimization.
We apply 75% mask ratio on the input image across all experiments, and discard 50% tokens after 0 th , 4 th , and 8 th layer respectively in default.
For more details please refer to the supplementary material.
[Table]
[Table]
[Table]
C OMPARISON WITH O THER M ETHODS
Fine-tuning on ImageNet-1k.
We follow the standard MLM evaluation protocol that fine-tunes the pre-trained model on the ILSVRC-2012 ImageNet for 100 epochs.
We follow the fine-tuning schedule and hyper-parameters in MAE
He et al.
out of fair comparison.
Under this setting, we compare PCAE with its main counterpart MAE and other self-supervised methods.
Referring to Table 1, PCAE largely outperforms MAE with the same number of epochs but saves around 35% computational cost.
For the same performance, PCAE only takes 1/8 overhead compared with MAE.
Although BEiT and CAE can benefit from DALL-E Ramesh et al.
(2021) which is pre-trained on 250M images, PCAE still reach a better or comparable fine-tuning performance.
Contrastive learning methods rely on complex data augmentation which badly increases the computational cost, especially for DINO which relies on the multi-crop operation.
As presented in Table 1, PCAE outperforms DINO in the fine-tuning benchmark at a much lower cost (around 1/12).
Object detection and segmentation.
Besides the classification task, we evaluate PCAE on object detection and segmentation to demonstrate its potential in dense prediction tasks.
In this setting, we fine-tune Mask R-CNN
He et al.
(2017) with the pre-trained parameters in an end-to-end manner on COCO Lin et al.
The ViT backbone is adapted for use with FPN Lin et al.
We report the box AP for object detection and the mask AP for instance segmentation in Table 2.
The results show that PCAE consistently outperforms other methods both in object detection and instance segmentation.
Throughput comparison.
As the training may be influenced by many factors, we show the throughput comparison between PCAE and MAE in Table 3.
The results show that PCAE is still much more [Table] [Table] Exp id
[Table]
×
[Table]
Exp id
[Table]
✓
[Table]
Exp
[Table]
✓
[Table]
Exp id
[Table]
✓
[Table]
Exp id =✓
[Table]
efficient than MAE even though MAE is already a very efficient method.
Specifically, PCAE at most accelerate 1.9 times throughput or save 2/3 memory usage compared with MAE.
We also adapt
PCAE to SimMIM, termed PCAE-SimMIM
[Table]
which accelerates the throughput 2.37 times without compromising performance.
For more details please refer to the supplementary.
4. D OWNSTREAM ACCELERATION
Model acceleration is critical for the practical deployment of vision transformers.
Due to the simplicity and flexibility of the proposed progressive token reduction strategy, it can be easily incorporated into the supervised setting and benefit the inference speed.
With the pre-trained model, we fine-tune it on ImageNet under different drop ratios for 100 epochs and report the Top-1 accuracy.
As presented in Table 6, PCAE balance the acceleration and performance well.
Specifically,
PCAE accelerates 57% throughput while the performance drop is within 0.6%.
Even for the extreme case where only 42 tokens are kept out of 196 tokens, the performance slightly drops 1.2% but the throughput improves 88%.
These results further demonstrate the effectiveness of PCAE as it manages to accelerate scenarios including self-supervised pre-training, supervised fine-tuning, and inference.
4. A BLATION S TUDY
We apply the same setting for all ablation studies unless other specific.
Specifically, we pre-train
PCAE on ImageNet for 100 epochs with the batch size of 1024 and fine-tune pre-trained parameters on ImageNet for 100 epochs with the default schedule of MAE for a fair comparison.
Initially, we highlight critical hyper-parameters of PCAE for clarity.
Referring to Eq 2, our discarding strategy is determined by I which decides where tokens are discarded and the drop ratio hyper-parameter.
As the value of I is too huge, we decompose it into two parameters, i.e., drop stage and drop case.
The drop stage indicates how many times we should discard tokens, i.e., card ( I ) .
Drop case indicates where we prefer to discard tokens, i.e., shallower layers, deeper layers, e.t.c.
(refer to Table 5 for more details).
In summary, drop stage, drop case, and drop ratio are three critical hyper-parameters for PCAE, and their effects are elaborated in the following.
The core operation in PCAE is discarding tokens during forwarding, so we elaborate this operation by answering two questions: how many tokens do we need to discard and where should we discard tokens with different settings of drop stage, drop ratio and drop case.
How many tokens should we discard?
We first fix the drop ratio and drop case, and ablate on the drop stage.
As presented in Table 4, the results show that removing redundant tokens from the prediction target significantly improves efficiency (exp 1 and others).
Among these experiments, we find that the reconstruction targets also prefer to high drop ratio, which is similar to the input end, and discarding 87.5% tokens (exp 4) achieves the best balance between performance and efficiency.
[Table]
We further the drop ratio equal to 87.5% and ablate on the drop stage, i
[Table]
e., we need to drop tokens in few times with a large drop ratio or more times with a smaller drop ratio. A large drop ratio improves efficiency but may suffer performance degradation, and vice versa. The empirical results in Table 7 show that dropping 3 times is the best choice to balance them.
[Table]
[Table]
[Table]
Then, we fix the drop stage and grid search for the best total drop ratio from 75% to 90%.
As presented in Table 8, totally discarding 80% tokens reaches the best performance.
However, compared with experiment 4, we find that discarding 50% tokens three times is more simple and efficient and only slightly degrades the performance.
Thus, we apply the configuration of exp 4 in default.
Where should we discard tokens?
We ablate on the drop case parameter here.
We propose five different cases to discard tokens and their details are illustrated at the right of Table 5.
The results in Table 5 show that evenly allocating the discarding operation among blocks performs better than continuously discarding tokens (compare exp 2, 3, 4 with exp 5, 6).
The drop case 1 performs best compared with case 2 and case 3 (exp 2, 3, 4), which indicates that token reduction supposes to be avoided at the last layer.
We also compare with the random drop in Table 5, and the result shows that the performance of randomly discarding tokens far falls behind that of our strategy (exp 1 and exp 2).
Besides, we provide more performance of random drop under different configurations in the supplementary, and results show the proposed strategy outperforms it in all cases.
C ONCLUSION
In this paper, we present PCAE which addresses the problem of the redundant reconstruction target in MIM methods.
The results show that PCAE outperforms MAE over multiple benchmarks in cases where PCAE even only reconstructs 18 tokens.
The lightweight reconstruction task accelerates training (at most 1.9 times throughput acceleration compared with MAE) without compromising performance.
We visualize the remained tokens and find that PCAE manages to retain semantic-abundant components in images and mainly discards tokens belonging to the background.
PCAE exhibits elegant symmetry at the input and output end as both them are compressed.
What beyond the proposed method is the insight that MIM-based methods benefit from more compact prediction targets.
We wish it can inspire future works.
A PPENDIX A D IVERSITY MEASUREMENT
Initially, we show how we measure the diversity.
We start by looking at the input token sequence from a patchify image.
Denote the number of token as N , and the dimension as D , the token sequence can be written as a matrix A ∈ R N × D .
Following the common practice in self-supervised learning Wang & Isola (2020); Ermolov et al.
(2021), we analyze the token representations on the unit sphere, i.e., we normalize the length of each token to one.
As the convex envelope of matrix rank, the nuclear norm is widely used to measure to diversity of matrix rows Cui et al.
Specifically, according to the theorem in Fazel (2002), ∥ A ∥ ∗ is the convex envelope to rank ( A )
when ∥
A ∥ F ≤ 1 .
Since ∥ A ∥ F = √ N (note that each row vector in A is unit vector), the convex envelope of rank ( A ) turn to ∥ A ∥ ∗ / √ N .
According to the property of the nuclear norm, ∥ A ∥ ∗ / √ N is bounded by √ k where k = min ( N, D ) .
In practical implementation, N is the length of token sequence and usually be 196, and D is the embedding dimension and usually be 768 or 512.
Thus, min ( N, D ) = N and ∥
A ∥ ∗ / √ N ≤ √ N .
Since the length of token sequences is various during forward process and the upper bound of ∥ A ∥ ∗ / √ N depends on N , we instead apply ∥ A ∥ ∗ /N which has a clear upper bound and get rid of the influence from N to measure the diversity in token sequences.
B I MPLEMENTATION D ETAILS
Architecture.
We study the ViT base architecture (12 transformer blocks with dimension 768) and follow the implementation of Dosovitskiy et al.
The teacher and the student branch share the same architecture.
We follow the implementation of He et al.
for the decoder part.
Specifically, we apply an FC (fully-connected layer) as the dimension adapter between the encoder and decoder and two transformer blocks to decode the masked tokens with the output of the student.
At last, we apply another FC to convert the decoder dimension into the encoder dimension for reconstruction.
What is different from MAE is that we only reconstruct tokens output by the teacher.
Hyper-parameters.
We follow the setting of hyper-parameters of MAE (He et al., 2022).
For the
EMA operation, we apply the fixed momentum equal to 0.9999.
C T HE N AIVE A PPROACH OF D ISCARDING P ART OF T HE T ARGETS .
[Table]
We present the naive approach of discarding part of the reconstruction targets of MAE here.
We discard patches with the proposed strategy and only reconstruct remanent patches with MAE.
We apply two settings for the base MAE.
In the first setting, we simply mask gradients from discarded patches while we only decode remanent tokens in the decoder in the second setting.
We observe that both settings suffer serious performance degradation as presented in Table 9.
This demonstrates what we stated in the main text that simply discarding part of the reconstruction target damages useful information as well as eliminates redundancy.
D A DAPT TO S IM MIM
For SimMIM, which forwards visible tokens together with the mask tokens at the encoder, we can regard the encoder in SimMIM as the decoder in MAE, and thereby speed up it via discarding most of the mask tokens.
Specifically, the length of the token sequence is 196 (78 visible tokens andmask tokens) in vanilla SimMIM, and we can speed up it by building a much more compact target for it.
Then, the length of the token sequence can be reduced to 96 (78 visible and 18 mask tokens), thereby improving the efficiency.
We adapt PCAE to SimMIM with minimum modification. Specifically, We run SimMIM for 100 epochs with the officially released code on a single machine with
[Figure]
[Table]
V-100 and record the practical GPU day.
We adapt PCAE to SimMIM with minimal modification and run PCAE-SimMIM for 100 epochs follow the configuration of SimMIM
[Table]
[Table]
Besides, we also report the throughput and memory usage in Table C by running on a single V-100.
As presented in Table 11, we find that
PCAE-SimMIM achieves comparable performance (PCAE-SimMIM 82.3 vs SimMIM 82.1 100 epoch) to SimMIM with 55% training hours (PCAE-SimMIM 11.1h vs SimMIM 20.3h on 8 V-100 machine).
Besides, we provide the throughput comparison between PCAESimMIM and SimMIM below, which shows PCAE still works well when transferred to other methods.
E C OMPARISON TO OTHER STRATEGY .
We consider two settings here.
In the first setting, we measure the token importance with the average attention from other tokens.
Specifically, we average the attention matrix of different heads in the self-attention block and calculate the mean of each column to get the importance of each token.
Then we discard less important tokens.
We follow the same implementation and settings with PCAE.
We term this setting as Token Importance.
In the second setting, we replace the momentum encoder of PCAE with a DINO pre-trained ViT-S model, and identify redundant tokens according to the saliency produced by pre-trained model or PCAE respectively for comparison.
We term this setting as Saliency-Pretrain or PCAE-Pretrain.
We run these settings for 100 epochs and report their finetuning accuracy in Table 12.
[Table]
The results shows that Token-Importance performs much worse compared with PCAE (80.7 vs 82.3).
We owe it to the lack of supervisory signals.
The network is hard to identify important tokens at the beginning stage where parameters are randomly initialized, and the worse beginning propagates errors to later phases, resulting in serious performance degradation ultimately.
We also find that PCAE-Pretrain performs better than Saliency-Pretrain.
It is caused by the diffculty in identifying redundant tokens at intermediate layers, especially at shallow layers for the saliency based method, and thus it may discard important tokens and causes performance degradation.
F M ORE COMPARISONS TO RANDOM DROP .
[Table]
RD
[Table]
RD 0.
RD
[Table]
RD 0. RD 0. We supplement more experiments here to show our advantage over randomly drop (RD).
For a fair comparison, we fix the total drop ratio equal to 87.5% and search for the best drop strategy for randomly discarding by the drop case and drop stage.
Specifically, we fix the drop stage and run the randomly drop strategy with different drop case (we select drop cases with better results by referring to Table 5).
The results in Table 13 show that the randomly drop strategy much worse than
PCAE.
Besides, we also fix the drop case and search on the drop stage, but the results still show that PCAE outperform the randomly drop strategy by a large margin.
G V ISUALIZATIONS UNDER D IFFERENT PATCH SIZES
We show the visualization to PCAE under different patch sizes.
The patch size determines the information amount of a token, and thus has influence on the discarding strategy in PCAE.
We pretrain models on ImageNet100 for 100 epochs under different patch size, and show their visualization in Figure 4.
The visualization indicates that PCAE discards similar content with different patch size. H M [Table]
[Table]
[Table]
[Table]
In this section, we present more ablation studies of PCAE.
Mask ratio.
MAE (He et al., 2022) suffers serious performance degradation once its mask ratio exceeds 75%, which prevents MAE keep fewer tokens in the encoder.
However, we surprisingly find that
PCAE is robust to extremely high mask ratios and even works well with them.
Specifically, as we presented in Table 14, the performance of PCAE is nearly not affected even when we mask 90% input patches.
We owe this merit to the relaxed reconstruction task in PCAE.
PCAE only reconstructs around 10% of patches while inputting 25% of patches, which makes the reconstruction task relatively simple, and thus the encoder could tolerate higher mask ratios.
The relaxed reconstruction task benefits both the encoder and the decoder in MAE.
Decoder depth.
We ablate the decoder depth of PCAE.
The results in Tableshow that the fine-tuning performance of
PCAE is robust to decoder depth but the performance in downstream tasks relies on deeper decoders, which is consistent with MAE (He et al., 2022).
Architectures.
We run PCAE with different architectures.
We apply the same configuration on hyper-parameters across different architectures and the results in Table 16 show that PCAE works well with different architectures, which demonstrates that the hyper-parameters are robust to model size.