\section{Conclusion and Future Work} \label{conclusion}
In this paper, we have introduced LIT, a hierarchical vision transformer which pays less attention in the early stages to ease the huge computational cost of self-attention modules over high-resolution representations. Specifically, 
LIT applies MLP blocks in the first two stages to focus on local patterns while employing standard Transformer blocks with sufficient heads in the later stages to handle long range dependencies.
Moreover, we have proposed a deformable token merging module, which is learned to adaptively merge informative patches to an output unit, with enhanced geometric transformations. 
Extensive experiments on ImageNet and COCO have demonstrated that LIT achieves better performance
compared with existing state-of-the-art HVT methods.
Future works may include finding better architectural configurations of LIT with neural architecture search (NAS) and improving MLP blocks in the early stages to enhance the capability of LIT to encode local patterns.
Besides, one may also consider applying efficient self-attention mechanisms in the later stages to achieve better efficiency, such as kernelization~\cite{peng2021random}, low-rank decomposition \cite{linformer}, memory \cite{rae2020compressive}, and sparsity \cite{sparse_attn} schemes, etc. 
Also note that LIT potentially brings some negative societal impacts, such as the unintended use for surveillance and the huge energy consumption and carbon emissions of training LIT models on ImageNet.

\section{Experiments} \label{experiments}



\subsection{ImageNet Classification}

We conduct experiments on ImageNet (ILSVRC2012) \cite{imagenet} dataset. ImageNet is a large-scale dataset which has $\sim$1.2M training images from 1K categories and 50K validation images.
We compare with CNN-based ResNet \cite{resnet} and Transformer-based models including DeiT~\cite{deit}, PVT~\cite{pvt} and Swin~\cite{swin}.
For simplicity, we denote them as ``Model-Ti/S/M/B'' to refer to their tiny, small, medium and base variants. Similarly, we define four variants of our LIT models: LIT-Ti, LIT-S, LIT-M and LIT-B. Detailed architecture settings are shown in Table~\ref{tab:model_arch}. For better comparison,
we design LIT-Ti as a counterpart to PVT-S, where both models adopt the absolute positional encoding.
Our LIT-S, LIT-M and LIT-B use the relative positional encoding, and these three models can be seen as competitors to Swin-Ti, Swin-S, Swin-B, respectively. 


\textbf{Implementation details.}
In general, all models are trained on ImageNet with 300 epochs and a total batch size of 1024. For all ImageNet experiments, training images are resized to $256 \times 256$, and $224 \times 224$ patches are randomly cropped from an image or its horizontal flip, with the per-pixel mean subtracted. We use the single-crop setting for testing. We use AdamW optimizer~\cite{adamw} with a cosine decay learning rate scheduler. The initial learning rate is $1\times10^{-3}$, and the weight decay is set to $5\times10^{-2}$. The initial values of learnable offsets in DTM are set to 0, and the initial learning rate 
for offset parameters is set to $1\times10^{-5}$.
For a fair comparison, we adopt the same training strategies as PVT and Swin when comparing our models to each of them. 

\begin{table}[]
\centering
\renewcommand\arraystretch{1.2}
\caption{Comparisons with state-of-the-art methods on ImageNet. All models are trained and evaluated with the input resolution of $224\times224$.}
\vspace{5pt}
\begin{tabular}{p{0.2\textwidth}|>{\centering}p{0.15\textwidth}>{\centering}p{0.15\textwidth}>{\centering\arraybackslash}p{0.16\textwidth}}
Method      & Params (M) & FLOPs (G) & Top-1 Acc. (\%) \\ \shline
ResNet-18~\cite{resnet}   & 12         & 1.8       & 68.5      \\
ResNet-50~    & 26         & 4.1       & 78.5      \\
ResNet-101   & 45         & 7.9       & 79.8      \\ \hline
DeiT-Ti~\cite{deit}      & 5          & 1.3       & 72.2      \\
DeiT-S      & 22         & 4.6       & 79.8      \\
DeiT-B      & 86         & 17.5      & 81.8      \\ \hline
PVT-Ti~\cite{pvt}      & 13         & 1.9       & 75.1      \\
PVT-S       & 25         & 3.8       & 79.8      \\
PVT-M       & 44         & 6.7       & 81.2      \\
PVT-L      & 61         & 9.8       & 81.7      \\ \hline
Swin-Ti~\cite{swin}     & 29         & 4.5       & 81.3      \\
Swin-S      & 50         & 8.7       & 83.0      \\
Swin-B      & 88         & 15.4      & 83.3      \\ \hline
LIT-Ti      & 19         & 3.6       & 81.1         \\
LIT-S       & 27         & 4.1       & 81.5         \\
LIT-M       & 48         & 8.6       & 83.0         \\
LIT-B       & 86         & 15.0      & 83.4         
\end{tabular}
\label{tab:main_table}
\end{table}


