Rethinking Vision Transformers for MobileNet Size and Speed

Introduction

The promising performance of Vision Transformers (ViTs) has inspired many follow-up works to further refine the model architecture and improve training strategies, leading to superior results on most computer vision benchmarks, such as classification , segmentation , detection , and image synthesis . As the essence of ViT, Multi Head Self Attention (MHSA) mechanism is proved to be effective in modeling spatial dependencies in $2$D images, enabling a global receptive field. In addition, MHSA learns second-order information with the attention heatmap as dynamic weights, which is a missing property in Convolution Neural Networks (CNNs) . However, the cost of MSHA is also obvious–quadratic computation complexity with respect to the number of tokens (resolution). Consequently, ViTs tend to be more computation intensive and have higher latency compared to widely adopted lightweight CNNs , especially on resource-constrained mobile devices, limiting their wide deployment in real-world applications.

Many research efforts are taken to alleviate this limitation. Among them, one direction is to reduce the quadratic computation complexity of the attention mechanism. Swin and following works propose window-based attention such that the receptive field is constrained to a pre-defined window size, which also inspires subsequent work to refine attention patterns . With the pre-defined span of attention, the computation complexity becomes linear to resolution. However, sophisticated attention patterns are generally difficult to support or accelerate on mobile devices because of intensive shape and index operations. Another track is to combine lightweight CNN and attention mechanism to form a hybrid architecture . The benefit comes two-fold. First, convolutions are shift invariant and are good at capturing local and detailed information, which can be considered as a good complement to ViTs . Second, by placing convolutions in the early stages while placing MHSA in the last several stages to model global dependency, we can naturally avoid performing MHSA on high resolution and save computations . Albeit achieving satisfactory performance, the latency and model size are still less competitive compared to lightweight CNNs. For instance, MobileViT achieves better performance than MobileNetV2 while being at least $5\times$ slower on iPhone 12. As applicable to CNNs, architecture search, pruning, and quantization techniques are also thoroughly investigated . Nevertheless, these models still emerge obvious weaknesses, e.g., EfficientFormer-L1 achieves comparable speed and better performance than MobileNetV2$\times 1.4$, while being $2\times$ larger. Thus, a practical yet challenging question arises:

Can we design a transformer-based model that is both light and fast, and preserves high performance?

In this work, we address the above question and propose a new family of mobile vision backbones. We consider three vital factors: number of parameters, latency, and model performance, as they reflect disk storage and mobile applications. First, we introduce novel architectural improvements to form a strong design paradigm. Second, we propose a fine-grained architecture search algorithm that jointly optimizes model size and speed for transformer models. With our network design and search method, we obtain a series of models under various constraints of model size and speed while maintaining high performance, named EfficientFormerV2. Under the exact same size and latency (on iPhone 12), EfficientFormerV2-S0 outperforms MobileNetV2 by $3.5\%$ higher top-1 accuracy on ImageNet-1K . Compared to EfficientFormer-L1 , EfficientFormerV2-S1 has similar performance while being $2\times$ smaller and $1.3\times$ faster (Tab. 2). We further demonstrate promising results in downstream tasks such as detection and segmentation (Tab. 3). Our contributions can be concluded as follows.

We comprehensively study mobile-friendly design choices and introduce novel changes, which is a practical guide to obtaining ultra-efficient vision transformer backbones.

We propose a novel fine-grained joint search algorithm that simultaneously optimizes model size and speed for transformer models, achieving superior Pareto optimality.

For the first time, we show that vision transformer models can be as small and fast as MobileNetV2 while obtaining much better performance. EfficientFormerV2 can serve as a strong backbone in various downstream tasks.

Related Work

Vaswani et al. propose attention mechanism to model sequences in NLP task, which forms transformer architecture. Transformers are later adopted to vision tasks by Dosovitskiy et al. and Carion et al. . DeiT improves ViT by training with distillation and achieves competitive performance against CNNs. Later research further improves ViTs by incorporating hierarchical design , injecting locality with the aid of convolutions , or exploring different types of token mixing such as local attention , spatial MLP mixer , and non-parameterized pool mixer . With appropriate changes, ViTs demonstrate strong performance in downstream vision tasks . To benefit from the advantageous performance, efficient deployment of ViTs has become a research hotspot, especially for mobile devices . For reducing the computation complexity of ViTs, many works propose new modules and architecture design , while others eliminate redundancies in attention mechanism . Similar to CNNs, architecture search , pruning , and quantization are also explored for ViTs.

