Boosting Vanilla Lightweight Vision Transformers via Re-parameterization
Authors: Zhentao Tan
Xiaodan Li
Yue Wu
Qi Chu
Le Lu
Nenghai Yu
Jieping Ye
Venue: ICLR 2024 poster (2024)
Keywords: Vision Transformers
Re-parameterization
Lightweight Models
Primary area: representation learning for computer vision, audio, language, and other modalities
TL;DR: We propose a linear-based re-parameterized module to improve the performance of vanilla lightweight ViTs.

Abstract: Large-scale Vision Transformers have achieved promising performance on downstream tasks through feature pre-training. However, the performance of vanilla lightweight Vision Transformers (ViTs) is still far from satisfactory compared to that of recent lightweight CNNs or hybrid networks. In this paper, we aim to unlock the potential of vanilla lightweight ViTs by exploring the adaptation of the widely-used re-parameterization technology to ViTs for improving learning ability during training without increasing the inference cost. The main challenge comes from the fact that CNNs perfectly complement with re-parameterization over convolution and batch normalization, while vanilla Transformer architectures are mainly comprised of linear and layer normalization layers. We propose to incorporate the nonlinear ensemble into linear layers by expanding the depth of the linear layers with batch normalization and fusing multiple linear features with hierarchical representation ability through a pyramid structure. We also discover and solve a new transformer-specific distribution rectification problem caused by multi-branch re-parameterization. Finally, we propose our Two-Dimensional Re-parameterized Linear module (TDRL) for ViTs. Under the popular self-supervised pre-training and supervised fine-tuning strategy, our TDRL can be used in these two stages to enhance both generic and task-specific representation. Experiments demonstrate that our proposed method not only boosts the performance of vanilla Vit-Tiny on various vision tasks to new state-of-the-art (SOTA) but also shows promising generality ability on other networks. Code will be available.

Published as a conference paper at ICLR 2024
BOOSTING VANILLA LIGHTWEIGHT VISION TRANSFORMERS VIA RE-PARAMETERIZATION
Zhentao Tan 1.3 Xiaodan Li Yue Wu', Qi Chuu, Le Lu², Nenghai Yu², Jieping Ye Alibaba Cloud', Alibaha Group?, University of Science and Technology of China", East China Normal University*
ftanzhentao. tzt, fiona. 1xd, matthew. wy e ibaba-inc. com [le.iu, yejieping. Sealeping.yee@alll yelbalibaba inc. com, {qchue, ynhō Jantfuuscc.eee edu. on
ABSTRACT
Large-scale Vision Transformers have achieved promising performance on down- stream tasks through feature pre-training. However, the performance of vanilla lightweight Vision Transformers (ViTs) is still far from satisfactory compared to that of recent lightweight CNNs hybrid networks. In this paper, we aim unlock the potential of vanilla lightweight ViTs by exploring the adaptation of the widely-used re- ee-aaaaeterization technology to ViTs for improving learning ability during training without increasing the inference cost. The main challenge comes from the fact that CNNs perfectly complement with re-pan re-parameterization over convolution and batch normalization, while vanilla Transformer architectures are mainly comprised of linear and layer normalization layers. We propose to in- corporate the linear ensemble into linear layers by expanding the depth of the linear layers with batch normalization and fusing multiple linear features with hi- erarchical representation ability through pyramid structure. We also discover and solve a new transformer-specifee distribution rectification problem caused by multi- re-parameterization. Finally, we propose our Two--Dieesioons Re- parameterized Linear module (TDRL) for ViTs. Under the popular self-supervised pre-training and supervised fine-tuning strategy, our TDRL can be used in these two stages to enhance both generic and task-specific representation. Experiments demonstrate that our proposed method not only boosts the performance of vanilla Vit- Vit-Tiny on various vision tasks to new state-of-the (SOTA) but also shows promising generality ability on other networks. Code will be available.
INTRODUCTION
Inspired by the remarkable success of Transformers in natural language processing (NLP), Vision Transformers (ViTs) (Dosovitskiy et al. (2020) have undergone significant advancements, especially when trained on large-scale datasets with self-supervised learning (e.g., contrastive learning (Chen) jet al. 202T) and masked image modeling (MIM) (He et al., (022)). These developments have led to the emergence of large-scale ViTs (Dehghani et al. (023). which are expected to promote performance mutations similar to Transformers in NLP and eventually become a unified framework for visual or even multimodal tasks (Li al.) 2023 Xu et al. 2023). However, the performance of lightweight VIT models is still far from satisfactory and even inferior to corresponding CNN counterparts (Howard et al. 2019 Woo et al. 2023). Lightweight deep models are still dominated by CNNs or carefully designed hybrid networks (Mehta & Rastegari] [022 Liu et al. (2023] Vasu et al.) 2023b). Given that large ViT models are progressively unifying multimodal feature representation, believe that it crucial to explore how unlock the potential of vanilla lightweight ViTs, thereby
achieving uniformity across different model scales.
Fortunately, recent research has recognized this issue and made efforts to take step forward (Wang et al. 2023 Huang et al.] 2023 Ren et al. 2023). MAE-Lite (Wang et al. 2023) gives detailed
* Thii work was supported by Anhui Provincial Science and Technology Major Project (No. 2023>020006) and the National Natural Science Foundation of China (No. 62272430).
Yue Wu and Jieping Ye are corresponding authors.

---

Published as a conference paper at ICLR 2024
analysis of the effects of MAE pre-training (He al. 2022) and finds that ViT-Tiny achieve promising classification performance with proper configuration: 1) increasing the number of heads to 12; 2) applying attention map distillation during pre-training; 3) carefully adjusting fine-tuning settings. These valuable insights focus on image classification, which may not be applicable to other tasks. TinyMIM (Ren et al. 2023) systematically studies different options in the distillation framework to take full advantage of MIM pretraining. Different from these two aforementioned methods which only focus on applying knowledge transfer during MAE pre-training, G2SD (Huang et al. 2023) proposes to benefit the learning of small/tiny ViT' from both MAE pre-training and downstream fine-tuning. While the generic-to-spect competitive performance for small models, its applicability may be limited due to the fine-tuning requirements of large-scale teacher models on downstream tasks. Considering that obtaining generic pre-trained large models (He et al. 2022] Radford et al. (02T) is relatively easier, it is more practical to perform knowledge distillation solely during the pre-training stage.
two-stage distillation approach does achieve
While previous methods focus more on exploring training recipes, we turn to enhance the lightweight ViTs itself: increase the model capacity during training while keeping the inference unchanged via re- parameterrizaiion technology (Ding al., 2021b a 2019: 2022). A typical re- parameterized module usually consists of bualt--banch networks primarily composed of convolu- tions and batch normalization (Toffe & Szegedy, 2015). Thanks to the particularity of batch nor- malization, these modules retain their adaptive normalization during training and can be merged into single convolution operation at inference. Consequently, the additional parameters within the re- re-parameterized module do not increase the inference cost. This approach has been success- fully employed in the CNN-related networks (Ding al. 202 a) and is even considered a default technique in recent lightweight network designs (Vasu et al. [023bb a). However, these convolution- based re- parameterized modules can not be directly applied to vanilla Transformers because of their non-convolution architecture. How to adapt this paradigm to vanilla ViTs remains unexplored.
In this paper, we systematically study the above issues and explore the linear-based re- parameterization of vanilla Vision Transformers, without incorporating any convolutional opera- tions. To enhance the training capacity of the linear layer, we design a delicate linear stack with batch normalization in between them to incorporate adaptive normalization. A pyramid multi-branch structure is further proposed to fuse hierarchical feature representations from linear- branches of different depths. Additionally, we discover the importance of distribution consistency along the depth dimension of deep networks for training stability, especially for self-attention in ViTs. To alleviate the distribution changes caused by the additive mechanism of re-parameterization, we in- corporate additional distribution rectification operations to normalize the outputs. Based on the above designs, we propose Two-Dimensional Re-parameterized Linear module (TDRL) for ViTs.
To validate the effectiveness of our TDRL, we follow the pre-training and line-tuning pipeline of MAE et al. (2022) and apply TDRL the ViT-Tiny model. It achieves new state-of-the-ar per- formance on various visual tasks, such as image classification, semantic segmentation, and object de- tection. Further experiments on more modees/architeetures including the relatively larger ViT-Smal network Dosovitskiy et al. (2020), generation network (Ho et al. 20200), and SOTA lightweight net- works (Vasu et al. 2023b) also validate the generality of our proposed re-parameterizzttoo method.
2 RELATED WORKS
Vision Transformers. ViTs (Dosovitskiy et al.
2020) have established the dominant position of Transformer architecture in the vision domain recently. It shows competitive performance on various downstream tasks (Touvron et al. [22T]] Li et al. 2022, Liu et al. 202T Zhang al., 2022a Peng
et al. 2023 Zhai et al., 2022] Zamir et al. 2022 Tan et al. 2023 Zhang et al. (022b). However, compared to CNN counterparts, ViT's perform poorly in limited model capacity or data scale due to the lack of inductive bias (Dosovitskiy et al. 2020 Touvron et al., 2021). Most lightweight ViT
works draw inspiration from CNN designs to build hybrid architecture (Mehta & Rastegari 2021
Wu et al. 2022b Chen et al., 2022 Liu et al. 2023) Vasu et al. 2023a). while few works tempt to improve the performance of vanilla ViTss (Wang et al., 22023] Huang et al. 20033] Ren et al.) 2023) In this paper, we focus on improving the vanilla lightweight VIT networks via re- parameterization.
Self-supervised Learning. Self-supervised learning is the mainstream powerful paradigm for rep- resentation modeling without the requirement of data labels (Balestriero et al. 2023). Among them,
2