\begin{small}
	\begin{table}[!ht]
		\begin{minipage}[t]{0.48\linewidth}
			\centering
            \renewcommand\arraystretch{1.2}
            \caption{Effect of our architecture design principle. * denotes the LIT model which adopts the same uniform patch merging strategies as in PVT-S or Swin-Ti.
            }
            \vspace{5pt}
            \scalebox{0.8}{
            \begin{tabular}{l|ccc}
            Model & Params & FLOPs & Top-1 Acc. (\%) \\ \shline
            PVT-S~\cite{pvt} & 25M & 3.8G & 79.8 \\
            \textbf{LIT-Ti*} & \bf{19M} & \textbf{3.6G} & \textbf{80.4} \\  \hline
            Swin-Ti~\cite{swin} & 28M & 4.5G & 81.3 \\
            \textbf{LIT-S*} & \bf{27M} & \textbf{4.1G} & 81.3 
            \end{tabular}
            }
            \label{tab:ablation_blk}
		\end{minipage}\hfill
		\begin{minipage}[t]{0.48\linewidth}
		
			\centering
            \renewcommand\arraystretch{1.2}
            \caption{Effect of the proposed deformable token merging module. We replace the uniform patch merging strategies in PVT-S and Swin-Ti with our DTM.}
            \vspace{5pt}
            \scalebox{0.78}{
            \begin{tabular}{l|ccc}
            Model & Params & FLOPs & Top-1 Acc. (\%) \\ \shline
            PVT-S~\cite{pvt} & 25M & 3.8G & 79.8 \\
            \textbf{PVT-S + DTM} & 25M & 3.8G & \textbf{80.5} \\  \hline
            Swin-Ti~\cite{swin} & 28M & 4.5G & 81.3 \\
            \textbf{Swin-Ti + DTM} & 28M & 4.5G & \textbf{81.6}
            \end{tabular}
            }
            \label{tab:ablation_dtm}
		\end{minipage}
	\end{table}
\end{small}



\textbf{Results on ImageNet.}
In Table~\ref{tab:main_table}, we compare LIT with several state-of-the-art methods on ImageNet. In general, all LIT models achieve on par or better performance than their counterparts while having fewer parameters and less FLOPs, without applying any existing efficient self-attention mechanisms. 
Specifically, our LIT-Ti outperforms PVT-S by 1.3\% on the Top-1 accuracy while the FLOPs is reduced by 0.2G. LIT-S surpasses Swin-Ti by 0.2\% with 0.4G less FLOPs. LIT-M achieves on par performance with Swin-S, whereas the FLOPs of LIT-M is reduced. LIT-B brings 0.1\% Top-1 accuracy increase over Swin-B while using 0.4G less FLOPs. 
For ResNet and DeiT, LIT demonstrates better performance when compared to them with the same magnitude of FLOPs and parameters (\textit{e}.\textit{g}., ResNet-50 v.s. LIT-S, DeiT-B v.s. LIT-B).


\subsection{Ablation Studies}
\textbf{Effect of the architecture design.}
To explore the effect of our architecture design principle in LIT, we conduct experiments on ImageNet and compare the architecture of LIT with two recent HVTs: PVT-S~\cite{pvt} and Swin-Ti~\cite{swin}. The results are shown in Table~\ref{tab:ablation_blk}.
In general, our architecture improves baseline PVT-S by 0.6\% in Top-1 accuracy while using less FLOPs (3.6G v.s. 3.8G). For Swin-Ti, our method reduces the FLOPs by 0.4G while achieving 
on par performance. It is also worth noting that the total number of parameters is reduced for both PVT-S and Swin-Ti. The overall performance demonstrates the effectiveness of the proposed architecture.  


\textbf{Effect of deformable token merging.}
To verify the effectiveness of our proposed DTM strategy, we replace the default patch merging scheme in PVT-S and Swin-Ti with DTM and train the models on ImageNet. The results are shown in Table~\ref{tab:ablation_dtm}. 
We 
observe that
for both models,
DTM introduces negligible FLOPs and parameters while 
improving
PVT-S and Swin-Ti by 0.7\% and 0.3\% in terms of the Top-1 accuracy, respectively.  Furthermore, we visualize the learned offsets in Figure~\ref{fig:dtm_vis}. As it shows, unlike 
previous uniform patch merging strategy where sampled locations are limited within the green rectangle, our DTM adaptively merges patches according to objects' scales and shapes
(\textit{e}.\textit{g}., koala leg, keyboard, cat tail). This again emphasizes the power of LIT to accommodate various geometric transformations. 


\begin{figure}[htp!]
	\centering
    \includegraphics[width=\textwidth]{figures/dpm_vis.pdf}
	\caption{Visualization of learned offsets by the proposed deformable token merging modules. Each image shows $4^3$ sampling locations (\textit{red dots}) in three DTM modules with $2\times2$ filter. The green rectangle outlines a $32\times32$ patch of the original image, which also indicates a regular sampling field of previous methods. Best viewed in color.}
	\label{fig:dtm_vis}
\end{figure}