We conclude two major drawbacks of the study in efficient ViT. First, many optimizations are not suitable for mobile deployment. For example, the quadratic computation complexity of the attention mechanism can be reduced to linear by regularizing the span or pattern of attention mechanism . Still, the sophisticated reshaping and indexing operations are not even supported on resource-constrained devices . It is crucial to rethink the mobile-friendly designs. Second, though recent hybrid designs and network search methods reveal efficient ViTs with strong performance , they mainly optimize the Pareto cure for one metric while being less competitive in others. For example, MobileViT is parameter efficient while being times slower than lightweight CNNs . EfficientFormer wields ultra-fast speed on mobile, but the model size is enormous. LeViT and MobileFormer achieve favorable FLOPs at the cost of redundant parameters.

Rethinking Hybrid Transformer Network

In this section, we study the design choices for efficient ViTs and introduce the changes that lead to the smaller size and faster speed without a performance drop. EfficientFormer-L1 is used as a baseline model given its superior performance on mobile devices.

Incorporating local information can improve performance and make ViTs more robust to the absence of explicit positional embedding . PoolFormer and EfficientFormer employ $3\times 3$ average pooling layers (Fig. 2(a)) as local token mixer. Replacing these layers with depth-wise convolutions (DWCONV) of the same kernel size does not introduce latency overhead, while the performance is improved by $0.6\%$ with negligible extra parameters ($0.02$M). Further, recent work suggest that it is also beneficial to inject local information modeling layers in the Feed Forward Network (FFN) in ViTs to boost performance with minor overhead. It is noteworthy that by placing extra depth wise $3\times 3$ convolutions in FFNs to capture local information, the functionality of original local mixer (pooling or convolution) is duplicated. Based on these observations, we remove the explicit residual-connected local token mixer and move the dept-wise $3\times 3$ CONV into the FFN, to get a unified FFN (Fig. 2(b)) with locality enabled. We apply the unified FFN to all stages of the network, as in Fig. 2(a,b). Such design modification simplifies the network architecture to only two types of blocks (local FFN and global attention), and boosts the accuracy to $80.3\%$ at the same latency (see Tab. 1) with minor overhead in parameters ($0.1$M). More importantly, this modification allows us to directly search the network depth with the exact number of modules in order to extract local and global information, especially at the late stages of the network, as discussed in Sec. 4.2.

2 Search Space Refinement

With the unified FFN and the deletion of residual-connected token mixer, we examine whether the search space from EfficientFormer is still sufficient, especially in terms of depth. We vary the network depth (number of blocks in each stage) and width (number of channels), and find that deeper and narrower network leads to better accuracy ($0.2\%$ improvement), less parameters ($0.13$M reduction), and lower latency ($0.1$ms acceleration), as in Tab. 1. Therefore, we set this network as a new baseline (accuracy $80.5\%$) to validate subsequent design modifications, and enable a deeper supernet for architecture search in Sec. 4.2.

In addition, 5-stage models with further down-sized spatial resolution ($\frac{1}{64}$) have been widely employed in efficient ViT arts . To justify whether we should search from a 5-stage supernet, we append an extra stage to current baseline network and verify the performance gain and overhead. It is noteworthy that though computation overhead is not a concern given the small feature resolution, the additional stage is parameter intensive. As a result, we need to shrink the network dimension (depth or width) to align parameters and latency to the baseline model for fair comparison. As seen in Tab. 1, the best performance of the 5-stage model surprisingly drops to $80.31\%$ with more parameters ($0.39$M) and latency overhead ($0.2$ms), despite the saving in MACs ($0.12$G). This aligns with our intuition that the fifth stage is computation efficient but parameter intensive. Given that 5-stage network can not introduce more potentials in our size and speed scope, we stick to 4-stage design. This analysis also explains why some ViTs offer an excellent Pareto curve in MACs-Accuracy, but tend to be quite redundant in size . As the most important takeaway, optimizing single metric is easily trapped, and the proposed joint search in Sec. 4.2 provides a feasible solution to this issue.