---

Published as a conference paper at ICLR 2024
Linear
Linear
Linear
L Linear
Conv
Conv
Conv
BN
ixN
Linear
Linear
BN
BN
BN
x
Intre Cran-Bracch Fusion (Eq.
Linear
Inter Branch Fusion (Eq.
xN
XN
Linear
Lincur
Original ViTs
TDRI.
Rectify
Pipeline
(a) Classical Architecture
(b) Re-parameterization for ViTs
(c) Re Re-parameterize Fusion
(a) Classical convolution-basel re- ce-parameterization architecture. (b) TDRL follows pyramid design of depths (i.e., I., n, n = {1, N).. Dashed lines with single linear layer indicate the skip connection. "Rectify" is the distribution rectification.
Figure 1:
(c) Re-parameeterzzttio
fusion: 1) merge each rep-branch into a single linear layer through Equation branches (each branch contains one linear layer) through Equation
2) merge multiple
masked image modeling (MIM) has achieved surprising performance on VITs (He al. 2022
Bao et al.) 2021). Taking raw pixel, semantic features, or discrete tokens reconstruction tar- gets, most methods explore performance upper bound by finding better supervisions or scaling up model capacity (Dehghani al. 2023; Wei al. 2022 Fang al.) 2023 Peng et al. (022). It
has been demonstrated that MIM technologies can benefit large models. But their performance on lightweight ViTs is always overlooked. MAE-Lite (Wang et al. 2023), TinyMIM (Ren al., 2023), and G2SD (Huang et al., (023)) are recent methods that investigate lightweight ViT's from the per- spectives of training configurations and knowledge distillation (Hinton et al.) 2015). Differently, develop a re-parameterized way for lightweight ViTs to take full advantage of MIM pre-training.
Structural Re-parameterization. Structural re- ee-arammetrizatiinn (or over-parameterization) means the technology to scale up the model capacity during training while keeping the inference un-
changed. very useful to train compact CNN networks (Ding et al. 2021b aj Guo et al. 2020 Cao et al. 2022 Ding et al., 2019). These modules are built with linear operations (c.g., K X K convolutions and average pooling) and batch normalization (Toffe & Szegedy 2015) enhance their training representation ability and keep efficient inference speed. However, these methods are de- signed over convolution which can not be directly applied to convolution- free vision Transformers. In this paper, we extend them to vanilla ViTs without any local convolutions.
3 METHODS
3.1 RE-VISITING STRUCTURAL RE-PARAMETERIZATTON
We first re-visit the p-parameterization of CNN networks (Ding et al. 2021b). As shown in Figure [ (a), typical module is a multi-branch additive architecture. Each branch consists of convolution and batch normalization. enhances the learning ability during training and can merged into single convolution layer for efficient inference. The merging process can be divided into two steps:
1) Entra-Brante Fusion; the fusion convolution and batch normalization within each branch. Let W Coot XC K and b E Cout denote the weight matrix and bias vector of convolution layer with x K kernel size, On input channels and Cout output channels. The scale, shift, mean, and variance of batch normalization denoted as Y, 3,8,, p. E IR Cout, respectively. The merged convolutional parameters are as follows:
b. (b. 1-pp + Br,
W
W
(1)
where is the output channel index. If converting the convolution to a linear layer, we can also easily fuse its weight and bias with batch normalization through Equation
3

---

Published as a conference paper at ICLR 2024
2) BiterBrracch Fusion: the multiple branches can be further merged to single convolution
N
N
n
b b=
Σ b
(2)
where N is the number of branches, W Th and b n are the fused weight and bias of the n-th branch.
3.2 TWO-DIMENSIONAL RE-PARAMETERIIEE LINEAR MODULE
CNN networks perfectly complement with re- ee-aarametrrization due the wide usage of batch nor- malization, for its intriguing calculation transformation characteristics during training and inference, which happens to the core of re-parameterization. However, ViT' are mainly comprised of lin- ear and layer normalization layers (Ba et al. 2016). Due to the fact that the mean and variance of layer normalization depend on the input during inference, layer normalization cannot be merged with other operations statically as the way of batch normalization. We can only re-parameterize linear layers in transformer networks, based on linear layers while keeping the re- parammeterizztion structure simple and mergeable in in- ference by incorporating Linear Ensemble to linear layers. What's more, we discover and solve a new transformer-speccffe Distribution Rectification problem with re-parameterization.
other words, we have to propose a new structure mainly
Linear Ensemble. The basic re- re-parameterized unit of CNNs instinctively incorporates statistical calculation characteristics with batch normalization. Replicating the unit into multiple branches im- proves the module capacity. Each branch of convolution incorporates explicit inductive bias through shared local kernels and padding for powerful spatial feature representation within a single layer. However, simple linear replicas suffer the homogeneous problem, with each replica updated with almost the same gradients during training. Thus, we propose to stack linear layers with batch nor- malization in-between them for the linear ensemble. The rationality is three folds: 1) Linear stacking with batch normalization similar to A..[[wicc is appropriate to transformers and plays an im- portant role to represent rich intra-token information (Dosovitskiy et al. 2020). Thus, it is inherently appropriate for transformer networks. 2) Although batch normalization is not as suitable as layer normalization for ViTs Yao et al. 2021). it still can be used in in-between layers for the linear en- semble while keeping the original layer normalization unchanged. 3) The stacked linear layers with batch normalization can be fused to single linear layer. The batch normalization can be fused with Let W e RGxX and b' RCi denote the weight and bias of the I-th = 1, 3, L) merged linear layer in the stack. Two adjacent layers (c.g., (I + 1)-th and (-th) can be merged as:
precedent linear layer via Equation
k= Ww k-1
Wi (1+,,1
wl - + +
(3)
where i 1, .... C++1 and 1,2, C1-1 are the channel indexes. Based on Equation 3 we can easily merge a sequence of linear layers of any length into one linear layer. Another effective re-parameterization strategy of CNN-based structure is use different kernel sizes (e.g., K K and 1 x 1) to vary the learning patterns of branches. Accordingly, we vary the depth of linear stacks build a pyramid multi-branch topology. The additive combination of features from these re-parameterization branches exhibits abstract representation ability from shallow to deep. It not only enhances the feature representation but also improves the diversity between branches.
Distribution Rectification. Numerous previous works have verified that distribution consistency along the depth dimension is critical to deep networks, with influential works like different ini- tialization and normalization methods (loffe & Szegedy 2015: Ba et al. 2016 He al.) 2015 Kumar 2017). The above proposed re-parameterization structure will change the distribution be- tween inputs and outputs due to its multi-branch additive mechanism, which will affect the training stability of Vision Transformers, especially for Multi-Heat Self Attention (MHSA).. The standard self-attention (Dosovitskiy et al., 2020) first calculates the attention map A E R M based on the pairwise similarity between query Q E M and key K E MMC, and then computes a weighted sum over all values V RR MxC;;
A
Attention (Q, K, V) so softmar(
VV,
A= QKT
VO
ML. stacks linear layers with GELJU and our stack multiple linear layers with batch normalization.
4

---

Published as a conference paper at ICLR 2024
Following Vaswani et al. 2017)) assuming that the components of Q and K are independent random k=
variables with mean 0 and variance the elements of attention map Ai.j Ex_1 QuaKKg have mean 0 and variance Ck. Typical Transformers perform Scaled Not-Product Attention which scale the dot product results by 1/vC& to ensure the final output variance back to 1. However, in case
Σ No Q", K = - m=1
Nx k the distribution of A
Q and K are re-parameterized as Q changes as follows:
n=
k=CL
No
NK
k=C. No NR ΣΣΣ
K O
Qi KK
(5)
k=1
1=1
m=1
k=1 m=1
The re-parameterized distribution changes are amplified through the attention mechanism, where the variance of elements in A increases to C& Nx. will increase the probability of extreme values of A and affecting the stability of training (as shown in Figure3). Considering that attention fragile to distribution and the distribution of Q, K much more complicated than the above assumption, we use an additional batch normalization to modulate Q' K' to rectify the distribution changes. In other components like FFN, this distribution change may also result the convergence bottleneck. Considering that layer normalization is already used between MHSA and FFN, and the variance change is not as large as in attention calculation we adopt the approach in Scaled Dot-Product Attention and re-scale the features with V rather than normalize
Main Architecture. As shown in Figure | (b), our proposed two-dimensional re-parameterized linear module consists of N re-parameterized branches (denoted as rep-branch) and an additional linear layer skip branch. The outputs of these branches are fused via element-wise addition. Each rep-branch consists of LN basic re-parameterized units (a linear layer followed by a batch normal- ization layer) and final linear layer. Let fn,r., (-) denote the n-th rep-branch with L,, basic units, X E R MxIn and Y E MxCau denote the input and output tensors where M is the sequence length. In this pyramid structure, we set the number of basic units in a rep-branch the same as the branch index. The whole calculation of TDRL can be formulated as:
Y = Recli I Linear X) Σ Jn,t, (X)), In T2,
where Rectify(-) is the distribution rectification operation mentioned before. It ensures that each rep-branch has different approximation abilities, thereby keeping their feature spaces away from each other. In the following, we use P-WNS to denote the detailed configuration of this Pyramid architecture with Width of N rep-branch and a Skip branch. In addition, we also design Regular version of TDRL whose Depth per branch is the same for comparison.
N
= Recti f (L(innarr X)
Σ (x)) n=1
Similarly, we denote this type of TDRL with N rep-branch, L basic units per branch, and skip branch as R-DLWNS. Compared to this regular version, we will show that the pyramid structure exhibits better performance and diversity under similar model parameter sizes.
Figure T (c) shows the merging of the proposed linear re- prraamtttrrizttion module for inference: 1) intra-branch fusion: merge batch normalization within each unit via Equation T and merge all linear layers in each rep-branch to a single linear layer via Equation 3 2) inter-branch fusion: merge all branches via Equation The proposed TDRL can replace arbitrary linear layers in ViTs.
4 EXPERIMENTS
4.1 EXPERIMENTAL SETTINGS
Datasets. Similar to previous MIM methods (He et al. we pre-train our lightweight ViT models on ImageNet (Degg et al. 2009) which contains about 1.2M
2022 Wang al., 2023 Huang al., 2023).
When extreme value exists A, the softmar(() )-) function will scale elements to close 0 1, which results in the extremely small gradients (since ły/ jo y(1 y). y = softmaz(z)).
'The variance is scaled to N following the same independent random assumption as hefore.
5