\textbf{Effect of MSA in each stage.}
To explore the effect of self-attention in recent HVTs, we train PVT-S on ImageNet and gradually remove self-attention layers at each stage. The results are presented in Table~\ref{tab:ablation_stages}. First, after replacing the SRA layers in PVT-S with standard MSA layers, we observe 1.1\% improvement on the Top-1 accuracy whereas the FLOPs is almost doubled. This indicates that PVT makes a trade-off between performance and efficiency. Next, by gradually removing MSA layers in the first two stages, the Top-1 accuracy only drops by 0.1\%, 0.5\%, respectively. It implies that the self-attention layers in the early stages of PVT contribute less than expected to the final performance, and they perform not much better than pure MLP layers. It can be attributed to the fact that shallow layers focus more on encoding local patterns. However, we observe a huge performance drop when removing self-attention layers in the last two stages. The results show that the self-attention layers play an important role in the later stages and capturing long range dependencies is essential for well-performed hierarchical vision Transformers.

\begin{table}[htp!]
\centering
\renewcommand\arraystretch{1.2}
\caption{Impact of the MSA layers at each stage. Note that PVT-S has four stages, which adopts SRA at all blocks instead of MSA. Here we denote ``w/ MSA'' as PVT-S with standard MSA layers. ``Removed Stage'' refers to the stages where we remove all self-attention layers. For example, ``1,2'' means we remove the self-attention layers in the first two stages. Note that ``0'' means a model without removing any self-attention layers.}
\scalebox{1.0}
{
\begin{tabular}{l|cccc}
Model & Removed Stage & Params (M) & FLOPs (G) & Top-1 Acc. (\%) \\ \shline
PVT-S~\cite{pvt} & 0 & 25 & 3.8 & 79.8 \\
PVT-S w/ MSA & 0 & 20 & 8.4 & 80.9 \\
PVT-S w/ MSA & 1 & 20 & 4.5 & 80.8 \\
PVT-S w/ MSA  & 1,2 & 19 & 3.6 & 80.4 \\
PVT-S w/ MSA  & 1,2,3 & 17 & 3.0 & 75.0 \\
PVT-S w/ MSA  & 1,2,3,4 & 14 & 2.8 & 66.8
\end{tabular}}
\label{tab:ablation_stages}
\vspace{5pt}
\end{table}



To better understand the phenomenon, we visualize the attention probabilities for PVT-S
without removing any MSA layers, which are depicted in Figure~\ref{fig:attn_prob_vis}. First, the attention map at the first stage shows that the query pixel almost pays no attention to other locations. At the second stage, the receptive field of the query pixel is slightly enlarged, but similar to the first stage. Considering that PVT-S only has one head at the first stage and two heads at the second stage, this strongly supports our hypothesis that very few heads in an MSA layer result in a smaller receptive field, such that a self-attention layer is almost equivalent to an FC layer. Furthermore, we observe relatively larger receptive fields from the attention maps of the last two stages. As a large receptive field usually helps to model longer dependencies, this explains the huge performance drop in Table~\ref{tab:ablation_stages} after we remove the MSA layers in the last two stages. 

\begin{figure}[htp!]
	\centering
    \includegraphics[width=0.9\textwidth]{figures/avg_attentions_all.pdf}
	\caption{Attention probabilities of PVT-S with standard MSA layers. For each stage, we visualize the attention map of each head (\textit{columns}) at each block (\textit{rows}).  All attention maps are averaged over 100 validation images. Each map shows the attention probabilities of a query pixel (\textit{green rectangle}) to other pixels. Darker color indicates higher attention probability and vice versa. Best viewed in color.}
	\label{fig:attn_prob_vis}
\end{figure}
\vspace{-1em}


\begin{table}[]
\centering
\renewcommand\arraystretch{1.2}
\caption{Object detection performance on the COCO \texttt{val2017} split using the RetinaNet framework.}
\vspace{5pt}
	\scalebox{0.9}
	{
\begin{tabular}{l|ccccccc}
\multirow{2}{*}{Backbone} & \multicolumn{7}{c}{RetinaNet}                                                                     \\ \cline{2-8} 
                          & \multicolumn{1}{c|}{Params (M)} & $\mathrm{AP}$      & $\mathrm{AP}_{50}$   & \multicolumn{1}{c|}{$\mathrm{AP}_{75}$}   & $\mathrm{AP}_{S}$   & $\mathrm{AP}_{M}$   & $\mathrm{AP}_{L}$   \\ \shline
ResNet-50~\cite{resnet}                  & \multicolumn{1}{c|}{38}         & 36.3    & 55.3 & \multicolumn{1}{c|}{38.6} & 19.3 & 40.0   & 48.8 \\
PVT-S~\cite{pvt}                     & \multicolumn{1}{c|}{34}         & 40.4    & 61.3 & \multicolumn{1}{c|}{43.0} & 25.0 & 42.9 & 55.7 \\
LIT-Ti                    & \multicolumn{1}{c|}{\bf{30}}         & \bf{41.6} &  \bf{62.8}    & \multicolumn{1}{c|}{\bf{44.7}}     & \bf{25.7}     &  \bf{44.4}    &  \bf{56.4}   \\ \hline
ResNet-101~\cite{resnet}                   & \multicolumn{1}{c|}{57}         & 38.5 & 57.8    & \multicolumn{1}{c|}{41.2}     & 21.4     &  42.6    &  51.1 \\
Swin-Ti~\cite{swin}                   & \multicolumn{1}{c|}{39}         & 41.5 & 62.1    & \multicolumn{1}{c|}{\bf{44.2}}     & 25.1     &  \bf{44.9}    &  55.5 \\
LIT-S               & \multicolumn{1}{c|}{39}       & \bf{41.6} & \bf{62.7}    & \multicolumn{1}{c|}{44.1}     &  \bf{25.6}    &  44.7    & \bf{56.5}  \\
\end{tabular}}
\label{tab:coco_retinanet}
\vspace{-10pt}
\end{table}