3 MHSA Improvements

We then study the techniques to improve the performance of attention modules without raising extra overhead in model size and latency. As shown in Fig. 2(c), we investigate two approaches for MHSA. First, we inject local information into the Value matrix ($V$) by adding a depth-wise $3\times 3$ CONV, which is also employed by . Second, we enable communications between attention heads by adding fully connected layers across head dimensions that are shown as Talking Head in Fig. 2(c). With these modifications, we further boost the performance to $80.8\%$ with similar parameters and latency compared to the baseline model.

4 Attention on Higher Resolution

Attention mechanism is beneficial to performance. However, applying it to high-resolution features harms mobile efficiency since it has quadratic time complexity corresponding to spatial resolution. We investigate strategies to efficiently apply MHSA to higher resolution (early stages). Recall that in the current baseline network obtained in Sec. 3.3, MHSA is only employed in the last stage with $\frac{1}{32}$ spatial resolution of the input images. We apply extra MHSA to the second last stage with $\frac{1}{16}$ feature size, and observe $0.9\%$ gain in accuracy. On the down side, the inference speed slows down by almost $2.7\times$. Thus, it is necessary to properly reduce complexity of the attention modules.

Although some work propose window-based attention , or downsampled Keys and Values to alleviate this problem, we find that they are not best-suited options for mobile deployment. Window-based attention is difficult to accelerate on mobile devices due to the sophisticated window partitioning and reordering. As for downsampling Keys ($K$) and Values ($V$) in , full resolution Queries ($Q$) are required to preserve the output resolution ($\mathbf{Out}$) after attention matrix multiplication:

where $B$, $H$, $N$, $C$ denotes batch size, number of heads, number of tokens, and channel dimension respectively. Based on our test, the latency of the model merely drops to $2.8$ms, which is still $2\times$ slower than the baseline model.

Therefore, to perform MHSA at the earlier stages of the network, we downsample all Query, Key, and Value to a fixed spatial resolution ($\frac{1}{32}$) and interpolate the outputs from the attention back to the original resolution to feed into the next layer, as shown in Fig. 2((d)&(e)). We refer to this method as Stride Attention. As in Tab. 1, this simple approximation significantly reduces the latency from $3.5$ms to $1.5$ms and preserves a competitive accuracy ($81.5\%$ vs. $81.7\%$).

5 Dual-Path Attention Downsampling

Most vision backbones utilize strided convolutions or pooling layers to perform a static and local downsampling and form a hierarchical structure. Some recent works start to explore attention downsampling. For instance, LeViT and UniNet propose to halve feature resolution via attention mechanism to enable context-aware downsampling with the global receptive field. Specifically, the number of tokens in Query is reduced by half so that the output from the attention module is downsampled:

However, it is nontrivial to decide how to reduce the number of tokens in Query. Graham et al. empirically use pooling to downsample Query , while Liu et al. propose to search for local or global approaches . To achieve acceptable inference speed on mobile devices, applying attention downsampling to early stages with high resolution is not favorable, restricting the values of existing works that search different downsampling approaches at higher-resolution.

Instead, we propose a combined strategy, which is dual-path attention downsampling, that wields both locality and global dependency, as in Fig. 2(f). To get downsampled Queries, we use pooling as static local downsampling, $3\times 3$ DWCONV as learnable local downsampling, and combine and project the results into Query dimension. In addition, the attention downsampling module is residual connected to a regular strided CONV to form a local-global manner, similar to the downsampling bottlenecks or inverted bottlenecks . As shown in Tab. 1, with slightly more parameters and latency overhead, we further improve the accuracy to $81.8\%$ with dual-path attention downsampling, which also has better performance than only using attention module for subsampling, i.e., attention downsampling.

EfficientFormerV2

As discussed, current arts merely focus on optimizing one metric, thus are either redundant in size or slow in inference . To find the most suitable vision backbones for mobile deployment, we propose to jointly optimize model size and speed. Furthermore, the network designs in Sec. 3 favor a deeper network architecture (Sec. 3.2) and more attentions (Sec. 3.4), calling for an improved search space and algorithm. In what follows, we present the supernet design of EfficientFormerV2 and its search algorithm.