---

Published as a conference paper at ICLR 2024
Table 1: Comparison with SOTA methods ImageNet validation. Teachers used by default for methods with pre-training, FT and P denote tine-tuning epochs and the size of backbone parameters respectively. fmeans performing distillation during fine-tuning.
Method
Network
Teacher
P (M)
Acc(%)
Without Pre-training
MobileNet-v3 (Howard et al.
2019)
CNNs
N/A
600
75.2
ConvNeXt-VI [Liu al.
20226
CNNs
N/A
600
77.5
VanillaNet-5 CChen et al. 2023)
CNNs
N/A
300
15.5
72.5
MobileViT-S (Mehta Rastegari
2021
Hybrid
N/A
300
78.3
EfficientVil-M33 (Liu et al. 2023
Hybrid
N/A
300
73.4
DeiT-Ti Touvron et al.
2021
ViTs
N/A
300
72.2
Manifold-T1 (Jia et al. 2021
VITs
CaiT-S24
6
75.1
MKD-Ti (Liu et al. [202aa]
ViTs
CaîT-S24
300
76.4
DeiT-Ti Touvron et al.
2021
VITs
RegNetY
300
6
74.5%
SSTA-11
Wu ct al.
2022a)
ViTs
DeÎT-S
300
75.2
ImageNet Pre-training
DMAE-Ti [Bai et al. 2023)
ViTs
viT-B
100
6
74.9
MAE-Lite Wang et al. 2023
ViTs
N/A
100
76.2
MAE-Tí He cl al. 2022)
VITs
N/A
200
6
75.2
TinyMIM-Ti RRe el al. 2023
VITs
TinyMIM-S
200
75.8
G2SD-Ti w/o S.D (Huang al. 2023
VITs
VIT-B
200
6
76.3
G2SD-Ti (Huang et al.
2023
VITs
VIT-B
200
77.0+
TDRL (ours)
VITS
VIT-B
200
78.3/78.61
MAE-Lite WWang et al.
2023
VITS
N/A
300
78.0
D-MAE-Lite (Wagg 2023)
VITS
VIT-B
300
6
78.4
TDRL (ours)
VITs
VIT-B
300
78.7/79.1
training images. We validate the performance downstream tasks including image classification on ImageNet(Deng et al., 2009), semantic image segmentation on ADE20K (Zhou et al.) 2019). object detection and instance segmentation on MS COCO (Lin et al. 2014).
Implementation Details. We mainly conduct experiments on applying the proposed TDRL to the projection layer of Q, K, v in multi- head self-attention (MHSA) and two linear layers in the feed-forward network (FFN). The configuration of TDRL remains consistent across all components. Considering both effectiveness and efficiency, the default recipe of TDRL is set P-W2S.
Following Wang et al. contains about 6M parameters. All blocks re-parameterized with TDRL. Due the recent SOTA methods (Wang et al. 2023 Huang et al. 2023) adopting MIM pipeline for taking full advantage of self-super learning, we also perform experiments based on them for a fair comparison. In pre- training. we use the MAE treatrind ViT-Base model as the teacher and perform generic distillation recipes as in Huang et al. (2023). More concretely, the student decoder contains Transformer blocks with an embedding dimension of 128. We align the outputs of the student decoder with the 4-th teacher decoder features including visible and masked patches to transfer generic knowledge. We use single linear layer to align the features of ViT-Tiny and its teacher model and ignore the class token in loss calculating. For optimization, we use AdamW optimizer (Loshchilov & Hutter) 2017) (with the initial learning rate 2.4e-3, weight decay 0.05, and batch size 4096) to train the model for 300 epochs. Images are randomly resized and cropped with a resolution 224x224.
2023). we conduct experiments on the classical vanilla ViT-Tiny which only
fine-tune pre-trained models for 200/300 epochs. For semantic seg- mentation, we replace the backbone of UperNet (Xiao et al. 2018) with ViT-Tiny and fine-tune the model for 160K iterations. We the BFiT(Bao et al. 202T semantic segmentation codehase. The input image resolution is 512x512. For object detection and instance segmentation tasks, we follow *The final projection layer in MHSA is ignored since its weight he merged with that of V.
For image classification,

---

Published as a conference paper at ICLR 2024
Table 2: Validation of semantic segmentation (ADE20K) and object detection tasks (MS COCO).
means that the resolution is 640x640 as in (Li et al. 2021 D. *means performing distillation during
fine-tuning. fmeans that the results based on Mask R-CNN (He et al. 2017
#Params (M)
Segmentation
Detection
Method
Seg/Det
mloU A Aboooe
A Apmask
Swin-T (Liu et al.
2021
59.9/47.8
44.5
46.0
41.6$
ConvNeXt-1 T
(Liu
2022b)
60.0/48.1
46.0
46.2
41.73
DINO-S Caron et al.
2021
42.0/44.5
44.0
49.1
43.3
iBOT-S Zhou et al.
2021
42.0/44.5
45.4
49.7
44.0
MAE-S (He et al.
2022)
42.0/44.5
41.1/44.97
45.3
40.8
MAE-Ti (e et al. 2022)
11.0/27.7
36.9/42.04
37.9/43.51
34.9/39.09
MAF-Lite Wang et al. 2023)
11.0/27.7
37.6
39.9*
35.4
D-MAE-Lite (Wagg et a..
2023
11.0/27.7
42.0
42.3*
37.4
G2SD-Ti (Huang et al. 2023)
11.0/27.7
41.4/44.5t
44.0/46.31
39.6/41.34
TDRL (ours)
11.0/27.7
42.5/45.24
46.5/47.47
41.5/42.11
the ViTDet (Li al. 2022) and use the detectron2 (Wu et al. 2019) codebase to train the model with 64 batch size for 100 epochs. The image resolution is 1024x 1024. If performing specific distillation in fine-tuning as Huang et al. (2023), the teacher is ViT- Base which achieves 83.6%. 48.1 mloU and 51.6 Abhooz on image classification, semantic segmentation and object detection. Re-parameterized architecture is used in the fine-tuning stage. More details are provided in the Appendix
4.2 COMPARISONS WITH SOTA METHODS
Image Classification. In Table L we summarize the detailed comparison of our TDRL with several types of SOTA methods image classification, including supervised methods (e.g., MobileNet- v3 (Howard et al., 2019) and ConvNetXt- V1-F (Liu et al. (0226)). self-supervised methods (e.g.. MAE-Ti (He et al. 2022 and some distillation methods with vanilla ViT-Tiny (e.g.. DMAE-Ti (Bai et al. 2023) and G2SD-Ti (Huang et al. 2023)). In general, our TDRL achieves the best classifica- tion accuracy under various epoch settings. For example, compared vanilla ViTs, it outperforms the best-performed one (e.g., G2SD and MAE-Lite) by 1.3% under 200 fine-tuning epochs and by 0.3% under 300 fine-tuning epochs. Compared carefully designed CNNs or hybrid networks, i.e. Mobile (Mehta & Rastegari 2021]. TDRL achieves 0.5% improvements. We also follow G2SD to perform specific distillation fine-tuning and find that our performance can be further improved to 79.1%.
Dense Prediction Tasks. Except for classification, we demonstrate the advance of pre-trained mod- els for dense prediction tasks, like segmentation, object detection and instance segmentation. As summarized in Table 2, TDRL achieves the best performance for dense prediction compared with other ViT-Tiny- A Apmask gains compared with MAE-Ti (He et al. 2022) and MAE-Lite (Wang et al. 2023). Com- pared with G2SD (Huang al.) 2023) which benefits from two-stage knowledge distil iistillation, TDRL achieves slightly better performance in object detection and instance segmentation without knowl- edge distillation during fine-tuning. What's more, compared with ViT-Small-based and elaborately designed CNNs/hybrid architectures, TDRL shows surprising performance to some extent. It out- performs MAE-S (He al. 2022) for both two tasks and shows superiority on all metrics (e.g.. 0.7 mloU) compared with Swin-T T (Liu et al.) (222I) which contains many inductive biases. We obverse that the performance without fine- fine-tuning distillation on ADE20K is not extraordinary as on MS COCO (i.e., its non-distillation performance on ADE20K is not close to G2SD with specific distil- lation), which may be caused by the data size. Insufficient data may require a well-learned teacher to guide the training, which results in the new best performance on ADE20K (i.e., 45.2 mloU). We provide an additional analysis on large-scale ImageNet in the Appendix B.I Compared to image classification, the improvement gaps from TDRL dense prediction tasks are considerably larger, which may attributed to task difficulty. TDRL can benefit better on complex dense prediction tasks than the elassification task by improving representation ability.
methods. Concretely, TDRL obtains more than 3.2 mloU, 3.9 1 pbbor and 3.1
7

---

Published as a conference paper at ICLR 2024
4.3 ABLATION STUDY AND ANALYSIS
In this section, systematically study the properties of the proposed TDRL.. Experiments mainly conducted on the ImageNet classification task. By default, we fine-tune the model for 100 epochs without teachers for efficiency. We give more experiments and analysis, such robustness evaluation and efficiency comparison, in the Appendix.
ImageNet Accuracy Gap
ImageNet Accuracy Gap (%)
Centered Kernel Alignment
ImageNet Accuracy Gup
(%)
(%)
1.2
0.8
0.5
0.6
0.8
0.4
0.6
0.4
-1.5
-1
0.4
100
1 101112
0.2
Parameters (M)
Blocks
R-DIW4 R-D4W1 R-D4W4
RR-DLWIS BOLLWSSRRDDWE WND
S.DIW3SSR-DIWE
PE
MHSA
FFN
= Skip
B awi Skip
P-WNS -P-NSS*
P-W2S
-P-W3S
= PP-22 FF-3S
(a)
(b)
(c)
(d)
Figure 2: (a) Effects of skip branch of TDRL in FFN. (b) Performance of various sizes for three
types of TDRL in FFN. The hybrid parameters of these variations are L and N. * means that keep the re- parameterized architecture in fine-tuning, (c) Comparison of CKA (Nguyen et al. 2020
similarity between rep-branches. (d) Comparison of embedding positions for ViTs.
Ablation on Re-parameterization. Here, we conduct the main ablation of TDRL when applying to Transformers, including the architecture design and re- -araameterized components in VITs.
1) Architecture Design. We first test the importance of skip branch through three variants in FFN, involving width and depth expansion. Skip branch performs shorter gradient propagation path compared to rep-branch to alleviate gradient vanishing. One can find that all variants suffer from non-negligible performance degradation without in Figure 2 (a). Then, validate the superiority of our pyramid structure compared to the regular version in terms of performance and module size. As shown in Figure (b), the pyramid-wise architecture (i.e., P-W NS) can achieve the best accuracy compared to other versions (e.g., R-DLWIS and R-DIW E-DIWNST NSS under similar parameter sizes. What's more, we compare the brter-branch diversity between our pyramid structure (P-WNS) and the width expansion version (R-DI W NS) through CKA similarity in Figure (c). It can be found that P-WNS shows a much richer representation ability than R-DIWIN NSW even under smaller number of rep- branches. We also compare the effects with or without TDRL in fine-tuning under P-WNS variants and find that keeping it can further stimulate the potential of the lightweight model. Considering both the effectiveness and efficiency, select P-W2S as the default settings.
2) Re-parameterized Components. Finally, we evaluate the effects of applying TDRL to different components of ViTs in Figure 2 (d). Except for Patch Embedding (PE), the introduction of TDRL in other components can bring significant improvements.
Attention Distribution Rectification. As analyzed before, ce-parameterization of self-attention will amplify the distribution changes, which may seriously affect the training stability. We ex- perimentally track the effects of distribution changes through maximum attention logits and their corresponding attention activation. In detail, we calculate the average maximum activation before and after softmax operation within each block ImageNet validation datasets (shown in Figure (a)-(b)). Without distribution rectification, the maximum value of their dot product is prone to ex- treme values for all the 12 Transformer blocks(over 1, 000), leading attention weights of near-zero entropy (i.e., almost one-hot attention map). In contrast, performing re-scale or normalization pre- vents divergence due to uncontrolled attention logit growth, which ensures that the behavior of the attention map is similar to that of the baseline. Accordingly, in Figure (c), we can clearly observe a rapid increase attention logits during the training process Compared with re-scale and normal-
'On experiments, we observe the NAN value after few thousand steps without normalization, which may he caused hy the data overflow.
8

---

Published as a conference paper at ICLR 2024
Logit Max
Attention Max
wio Rectify
Accuracy(%
1000
5000
0.8
76.8
76.79
100
4000
0.6
76.7
10
1000
76.5
76.48
76.4
1000
0.0
76
76.2
6
1011
6
8
10
11
1200
2000
1600
4800
76.1
Blocks
Blocks
76
Initial
Raseline
-nnital
Roseliie
Roctify
Ierations
75.9
-Nom
Scale
-Norm
Scall
Allention Max Logii Mas
Buseline Scale II Norm
(a)
(b)
(c)
(d)
Figure 3: (a)-(b) Effects of feature distributions Q/K re-parameterization with the recipe of R-DIW4. The horizontal axis represents the index of blocks. Initial represents the random initial- ization. The results without rectification come from the last checkpoint before the collapse, while others come from the 300-epoch pre-trained checkpoint. The logit value has been scaled by 1/ L/vCa. (c) The trend of logits maximum (blue) and attention maximum (orange) during training. The hori- zontal axis represents the steps. (d) The Top-1 ImageNet classification accuracy of different settings.
ization, we find that normalization shows superiority in terms of ImageNet classification accuracy (in Figure (d)). We provide more analysis about FFN in the Appendix B.2
Generality of TDRL. In addition to applying TDRL in vanilla ViT-Tiny, also apply it other networks to show its generic ability. As summarized in Table 3 we first validate the effect of TDRL in a slightly larger model (i.e., ViT-Small and Swin-Ti (Liu et al. 202I)). then expand the experiments on the recent lightweight CNN, hybrid networks (VanillaNet-5 (Chen et al. (023) and Mobileone-SO (Vasu al., 2023b)), and image generation models (e.g.. DDPM (Ho et al.) 2020 ) It can be found that all these methods benefit from the proposed TDRL, indicating that our proposed TDRL suitable for various networks on different tasks.
Table 3: Applications of TDRL on various networks and different tasks. For ViT-Small, we fol- low the same pre-training recipe and fine-tune it for 100 epochs without distillation. The model is re-parameterized fine-tuning. For other networks, we use official codes and replace the corre- sponding linear layers with the proposed TDRL. For image generation, we conduct experiments on Cifarl0 (Krizhevsky et al. 2009).
Classification Accuracy (%) f Swin-Ti
Image Generation FID + DDPM
TDRL
ViT-Small
Mobileone-S5
VanillaNet-5
80.8
76.2 78.2 (+2.0)
71.3 75.1 (+3.8)
71.1
10.4
81.3 (+0.5)
71.5 (+0.4)
9.2 (+1.2)
CONCLUSION
This paper explores the potential of boosting vanilla lightweight ViTs via ce-parametrrzaatonn. To enhance the representation ability linear layers in ViTs, we propose multi-branch pyramid archi- tecture (TDRL) with branches consisting of various depths of linear layers and batch normalization. What's more, we discover and alleviate the distribution explosion problem when applying the pro- posed TDRL to Vision Transformers by distribution rectification. Experiments show that our TDRL can efficiently improve the performance lightweight ViTs as well as other transformer or hybrid networks.
Limitations and societal impact. Similar to previous re- re-parametrrized methods, our TDRL im- proves performance without compromising inference efficiency. But it results in extra training costs for larger capacity. However, existing models can still benefit from TDRL under similar training costs (see B.6). We hope work can promote the research lightweight VIT' in the future.

---

Published as a conference paper at ICLR 2024
REPRODUCIBILITY STATEMENT
Our proposed method, TDRL, is lightweight module whose PyTorch-style implementation is pro- vided in the supplementary materials, The pre-training and fine-tuning settings can be found in the Appendix A. A.2
REFERENCES
Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E Hinton. Layer normalization. arXiv preprint (6l.:6607.06400, 2016.
Yutong Bai, Zeyu Wang, Junfei Xiao, Chen Wei, Huiyu Wang, Alan L Yuille. Yuyin Zhou, and Cihang Xie. Masked autoencoders enable efficient knowledge distillers. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 24256-2444 24265, 2023.
Randall Balestriero, Mark Ibrahim, Vlad Sobal, Ari Morcos, Shashank Shekhar, Tom Goldstein, Florian Bordes, Adrien Bardes, Gregoire Mialon, Yuandong Tian, et al. A cookbook of self- supervised learning. arXiv preprint ar Xrvv::2304. /22/0, 2023.
Hangbo Bao, Li Dong, Songhao Piao, and Furu Wei. Beit: Bert pre-training of image transformers. arXiv preprint arXivv2106.008244 2021.
Jinming Cao, Yangyan Li, Mingchao Sun, Ying Chen, Dani Lischinski, Daniel Cohen-Or, Baoquan Chen, and Changhe Tu. Do-conv: Depthwise over- wer-parameterized convolutional layer. IEEE Transactions on Image Processing, 31 11:3726-3736, 2022.
Mathilde Caron, Hugo Touvron, Ishan Misra, Hervé Jégou, Julien Mairal, Piotr Bojanowski, and Armand Joulin. Emerging properties in self- vision transformers. In Proceedings of the IEEE/CVF international conference on computer vision, pp. 9650-9660, 2021.
Hanting Chen, Yunhe Wang. Jianyuan Guo, and Dacheng Tao. Vanillanet: the power of minimalism in deep learning. arXiv preprint arXiv:2305. /2972, 2023.
Xinlei Chen, Saining Xie, and Kaiming He. empirical study of training ceff-supervied vision transformers. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pp- 9640-9649, 2021.
Yinpeng Chen, Xiyang Dai, Dongdong Chen, Mengchen Liu, Xiaoyi Dong, Lu Yuan, and Zicheng Liu. Mobile-former: Bridging mobilenet and transformer. In Proceedings of the TEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 5270-5279, 2022.
Ekin D Cubuk, Barret Zoph, Jonathon Shlens, and Quoc V Le. Randaugment: Practical automated data augmentation with reduced search space. In Proceedings of the IEEE/CVE conference on computer vision and pattern recognition workshops, pp. 702-703, 2020.
Mostafa Dehghani, Josip Djolonga, Basil Mustafa, Piotr Padlewski. Jonathan Heek, Justin Gilmer, Andreas Steiner, Mathilde Caron, Robert Geirhos, Ibrahim Alabdulmohsin, al. Scaling vision transformers to 22 billion parameters. arXiv preprint 230.:23300004422 2023.
Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A large-scale hi- erarchical image database. In 2009 IEEE conference computer vision and pattern recognition, pp. 248 255. leee, 2009.
Xiaohan Ding, Yuchen Guo, Guiguang Ding, and Jungong Han. Acnet: Strengthening the kernel skeletons for powerful enn via asymmetric convolution blocks. In Proceedings of the IEEE/CVF international conference on computer vision, pp. 1911-1920, 2019.
Xiaohan Ding, Xiangyu Zhang, Jungong Han. and Guiguang Ding. Diverse branch block: Building a convolution as an inception like unit. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 10886-10895. 2021a.
10

---

Published as a conference paper at ICLR 2024
Xiaohan Ding, Xiangyu Zhang, Ningning Ma. Jungong Han, Guiguang Ding, and Jian Sun. Repvge: Making vgg-style convnets great again. In Proceedings of the IEEE/CVE conference on computer vision and pattern recognition, pp.
2021b.
Xiaohan Ding, Xiangyu Zhang, Jungong Han, and Guiguang Ding. Scaling up your kernels to 31x31: Revisiting large kernel design in cmns. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 11963-11975, 2022.
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenbomn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al. An image is worth 16x16 words: Transformers for image recognition at scale. arXiv preprint arXiv:2010.1 /929, 2020.
Yuxin Fang, Wen Wang, Binhui Xie, Quan Sun, Ledell Wu, Xinggang Wang, Tiejun Huang, Xinlong Wang, and Yue Cao. Eva: Exploring the limits of masked visual representation learning at scale. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp- 2023.
Shuxuan Guo, Jose M Alvarez, and Mathieu Salzmann. Expandnets: Linear over-paramettriizaton to train compact convolutional networks. Advances in Neural Information Processing Systems, 33:1298-1310 2020.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Delving deep into rectifiers: Surpassing human-level performance imagenet classification. In Proceedings of the IEEE international conference on computer vision, pp. 1026-1034, 2015.
Kaiming He, Georgia Gkioxari, Piotr Dollár, and Ross Girshick. Mask r-enn. In Proceedings of the IEEE international conference on computer vision, pp. 2961-2969, 2017.
Kaiming He, Xinlei Chen, Saining Xie, Yanghao Li, Piotr Dollár, and Ross Girshick. Masked au- toencoders are scalable vision learners. In Proceedings of the IEEE/CVF Conference Computer Vision and Pattern Recognition, pp. 16000-16009, 2022.
Dan Hendrycks, Steven Basart, Norman Mu, Saurav Kadavath, Frank Wang, Evan Dorundo, Rahul Desai, Tyler Zhu, Samyak Parajuli, Mike Guo, al. The many faces of robustness: A critical analysis of out-of-distributtion generalization. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pp. 8340-8349 2021a.
Dan Hendrycks, Kevin Zhao, Steven Basart, Jacob Steinhardt, and Dawn Song. Natural adver- sarial examples. In Proceedings of the IEEE/CVF Conference Computer Vision and Pattern Recognition, pp. 15262-15271 2021b.
Geoffrey Hinton, Oriol Vinyals, and Jeff Dean. Distilling the knowledge in neural network. arXiv preprint arxiv:1503 550.00233, 2015.
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in neural information processing systems, 33:6840-688 6851, 2020.
Andrew Howard. Mark Sandler, Grace Chu, Liang-Chice Chen. Bo Chen, Mingxing Tan, Weijun Wang, Yukun Zhu. Ruoming Pang, Vijay Vasudevan, et al. Searching for mobilenetv3. In Pro- ceedings of the IEEE/CVE international conference on computer vision, pp. 1314-1324, 2019.
Wei Huang, Zhiliang Peng, Li Dong, Furu Wei, Jianbin Jiao, and Qixiang Ye. Generic-to sp-seccf distillation of masked autoencoders. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 15996-16005, 2023.
Sergey loffe and Christian Szegedy. Batch normalization: Accelerating deep network training by reducing internal covariate shift. In International conference on machine learning, pp. 448-456 pmlr, 2015.
Ding Jia, Kai Han, Yunhe Wang, Yehui Tang, Jianyuan Guo, Chao Zhang, and Dacheng Tao. Ef- ficient vision transformers via fine- aeee-graned manifold distillation. arXiv e-prints, pp. arXiv-2107, 2021.
11

---

Published as a conference paper at ICLR 2024
Alex Krizhevsky, Geoffrey Hinton, et al. Learning multiple layers of features from tiny images. 2009.
Siddharth Krishna Kumar. On weight initialization in deep neural networks. arXiv preprint arXiv:1704.0886. 088863.. 2017.
Junnan Li, Dongxu Li, Silvio Savarese, and Steven Hoi. Blip-2: Bootstrapping language- image pre-training with frozen image encoders and large language models. arXiv preprint arXiv:2301. 12009 2023.
Yanghao Li, Saining Xie, Xinlei Chen, Piotr Dollar, Kaiming He, and Ross Girshick. Benchmarking detection transfer learning with vision transformers. arXiv preprini arXiv:2111.11429) 2021.
Yanghao Li, Hanzi Mao, Ross Girshick, and Kaiming He. Exploring plain vision transformer back- bones for object detection. In Computer Vision-ECCC 2022: 17th European Conference, Tel Avix, Israel October 23-27, 2022, Proceedings, Part IX, pp. 280-296. Springer, 2022.
Tsung-Yi Yi Lin, Michael Maire, Serge Belongie, James Hays, Pietro Perona, Deva Ramanan, Piotr Dollár, and C Lawrence Zitnick. Microsoft coco: Common objects context. In Computer EssionEECV 2014: 13th European Conference, Zurich, Switzerland, September 6-12, 2014, Proceedings, Part 13, pp. 740-755. Springer, 2014.
Jihao Liu, Boxiao Liu, Hongsheng Li, and Yu Liu. Meta knowledge distillation. arXiv preprint arXiv:2202.07940, 2022a.
Xinyu Liu, Houwen Peng. Ningxin Zheng, Yuqing Yang, Han Hu, and Yixuan Yuan. Efficientvit: Memory efficient vision transformer with cascaded group attention. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 14420-14430, 2023.
Ze Liu, Yutong Lin, Yue Cao, Han Hu, Yixuan Wei, Zheng Zhang, Stephen Lin, and Baining Guo. Swin transformer: Hierarchical vision transformer using shifted windows. In Proceedings of the IEEE/CVF international conference on computer vision, pp. 10012-10022, 2021.
Zhuang Liu, Hanzi Mao, Chao-Yuse Wu, Christoph Feichtenhofer, Trevor Darrell, and Saining Xie. A convnet for the 2020s. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pp. 11976-11986, 2022b.
Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. arXiv preprint 17Xiv:1111.05101, 2017.
Sachin Mehta and Mohammad Rastegari. Mobilevit: lighh- general-purpose and mobile- friendly vision transformer. arXiv preprint arXiv:2110.02178, 2021.
Thao Nguyen, Maithra Raghu, and Simon Kornblith. Do wide and deep networks learn the same things? uncovering how neural network representations vary with width and depth. arXiv preprint N0N.:20000..322,, 2020.
Zhiliang Peng, Li Dong, Hangbo Bao, Qixiang Ye, and Furu Wei. Beit v2: Masked image modeling with vector-quantized visual tokenizers. arXiv preprint arXiv:2208.06688 06366, 2022.
Zhiliang Peng, Zonghao Guo, Wei Huang, Yaowei Wang, Lingxi Xie, Jianbin Jiao, Qi Tian, and Qix- iang Ye. Conformer: Local features coupling global representations for recognition and detection. IEEE Transactions on Puttern Analysis and Machine Intelligence, 2023.
Alec Radford, Jong Wook Kim, Chris Hallacy, Aditya Ramesh, Gabriel Goh, Sandhini Agarwal, Girish Sastry, Amanda Askell, Pamela Mishkin, Jack Clark, et al. Learning transferable visual models from natural language supervision. In International conference machine learning, pp- 8748-8763, PMLR, 2021.
Benjamin Recht, Rebecca Roelofs, Ludwig Schmidt, and Vaishaal Shankar. Do imagenet classifiers generalize to imagenet? In International conference on machine learning, pp. 5389-548840 5400. PMLR, 2019.
12

---

Published as a conference paper at ICLR 2024
Sucheng Ren, Fangyun Wei, Zheng Zhang, and Han Hu. Tinymim: An empirical study of distilling mim pre- pre-trained models. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 3687-3697, 2023.
Zhentao Tan, Yue Wu, Qiankun Liu, Qi Chu, Le Lu, Jieping Ye, and Nenghai Yu. Exploring the application of large-scale pre-trained models adverse weather removal. arXiv preprint arXiv:2306.00000 2023.
Hugo Touvron, Matthieu Cord, Matthijs Douze, Francisco Massa, Alexandre Sablayrolles, and Hervé Jégou. Training data-efficient image transformers & distillation through attention. In International conference on machine learning, pp. 10347-10357. PMLR, 2021.
Pavan Kumar Anasosalu Vasu, James Gabriel, Jeff Zhu, Oncel Tuzel, and Anurag Ranjan. Fastvit: A fast hybrid vision transformer using structural reparameterization. arXiv preprint
arXiv:2303. 14189, 2023a.
Pavan Kumar Anasosalu Vasu, James Gabriel, Jeff Zhu, Oncel Tuzel, and Anurag Ranjan. Mo- bileone: An improved one millisecond mobile backbone. In Proceedings of the IEEE/CVF Con- ference Computer Vision and Pattern Recognition, pp. 7907-7917, 2023b.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural informa- tion processing systems, 30, 2017.
Haohan Wang, Songwei Ge, Zachary Lipton, and Eric P Xing. Learning robust global representa- tions by penalizing local predictive power. Advances in Neural Information Processing Systems, 32, 2019.
Shaoru Wang, Jin Gao, Zeming Li, Xiaoqin Zhang, and Weiming Hu. A closer look at self- supervised lightweight vision transformers. In International Conference on Machine Learning, pp. 35624-35641. PMLR, 2023.
Chen Wei, Haoqi Fan, Saining Xie, Chao- Wu, Alan Yuille, and Christoph Feichten- hofer. Masked feature prediction for self-supervised visual pre-training. In Proceedings of the IEEE/CVF Conference Computer Vision and Pattern Recognition, pp. 14668-1680 14678, 2022.
Sanghyun Woo, Shoubhik Debnath, Ronghang Hu, Xinlei Chen, Zhuang Liu, So Kweon, and Saining Xie. Convnext v2: Co-designing and scaling convnets with masked autoencoders. arXiv preprint arXiv:2301.00800 2023.
Haiyan Wu, Yuting Gao, Yinqi Zhang, Shaohui Lin, Yuan Xie, Xing Sun, and Ke Li. Self-supervised models are good teaching assistants for vision transformers. International Conference on Ma- chine Learning, pp. 24031-24042. PMLR, 2022a.
Kan WU, Jinnian Zhang, Houwen Peng. Mengchen Liu, Bin Xiao, Jianlong Fu, and Lu Yuan. Tinyvit: Fast pretraining distillation for small vision transformers. In European Conference Computer Vision, pp. 68-85. Springer, 2022b.
Yuxin Wu, Alexander Kirillov, Francisco Massa, Wan Yen Lo, and Ross Girshick. Detectron2. 2019.
Tete Xiao, Yingeheng Liu, Bolei Zhou, Yuning Jiang, and Jian Sun. Unified perceptual parsing for scene understanding. In Proceedings the the European conference on computer vision (ECCV), pp. 418-434, 2018.
Haiyang Xu, Qinghao Ye, Ming Yan, Yaya Shi, Jiabo Ye, Yuanhong Xu, Chenliang Li, Bin Bi, Qi Qian, Wei Wang, al. mplug-2: A modularized multi-moda. foundation model across text, image and video. arXiv preprini arXiv:2302.004400 2023.
Zhuliang Yao, Yue Cao, Yutong Lin, Ze Liu, Zheng Zhang, and Han Hu. Leveraging batch nor- malization for vision transformers. In Proceedings of the IEEE/CVF International Conference on Computer Vision pp. 413-422, 2021.
13

---

Published as a conference paper at ICLR 2024
Syed Waqas Zamir, Aditya Arora, Salman Khan, Munawar Hayat, Fahad Shahbaz Khan, and Ming- Hsuan Yang. Restormer: Efficient transformer for high-resolution image restoration. Proceed- ings the IEEE/CVF conference computer vision and pattern recognition, pp. 5728-5739, 2022.
Xiaohua Zhai, Alexander Kolesnikov, Neil Houlsby, and Lucas Beyer. Scaling vision transformers. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp- 12104-12113, 2022.
Xiaosong Zhang. Feng Liu, Zhiliang Peng, Zonghao Guo, Fang Wan, Xiangyang Ji, and Qixiang Ye. Integral migrating pre-trained transformer encoder-decoders for visual object detection. arXiv preprint ar arXiv:2205.09660 2022a.
Xinyu Zhang, Jiahui Chen, Junkun Yuan, Qiang Chen, Jian Wang, Xiaodi Wang, Shumin Han, Xiaokang Chen, Jimin Pi, Kun Yao, et al. Cae v2: Context autoencoder with clip target. arXiv preprint arXiv:2211.0977 2022b.
Bolei Zhou, Hang Zhao, Xavier Puig, Tete Xiao, Sanja Fidler, Adela Barriuso, and Antonio Torralba. Semantic understanding of scenes through the ade20k dataset. International Journal of Computer Vision, 127:302-321, 2019.
Jinghao Zhou, Chen Wei, Huiyu Wang, Wei Shen, Cihang Xie, Alan Yuille, and Tao Kong. ibot: Image bert pre-training with online tokenizer. arXiv preprint arXiv:2111008332 2021.
14

---

Published as a conference paper at ICLR 2024
A EXPERIMENTAL DETAILS
A.1 PRE-TRAINING DETAILS
Our self- sefffsuperviee pre-traiin pre-training strategy follows the recent popular MAE (He al. 2022). including the optimizer, learning rate, batch size, mask ratio, etc. As for models, we use ViT-TTny as the encoder and replace its linear layers in MHSA/FFN with the proposed TDRL. By default, we re- parameterize all blocks of ViT-Tiny. Following MAE-Lite (Wang et al. 2023), we set the number of heads in ViT-Tiny as 12. In the decoder, we use blocks with an embedding dimension of 128. The teacher model is MAE pre-trained ViT-Base provided by the official repository We use an additional linear layer to align the last decoder features from the student and the 4-th decoder features from the teacher and calculate the loss for both visible and invisible patches.
A.2 FINE-TUNING DETAILS
To evaluate the effectiveness of the proposed TDRL, we fine-tune the pre-trained models on three mainstream tasks, including classification, semantic segmentation, object detection and instance segmentation tasks.
Table 4: Fine-tuning settings of ViT-Tiny for ImageNet classification.
Config
Value (w/o distillation)
Value (w distillation)
Teacher
N/A {5. 5, 10} {100, 200, 300}
ViT-Base
Warmup epochs
{5. 5, 10) (100, 200, 300)
Training epochs
Layer-wise Ir decay
0.85 (w/o TDRL), 0.65 (TDRL)
0.75 (w/o TDRL), 0.65 (TDRL)
Optimizer
AdamW
Base learning rate
le-3
Weight decay
0.05
Optimizer momentum
31, B2 0.9, 0.999
Batch size
1024
Learning rate schedule
Cosine decay
Augmentation
RandAug(10, 0.5) CCbuu et al. 2020
Colorjitter
0.3
Label smoothing
0
Mixup, Cutmix
0.2, 0
Drop path
0
Image Classification. We follow previous work (Huang et al. 2023 Wang et al. 2023) set the fine-tuning recipes and summarize them in Table 4 The difference between using TDRL and merging TDRL in the fine-tuming stage comes from the layer decay (e.g., 0.65 vs. 0.85). To evaluate the effect of these two settings for the baseline, we also fine-tune it with 0.65 layer decay and find that the performance similar to the original one. Thus, the improvements indeed come from our proposed TDRL, rather than the fine-tuning recipes. When using ViT-Small as the target, the augmentation recipes are the same as Huang et al. (0023).
Semantic Segmentation. In this experiment, 2023) and follow its settings. Differently, we also change the layer decay when using TDRL in fine-tuning. In detail, we set layer decay to 0.80 when using specific distillation, otherwise set it to 0.75.
codebase provided by G2SD (Huang et al.
Object Detection and Instance Segmentation. Following G2SD (Huang et al. 2023). fine-tune the model for 100 epochs with a batch size of 64. The layer decay is set to 0.7 by default. We set the learning rate 3c-f if distillation is not used, otherwise set it to le-4.
15

---

Published as a conference paper at ICLR 2024
A.3 MORE IMPLEMENTATION DETAILS OF TDRL
When applying TDRL to the Patch Embedding, we should make some modifications. Formally, PE performs k k convolution with & stride encode image patches ( with size of k k) independently. To combine with the proposed TDRL, we replace the first linear layer of each branch (containing rep-branch and skip branch) with a convolution layer. At inference time, batch normalization can be converted into convolution followed by Ding et al. (202Tb) And we thereafter merge convolution and linear as follows:
Coul
Cout
W JJn
W W 1,k W.gin by
Σ
W 1g +
where W E R. X,, X 6 E R W' E RC. and b' E RoCt are the weights and biases of convolution and linear. And Cout C is the prerequisite. By the way, this way can be used when combining our TDRL with other convolutions.
A.4 MORE DETAILS OF APPLYING TDRL TO OTHER MODELS
Here, we give more details when applying TDRL to different models which are summarized in Table 3 For ViT-Small, we follow the same settings as Vin-Tiny. TDRL is used in both FFN and MHSA The fine-tuning settings are the same as MAE (He et al. 2022). For Swin-Ti (Liu et al. 2021 we replace the linear layers with the proposed TDRL in the FFN. We train the Swin Ti with or without our TDRL on ImageNet directly for 100 epochs. For Mobileone (Vasu et al. 2023b 8
which has already convolution-basee re-parameterized modules with our TDRL and also combine TDRL with 3 x 3 convolution re-parameterized modules. To fuse 3 x 3 convolution and linear layer, we can follow the Equation 8 For VanillaNet Chen et al. 2023
and 1 x 1 convolution- based re- paraamterrization, we replace its
replace its two sequential 1 x convolutions with the proposed TDRL. The batch size is set to 512. For DDPM (Ho et al. 2020 10
we replace its 1 x 1 convolutions within all the self-attention blocks with our TDRL. All the results summarized in Table 3 Jrr reproduced by ourselves. In these experiments, the default setting of TDRL is P- P-W2S.
B MORE ANALYSIS
B. DISTILLATION: ONE-STAGE VS. TWO-STAGE
In Table 2 we find that TDRL may still need a well-learned teacher in fine-tuning to achieve the SOTA performance when data is insufficient (e.g., ADE20K (Zhou et al. 2019)). which can be alleviated by increasing the training data (e.g., MS COCO (Lin et al. 20144). Here, we further com- pare the ote-stage distillation and the two-stage distillation a large-scale ImageNet classification dataset (Deng et al. 2009) in Table One can see that our proposed TDRL shows superiority com- pared the baseline either with specific distillation or without specific distillation. More concretely, the improvement in terms of accuracy is at least larger than 0.97%. Compared to the baseline with specific distillation, our TDRL still outperforms 0.62% without the specific distillation, indicating its advantage can be stimulated by enough data.
B.2 MORE ANALYSIS FOR DISTRIBUTION RECTIFICATION
We have analyzed the impacts of distribution for Attention calculation before. Here, we give more discussion about it for other components of VITs. Due to the pre-normalization mechanism, the dis- tribution changes of features will accumulate through the skip connection within FFN and MHSA. That is, distribution rectification may also be important when applying TDRL FFN the V pro- jection in MHSA. To evaluate it, we test the accuracy gap between the models with and without distribution rectification. As shown in Figure 4 the performance gap increases from around zero to
"https//gituu..commmiiii http://hww...oooiooooooooooanaae
16

---

Published as a conference paper at ICLR 2024
Table 5: Comparison ImageNet with without specific distillation in fine-tuning. All models are fine-tuned with 100 epochs. Re-parameterized architecture is kept in fine-tuning.
Method
Specific Distillation
Accuracy (%)
A (%)
Baseline
76.24
TDRL (ours)
X
77.39
+1.15
Baseline
76.77
TDRL (ours)
77.74
+0.97
O. 1818 with the increasing of rep-branchs in FFN. It indicates that the greater the change in distri- bution, the more significant the corrective effect. In addition, we compare the difference between applying TDRL only in FFN and both in FFN and MHSA (0.13% 0.21%). It can be found that the effect distribution rectification also proportional to the number of layers applied to TDRL.
ImageNel Accurucy Gup (%)
0.25
0.15
0.1
-0.05
2 Number of rep-brunch N' for P-WNS OFN OFN+*MHA
Figure 4: Comparison of distribution rectification. The values are the accuracy gap between the models with and without distribution rectification. We use P-WNS as the configuration of TDRL. All models are fine-tuned for 100 epochs with re- parameterization.
Table 6: Robustness comparison. "IN" is short for ImageNet.
Method
IN
IN-A
IN-R
IN-S
IN-V2-F
IN-V2-Thr
IN-V2-Top
G2SD-Ti
77.0
12.9
39.0
25.9
65.6
D-MAE-Lite
78.4
13.9
40.6
28.0
66.7
74.9
80.1
TDRL (ours)
78.7
14.7
41.4
28.1
67.1
75.5
80.3
B.3 ROBUSTNESS EVALUATION
We evaluate the robustness by directly testing these ImageNet-trained methods on several ImageNet variants, including ImageNet-A (Hendrycks et al. 2021b), ImageNet-R (Hendrycks al. 202 ImageNet-S (Wang et al. 2019) and ImageNet- V2 (Recht et al., 2019). In Table [6] we can see that TDRL outperforms other methods on all test sets, which implies that our method hold the generalization capability while boosting the downstream task performances.
B.4 COMPARISON WITHOUT PRETRAINING
Here, we evaluate the efficiency of our proposed TDRL without MIM pertaining. Specifically, directly train ViT-Tiny and our TDRL ImageNet for 100 epochs. ViT-Tiny achieves 63.5% accuracy, and our TDRL increases the accuracy to 65.9%. This proves that our TDRL does not rely on pretraining and distillation.
17

---

Published as a conference paper at ICLR 2024
Table 7: Comparison of our TDRL under various fine-tuning epochs.
Fine-tuning Epochs
100
200
1,000
Accuracy (%)
77.7
78.6
79.9
Table 8: Efficiency comparison between pre-training and inference on V100 GPUs. the pre- training. the batch size per GPS is set to 256. And the inference batch size is 128. P is the learnable parameters and FLOPs denotes the computational complexity. Values in the (-) denote the propor- tion of increase compared to the baseline (i.e., G2SD-Ti (Huang et al. 2023
Pre-training Epoch Times (s)
Inference Speed (s/iteration)
Method
Memory (G)
P (M)
FLOPs (G)
Baseline
12.57
326 462 (+41.7%)
5.72
1.26
0.35 0.35 (+0%)
TDRL (ours)
18.22 (+44.9%)
48.86
9.72
B.5 LONG EPOCHS OR LARGE MODELS
To evaluate the effect of the training epoch, we fine-tune our TDRL for different epochs. As sum- marized in Table 7 our proposed TDRL can be beneficial for larger epochs. We further explore the potential of our TDRL for large models, such as DeiT-B [Touvron et al. 2021). In detail, we directly train DeiT-B with or without TDRL on ImageNet for 100 epochs and find that TDRL can still improve the accuracy by 0.6% for large models. Note that we only adopt TDRL in Q, K V of DeîT-B to reduce training costs. In addition, we find that applying our TDRL to a larger network may require stronger regular constraints (e.g., weight decay) during training.
B.6 EFFICIENCY COMPARISON
Although TDRL improves the model capacity of ViT', it brings additional optimized parameters. Here, we summarize the training cost and inference speed in Table It can be found that the cost increase of memory and training times does not exceed 50%. And the improvement in training pa- rameters and computational complexity is relatively significant. In the inference stage, our inference speed is fast the baseline. For fair comparison, the total pre- pre-training time the same between the baseline and our TDRL. As summarized in Table our TDRL still outperforms the baseline (i.e., G2SD) by 0.59% terms of image classification ac- curacy, which indicates our superiority. In addition, we can flexibly select the recipes of TDRI. in terms of module size and replacement places to balance the training cost and test performance. To validate it, compare the trend of image classification accuracy and pre-training times in Figure As the pre-training cost increases, we can efficiently improve the performance in classification.
reduce the pre-training epochs to keep
B.7 MORE COMPARISON WITH NAIVE DESIGNS
Here, compare proposed TDRL with the naive version that directly converts convolutions to linear layers in Figure [ (a). As summarized in Table IO we can find that our proposed TDRL shows better performance than the naive version under similar parameter conditions. For example, TDRL outperforms the naive version by 0.64% at around 30 MB training parameters. It further demonstrates the superiority of our proposed method.
B.8 OVERFITTING ANALYSIS
As our TDRL is much larger than a single linear layer, one of the concerns of practical application may be listed in the overfitting, In experiments, have not observed overfitting issues. The reasons may be as follows: 1) The datasets we used are relatively big for ViT-Tiny, even with our proposed TDRL; 2) Batch normalization in TDRL not only improves the training representation but
18

---

Published as a conference paper at ICLR 2024
Table 9: Comparison of our TDRL and baseline under the similar pre-training time. The configura- tion of TDRL is P-WIS applied in FFN. The fine-tuning efficiency of the bascline and our method the same since we merge the re-parameterized architecture after pre-training.
Method
Pre-train time (hours) Pre-train epoch
Fine-tune epoch
Accuracy (%)
Baseline (G2SD)
32.33
300
100
76.24
TDRL (ours)
32.02
220
100
76.73 (+0.59)
Image Chassifictioon Accuracy (%)
Image Classification Accuracy (%)
Image Classification Accuracy vs. Times iteration (6)
vs, Parameters (M)
vs. Training Memory (G)
76.5
76.5
76.5
76.4
76.4
76.4
76.3
76.3
76.3
76.2
76.2
76.2
5.5
6.5
6.9
6,95
7.05
7.1
0.3
0.35
0-4
0.45
6 Bascline
P-WIS-V-3
Baseline
P-WIS-V.3
&P-WIS.V.6
P-11S-V-12
0 P.W1S--
P-W1S-V-12
P-WIS PPWISS-V-1
Figure 5: Comparison of image classification accuracy and pre-training efficiency. " "-V-N" repre- sents that the TDRL is applied in V within MHSA for the first N blocks. The pre-training epoch is set to 300, while all models are fine-tuned for 100 epochs without re- ee-paammtterieed architectures.
also reduces the risk of overfitting: 3) The feature dimension along the network is not changed (fea- tures are limited to the original dimension before outputting from TDRL), resulting in the increased intrinsic dimension. This also reduces the risk of overfitting compared to directly increasing the depth of the network the feature dimension. We further train ViT-Tiny with without TDRL a small dataset (i.e., Cifar10 (Krizhevsky et al.) 2009)) which contains 50,000 training images. The results indicate that overfitting still does not occur. TDRL still improves the performance of ViT- Tiny by 1.0% (70.9% vs. 69.9%). When the size of datasets is too small for the network, we may face overfitting issues. However, considering the rapid development of data size and our lightweight model research targets, the probability of overfitting in practical applications is very low.
B.9 MORE ANALYSIS FOR DENSE PREDICTION TASKS.
To demonstrate the gains resulting from our re- ce-parameerrizatoo structure rather than the additional batch normalization layers, conduct experiments normalization to the MHSA and FFN. We observe that adding only the batch normalization layer does not bring any effective improvement (37.8 vs. 41.4 in terms of mloU). This experiment further validates the effectiveness of our re- e-parameterizationn structure.
semantic segmentation that apply batch
19

---

Published as a conference paper at ICLR 2024
Table 10: Comparison of our TDRL and naive version that directly converts convolutions to linear layers in the typical CNN-based re- --paaameterized module. The re-parameterized modules are ap- plied in FFN. The naive version contains 8 branches, while TDRL set to R-D2WIS for similar parameters.
Method
Trainin Parameters (M) Accuracy (%)
Baseline
5.72
76.24
Naive
30.75
76.26 (+0.02)
TDRL
31.03
76.90 (+0.66)
20


Reviews:
[{"review_id": "d9L2JESyqd", "reviewer": "Reviewer_HAf5", "summary": "This paper aims to improve the potential of vanilla lightweight ViTs by exploring the adaptation of the re-parameterization technology to ViTs for improving learning ability during training without increasing the inference cost. The authors propose to incorporate the nonlinear  ensemble into linear layers. They increase the depth of the linear layers during the re-parameterization training stage. And they re-scale the features with $\\sqrt(N)$ rather than normalize it in Scaled Dot-Product Attention to solve a transformer-specific distribution rectification.", "questions": "In Sec 3.2, the authors claim that \"The linear operation in transformers focuses more on intra-token semantic mining rather than inter-token spatial exploration. Linear stacking is inherently appropriate to transformer networks.\" I think the authors should explain this in more details. I am confused why Linear stacking is inherently appropriate to transformer networks.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "-1- The motivation of this paper is clear.\n\n-2- The pyramid multi-branch topology, i.e., n-th rep-branch with $L_n$ basic units is interesting.", "weaknesses": "-1- In the Sec 3.2 Main Architecture, authors use P-WNS to denote the architecture with Width of N rep-branch, what about the different $L_N$ in every rep-branch in the P-WNS? And there is only 1$\\times$1 (without K$\\times$K) in linear projection, the accuracy gaps in different configurations of \"pyramid multi-branch topology\" seems minor (Figure 2a).\n\n-2- About the baseline performance. Table 1 reports the performance improvement on ViT-B, with the distillation phase in pre-training. Why the proposed TDRL should be combined with distillation? Why not use a stronger and simpler baseline such as DeiT-B (can obtain 81.8 top-1 accuracy on ImageNet)?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "7mvPoiDUtd", "reviewer": "Reviewer_WZef", "summary": "This paper addresses the problem when using re-parameterization to boost vanilla lightweight ViTs. The main problem is that vanilla Transformer architectures are mainly comprised of linear and layer normalization layers, which do not complement with re-parameterization. This paper incorporates the nonlinear ensemble into linear layers. It also discovers and solves a new transformer-specific distribution rectification problem. It performs lots of experiments in order to verify its effectiveness.", "questions": "See weakness.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The paper is generally well-written, clearly structured, and quite easy to follow.\nThe target issues of the paper are meaningful and worth exploring. The motivation is clear.\nThis submission gives a valuable implementation of such an idea.", "weaknesses": "1. To some extent, this paper lacks novelty. The intra-branch fusion and inter-branch fusion are commonly seen techniques in the re-parameterization of CNN networks, as shown in Sec 3.1. This paper only applies this technique to ViTs. Although the targeted module is changed from convolutions to linear projections, this modification is straightforward. This paper uses an additional batch normalization to modulate Q^{'} and K^{'} and names this operation a distribution rectification operation. This operation is quite simple.\n\n2. In Table 1, TDRL only outperforms D-MAE-Lite for 0.3 (78.7 vs. 78.4). The advantage does seem not significant.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "v7NiPtQRrC", "reviewer": "Reviewer_wWZH", "summary": "The paper introduced a new reparametrization approach for ViTs. Unlike standard reparametrization methods only considering convolution layers with kernel sizes larger than 1, this paper studied how to reparametrize the linear layer, the fundamental building block for Vision Transformers. By combining features from multiple branches with different depths, the proposed method exhibited improved performance for ViT-Ti and some other tiny ConvNet-based architectures on ImageNet classification and MSCOCO object detection benchmarks.", "questions": "See the weakness.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "1. To the best of my knowledge, the proposed method, TDRL, is the first one tailoring for the reparametrization of the linear layer, and thus for ViTs. This is important, particularly for the tiny models, because pure ViTs lag behind the ConvNets and hybrid models due to the lack of inductive bias and inadequate capacity. Reparametrization, training with extra parameters but shrinking back at the inference stage, is a good way to fit more knowledge into the models without incurring extra inference costs.\n\n2. TDRL seems to be a generic method and can be applied to many architectures, e.g., VanillaNet, and different tasks, e.g., DDPM, with linear layers.\n\n3. With the help of varied-depth branches and feature rectification, TDRL achieves promising results on classification and dense prediction tasks, especially the latter.", "weaknesses": "1. All the experiments are conducted with tiny or small models and a few pre-training epochs. I am curious whether TDRL would still be beneficial for larger models and when using larger epochs, since reparametrization works for ConvNets with various capacities. This might further strengthen the paper.\n\n2. I noticed that the gain from TDRL is considerably larger on dense prediction tasks. Do the authors have any explanation for it? Is it due to the use of BN in TDRL since some dense prediction tasks favor architectures with BN? If so, can the authors isolate the contributions of reparametrization and BN?\n\n3. Certain parts of the presentation could be further improved. For example, Fig. 1c does not help grasp how to do intra/inter-branch merging.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "W6QR3Zh6Z6", "reviewer": "Reviewer_GsFu", "summary": "This paper proposes a reparameterization technique to improve performance of lightweight vanilla ViT. The reparameterization technique is used in CNNs but surprisingly unexplored in ViTs. The paper presents a linear-based re-parameterization for vanilla Vision Transformers, without any convolutional operations. The authors introduce a specific linear stack with batch normalization for added nonlinearity and a pyramid multi-branch structure for feature fusion. They also highlight the importance of distribution consistency in deep networks, especially for ViTs, and introduce a Two-Dimensional Re-parameterized Linear module (TDRL) to address this. To enhance re-parameterization, different depths of linear stacks are used to create a pyramid multi-branch structure. Since this multi-branch structure can lead to instability in training,  authors suggest 1) an additional batch normalization to modulate the distributions of Q and K in the attention mechanism to rectify the distribution changes. 2) For other components like Feed Forward Networks (FFN), where the variance change isn't as drastic as in attention calculation, features are re-scaled using the square root of N instead of being normalized.", "questions": "1. Can you discuss on the computational complexity of the TDRL architecture?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 4, "strengths": "1. The paper looks into reparameterization technique for ViTs performance improvement. The technique that is common in CNN but overlooked in ViTs\n2. The proposed TDRL architecture seems to ensure diverse feature spaces among branches, enhancing the model's representation capabilities. TDRL seems to me a modular design that can replace any linear layer in ViTs", "weaknesses": "1. The authors experiment on the vanilla ViT-Tiny model. It is not clear what would be the impact of this TDRL module if applied in other tiny transformers such as Swin-tiny etc.\n2. The introduction of multiple branches and added nonlinearity might lead to overfitting, especially on smaller datasets. I don't see the authors discussing anything around it", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "oofbE5f2yY", "reviewer": "Reviewer_Yszy", "summary": "The paper proposes a new re-parameterization technique for tiny vision transformers. They stack multiple branches with different numbers of linear layers with BN, and the whole branches can be merged into a single linear layer during inference. The experimental results demonstrate the effectiveness of the proposed method.", "questions": "See weaknesses above.", "limitations": "", "rating": 6, "confidence": 5, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "The paper is technically sound and easy to understand. The method of reparameterization on the linear layer is shown to be useful and easy to implement.\n\nThe experimental results are rich and show the effectiveness of the proposed method.", "weaknesses": "The way of merging a stack of linear layers with BN has already been studied by [1].\n\nYou mentioned several times that BN can be used in between layers for nonlinearity, and the proposed re-parameterization method is nonlinear. Why does BN have nonlinearity? Isn't the proposed method a linear ensemble re-parameterization structure?\n\nWhy is the experimental setting based on MIM? What is the performance of training ViT with the proposed method on ImageNet directly? What is the performance of training vanilla ViT with MIM?\n\n[1] VanillaNet: the Power of Minimalism in Deep Learning.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (poster)