\begin{table}[]
\centering
\renewcommand\arraystretch{1.2}
\caption{Object detection and instance segmentation performance on the COCO \texttt{val2017} split using the Mask R-CNN framework. $\mathrm{AP}^b$ and $\mathrm{AP}^m$ denote the bounding box AP and mask AP, respectively.}
\vspace{4pt}
\scalebox{0.9}
{
\begin{tabular}{l|ccccccc}
\multirow{2}{*}{Backbone} & \multicolumn{7}{c}{Mask R-CNN}                                                                    \\ \cline{2-8} 
                          & \multicolumn{1}{c|}{Params (M)} & $\mathrm{AP}^{b}$       & $\mathrm{AP}^{b}_{50}$    & \multicolumn{1}{c|}{$\mathrm{AP}^{b}_{75}$}   & $\mathrm{AP}^{m}$      & $\mathrm{AP}^{m}_{50}$      & $\mathrm{AP}^{m}_{75}$   \\ \shline
ResNet-50~\cite{resnet}                  & \multicolumn{1}{c|}{44}         & 38.0      & 58.6 & \multicolumn{1}{c|}{41.4} & 34.4 & 55.1 & 36.7 \\
PVT-S~\cite{pvt}                     & \multicolumn{1}{c|}{44}         & 40.4    & 62.9 & \multicolumn{1}{c|}{43.8} & 37.8 & 60.1 & 40.3 \\
LIT-Ti                    & \multicolumn{1}{c|}{\bf{40}}         & \bf{42.0} & \bf{64.9}      & \multicolumn{1}{c|}{\bf{45.6}}     &   \bf{39.1}   & \bf{61.9}     &  \bf{41.9}  \\ \hline
ResNet-101~\cite{resnet}                   & \multicolumn{1}{c|}{63}         & 40.4 & 61.1    & \multicolumn{1}{c|}{44.2}     & 36.4     &  57.7    &  38.8 \\
Swin-Ti~\cite{swin}                   & \multicolumn{1}{c|}{48}         & 42.2 & 64.6    & \multicolumn{1}{c|}{46.2}     & 39.1     &  61.6    &  42.0 \\
LIT-S               & \multicolumn{1}{c|}{48}       & \bf{42.9} &    \bf{65.6} & \multicolumn{1}{c|}{\bf{46.9}}     &   \bf{39.6}   &  \bf{62.3}    &  \bf{42.4} \\
\end{tabular}}
\label{tab:coco_maskrcnn}
\vspace{-10pt}
\end{table}

\subsection{Object Detection and Instance Segmentation on COCO}
In this section, we conduct experiments on COCO 2017~\cite{coco} dataset to show the performance of LIT on object detection and instance segmentation. COCO is a large-scale dataset which contains $\sim$118K images for the training set and $\sim$5K images for the validation set. For a fair comparison, we evaluate LIT-Ti and LIT-S on two base detectors: RetinaNet~\cite{retinanet} and Mask R-CNN~\cite{maskrcnn}.
For the experiments with both detectors, we consider 
CNN-based ResNet~\cite{resnet} and Transformer-based models including PVT-S~\cite{pvt} and Swin-Ti~\cite{swin}.
Following common practice~\cite{detr,pvt}, we measure the performance of all models by Average Precision (AP) in COCO.


\textbf{Implementation details.}
All models are trained on 8 V100 GPUs, with $1\times$ schedule (12 epochs) and a total batch size of 16. We use AdamW~\cite{adamw} optimizer with a step decay learning rate scheduler. Following PVT~\cite{pvt}, the initial learning rates are set to $1\times10^{-4}$ and $2\times10^{-4}$ for RetinaNet and Mask R-CNN, respectively.
The weight decay is set to $1\times10^{-4}$ for all models. Results of Swin-Ti based detectors are adopted from Chu~\textit{et al}. ~\cite{chu2021Twins}. 
At the training stage, we initialize the backbone with the pretrained weights on ImageNet. The training images are resized to the shorter size of 800 pixels, and the longer size is at most 1333 pixels. During inference, we fix the shorter side of an image to 800 pixels.


\textbf{Results on COCO.}
Table~\ref{tab:coco_retinanet} shows the comparisons of different backbones on object detection based on RetinaNet. 
By comparing LIT with PVT and ResNet counterparts, we find that our model outperforms both backbones on object detection in almost all metrics. Similar results can be found in Table~\ref{tab:coco_maskrcnn}, where LIT again surpasses compared methods on object detection and instance segmentation using the Mask R-CNN framework. The consistent improvement demonstrates the effectiveness of LIT for dense prediction tasks. 




\section{Introduction}