As discussed in Sec. 3.2, we employ a 4-stage hierarchical design which obtains feature sizes in $\{\frac{1}{4},\frac{1}{8},\frac{1}{16},\frac{1}{32}\}$ of the input resolution. Similar to its predecessor , EfficientFormerV2 starts with a small kernel convolution stem to embed input image instead of using inefficient embedding of non-overlapping patches,

where $\texttt{S}_{i,j}$ is a learnable layer scale and the FFN is constructed by two properties: stage width $C_{j}$ and a per-block expansion ratio $E_{i,j}$. Note that each FFN is residual connected. In the last two stages, both local FFN and global MHSA blocks are used. Therefore, on top of Eqn. 4, global blocks are defined as:

with ab as a learnable attention bias for position encoding.

2 Jointly Optimizing Model Size and Speed

Though the baseline network EfficientFormer is found by latency-driven search and wields fast inference speed on mobile, there are two major drawbacks for the search algorithm. First, the search process is merely constrained by speed, resulting in the final models being parameter redundant, as in Fig. 1. Second, it only searches for depth (number of blocks $N_{j}$ per stage) and stage width $C_{j}$, which is in a coarse-grained manner. In fact, the majority of computations and parameters of the network are in FFNs, and the parameter and computation complexity are linearly related to its expansion ratio $E_{i,j}$. $E_{i,j}$ can be specified independently for each FFN without the necessity to be identical. Thus, searching $E_{i,j}$ enables a more fine-grained search space where the computations and parameters can distribute flexibly and non-uniformly within each stage. This is a missing property in most recent ViT NAS arts , where $E_{i,j}$ remains identical per stage. We propose a search algorithm that enables a flexible per-block configuration, with joint constraints on size and speed, and finds vision backbones best suited for mobile devices.

First, we introduce the metric guiding our joint search algorithm. Given the fact that the size and latency of a network all matter when evaluating mobile-friendly models, we consider a generic and fair metric that better understands the performance of a network on mobile devices. Without loss of generality, we define a Mobile Efficiency Score (MES):

where $i\in\{{size},latency,...\}$ and $\alpha_{i}\in(0,1]$ indicating the corresponding importance. $M_{i}$, and $U_{i}$ represent the metric and its unit. $Score$ is a pre-defined base score set as $100$ for simplicity. Model size is calculated by the number of parameters, and latency is measured as running time when deploying models on devices. Since we focus on mobile deployment, the size and speed of MobileNetV2 are used as the unit. Specifically, we define $U_{size}=3\text{M}$, and $U_{latency}$ as $1$ms latency on iPhone 12 (iOS 16) deployed with CoreMLTools . To emphasize speed, we set $\alpha_{latency}=1.0$ and $\alpha_{size}=0.5$. Decreasing size and latency can lead to a higher MES, and we search for Pareto optimality on MES-Accuracy. The form of MES is general and can be extended to other metrics of interest, such as inference-time memory footprint and energy consumption. Furthermore, the importance of each metric is easily adjustable by appropriately defining $\alpha_{i}$.

2.2 Search Space and SuperNet

Search space consists of: (i) the depth of the network, measured by the number of blocks $N_{j}$ per stage, (ii) the width of the network, i.e., the channel dimension $C_{j}$ per stage, and (iii) expansion ratio $E_{i,j}$ of each FFN. The amount of MHSA can be seamlessly determined during depth search, which controls the preservation or deletion of a block in the supernet. Thus, we set every block as MHSA followed by FFN in the last two stages of the supernet and obtain subnetworks with the desired number of global MHSA by depth search.

Supernet is constructed by using a slimmable network that executes at elastic depth and width to enable a pure evaluation-based search algorithm. Elastic depth can be naturally implemented through stochastic drop path augmentation . As for width and expansion ratio, we follow Yu et al. to construct switchable layers with shared weights but independent normalization layers, such that the corresponding layer can execute at different channel numbers from a predefined set, i.e., multiples of $16$ or $32$. Specifically, the expansion ratio $E_{i,j}$ is determined by the channels of the depth-wise $3\times 3$ Conv in each FFN, and stage width $C_{j}$ is determined by aligning the output channels of the last projection ($1\times 1$ Conv) of FFN and MHSA blocks. The switchable execution can be expressed as:

where $w^{:c}$ refers to slicing the first $c$ filters of the weight matrix to obtain a subset of output, and $\gamma_{c}$, $\beta_{c}$, $\mu_{c}$, and $\sigma_{c}$ are the parameters and statistics of the normalization layer designated for width $c$. The supernet is pre-trained with Sandwich Rule by training the largest, the smallest, and randomly sampled two subnets at each iteration (we denote these subnets as max, min, rand-1, and rand-2 in Alg. 1).

Discussion. The pruning of our supernet is partially inspired by the slimmable network . However, the differences are also significant. First, the search objective is different. We apply the introduced joint search objective for optimizing model size and efficiency (Sec. 4.2.1). Second, the search actions are different. Depth is pruned through the reduction of each block, which is possible since we unify the design and only adopt two blocks: Feed Forward Network (Sec. 3.1) and attention block. The way of pruning the depth is different from the slimmable network. Unifying all the flexible search actions under one joint objective has not been studied for transformers before.

2.3 Search Algorithm

Now that search objective, search space, and supernet are formulated, we present the search algorithm. Since the supernet is executable at elastic depth and switchable width, we can search the subnetworks with the best Pareto curve by analyzing the efficiency gain and accuracy drop with respect to each slimming action. We define the action pool as:

where $A_{N[i,j]}$ denotes slimming each block, $A_{C[j]}$ refers to shrinking the width of a stage, and $A_{E[i,j]}$ denotes slimming each FFN to a smaller expansion. Initializing the state with full depth and width (largest subnet), we evaluate the accuracy outcome ($\Delta\texttt{Acc}$) of each frontier action on a validation partition of ImageNet-1K, which only takes about 4 GPU-minutes. Meanwhile, parameter reduction ($\Delta\texttt{Params}$) can be directly calculated from layer properties, i.e., kernel size, in-channels, and out-channels. We obtain the latency reduction ($\Delta\texttt{Latency}$) through a pre-built latency look-up table measured on iPhone 12 with CoreMLTools. With the metrics in hand, we can compute $\Delta\texttt{MES}$ through $\Delta\texttt{Params}$ and $\Delta\texttt{Latency}$, and choose the action with the minimum per-MES accuracy drop: $\hat{A}\leftarrow{\arg\min}_{A}\frac{\Delta\texttt{Acc}}{\Delta\texttt{MES}}$. It is noteworthy that though the action combination is enormous, we only need to evaluate the frontier one at each step, which is linear in complexity. Details can be found in Alg. 1.

Experiments

Implementation Details. We implement the model through PyTorch 1.12 and Timm library , and use $16$ NVIDIA A100 GPUs to train our models. We train the models from scratch by $300$ and $450$ epochs on ImageNet-1K , with AdamW optimizer. Learning rate is set to $10^{-3}$ per $1,024$ batch size with cosine decay. We use a standard image resolution, i.e., $224\times 224$, for both training and testing. Similar to DeiT , we use RegNetY-16GF with $82.9\%$ top-1 accuracy as the teacher model for hard distillation. We use three testbeds to benchmark the latency:

iPhone 12 - NPU. We get the latency on iPhone 12 (iOS 16) by running the models on Neural Engine (NPU). The models (batch size of $1$) are compiled with CoreML .

Pixel 6 - CPU. We test model latency on Pixel 6 (Android) CPU. To obtain the latency for most works under comparison, we replace the activation from all models to ReLU to get fair comparisons. The models (batch size of $1$) are compiled with XNNPACK .

Nvidia GPU. We also provide the latency on a high-end GPU–Nvidia A100. The models (batch size of $64$) are deployed in ONNX and executed by TensorRT .