Transformers have made substantial strides in natural language processing (NLP) (\textit{e}.\textit{g}., \cite{transformer, devlin2019bert}) and recently in the computer vision (CV) field (\textit{e}.\textit{g}., \cite{vit, deit}). Inspired by the pyramid design~\cite{fpn} in CNNs, recent hierarchical vision Transformers (HVTs)~\cite{pvt,swin,cvt,contnet} divide transformer blocks into several stages and progressively shrink feature maps as the network goes deeper. However, high-resolution feature maps in the early stages result in long token sequences, which brings huge computational cost and memory consumption due to the quadratic complexity of self-attention. For instance, a feature map of size $56 \times 56 \times 96$ 
costs 2.0G 
FLOPs
in one Multi-head Self-Attention (MSA)~\cite{transformer}, while the entire model of ResNet-18~\cite{resnet} only requires 1.8G FLOPs. Such a huge computing cost makes it difficult to apply Transformers into broad computer vision tasks. 


Several emerging efforts have been made to reduce the computational cost in the early stages of HVTs. For example, some works~\cite{cvt, halonet} reduce the number of self-attention heads in an MSA layer or further decreasing the number of Transformer blocks~\cite{vit}. Another line of works proposes to
trade-off accuracy and efficiency for MSA
via heuristic approximation, such as spatial reduction attention (SRA)~\cite{pvt} and shifted window based multi-head self-attention (SW-MSA)~\cite{swin}. There are also studies simply employ
convolutional layers~\cite{botnet, levit} when the resolution of feature maps are considerably large. However, how much the early adoption of self-attention layers really contributes to the final performance remains unclear.


In this paper, we present a \textbf{L}ess attention v\textbf{I}sion \textbf{T}ransformer (\textbf{LIT}) to address the aforementioned problem for HVTs. Specifically, we propose to exclusively use MLP layers to capture local patterns in the early stages while introducing MSA layers with sufficient number of heads to handle long range dependencies in the later stages. 
The motivation comes from two aspects. 
First, previous studies in CNNs have shown that shallow layers focus on local patterns and deeper layers tend to produce more discriminative features~\cite{WuSH19,HouCHBTT19,zfnet}.
Second, from the theoretical perspective, a self-attention layer with sufficient heads applied to images can express any convolutional layer~\cite{selfconv}, and the size of the receptive field for an MSA layer is positively correlated with the number of heads. Therefore, fewer heads in an MSA layer theoretically result in a smaller receptive field, where the extreme case is as expressive as a $1 \times 1$ convolution that can be viewed as a standard 
FC layer applied to each pixel independently. 
To be emphasized, by exploiting MLP blocks in the early stages, the model avoids the huge computational cost and memory footprint
arising from
self-attention on high-resolution feature maps. Moreover, applying self-attention in the later stages to capture long range dependencies is quite efficient due to the progressive shrinking pyramid.
Our comprehensive results show that 
such a simple architecture design
brings a sweet spot between model performance and efficiency.


Furthermore, recent HVTs either adopt a standard convolutional layer or a linear projection layer to merge nearby tokens~\cite{pvt, swin}, aiming to control the scale of feature maps.
However, such methods hinder the representational power for vision Transformers to model geometric transformations, 
considering that not every pixel equally contributes to an output unit~\cite{LuoLUZ16}. To this end, we propose a Deformable Token Merging (DTM) module, inspired by deformable convolutions~\cite{dconv1,dconv2}, where we learn a grid of offsets to adaptively augment the spatial sampling locations for merging neighboring patches from a sub-window in a feature map. In this way, we can obtain more informative downsampled tokens for subsequent processing.



Our contributions can be summarized as follows. First, 
we propose a simple HVT structure with pure MLP blocks in the early stages, making it tractable for Transformers to process high-resolution images, especially in dense prediction tasks.
Second, we propose a deformable token merging module to adaptively merge more informative patches to deliver hierarchical representations, with enhanced transformation modeling capability.
Finally, we conduct extensive experiments to show 
that the proposed LIT performs favorably against other state-of-the-art vision Transformers with similar or even reduced computational complexity and memory consumption.



\section{Proposed Method} \label{method}



\subsection{Overall Architecture}
\label{overall}
The overall architecture of LIT is illustrated in Figure~\ref{fig:model_arch}. Let ${\mathcal I} \in \mathbb{R}^{H \times W \times 3}$ be an input RGB image, where $H$ and $W$ represent the height and width, respectively. We first split ${\mathcal I}$ into non-overlapping patches with a patch size of $4\times4$, and thus the initial feature dimension of each patch is $4 \times 4 \times 3 = 48$. Next, a linear embedding layer is exploited to project each patch into dimension $C_1$, serving as the initial input for the following pipeline. The entire model is divided into 4 stages. Letting $s\in [1,2,3,4]$ be the index of a stage, we employ $L_s$ blocks at each stage, where the first two stages solely utilise MLP blocks 
to encode local representations and the last two stages employ standard Transformer blocks~\cite{vit} to handle longer dependencies. 
At each stage, we scale the input feature maps into $\frac{H_{s-1}}{P_s} \times \frac{W_{s-1}}{P_s} \times C_s$, where $P_s$ and $C_s$ represent the patch size and the hidden dimension at the $s$-th stage, respectively. 
For the last two stages, we set $N_s$ self-attention heads in each Transformer block.