Evaluation on Single Metric. We show the comparison results in Tab. 2, which includes the most recent and representative works on vision transformers and CNNs. The works that do not have public models or are not compatible with mobile devices are not contained in Tab. 2. EfficientFormerV2 series achieve the SOTA results on a single metric, i.e., number of parameters or latency. For model size, EfficientFormerV2-S0 outperforms EdgeViT-XXS by $1.3\%$ top-1 accuracy with even $0.6$M fewer parameters and MobileNetV2$\times 1.0$ by $3.5\%$ top-1 with similar number of parameters. For large models, EfficientFormerV2-L model achieves identical accuracy to recent EfficientFormer-L7 while being $3.1\times$ smaller. As for speed, with comparable or lower latency, EfficientFormerV2-S2 outperforms UniNet-B1 , EdgeViT-S , and EfficientFormer-L1 by $0.8\%$, $0.6\%$ and $2.4\%$ top-1 accuracy, respectively. We hope the results can provide practical insight to inspire future architecture design: modern deep neural networks are robust to architecture permutation, optimizing the architecture with joint constraints, such as latency and model size, will not harm individual metrics.

Jointly Optimized Size and Speed. Further, we demonstrate the superior performance of EfficientFormerV2 when considering both model size and speed. Here we use MES as a more practical metric to assess mobile efficiency than using size or latency alone. EfficientFormerV2-S1 outperforms MobileViT-XS , EdgeViT-XXS , and EdgeViT-XS by $4.2\%$, $4.6\%$, and $1.5\%$ top-1, respectively, with far higher MES. With $1.8\times$ higher MES, EfficientFormerV2-L outperforms MobileFormer-508M by $4.0\%$ top-1 accuracy. The evaluation results answer the central question raised at the beginning: with the proposed mobile efficiency benchmark (Sec. 4.2.1), we can avoid entering a pitfall achieving seemingly good performance on one metric while sacrificing too much for others. Instead, we can obtain efficient mobile ViT backbones that are both light and fast.

2 Downstream Tasks

Object Detection and Instance Segmentation. We apply EfficientFormerV2 as backbone in Mask-RCNN pipeline and experiment on MS COCO 2017 . The model is initialized with ImageNet-1K pretrained weights. We use AdamW optimizer with an initial learning rate as $2\times 10^{-4}$ and conduct training for $12$ epochs with resolution as $1333\times 800$. Following Li et al. , we apply a weight decay as $0.05$ and freeze the normalization layers in the backbone. As in Tab. 3, with similar model size, our EfficientFormerV2-S2 outperform PoolFormer-S12 by $6.1$ APbox and $4.9$ APmask. EfficientFormerV2-L outperforms EfficientFormer-L3 by $3.3$ APbox and $2.3$ APmask.

Semantic Segmentation. We perform experiments on ADE20K , a challenging scene segmentation dataset with $150$ categories. Our model is integrated as a feature encoder in Semantic FPN pipeline, with ImageNet-1K pretrained weights. We train our model on ADE20K for $40$K iterations with batch size as $32$ and learning rate as $2\times 10^{-4}$ with a poly decay by the power of $0.9$. We apply weight decay as $10^{-4}$ and freeze the normalization layers. Training resolution is $512\times 512$, and we employ a single scale testing on the validation set. As in Tab. 3, EfficientFormerV2-S2 outperforms PoolFormer-S12 and EfficientFormer-L1 by $5.2$ and $3.5$ mIoU, respectively.

3 Ablation Analysis on Search Algorithm

We compare the proposed search algorithm with the random search and the one from EfficientFormer . As seen in Tab. LABEL:tab:ablate_search, our search algorithm obtains models with much better performance than random search, i.e., Random 1 and Random 2. Compared with EfficientFormer , we achieve higher accuracy under similar parameters and latency, demonstrating the effectiveness of fine-grained search and joint optimization of latency and size.

Discussion and Conclusion

In this work, we comprehensively study transformer backbones, identify inefficient designs, and introduce mobile-friendly novel architectural changes. We further propose a fine-grained joint search on size and speed and obtain the EfficientFormerV2 model family. We extensively benchmark and compare our work with existing studies on different hardware and demonstrate that EfficientFormerV2 is both lightweight, ultra-fast in inference speed and high performance. Since we focus on size and speed, one future direction is to apply the joint optimization methodology to subsequent research exploring other critical metrics, such as memory footprint and CO2 emission.

Acknowledgements

This work is partly supported by the Army Research Office/Army Research Laboratory via grant W911-NF-20-1-0167 to Northeastern University, CNS1909172.

References