\subsection{Block Design in LIT} \label{blockdesign}
As shown in Figure~\ref{fig:model_arch}, LIT employs two types of blocks: MLP blocks and Transformer blocks. 
In the early stages, we apply MLP blocks.
Concretely, an MLP block is built upon an MLP which 
consists of two FC layers with GELU~\cite{gelu} non-linearity in between. For each MLP at the $s$-th stage, an expansion ratio of $E_s$ is used. Specifically, the first FC layer expands the dimension of a token from $C_s$ to $E_s \times C_s$, and the other FC layer reduces the dimension back to $C_s$. Formally, letting ${\bf X} \in \mathbb{R}^{(\frac{H_{s-1}}{P_s} \times \frac{W_{s-1}}{P_s}) \times C_s}$ be the input of the $s$-th stage and $l$ be the index of a block, an MLP block can be formulated as
\begin{equation}
    {\bf X}_l = {\bf X}_{l-1} + \mathrm{MLP}(\mathrm{LN}({\bf X}_{l-1})),
\end{equation}
where $\mathrm{LN}$ indicates the layer normalization~\cite{layernorm} and $\mathrm{MLP}$ denotes an MLP.
In the last stages, a Transformer block as described in ViT~\cite{vit} contains an MSA layer and an MLP, which can be expressed as
\begin{align}
    \label{eq:msa}
    {\bf X}^{'}_{l-1} &= {\bf X}_{l-1} + \mathrm{MSA}(\mathrm{LN}({\bf X}_{l-1})), \\
    \label{eq:mlp}
    {\bf X}_l &= {\bf X}^{'}_{l-1} + \mathrm{MLP}(\mathrm{LN}({\bf X}^{'}_{l-1})).
\end{align}


With this architecture, our model benefits from two main advantages: First, we 
avoid the
huge computational costs and memory footprint that are introduced by
long sequences
in the early stages. 
Second, unlike recent works that shrink the attention maps using sub-windows~\cite{swin} or reduce the spatial dimensions of the key and value matrices, 
we keep standard MSA layers in the last two stages so as to maintain the capability of LIT to handle long range dependencies while keeping mild FLOPs due to the pyramid structure. We will show in the ablation study that our simple architecture design outperforms the state-of-the-art hierarchical ViT variants on ImageNet with comparable FLOPs.

\paragraph{Remark.}
Here we justify the rationality of applying pure MLP blocks in the early stages by considering the relationship among a convolutional layer, an FC layer and an MSA layer.
Firstly, we begin with a review of a standard convolutional layer. Let ${\bf X} \in \mathbb{R}^{H\times W \times C_{in}}$ be the input feature map, and let ${\bf W} \in \mathbb{R}^{K\times K \times C_{in} \times C_{out} }$ be the convolutional weight tensor, where $K$ is the kernel size, $C_{in}$ and $C_{out}$ are the input and output channel dimensions, respectively. For simplicity, we omit the bias term and use ${\bf X}_{p,:}$ to represent  ${\bf X}_{i,j,:,:}$, where $(i,j)$ denotes the pixel index and $p \in [H]\times[W]$. Given a convolutional kernel of $K\times K$ sampling locations, the output for a pixel $p$ can be formulated as
\begin{equation} \label{conv}
    \mathrm{Conv}({\bf X})_{p,:} = \sum_{k \in [K\times K]} {\bf X}_{p + g(k),:}{\bf W}_{g(k),:,:},
\end{equation}
where $g : [K \times K]\rightarrow\Delta_K$ is a bijective mapping of sampling indexes onto the pre-specified offsets $\Delta_K$. For example, let $\Delta_K = \{(-1,-1), (-1,0), ..., (0,1), (1,1)\}$ be a $3 \times 3$ kernel with dilation 1, then $g(0) = (-1, -1)$ represents the first sampling offset.

When $K = 1$, the weight tensor ${\bf W}$ is equivalent to a matrix, such that ${\bf W} \in \mathbb{R}^{C_{in}\times C_{out}}$. In this case, Eq.~(\ref{conv}) can express an FC layer, and the output for a pixel $p$ is defined by
\begin{equation} \label{fc}
    \mathrm{FC}({\bf X})_{p,:}  = {\bf X}_{p,:}{\bf W}_{:,:}.
\end{equation}
Last, let $N_h$ be the number of heads in an MSA layer and ${\bf W}^{(h)} \in \mathbb{R}^{C_{in}\times C_{out}}$ be learnable parameters of the $h$-th head. Under mild conditions,
Cordonnier~\textit{et al}. ~\cite{selfconv} prove that the output from an MSA layer at pixel $p$ is identical to 
\begin{equation} \label{re_msa}
    \mathrm{MSA}({\bf X})_p = \sum_{h \in [N_h]}{\bf X}_{p+f(h),:} {\bf W}^{(h)},
\end{equation}
where $f : [N_h]  \rightarrow \Delta_K$ is a bijective mapping of heads onto pixel shifts. In that case,
Eq.~(\ref{re_msa}) can be seen to be equivalent to a convolutional layer with a kernel size of $\sqrt{N_h}\times\sqrt{N_h}$. We refer detailed explanations to \cite{selfconv}.

From Eqs.~(\ref{conv})-(\ref{re_msa}), we observe almost equivalent mathematical formulations among a convolutional layer, an FC layer and an MSA layer. In fact, the main difference lies in the size of their receptive fields. Specifically, the size of the receptive field in a convolutional layer is $K\times K$. For an FC layer, $K$ is always $1$. For an MSA layer, $K$ is equivalent to $\sqrt{N_h}$. The observation tells us that when $N_h$ is pretty small (\textit{e}.\textit{g}., 1 or 2), an MSA layer theoretically has a limited receptive field that is approximately or equal to an FC layer. Considering most recent HVTs~\cite{pvt,cvt,contnet} adopt very few heads in the early stages, this justifies our method of applying pure MLP blocks in the first two stages without self-attention.

\subsection{Deformable Token Merging}
Previous works on HVTs~\cite{pvt,swin} rely on patch merging to achieve pyramid feature representations. However, they merge patches from a regular grid and neglect the fact that not every patch contributes equally to an output unit~\cite{LuoLUZ16}. Inspired by deformable convolutions~\cite{dconv1,dconv2}, we propose a deformable token merging module to learn a grid of offsets to adaptively sample more informative patches. 
Formally, a deformable convolution is formulated as 
\begin{equation} \label{dcn}
    \mathrm{DC}({\bf X})_{p,:} = \sum_{k \in [K\times K]} {\bf X}_{p + g(k) + \Delta g(k),:}{\bf W}_{g(k),:,:}.
\end{equation}
Compared to a standard convolution operation as in Eq.~(\ref{conv}), DC learns an offset $\Delta g(k)$ for each pre-specified offset $g(k)$. Learning $\Delta g(k)$ requires a separate convolutional layer, which is also applied over the input feature map ${\bf X}$. To merge patches in an adaptive manner, we adopt one DC layer in a DTM module, which can be formulated by 
\begin{equation}
    \mathrm{DTM}({\bf X}) =  \mathrm{GELU}( \mathrm{BN}( \mathrm{DC}({\bf X}))),
\end{equation}
where $\mathrm{BN}$ denotes the batch normalization~\cite{batchnorm} and we employ the GELU non-linearity.
We will show in the ablation study that the sampling locations in DTM are adaptively adjusted when objects' scales and shapes change, benefiting from the learned offsets. 
Also note that our light-weight DTM introduces negligible FLOPs and parameters compared to regular grid sampling in baselines, thus making it a plug-and-play module for recent HVTs.

\begin{table}[]
\renewcommand\arraystretch{1.2}
\caption{Architecture specifications of LIT. $P$ denotes the patch size and $C$ is the channel dimension. $L$ refers to the number of blocks applied at each stage. $N$ is the number of self-attention heads. As the first two stages do not have self-attention layers, we only show the numbers of $N_3$ and $N_4$. The expansion ratios of LIT-Ti at each stage are [8, 8, 4, 4]. We use an expansion ratio of 4 for all MLP layers in LIT-S, LIT-M and LIT-B. The numbers denoted under each stage index (\textit{e}.\textit{g}., $28 \times 28$) represents the size of the feature maps at that stage, based on the input resolution of $224 \times 224$.}
\vspace{5pt}
\begin{tabular}{c|c|l|l|l|l}
 Stage &
  Layer Name &
  \multicolumn{1}{c|}{LIT-Ti} &
  \multicolumn{1}{c|}{LIT-S} &
  \multicolumn{1}{c|}{LIT-M} &
  \multicolumn{1}{c}{LIT-B} \\ \shline
\multirow{2}{*}{\begin{tabular}[c]{@{}c@{}}Stage 1\\ $56\times56$\end{tabular}} &
  Patch Embedding &
  $\begin{array}{l}P_1=4\\C_1 = 64\end{array}$ &
  $\begin{array}{l}P_1=4\\C_1 = 96\end{array}$ &
  $\begin{array}{l}P_1=4\\C_1 = 96\end{array}$ &
  $\begin{array}{l}P_1=4\\C_1 = 128\end{array}$ \\ \cline{2-6} 
 &
  MLP Block &
  $\begin{array}{l}L_1 = 3\end{array}$ &
  $\begin{array}{l}L_1 = 2\end{array}$ &
  $\begin{array}{l}L_1 = 2\end{array}$ &
  $\begin{array}{l}L_1 = 2\end{array}$ \\ \hline
\multirow{2}{*}{\begin{tabular}[c]{@{}c@{}}Stage 2\\ $28\times28$\end{tabular}} &
  DTM &
  $\begin{array}{l}P_2=2\\C_2 = 128\end{array}$ &
  $\begin{array}{l}P_2=2\\C_2 = 192\end{array}$ &
  $\begin{array}{l}P_2=2\\C_2 = 192\end{array}$ &
  $\begin{array}{l}P_2=2\\C_2 = 256\end{array}$ \\ \cline{2-6} 
 &
  MLP Block &
  $\begin{array}{l}L_2 = 4\end{array}$ &
  $\begin{array}{l}L_2 = 2\end{array}$ &
  $\begin{array}{l}L_2 = 2\end{array}$ &
  $\begin{array}{l}L_2 = 2\end{array}$ \\ \hline
\multirow{3}{*}{\begin{tabular}[c]{@{}c@{}}Stage 3\\ $14\times14$\end{tabular}} &
  DTM &
  $\begin{array}{l}P_3=2\\C_3 = 320\end{array}$ &
  $\begin{array}{l}P_3=2\\C_3 = 384\end{array}$ &
  $\begin{array}{l}P_3=2\\C_3 = 384\end{array}$ &
  $\begin{array}{l}P_3=2\\C_3 = 512\end{array}$ \\ \cline{2-6} 
 &
  Transformer Block &
  $\begin{array}{l}N_3 = 5\\L_3 = 6\end{array}$ &
  $\begin{array}{l}N_3 = 12\\L_3 = 6\end{array}$ &
  $\begin{array}{l}N_3 = 12\\L_3 = 18\end{array}$ &
  $\begin{array}{l}N_3 = 16\\L_3 = 18\end{array}$ \\ \hline
\multirow{3}{*}{\begin{tabular}[c]{@{}c@{}}Stage 4\\ $7\times7$\end{tabular}} &
  DTM &
  $\begin{array}{l}P_4=2\\C_4 = 512\end{array}$ &
  $\begin{array}{l}P_4=2\\C_4 = 768\end{array}$ &
  $\begin{array}{l}P_4=2\\ C_4 = 768\end{array}$ &
  $\begin{array}{l}P_4=2\\C_4 = 1024\end{array}$ \\ \cline{2-6} 
 &
  Transformer Block &
  $\begin{array}{l}N_4 = 8\\L_4 = 3\end{array}$ &
  $\begin{array}{l}N_4 = 24\\L_4 = 2\end{array}$ &
  $\begin{array}{l}N_4 = 24\\L_4 = 2\end{array}$ &
  $\begin{array}{l}N_4 = 32\\L_4 = 2\end{array}$ \\
\end{tabular}
\label{tab:model_arch}
\end{table}





\section{Related Work}

\paragraph{Vision Transformers.} 
Vision Transformers are models which adopt the self-attention mechanism ~\cite{transformer} into CV tasks. Recent works towards Vision Transformers either follow a hybrid architecture that combines convolution and self-attention~\cite{detr,botnet}, or design a pure self-attention architecture without convolution~\cite{sasa,lrnet}.
More recently, Dosovitskiy~\textit{et al}. ~\cite{vit} propose a Vision Transformer (ViT) which achieves promising results on ImageNet. 
Since then, a few subsequent works have been proposed to improve ViT from different aspects. 
For example, some works~\cite{t2t,localvit} seek to bring locality into ViT as they find ViT failed to model the important local structures (\textit{e}.\textit{g}., edges, lines). 
Another line of works aims to explore deeper architectures~\cite{deepvit, godeeper} by stacking more Transformer blocks. 


There is also a prevailing trend to introduce hierarchical representations into ViT~\cite{hvt,ceit,pvt,swin,cvt,pit,halonet}. 
To do so, these works divide Transformer blocks into several stages and downsample feature maps as the network goes deeper.
However, high-resolution feature maps in the early stages inevitably result in high computational and memory costs due to the quadratic complexity of the self-attention module. 
Targeting at this problem, Wang~\textit{et al}. ~\cite{pvt} propose to reduce the spatial
dimensions of attention's key and value matrices. Liu~\textit{et al}. ~\cite{swin} propose to limit self-attention in non-overlapped local windows. However, these replacements seek to shrink the global attention maps for efficiency.
In this paper, we elaborately design the shallow layers with pure MLPs, that are powerful enough to encode local patterns. 
This neat architecture design keeps the capability for modelling global dependencies in the later stages while easing the prohibitively expensive complexity introduced by high-resolution feature maps, especially in the dense prediction tasks.

\paragraph{Deformable Convolutions.} 
Deformable convolutions (DC) are initially proposed by Dai~\textit{et al}. ~\cite{dconv1} in  object detection and semantic segmentation tasks. Different from the regular sampling grid of a standard convolution, DC adaptively augments the spatial sampling locations with learned offsets. One following work by Zhu~\textit{et al}. ~\cite{dconv2} improves DC with a modulation mechanism, which modulates the input feature amplitudes from different spatial locations. 
With the advantage on modeling geometric transformations, many works adopt DC to target various CV problems. For example,
Zhu~\textit{et al}. ~\cite{deformable_detr} propose a deformable attention module for object detection. Shim~\textit{et al}. ~\cite{refsr} construct a similarity search and extraction network built upon DC layers for single image super-resolution. Thomas~\textit{et al}. ~\cite{kpconv} introduce a deformable kernel point convolution operator for point clouds.
In this paper, we propose a deformable token merging module to adaptively merge more informative image tokens. Unlike previous works that merge tokens from a regular grid, DTM introduces better transformation modeling capability for HVTs.


\begin{figure}
	\centering
    \includegraphics[width=\textwidth]{figures/framework.pdf}
	\caption{\textbf{Overall architecture of LIT.} The model is divided into four stages, where we apply MLP blocks in the first two stages and employ standard Transformer blocks in the last two stages. ``DTM'' denotes our proposed deformable token merging module.}
	\label{fig:model_arch}
\end{figure}
