# Multi-Scale VMamba: Hierarchy in Hierarchy
Visual State Space Model

Yuheng Shi
City University of Hong Kong
<EMAIL_ADDRESS>
&Minjing Dong
City University of Hong Kong
<EMAIL_ADDRESS>
&Chang Xu
University of Sydney
<EMAIL_ADDRESS>

###### Abstract

Despite the significant achievements of Vision Transformers (ViTs) in various
vision tasks, they are constrained by the quadratic complexity. Recently,
State Space Models (SSMs) have garnered widespread attention due to their
global receptive field and linear complexity with respect to the input length,
demonstrating substantial potential across fields including natural language
processing and computer vision. To improve the performance of SSMs in vision
tasks, a multi-scan strategy is widely adopted, which leads to significant
redundancy of SSMs. For a better trade-off between efficiency and performance,
we analyze the underlying reasons behind the success of the multi-scan
strategy, where long-range dependency plays an important role. Based on the
analysis, we introduce Multi-Scale Vision Mamba (MSVMamba) to preserve the
superiority of SSMs in vision tasks with limited parameters. It employs a
multi-scale 2D scanning technique on both original and downsampled feature
maps, which not only benefits long-range dependency learning but also reduces
computational costs. Additionally, we integrate a Convolutional Feed-Forward
Network (ConvFFN) to address the lack of channel mixing. Our experiments
demonstrate that MSVMamba is highly competitive, with the MSVMamba-Tiny model
achieving 82.8% top-1 accuracy on ImageNet, 46.9% box mAP, and 42.2% instance
mAP with the Mask R-CNN framework, 1x training schedule on COCO, and 47.6%
mIoU with single-scale testing on ADE20K. Code is available at
https://github.com/YuHengsss/MSVMamba.

##  1 Introduction

In the domain of computer vision, the extraction of features plays a pivotal
role in the performance of various tasks, ranging from image classification to
more complex applications like detection and segmentation. Traditionally,
Convolutional Neural Networks (CNNs) [26, 41, 19, 22, 34] have been the
backbone of feature extraction methodologies, prized for their linear scaling
complexity and proficiency in capturing local patterns. However, CNNs often
fall short in encapsulating global context, a limitation that becomes
increasingly apparent in tasks requiring a comprehensive understanding of the
entire visual field. In contrast, Vision Transformers (ViTs) [6, 33, 51, 44]
have emerged as a compelling alternative, boasting an inherent global
receptive field that allows for the direct capture of long-range dependencies
within an image. Despite their advantages, ViTs are hampered by their
quadratic scaling complexity concerning the input size, which significantly
constrains their applicability to downstream tasks such as object detection
and segmentation, where efficiency is paramount. Recently, State Space Model
(SSM)-based approaches [12, 42, 8] have garnered attention for their ability
to combine the best of both worlds: a global receptive field and linear
scaling complexity. Notably, Mamba [10] introduces a hardware-aware and input-
dependent algorithm that significantly enhances the performance and efficiency
of SSMs. Inspired by Mamba’s success, a burgeoning body of work has sought to
leverage its advantages for vision tasks, pioneering efforts such as ViM [61]
and VMamba [31].

Figure 1: Comparison on ImageNet.

The S6 block, developed by Mamba [10], was originally designed for NLP tasks.
To adapt S6 for vision tasks, images are first divided into patches and then
flattened into a patch sequence along the scanning path. To accommodate the
non-causal nature of image data, the multi-scan strategy is widely adopted for
vision tasks, such as ViM [61] which enhances the sequence by summing it in
both forward and reverse directions, and VMamba [31] which integrates
horizontal and vertical scans. However, unlike NLP models, which can have up
to billions of parameters, current vision backbones always take computational
costs into consideration, _i.e._ , the trade-off between accuracy and
efficiency. This constraint on model size inherently limits the long-range
modeling capabilities of SSMs in vision tasks. Taking ViM-Tiny [61] as an
example, placing the cls token in the middle of the sequence yields markedly
better results than positioning it at the ends. This suggests that central
placement compensates for the model’s limited ability to integrate distant
information, highlighting the difficulties of handling long-range dependencies
in parameter-constrained vision models. We refer to this as the long-range
forgetting problem. In this work, we analyze how the multi-scan strategy in
[61, 31] helps to alleviate this problem. Compared to the single-scan
strategy, the multi-scan one allows long-range decay to manifest in various
directions within 2D images. However, the increased scanning routes bring
multiples of computations, significantly increasing redundancy and limiting
efficiency. Thus, we aim to pursue a better trade-off between the performance
and efficiency of Mamba in vision tasks.

The most direct and effective method to address the long-range forgetting
problem is to shorten the sequence length, which can be achieved by
downsampling the feature map. However, placing all scanning routes on a
downsampled feature map could result in the loss of fine-grained features and
the downstream task performance. Through the visualization of different scans,
we show that the decay rates could vary for different scanning routes, which
motivates us to develop a hierarchical design of multi-scan. In this work, we
propose a Multi-Scale 2D (MS2D) scanning strategy to alleviate the long-range
forgetting problem with limited computational costs. Specifically, we divide
the scanning directions of SS2D [31] into two groups: one retains the original
resolution and is processed by the S6 block, while the others are downsampled,
processed by the S6 block, and then upsampled, which not only shortens the
sequence length for long-range dependencies learning but also alleviates
redundancy. Building on the VMamba with its hierarchical architecture, we
incorporate another hierarchical design within the block, creating a hierarchy
within a hierarchy. Furthermore, we introduce a Convolutional Feed-Forward
Network (ConvFFN) within each block to bolster the model’s capability for
channel-wise information exchange and local feature capture.

We conduct extensive experiments to validate the effectiveness of MSVmamba
across a spectrum of tasks, including image classification, object detection,
and semantic segmentation. Detailed comparisons on the ImageNet-1K [3] dataset
are illustrated in Fig. 1. Specifically, MSVMamba-Tiny achieves a notable
improvement over VMamba-Tiny, enhancing the Top-1 accuracy by 0.6% while
concurrently reducing the computational demand by 17% in terms of GFLOPs.
Furthermore, our MSVMamba model also surpasses SOTA SSM-based models in
performance on the COCO [29] and ADE20K [59] datasets for object detection and
semantic segmentation tasks, respectively.

##  2 Related work

###  2.1 Generic Vision Backbone

The evolution of generic vision backbones has significantly shaped the
landscape of computer vision, transitioning from CNNs [26, 41, 19, 55, 22, 20,
43, 34] to ViTs [6, 33, 32, 51, 52, 5, 44, 58]. CNNs have been the cornerstone
of vision-based models, dominating vision tasks in the early era of deep
learning. The classic CNNs such as AlexNet [26], VGG [41], and ResNet [19],
have paved the way for numerous subsequent innovations [55, 22, 20, 43, 21,
34, 4, 50]. These designs have significantly improved performance on a wide
range of vision tasks by enhancing the network’s ability to capture complex
patterns and features from visual data. The Vision Transformer [6], drawing
inspiration from the success of transformers [48] in natural language
processing, has emerged as a formidable contender to conventional CNNs for
vision-related tasks. ViT reimagines image processing by segmenting an image
into patches and employing self-attention mechanisms to process these
segments. This innovative approach enables the model to discern global
dependencies across the entire image, a significant leap forward in
understanding complex visual data. However, the ViT architecture demands
considerable computational resources and extensive datasets for effective
training. Moreover, its performance is intricately tied to the input sequence
length, exhibiting a quadratic complexity that can escalate processing costs.
In response, subsequent research has focused on developing more efficient
training strategies [44, 46, 25], hierarchical network structures [33, 32, 5,
51, 52, 16], refined spatial attention mechanisms [58, 60, 53, 47, 17, 5] and
convolution-based design [2, 30, 15, 37] to address these issues.

###  2.2 State Space Models

State Space Models (SSMs) [11, 13, 12, 42, 8] have garnered increasing
attention from researchers due to their computational complexity, which grows
linearly with the length of the input sequence, and their inherent global
awareness properties. To reduce the computational resource consumption of
SSMs, S4 [12] introduced a diagonal structure and combined it with a diagonal
plus low-rank approach to construct structured SSMs. Subsequently, S5 [42] and
H3 [8] further enhanced the efficiency of SSM-based models by introducing
parallel scanning techniques and improving hardware utilization. Mamba [10]
then introduced the S6 block, incorporating data-dependent parameters to amend
the Linear Time Invariant(LTI) characteristics of previous SSM models,
demonstrating superior performance over transformers on large-scale datasets.
In the realm of vision tasks, S4ND [36] pioneered the application of SSM
models in vision tasks by treating visual data as 1D, 2D, and 3D signals.
U-Mamba [35] combined CNNs with SSMs for medical image segmentation. ViM [61]
and VMamba [31] integrated the S6 block into vision backbone design, employing
multiple scanning directions to accommodate the non-casual nature of image
data, achieving competitive results against ViTs and CNNs. Motivated by the
success of ViM and VMamba, a plethora of Mamba-based works [38, 24, 7, 56, 1,
27, 57, 40] have emerged across various vision tasks, including vision
backbone design [38, 24, 56], medical image segmentation [40, 28], and video
classification [27], showcasing the potential of SSM-based approaches in
advancing the field of computer vision.

##  3 Method

In this section, we first summarize the state space model in Section 3.1.
Subsequently, in Section 3.2, we provide an in-depth analysis of the multi-
scan strategy in existing vision Mamba models. Following the analysis, Section
3.3 tackles the redundancy and long-range dependency issue by introducing a
Multi-Scale 2D (MS2D) scanning strategy. Finally, Section 3.4 details the
integration of the Multi-Scale State Space (MS3) block, which incorporates the
MS2D technique alongside a ConvFFN. Building upon the MS3 block, various model
configurations are developed across different scales, illustrating the
adaptability and scalability of our proposed approach.

###  3.1 Preliminaries

State Space Models. Classical State Space Models (SSMs) represent a continuous
system that maps an input sequence $x(t)\in\mathbb{R}^{L}$ to a latent space
representation $h(t)\in\mathbb{R}^{N}$ and subsequently predicts an output
sequence $y(t)\in\mathbb{R}^{L}$ based on this representation. Mathematically,
an SSM can be described as follows:

$\begin{array}[]{l}h^{\prime}(t)=\mathbf{A}h(t)+\mathbf{B}x(t),\;y(t)=\mathbf{C}h(t),\end{array}$
(1)

where $\mathbf{A}\in\mathbb{R}^{N\times N}$, $\mathbf{B}\in\mathbb{R}^{N\times
1}$ and $\mathbf{C}\in\mathbb{R}^{1\times N}$ are learnable parameters.

(a) Horizontal Scan

(b) Vertical Scan

(c) Decay Ratio

(d) Binary Decay Ratio

Figure 2: Illustration of decay along horizontal, vertical scanning routes and
their ratio.

Discretization. To adapt continuous State Space Models (SSMs) for use within
deep learning frameworks, it is crucial to implement discretization
operations. By incorporating a timescale parameter
$\mathbf{\Delta}\in\mathbb{R}$ and employing the widely utilized zero-order
hold (ZOH) as the discretization rule, the discretized versions of
$\mathbf{A}$ and $\mathbf{B}$ (denoted as $\overline{\textbf{A}}$ and
$\overline{\mathbf{B}}$, respectively) can be derived, with which, Eq. 1 can
be reformulated into a discretized manner as:

$\begin{array}[]{l}h(t)=\overline{\mathbf{A}}h(t-1)+\overline{\mathbf{B}}x(t),\;y(t)=\mathbf{C}h(t),\\\
\text{where
}\overline{\mathbf{A}}=e^{\mathbf{\Delta}\mathbf{A}},\;\overline{\mathbf{B}}=(\mathbf{\Delta}\mathbf{A})^{-1}(e^{\mathbf{\Delta}\mathbf{A}}-\mathbf{I})\mathbf{\Delta}\mathbf{B}\approx\mathbf{\Delta}\mathbf{B},\\\
\end{array}$ (2)

where $\mathbf{I}$ denotes the identity matrix. Afterward, the process of Eq.
2 could be implemented in a global convolution manner as:

$\begin{array}[]{l}y=x\odot\overline{\mathbf{K}},\;\overline{\mathbf{K}}=\left(\mathbf{C}\overline{\mathbf{B}},\mathbf{C}\overline{\mathbf{AB}},\ldots,\mathbf{C}\overline{\mathbf{A}}^{L-1}\overline{\mathbf{B}}\right),\end{array}$
(3)

where $\overline{\mathbf{K}}\in\mathbb{R}^{L}$ is the convolution kernel.

Selective State Space Models. The Selective State Space (S6) mechanism,
introduced by Mamba [10], renders the parameters
$\overline{\mathbf{B}},\overline{\mathbf{C}}$, and $\mathbf{\Delta}$ input-
dependent, thereby enhancing the performance of SSM-based models. After making
$\overline{\mathbf{B}},\overline{\mathbf{C}}$, and $\mathbf{\Delta}$ input-
dependent, the global convolution kernel in Eq. 3 could be rewritten as:

$\overline{\mathbf{K}}=\left(\mathbf{C}_{L}\overline{\mathbf{B}}_{L},\mathbf{C}_{L}\overline{\mathbf{A}}_{L-1}\overline{\mathbf{B}}_{L-1},\ldots,\mathbf{C}_{L}\prod_{i=1}^{L-1}\mathbf{\overline{A}}_{i}\overline{\mathbf{B}}_{1}\right).$
(4)

###  3.2 Analysis of Multi-Scan Strategy

When processing image data using the S6 block, the 2D feature map
$\mathbf{Z}\in\mathbb{R}^{H\times W\times D}$ is flattened into a 1D sequence
of image tokens, denoted as $\mathbf{X}\in\mathbb{R}^{L\times D}$. According
to Eq. 4 and Eq. 2, the contribution of the $m_{th}$ token to the construction
of the $n_{th}$ token ($m<n$) in S6 can be expressed as:

$\displaystyle\mathbf{C}_{n}\prod_{i=m}^{n}\mathbf{\overline{\mathbf{A}}}_{i}\overline{\mathbf{B}}_{m}=\mathbf{C}_{n}\overline{\mathbf{A}}_{(m\rightarrow
n)}\overline{\mathbf{B}}_{m},\text{ where
}\overline{\mathbf{A}}_{(m\rightarrow
n)}=e^{\sum_{i=m}^{n}\mathbf{\Delta}_{i}\mathbf{A}}.\vspace{-3mm}$ (5)

Typically, the learned $\mathbf{\Delta}_{i}\mathbf{A}$ is negative, which
biases the model towards prioritizing recent tokens’ information.
Consequently, as the sequence length increases, the exponential term
$e^{\sum_{i=m}^{n}\mathbf{\Delta}_{i}\mathbf{A}}$ in Eq. 5 decays
significantly, resulting in minimal contributions from distant tokens. We
refer to it as the long-range forgetting issue, which has also been observed
in recent studies applying S6 to vision tasks [14]. Although this problem can
be mitigated by increasing the number of parameters and the depth of the
model, such adjustments introduce additional computational costs. Furthermore,
the causal property of the S6 block ensures that information can only
propagate in a unidirectional manner between tokens, preventing earlier tokens
from accessing information from subsequent tokens.

The inherent non-causal nature of images renders the direct application of the
S6 block to vision-related tasks less than optimal, as identified by ViM [61].
To mitigate this limitation, ViM [61] and VMamba [31] have introduced
methodologies that entail scanning image features across various directions
and then integrating these features. Generally, the updated token along one of
the scanning routes, denoted as $Scan(\mathbf{Z}_{(p,q)})$, where $(p,q)$
indicates the coordinate, could be obtained by:

$\begin{array}[]{c}Scan(\mathbf{Z}_{(p,q)})=\mathbf{C}_{\alpha}\sum_{i=1}^{\alpha}\overline{\mathbf{A}}_{(i\rightarrow\alpha)}\overline{\mathbf{B}}_{i}\sigma(\mathbf{Z})_{i}.\\\
\end{array}$ (6)

In Eq. LABEL:eq:general_scan, $\sigma$ represents the transformation that
converts a 2D feature map into a 1D sequence, and $\alpha$ denotes the
corresponding index of $\mathbf{Z}{(p,q)}$ in the transformed 1D sequence.
Afterward, results from multi-scan routes are added together to produce
enhanced feature $\mathbf{Z}^{\prime}{(p,q)}$, which can be denoted as:

$\mathbf{Z^{\prime}}_{(p,q)}=\sum_{k}Scan_{k}(\mathbf{Z}_{(p,q)})=\sum_{k}\mathbf{C}_{\alpha_{k}}\sum_{i=1}^{\alpha_{k}}\overline{\mathbf{A}}_{(i\rightarrow\alpha_{k})}\overline{\mathbf{B}}_{i}\sigma_{k}(\mathbf{Z})_{i}.\\\
$ (7)

This multi-scan strategy allows tokens to access information from each other.
In ViM [61], two distinct scanning routes correspond to two different
transformations in Eq. LABEL:eq:general_scan, specifically, the flatten and
the flatten with flip transformations. Similarly, VMamba [31] extends the
basic bidirectional scanning by incorporating both horizontal and vertical
scanning directions, yielding four distinct scanning routes. Besides, the
multi-scan strategy also alleviates the long-range forgetting problem by
minimizing the effective distance between tokens. For tokens at coordinates
$(p_{1},q_{1})$ and $(p_{2},q_{2})$, the strategy employs multiple scanning
routes, each potentially altering their relative positions. The minimum
distance across these routes is given by
$\min_{k}d_{k}((p_{1},q_{1})\rightarrow(p_{2},q_{2}))$, where $d_{k}$
represents the distance between the tokens in the sequence generated by the
$k$-th scan. By reducing this distance, the multi-scan strategy reduces the
decay of influence between distant tokens, thereby enhancing the model’s
ability to maintain and utilize long-range information.

To more intuitively demonstrate the relationship between the multi-scan
strategy and long-range decay, we visualize the exponential term
$e^{\sum_{i=1}^{\alpha}\mathbf{\Delta}_{i}\mathbf{A}}$ along the horizontal
and vertical scanning directions in VMamba-Tiny with respect to the central
token in Fig. 2(a) and Fig. 2(b). Specifically, we randomly select 50 images
from the ImageNet [3] validation set and compute the average decay along
scanning routes at the last layer of the final stage across these images and
feature dimensions. We use a higher input resolution to enhance the quality of
the visualization.

According to these observations, the success of the multi-scan strategy in
VMamba can be attributed to its mitigation of the non-causal properties of
image data and alleviation of the long-range forgetting problem. However, as
the number of scanning routes increases, the computational cost also rises
linearly, introducing computational redundancy. In Fig. 2(c), we illustrate
the maximum ratio of Fig. 2(a) to Fig. 2(b) and vice versa. While in Fig.
2(d), we present a binarized version of Fig. 2(c), applying a threshold of 10,
which covers more than 40% of the entire figure. This phenomenon indicates
that the varying decay rates across different scanning routes lead to certain
routes dominating the decay dynamics, which can also be attributed to the
long-range forgetting problem. The existence of dominant scanning routes
implies that some scans contribute significantly more to information retention
than others, leading to computational redundancy in the multi-scan strategy.

Figure 3: Illustration of the Multi-Scale 2D-Selective-Scan on an image

###  3.3 Multi-Scale 2D Scanning

As discussed in the last subsection, the contribution of tokens decays with
increasing scanning distance. The most effective and direct way to alleviate
the long-range forgetting problem is to reduce the number of tokens.
Simultaneously, since the computational complexity of the S6 block is linearly
dependent on the number of tokens, reducing the token count also enhances
efficiency. Thus, an alternative approach to address the aforementioned issue
is to apply the multi-scan strategy on a downsampled feature map. However,
setting all scans on a downsampled feature map will ignore fine-grained
features and result in unavoidable information loss. Thus, scanning along the
full-resolution feature map is also essential.

Motivated by these considerations, we introduce a simple yet effective Multi-
Scale 2D scanning(MS2D) strategy, as depicted in Fig. 3. Our approach
commences with the generation of hierarchical feature maps at varying scales,
achieved through the application of Depthwise Convolution (DWConv) with
distinct stride values. These multi-scale feature maps are then processed
through four distinct scanning routes within VMamba. Specifically, we uitlize
DWConvs with strides of 1 and $s$ to obtain feature map
$\mathbf{Z}_{1}\in\mathbb{R}^{H\times W\times D}$ and
$\mathbf{Z}_{2}\in\mathbb{R}^{\frac{H}{s}\times\frac{W}{s}\times D}$,
respectively. Afterwards, $\mathbf{Z}_{1}$ and $\mathbf{Z}_{2}$ are processed
by two S6 blocks as:

$\displaystyle\mathbf{Y}_{1}$ $\displaystyle=S6(\sigma_{1}(\mathbf{Z}_{1})),$
(8) $\displaystyle[\mathbf{Y}_{2},\mathbf{Y}_{3},\mathbf{Y}_{4}]$
$\displaystyle=S6([\sigma_{2}(\mathbf{Z}_{2}),\sigma_{3}(\mathbf{Z}_{2}),\sigma_{4}(\mathbf{Z}_{2})]),$
(9)

where $\sigma$ is transformation that convert 2D feature maps into 1D
sequences used in SS2D, and $\mathbf{Y}$ is the processed sequence. These
processed sequences are converted back into 2D feature maps, and the
downsampled feature maps are interpolated for merging:

$\displaystyle\mathbf{Z}^{\prime}_{i}$
$\displaystyle=\gamma_{i}(\mathbf{Y}_{i}),i\in\\{1,2,3,4\\},$ (10)
$\displaystyle\mathbf{Z}^{\prime}$
$\displaystyle=\mathbf{Z}^{\prime}_{1}+\text{Interpolate}(\sum(\mathbf{Z}^{\prime}_{j})),j\in\\{2,3,4\\},$
(11)

where $\gamma$ is the inverse transformation of $\sigma$ and
$\mathbf{Z}^{\prime}$ is the feature map enhanced by MS2D.

(a) SS2D

(b) MS2D

Figure 4: Illustration of the decay with different scanning routes in SS2D and
MS2D.

The downsampling operation reduces the sequence length by a factor of $s^{2}$,
which also shortens the distance between tokens in Eq. 5 by a factor of
$s^{2}$, thereby alleviating the long-range forgetting problem. As the
computational complexity of a single S6 block is $O(9LDN)$ [10], where $N$
denotes the SSM dimension, replacing SS2D with MS2D reduces the total sequence
length across four scans from $4L$ to $(1+3/s^{2})L$, thereby improves the
efficiency. Practically, the downsampling rate is set to 2. It is worth noting
that sequences from $\mathbf{Z}_{2}$ are processed by the same S6 block. This
approach maintains the same accuracy as using multiple S6 blocks for different
scanning routes while effectively reducing the number of parameters.

To better illustrate the alleviation of the long-range forgetting problem, we
also provide empirical evidence, as shown un Fig. 4. We compare the decay
along scanning routes in the SS2D of VMamba and our MS2D, focusing on the last
token with the same configuration as Fig. 2. The decay maps in downsampled
features are interpolated back. As observed, the decay rate along scanning
routes in downsampled maps is significantly alleviated, enhancing the
capability to capture global information.

###  3.4 Overall Model Architecture

Figure 5: Detailed architecture of Multi-Scale State Space (MS3) block,
consisting of a Multi-Scale Vision Space State (MSVSS) block and a
Convolutional Feed-Forward Network (ConvFFN) block.

In this study, we extend the capabilities of the VMamba framework by
substituting its VSS block with our Multi-Scale State Space (MS3) block. The
architectural framework of the MS3 block is delineated in Fig. 5, comprising a
Multi-Scale Vision Space State(MSVSS) component and a Convolutional Feed-
Forward Network (ConvFFN). The MSVSS component is devised by adapting the
vision state space framework in VMamba, substituting the SS2D with an MS2D to
further introduce a hierarchy design in a single layer. Additionally, a
Squeeze-Excitation (SE) [21] block is integrated subsequent to the multi-scale
2D scanning, as informed by related literature [24, 38]. Diverging from the
conventional focus on token mixing in prior vision Mamba architectures [61,
31, 24], our design introduces a channel mixer to augment the flow of
information across different channels, aligning with the structural paradigms
of typical vision transformers. In concordance with preceding studies [53, 49,
15, 52], the ConvFFN which consists of a depth-wise convolution and two fully
connected layers is employed as the channel mixer. To ensure that the ConvFFN
and MSVSS block consume approximately the same FLOPs following LeViT [9], we
adopt a modest channel expansion ratio of 2. Upon the amalgamation of MSVSS
and ConvFFN within the MS3 block, meticulous adjustments are made to the
number of blocks to ensure a comparable computational budget, facilitating a
fair comparison.

To empirically validate the efficacy of our proposed modifications, we
introduce model variants across different scales, the specifications of which
are detailed in our appendix Network Architecture. These variants, namely
Nano, Micro, and Tiny, are characterized by their parameter counts of 6.9M,
11.9M, and 33.0M, respectively. In terms of computational expenditure, these
models require 0.9, 1.5, and 4.6 GFLOPs correspondingly, demonstrating a
scalable approach to model design that accommodates varying computational
constraints.

##  4 Experimental Validation

###  4.1 ImageNet Classification

Table 1: Accuracy comparison across various models on ImageNet-1K. Method | #param. | FLOPs |  | Top-1
---
Acc(%).
RegNetY-800M [39] | 6M | 0.8G | 76.3
RegNetY-1.6G [39] | 11M | 1.6G | 78.0
RegNetY-4G [39] | 21M | 4.0G | 80.0
DeiT-S [44] | 22M | 4.6G | 79.8
DeiT-B [44] | 86M | 17.5G | 81.8
Swin-T [33] | 29M | 4.5G | 81.3
Swin-S [33] | 50M | 8.7G | 83.0
Swin-B [33] | 88M | 15.4G | 83.5
ViM-T [61] | 7M | 1.5G | 76.1
ViM-S [61] | 26M | 5.1G | 80.5
VMamba-T [31] | 22M | 5.6G | 82.2
VMamba-B [31] | 44M | 11.2G | 83.5
LocalVMamba-T [24] | 26M | 5.7G | 82.7
LocalVMamba-S [24] | 50M | 11.4G | 83.7
EffVMamba-T [38] | 6M | 0.8G | 76.5
EffVMamba-S [38] | 11M | 1.3G | 78.7
EffVMamba-B [38] | 33M | 4.0G | 81.8
MSVMamba-N | 7M | 0.9G | 77.3
MSVMamba-M | 12M | 1.5G | 79.8
MSVMamba-T | 33M | 4.6G | 82.8

Settings. Our models are trained and tested on the ImageNet-1K dataset [3]. In
alignment with previous works [33, 31, 24], all models undergo training for
300 epochs, with the initial 20 epochs dedicated to warming up. The training
utilizes a batch size of 1024 across 8 GPUs. We employ the AdamW optimizer,
setting the betas to (0.9, 0.999) and momentum to 0.9. The learning rate is
managed through a cosine decay scheduler, starting from an initial rate of
0.001, coupled with a weight decay of 0.05. Additionally, we leverage the
exponential moving average (EMA) and implement label smoothing with a factor
of 0.1 to enhance model performance and generalization. During testing, images
are center cropped with the size of 224$\times$224\. When dealing with the
routes for multi-scale scanning, we empirically select top-left to the bottom-
right for dealing with the full-resolution feature map while the other three
scans are responsible for scanning the downsampled feature map.

Results. Tab. 1 showcases our MSVMamba models against established CNNs, ViTs,
and SSM-based models on ImageNet-1K. MSVMamba-N, with 7M parameters and 0.9G
FLOPs, achieves 77.3% top-1 accuracy, outperforming similar-cost Conv-based
RegNetY-800M and SSM-based EfficientVMamba-T. MSVMamba-M, with 12M parameters
and 1.5G FLOPs, reaches 79.8% accuracy, surpassing ViM-T by 3.7%. The
MSVMamba-T model attains 82.8% accuracy with 33M parameters and 4.6G FLOPs,
exceeding VMamba-T by 0.6% with much less computational cost. These results
highlight MSVMamba’s efficiency and scalability, offering a robust option for
high-accuracy, resource-efficient model design.

###  4.2 Object Detection

Table 2: Object detection and instance segmentation with Mask R-CNN on COCO. The FLOPs are computed for an input size of $1280\times 800$. Multi-scale training is exclusively implemented in the 3× schedule. All backbones are pre-trained on the ImageNet-1K dataset. Backbone | Mask R-CNN 1x Schedule | Mask R-CNN 3x Schedule | #param. | FLOPs
---|---|---|---|---
AP${}^{\text{b}}$ |  AP${}^{\text{b}}_{\text{50}}$ |  AP${}^{\text{b}}_{\text{75}}$ |  AP${}^{\text{m}}$ |  AP${}^{\text{m}}_{\text{50}}$ |  AP${}^{\text{m}}_{\text{75}}$ |  AP${}^{\text{b}}$ |  AP${}^{\text{b}}_{\text{50}}$ |  AP${}^{\text{b}}_{\text{75}}$ |  AP${}^{\text{m}}$ |  AP${}^{\text{m}}_{\text{50}}$ |  AP${}^{\text{m}}_{\text{75}}$
PVT-T [51] | 36.7 | 59.2 | 39.3 | 35.1 | 56.7 | 37.3 | 39.8 | 62.2 | 43.0 | 37.4 | 59.3 | 39.9 | 33M | 208G
LightViT-T [23] | 37.8 | 60.7 | 40.4 | 35.9 | 57.8 | 38.0 | 41.5 | 64.4 | 45.1 | 38.4 | 61.2 | 40.8 | 28M | 187G
EffVMamba-S [38] | 39.3 | 61.8 | 42.8 | 36.7 | 58.9 | 39.2 | 41.6 | 63.9 | 45.6 | 38.2 | 60.8 | 40.7 | 31M | 197G
MSVMamba-M | 43.8 | 65.8 | 47.7 | 39.9 | 62.9 | 42.9 | 46.3 | 68.1 | 50.8 | 41.8 | 65.1 | 44.9 | 32M | 201G
Swin-T [33] | 42.7 | 65.2 | 46.8 | 39.3 | 62.2 | 42.2 | 46.0 | 68.1 | 50.3 | 41.6 | 65.1 | 44.9 | 48M | 267G
ConvNeXt-T [34] | 44.2 | 66.6 | 48.3 | 40.1 | 63.3 | 42.8 | 46.2 | 67.9 | 50.8 | 41.7 | 65.0 | 44.9 | 48M | 262G
VMamba-T [31] | 46.5 | 68.5 | 50.7 | 42.1 | 65.5 | 45.3 | 48.5 | 69.9 | 52.9 | 43.2 | 66.8 | 46.3 | 42M | 286G
LocalVMamba-T [24] | 46.7 | 68.7 | 50.8 | 42.2 | 65.7 | 45.5 | 48.7 | 70.1 | 53.0 | 43.4 | 67.0 | 46.4 | 45M | 291G
EffVMamba-B [38] | 43.7 | 66.2 | 47.9 | 40.2 | 63.3 | 42.9 | 45.0 | 66.9 | 49.2 | 40.8 | 64.1 | 43.7 | 53M | 252G
MSVMamba-T | 46.9 | 68.8 | 51.4 | 42.2 | 65.6 | 45.4 | 48.3 | 69.5 | 53.0 | 43.2 | 66.8 | 46.9 | 53M | 252G

Setup. We evaluate our MSVMamba on the MSCOCO [29] dataset using the Mask
R-CNN [18] framework for object detection and instance segmentation tasks.
Following previous works [33, 31], we utilize backbones pretrained on
ImageNet-1K for initialization. We employ standard training strategies of
1$\times$ (12 epochs) and 3$\times$ (36 epochs) with Multi-Scale (MS) training
for a fair comparison.

Results. Tab. 2 presents a performance comparison of our method against CNNs,
ViTs, and SSM-based models. Our model consistently outperforms others across
various variants and training settings. Specifically, MSVMamba-T outperforms
Swin-T by +4.2 box AP and +2.9 mask AP under the 1$\times$ schedule and also
shows improvements in both box AP and mask AP under the 3$\times$ schedule.

###  4.3 Semantic Segmentation

Table 3: We present the results of semantic segmentation on the ADE20K dataset [59] using the UperNet framework [54]. The computational complexity, measured in FLOPs, is calculated for input dimensions of $512\times 2048$. The abbreviations "SS" and "MS" refer to single-scale and multi-scale testing, respectively. Method |  | mIoU
---
SS
| mIoU
---
MS
#param. | FLOPs
ResNet-50 [19] | 42.1 | 42.8 | 67M | 953G
DeiT-S+MLN [45] | 43.8 | 45.1 | 58M | 1217G
Swin-T [33] | 44.4 | 45.8 | 60M | 945G
ConvNeXt-T [34] | 46.0 | 46.7 | 60M | 939G
VMamba-T [31] | 47.3 | 48.3 | 55M | 964G
LocalVMamba-T [24] | 47.9 | 49.1 | 57M | 970G
EffVMamba-S [38] | 41.5 | 42.1 | 29M | 505G
EffVMamba-B [38] | 46.5 | 47.3 | 65M | 930G
MSVMamba-M | 45.1 | 45.4 | 42M | 875G
MSVMamba-T | 47.6 | 48.5 | 65M | 942G

Setup. Consistent with the methodologies used in Swin [33] and VMamba [31], we
utilize the UperHead [54] framework atop an ImageNet pre-trained MSVMamba
backbone. The training process is conducted over 160K iterations with a batch
size of 16. We employ the AdamW optimizer with a learning rate set at $6\times
10^{-5}$. Our experiments are primarily conducted using a default input
resolution of $512\times 512$. Additionally, we also incorporate Multi-Scale
(MS) testing to assess performance variations.

Results. We present the detailed results of our model and other competitors in
Tab. 3, which includes both single-scale and multi-scale testing. Our MSVMamba
consistently outperforms the Swin, ConNeXt, and VMamba models in the tiny
variant by margins of +2.2, +1.6, and +0.3 mIoU, respectively.

###  4.4 Ablation Study

Table 4: Evolutionary trajectory from VMamba to MSVMamba. When integrating the ConvFFN and setting the state space dimension $N=1$, we meticulously calibrate the quantity of blocks to maintain a nearly constant computational cost, measured in GFLOPs. Model | MS2D | SE | ConvFFN | $N=1$ | #param. | FLOPs |  | Top-1
---
Acc(%).
VMamba-Nano |  |  |  |  | 4.4M | 0.87G | 69.6
MSVMamba-Nano | ✓ |  |  |  | 4.8M | 0.89G | $71.9_{\uparrow 2.3}$
✓ | ✓ |  |  | 5.3M | 0.89G | $72.4_{\uparrow 2.8}$
✓ | ✓ | ✓ |  | 6.6M | 0.94G | $74.4_{\uparrow 4.8}$
✓ | ✓ | ✓ | ✓ | 6.9M | 0.88G | $75.1_{\uparrow 5.5}$

To validate the effectiveness of the proposed modules, we conducted a
comprehensive ablation study. Specifically, we scaled the VMamba-Tiny model by
setting its embedding dimension $d$ to 48, the state space dimension $N$ to 8,
and the number of blocks in the four different stages to $[1,2,4,2]$. The
scaled model, referred to as VMamba-Nano, has parameters and computational
costs of 4.4M and 0.87GFLOPs, respectively. This model serves as the baseline
for our ablation experiments. Models in ablation study are conducted over a
training schedule of 100 epochs, except where noted otherwise, to reduce
training time.

On Multi-Scale 2D Scan. After replacing the SS2D in VMamba with our MS2D,
while maintaining a roughly equivalent FLOP count, the accuracy increased from
69.6$\%$ to 71.9$\%$, as demonstrated in the Tab. 4. Furthermore, we conducted
an ablation on the number of scans in the multi-scale scan, considering both
full-resolution and half-resolution branches. The results are shown in Tab. 5.
Placing all scans in the half-resolution branch led to a significant loss of
fine-grained features, resulting in a substantial decrease in model accuracy.
Positioning two or three scans in the full resolution branch, compared to just
one, resulted in accuracy improvements of 0.1$\%$ and 0.6$\%$, respectively,
but introduced an additional computational cost of approximately 12$\%$ and
25$\%$. Allocating four scans to the full resolution branch, effectively
reverting to the SS2D method, increased the computational cost by 34$\%$ while
only improving accuracy by 0.4$\%$.

Table 5: Ablation study on Multi-Scale 2D Scan. Full | Half | #param | FLOPs |  | Top-1
---
Acc(%).
0 | 4 | 5.1M | 0.74G | 63.1
1 | 3 | 4.8M | 0.89G | 71.9
2 | 2 | 5.0M | 1.00G | 72.0
3 | 1 | 5.3M | 1.11G | 72.5
4 | 0 | 5.1M | 1.19G | 72.3

For an optimal trade-off between computational cost and accuracy, we select
one scan in the full-resolution branch as the default setting. Building upon
the MS2D foundation, we introduce an SE block following EfficientVMamba [38],
which further enhanced accuracy by 0.5$\%$ with minimal additional
computational cost.

On ConvFFN. Upon replacing SS2D with MS2D and incorporating a SE block, we
constructed a model that utilizes ConvFFN as a channel mixer. When only using
SSM, the model exhibited insufficient information exchange between channels.
The integration of ConvFFN as a channel mixer significantly enhanced the
model’s capability for inter-channel information interaction. As indicated in
Tab. 4, the addition of ConvFFN resulted in an additional accuracy improvement
of 2.0$\%$. To maintain a roughly equivalent computational cost, we adjusted
the number of blocks within the model. Besides, we set the state space
dimension $N=1$ and stacked one more block to further enhance the capability
of capturing long-range information while maintaining a roughly constant
computational cost. This operation resulted in an additional accuracy
improvement of 0.7$\%$, as shown in Tab. 4.

##  5 Limitations

The design of multi-scale VMamba aims to tackle the long-range forgetting
problem of Mamba models with limited parameters on vision tasks. Although the
proposed model has proven to be effective, its scalability remains to be
explored since this issue can also be alleviated by increasing model sizes. In
such cases, the multi-scale design may have only a marginal improvement.

##  6 Conclusion

In this paper, we introduced Multi-Scale VMamba (MSVMamba), an SSM-based
vision backbone that leverages the advantages of linear complexity and global
receptive field. We developed the Multi-Scale 2D (MS2D) scanning technique to
minimize computational redundancy and alleviate the long-range forgetting
problem in parameter-limited vision models. Additionally, we incorporated the
Convolutional Feed-Forward Network (ConvFFN) to enhance the exchange of
information between channels, thereby significantly improving the performance
of our model. Our experiments demonstrate that MSVMamba consistently
outperforms popular models from various architectures, including ConvNeXt,
Swin Transformer, and VMamba, in image classification and downstream tasks.

## References

  * [1] Tianxiang Chen, Zhentao Tan, Tao Gong, Qi Chu, Yue Wu, Bin Liu, Jieping Ye, and Nenghai Yu.  Mim-istd: Mamba-in-mamba for efficient infrared small target detection.  arXiv preprint arXiv:2403.02148, 2024.
  * [2] Zihang Dai, Hanxiao Liu, Quoc V Le, and Mingxing Tan.  Coatnet: Marrying convolution and attention for all data sizes.  NeurIPS, 2021.
  * [3] Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei.  Imagenet: A large-scale hierarchical image database.  In CVPR, 2009.
  * [4] Xiaohan Ding, Xiangyu Zhang, Yizhuang Zhou, Jungong Han, Guiguang Ding, and Jian Sun.  Scaling up your kernels to 31x31: Revisiting large kernel design in cnns.  In CVPR, 2022.
  * [5] Xiaoyi Dong, Jianmin Bao, Dongdong Chen, Weiming Zhang, Nenghai Yu, Lu Yuan, Dong Chen, and Baining Guo.  Cswin transformer: A general vision transformer backbone with cross-shaped windows.  In CVPR, 2022.
  * [6] Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, 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.  In ICLR, 2020.
  * [7] Chengbin Du, Yanxi Li, and Chang Xu.  Understanding robustness of visual state space models for image classification.  arXiv preprint arXiv:2403.10935, 2024.
  * [8] Daniel Y Fu, Tri Dao, Khaled K Saab, Armin W Thomas, Atri Rudra, and Christopher Ré.  Hungry hungry hippos: Towards language modeling with state space models.  arXiv preprint arXiv:2212.14052, 2022.
  * [9] Benjamin Graham, Alaaeldin El-Nouby, Hugo Touvron, Pierre Stock, Armand Joulin, Hervé Jégou, and Matthijs Douze.  Levit: a vision transformer in convnet’s clothing for faster inference.  In ICCV, 2021.
  * [10] Albert Gu and Tri Dao.  Mamba: Linear-time sequence modeling with selective state spaces.  arXiv preprint arXiv:2312.00752, 2023.
  * [11] Albert Gu, Tri Dao, Stefano Ermon, Atri Rudra, and Christopher Ré.  Hippo: Recurrent memory with optimal polynomial projections.  NeurIPS, 2020.
  * [12] Albert Gu, Karan Goel, and Christopher Ré.  Efficiently modeling long sequences with structured state spaces.  arXiv preprint arXiv:2111.00396, 2021.
  * [13] Albert Gu, Isys Johnson, Karan Goel, Khaled Saab, Tri Dao, Atri Rudra, and Christopher Ré.  Combining recurrent, convolutional, and continuous-time models with linear state space layers.  NeurIPS, 2021.
  * [14] Hang Guo, Jinmin Li, Tao Dai, Zhihao Ouyang, Xudong Ren, and Shu-Tao Xia.  Mambair: A simple baseline for image restoration with state-space model.  arXiv preprint arXiv:2402.15648, 2024.
  * [15] Jianyuan Guo, Kai Han, Han Wu, Yehui Tang, Xinghao Chen, Yunhe Wang, and Chang Xu.  Cmt: Convolutional neural networks meet vision transformers.  In CVPR, pages 12175–12185, 2022.
  * [16] Kai Han, An Xiao, Enhua Wu, Jianyuan Guo, Chunjing Xu, and Yunhe Wang.  Transformer in transformer.  In NeurIPS, 2021.
  * [17] Ali Hassani, Steven Walton, Jiachen Li, Shen Li, and Humphrey Shi.  Neighborhood attention transformer.  In CVPR, 2023.
  * [18] Kaiming He, Georgia Gkioxari, Piotr Dollár, and Ross Girshick.  Mask r-cnn.  In ICCV, 2017.
  * [19] Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun.  Deep residual learning for image recognition.  In CVPR, 2016.
  * [20] Andrew G Howard, Menglong Zhu, Bo Chen, Dmitry Kalenichenko, Weijun Wang, Tobias Weyand, Marco Andreetto, and Hartwig Adam.  Mobilenets: Efficient convolutional neural networks for mobile vision applications.  arXiv preprint arXiv:1704.04861, 2017.
  * [21] Jie Hu, Li Shen, and Gang Sun.  Squeeze-and-excitation networks.  In CVPR, 2018.
  * [22] Gao Huang, Zhuang Liu, Geoff Pleiss, Laurens Van Der Maaten, and Kilian Weinberger.  Convolutional networks with dense connectivity.  IEEE TPAMI, 2019.
  * [23] Tao Huang, Lang Huang, Shan You, Fei Wang, Chen Qian, and Chang Xu.  Lightvit: Towards light-weight convolution-free vision transformers.  arXiv preprint arXiv:2207.05557, 2022.
  * [24] Tao Huang, Xiaohuan Pei, Shan You, Fei Wang, Chen Qian, and Chang Xu.  Localmamba: Visual state space model with windowed selective scan.  arXiv preprint arXiv:2403.09338, 2024.
  * [25] Zihang Jiang, Qibin Hou, Li Yuan, Daquan Zhou, Xiaojie Jin, Anran Wang, and Jiashi Feng.  Token labeling: Training a 85.5% top-1 accuracy vision transformer with 56m parameters on imagenet.  arXiv preprint arXiv:2104.10858, 2021.
  * [26] Alex Krizhevsky, Ilya Sutskever, and Geoffrey E Hinton.  Imagenet classification with deep convolutional neural networks.  NeurIPS, 2012.
  * [27] Kunchang Li, Xinhao Li, Yi Wang, Yinan He, Yali Wang, Limin Wang, and Yu Qiao.  Videomamba: State space model for efficient video understanding.  arXiv preprint arXiv:2403.06977, 2024.
  * [28] Weibin Liao, Yinghao Zhu, Xinyuan Wang, Cehngwei Pan, Yasha Wang, and Liantao Ma.  Lightm-unet: Mamba assists in lightweight unet for medical image segmentation.  arXiv preprint arXiv:2403.05246, 2024.
  * [29] Tsung-Yi Lin, Michael Maire, Serge Belongie, James Hays, Pietro Perona, Deva Ramanan, Piotr Dollár, and C Lawrence Zitnick.  Microsoft coco: Common objects in context.  In ECCV, 2014.
  * [30] Weifeng Lin, Ziheng Wu, Jiayu Chen, Jun Huang, and Lianwen Jin.  Scale-aware modulation meet transformer.  In ICCV, 2023.
  * [31] Yue Liu, Yunjie Tian, Yuzhong Zhao, Hongtian Yu, Lingxi Xie, Yaowei Wang, Qixiang Ye, and Yunfan Liu.  Vmamba: Visual state space model.  arXiv preprint arXiv:2401.10166, 2024.
  * [32] Ze Liu, Han Hu, Yutong Lin, Zhuliang Yao, Zhenda Xie, Yixuan Wei, Jia Ning, Yue Cao, Zheng Zhang, Li Dong, et al.  Swin transformer v2: Scaling up capacity and resolution.  In CVPR, 2022.
  * [33] 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 ICCV, 2021.
  * [34] Zhuang Liu, Hanzi Mao, Chao-Yuan Wu, Christoph Feichtenhofer, Trevor Darrell, and Saining Xie.  A convnet for the 2020s.  In CVPR, 2022.
  * [35] Jun Ma, Feifei Li, and Bo Wang.  U-mamba: Enhancing long-range dependency for biomedical image segmentation.  arXiv preprint arXiv:2401.04722, 2024.
  * [36] Eric Nguyen, Karan Goel, Albert Gu, Gordon Downs, Preey Shah, Tri Dao, Stephen Baccus, and Christopher Ré.  S4nd: Modeling images and videos as multidimensional signals with state spaces.  NeurIPS, 2022.
  * [37] Xuran Pan, Chunjiang Ge, Rui Lu, Shiji Song, Guanfu Chen, Zeyi Huang, and Gao Huang.  On the integration of self-attention and convolution.  In CVPR, 2022.
  * [38] Xiaohuan Pei, Tao Huang, and Chang Xu.  Efficientvmamba: Atrous selective scan for light weight visual mamba.  arXiv preprint arXiv:2403.09977, 2024.
  * [39] Ilija Radosavovic, Raj Prateek Kosaraju, Ross Girshick, Kaiming He, and Piotr Dollár.  Designing network design spaces.  In CVPR, 2020.
  * [40] Jiacheng Ruan and Suncheng Xiang.  Vm-unet: Vision mamba unet for medical image segmentation.  arXiv preprint arXiv:2402.02491, 2024.
  * [41] Karen Simonyan and Andrew Zisserman.  Very deep convolutional networks for large-scale image recognition.  arXiv preprint arXiv:1409.1556, 2014.
  * [42] Jimmy TH Smith, Andrew Warrington, and Scott W Linderman.  Simplified state space layers for sequence modeling.  arXiv preprint arXiv:2208.04933, 2022.
  * [43] Mingxing Tan and Quoc Le.  Efficientnet: Rethinking model scaling for convolutional neural networks.  In ICML, 2019.
  * [44] Hugo Touvron, Matthieu Cord, Matthijs Douze, Francisco Massa, Alexandre Sablayrolles, and Hervé Jégou.  Training data-efficient image transformers & distillation through attention.  In ICML, 2021.
  * [45] Hugo Touvron, Matthieu Cord, and Hervé Jégou.  Deit iii: Revenge of the vit.  In ECCV, 2022.
  * [46] Hugo Touvron, Matthieu Cord, Alexandre Sablayrolles, Gabriel Synnaeve, and Hervé Jégou.  Going deeper with image transformers.  In ICCV, 2021.
  * [47] Zhengzhong Tu, Hossein Talebi, Han Zhang, Feng Yang, Peyman Milanfar, Alan Bovik, and Yinxiao Li.  Maxvit: Multi-axis vision transformer.  In ECCV, 2022.
  * [48] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin.  Attention is all you need.  In NeurIPS, 2017.
  * [49] A Wang, H Chen, Z Lin, H Pu, and G Ding.  Repvit: Revisiting mobile cnn from vit perspective. arxiv 2023.  arXiv preprint arXiv:2307.09283, 2023.
  * [50] Wenhai Wang, Jifeng Dai, Zhe Chen, Zhenhang Huang, Zhiqi Li, Xizhou Zhu, Xiaowei Hu, Tong Lu, Lewei Lu, Hongsheng Li, et al.  Internimage: Exploring large-scale vision foundation models with deformable convolutions.  arXiv preprint arXiv:2211.05778, 2022.
  * [51] Wenhai Wang, Enze Xie, Xiang Li, Deng-Ping Fan, Kaitao Song, Ding Liang, Tong Lu, Ping Luo, and Ling Shao.  Pyramid vision transformer: A versatile backbone for dense prediction without convolutions.  In ICCV, 2021.
  * [52] Wenhai Wang, Enze Xie, Xiang Li, Deng-Ping Fan, Kaitao Song, Ding Liang, Tong Lu, Ping Luo, and Ling Shao.  Pvt v2: Improved baselines with pyramid vision transformer.  Computational Visual Media, 2022.
  * [53] Zhuofan Xia, Xuran Pan, Shiji Song, Li Erran Li, and Gao Huang.  Dat++: Spatially dynamic vision transformer with deformable attention.  arXiv preprint arXiv:2309.01430, 2023.
  * [54] Tete Xiao, Yingcheng Liu, Bolei Zhou, Yuning Jiang, and Jian Sun.  Unified perceptual parsing for scene understanding.  In ECCV, 2018.
  * [55] Saining Xie, Ross Girshick, Piotr Dollár, Zhuowen Tu, and Kaiming He.  Aggregated residual transformations for deep neural networks.  In CVPR, 2017.
  * [56] Chenhongyi Yang, Zehui Chen, Miguel Espinosa, Linus Ericsson, Zhenyu Wang, Jiaming Liu, and Elliot J Crowley.  Plainmamba: Improving non-hierarchical mamba in visual recognition.  arXiv preprint arXiv:2403.17695, 2024.
  * [57] Yuhuan Yang, Chaofan Ma, Jiangchao Yao, Zhun Zhong, Ya Zhang, and Yanfeng Wang.  Remamber: Referring image segmentation with mamba twister.  arXiv preprint arXiv:2403.17839, 2024.
  * [58] Weihao Yu, Mi Luo, Pan Zhou, Chenyang Si, Yichen Zhou, Xinchao Wang, Jiashi Feng, and Shuicheng Yan.  Metaformer is actually what you need for vision.  In CVPR, 2022.
  * [59] Bolei Zhou, Hang Zhao, Xavier Puig, Sanja Fidler, Adela Barriuso, and Antonio Torralba.  Scene parsing through ade20k dataset.  In CVPR, 2017.
  * [60] Lei Zhu, Xinjiang Wang, Zhanghan Ke, Wayne Zhang, and Rynson Lau.  Biformer: Vision transformer with bi-level routing attention.  In CVPR, 2023.
  * [61] Lianghui Zhu, Bencheng Liao, Qian Zhang, Xinlong Wang, Wenyu Liu, and Xinggang Wang.  Vision mamba: Efficient visual representation learning with bidirectional state space model.  arXiv preprint arXiv:2401.09417, 2024.

## Appendix

### Throughput Comparison

We present the throughput comparison of our MSVMamba variants, and the
baseline VMamba-Tiny on the ImageNet-1K dataset in Tab. 6. The throughput is
measured in images per second, processed on a NVIDIA GeForce RTX 2080 Ti GPU.
The image size for all tests is $224\times 224$, and the batch size is set to
128. Under similar GFLOPs, our MSVMamba-Tiny model nearly doubles the
throughput of the baseline VMamba-Tiny.

Table 6: Throughput comparison on ImageNet-1K. Model | Throughput (img/sec)
---|---
VMamba-Tiny | 218.86
MSVMamba-Nano | 1165.71
MSVMamba-Micro | 887.14
MSVMamba-Tiny | 421.80

### Qualitative Analysis

We provide visualizations in Fig. 6 comparing the proposed MS2D and SS2D
configurations in VMamba. These visualizations are generated by converting the
S6 layer into an attention format, as demonstrated by VMamba [31]. The results
clearly show that the full-resolution in MS2D scan captures more detailed
features, whereas the scans at half resolution primarily focus on broader
architectural details, compared to SS2D. The proposed hierarchical scanning
pattern facilitates the current layer’s ability to discern and amalgamate
features across various levels of abstraction.

Figure 6: Attention maps from four distinct scanning directions, generated by
SS2D and our MS2D in the last layer of the second stage. In the second row,
full-resolution scan (first scan) captures fine-grained features, whereas
scans at half resolution capture coarse-grained features. Maps are rendered at
a higher resolution to enhance visualization quality.

### Network Architecture

Table 7:  Architectural overview of the MSVMamba series.

layer name | output size | Nano | Micro | Tiny
---|---|---|---|---
stem | 56$\times$56 | $\begin{array}[]{c}\text{conv 4$\times$4, 48, stride 4}\\\ \end{array}$ | $\begin{array}[]{c}\text{conv 4$\times$4, 64, stride 4}\\\ \end{array}$ | $\begin{array}[]{c}\text{conv 4$\times$4, 96, stride 4}\\\ \end{array}$
stage 1 | 28$\times$28 | $\begin{array}[]{c}\text{MS3 $\times$1}\\\ \text{conv 2$\times$2, 96, stride 2}\\\ \end{array}$ | $\begin{array}[]{c}\text{MS3 $\times$1}\\\ \text{conv 2$\times$2, 128, stride 2}\\\ \end{array}$ | $\begin{array}[]{c}\text{MS3 $\times$1}\\\ \text{conv 2$\times$2, 192, stride 2}\\\ \end{array}$
stage 2 | 14$\times$14 | $\begin{array}[]{c}\text{MS3 $\times$2}\\\ \text{conv 2$\times$2, 192, stride 2}\\\ \end{array}$ | $\begin{array}[]{c}\text{MS3 $\times$2}\\\ \text{conv 2$\times$2, 256, stride 2}\\\ \end{array}$ | $\begin{array}[]{c}\text{MS3 $\times$2}\\\ \text{conv 2$\times$2, 384, stride 2}\\\ \end{array}$
stage 3 | 7$\times$7 | $\begin{array}[]{c}\text{MS3 $\times$5}\\\ \text{conv 2$\times$2, 384, stride 2}\\\ \end{array}$ | $\begin{array}[]{c}\text{MS3 $\times$5}\\\ \text{conv 2$\times$2, 512, stride 2}\\\ \end{array}$ | $\begin{array}[]{c}\text{MS3 $\times$9}\\\ \text{conv 2$\times$2, 768, stride 2}\\\ \end{array}$
stage 4 | 7$\times$7 | MS3 $\times$ 2 | MS3 $\times$ 2 | MS3 $\times$ 2
| 1$\times$1 | average pool, 1000-d fc, softmax
Param. (M) | 6.9 | 11.9 | 33.0
FLOPs | 0.9$\times 10^{9}$ | 1.5$\times 10^{9}$ | 4.6$\times 10^{9}$

In Tab. 7, we present the detailed architecture of our model variants,
including the Nano, Micro, and Tiny versions, each with varying channels and
block numbers.


# MIA-Former: Efficient and Robust Vision Transformers
via Multi-grained Input Adaptation

Zhongzhi Yu1, Yonggan Fu1, Sicheng Li2, Chaojian Li1, Yingyan Lin1

###### Abstract

Vision transformers (ViTs) have recently demonstrated great success in various
computer vision tasks, motivating a tremendously increased interest in their
deployment into many real-world IoT applications. However, powerful ViTs are
often too computationally expensive to be fitted onto real-world resource-
constrained devices, due to (1) their quadratically increased complexity with
the number of input tokens and (2) their overparameterized self-attention
heads and model depth. In parallel, different images are of varied complexity
and their different regions can contain various levels of visual information,
e.g., a sky background is not as informative as a foreground object in object
classification tasks, indicating that treating all regions/tokens equally in
terms of model complexity is unnecessary while such opportunities for trimming
down ViTs’ complexity have not been fully explored. To this end, we propose a
Multi-grained Input-Adaptive Vision TransFormer framework dubbed MIA-Former
that can input-adaptively adjust the structure of ViTs at three coarse-to-
fine-grained granularities (i.e., model depth and the number of model
heads/tokens). In particular, our MIA-Former adopts a low-cost network trained
with a hybrid supervised and reinforcement training method to skip unnecessary
layers, heads, and tokens in an input adaptive manner, reducing the overall
computational cost. Furthermore, an interesting side effect of our MIA-Former
is that its resulting ViTs are naturally equipped with improved robustness
against adversarial attacks over their static counterparts, because MIA-
Former’s multi-grained dynamic control improves the model diversity similar to
the effect of ensemble and thus increases the difficulty of adversarial
attacks against all its sub-models. Extensive experiments and ablation studies
validate that the proposed MIA-Former framework can (1) effectively allocate
computation budgets adaptive to the difficulty of input images, achieving
state-of-the-art (SOTA) accuracy-efficiency trade-offs, e.g., 20% computation
savings with the same or even a higher accuracy compared with SOTA dynamic
transformer models, and (2) boost ViTs’ robustness accuracy under various
adversarial attacks over their vanilla counterparts by 2.4% and 3.0%,
respectively.

##  1 Introduction

Vision transformers (ViTs) have been proven to be a powerful architecture on
various computer vision tasks (Dosovitskiy et al. 2020; Touvron et al. 2021;
Chen et al. 2021a; Caron et al. 2021; Strudel et al. 2021), especially when
being scaled up with larger model sizes and more training data (Dosovitskiy et
al. 2020; Steiner et al. 2021; Ridnik et al. 2021). However, powerful ViTs
often come with prohibitive computational overhead. Specifically, (1) the
tokens consisting of merely background information and (2) redundant heads
within the multi-head self-attention (MSA) module of ViTs which learn similar
features can lead to unnecessary yet non-negligible inference costs. Taking
the widely used DeiT-Small model (Touvron et al. 2021) as an example, running
inference on a single image with a resolution of $224\times 224$ requires over
4.6 Giga floating point operations (GFLOPs), making it challenging to deploy
ViTs onto many real-world resource-constrained devices for supporting
intelligent internet of things (IoT) applications. Thus, there is an urgent
need to reduce the computational cost of ViTs.

On the other hand, in real-world applications, the complexity of images can
vary significantly. As such, processing all the images with the same model
complexity of ViTs could be overcooked. For example, for most of the time
during video surveillance, the video may be just staring at an empty
background, and the corresponding images can be processed with a naively
simple model, saving a large portion of computational cost while still
achieving satisfying accuracy. Thus, a straightforward solution for trimming
down ViTs’ complexity is to perform input-adaptive dynamic inference. Although
dynamic inference has been extensively explored for convolutional neural
networks (CNNs) through various dynamic dimensions (e.g., model depth, channel
number, and model bit-width) (Hu et al. 2020; Wang et al. 2018; Shen et al.
2020; Wang et al. 2020), only a few pioneering works have considered this
aspect for ViTs (Rao et al. 2021; Wang et al. 2021b), which yet merely focus
on reducing the computational budget by adaptively adjusting the number of
input tokens. However, as suggested in (Zhou et al. 2021), the similarity
between heads and feature maps can increase significantly in deeper ViT
layers, implying that the token dimension is not the only source of
redundancy, and the unexplored depth and head dimensions could lead to more
efficient ViTs.

To this end, we target a multi-grained ViT framework that can fully explore
the redundancy in ViTs and make the following contributions:

  * •

We propose a Multi-grained Input-Adaptive vision transFormer framework, dubbed
MIA-Former, in order to trim down the redundancy of ViTs from multiple
dimensions at three coarse-to-fine-grained granularities.

  * •

We propose a low-cost MIA-Controller to make input-adaptive decisions, which
is jointly trained with the ViT models via a hybrid supervised and
reinforcement learning (RL) scheme.

  * •

We empirically find that thanks to the proposed hybrid supervised and
reinforcement training method, MIA-Former is equipped with improved robustness
to various types of adversarial attacks, achieving a win-win in both
robustness and efficiency.

  * •

Extensive experiments and ablation studies based on both DeiT-based (Touvron
et al. 2021) and LeViT-based (Graham et al. 2021) models show that the
proposed MIA-Former can be used as a plug-in module on top of a wide range of
ViTs to achieve better accuracy-efficiency trade-offs and boosted adversarial
robustness, compared with state-of-the-art (SOTA) vanilla ViTs, input-adaptive
ViTs, as well as CNNs. Specifically, MIA-Former achieves a 20.1% FLOPs
reduction and 2.4% higher robustness accuracy under Projected Gradient Descent
(PGD) (Kurakin et al. 2016) attacks together with the same natural accuracy,
compared with the original DeiT-Small model.

##  2 Related Works

Vision Transformers. Transformers are first introduced to natural language
processing tasks in (Vaswani et al. 2017) and achieve impressive performance
in language modeling, machine translation, and other downstream tasks. It has
been shown that the self-attention module in transformers can serve as an
effective way to model the token-wise relationship of sentences. (Dosovitskiy
et al. 2020) then proposes ViTs, which is a pioneering work to extend
transformer architectures to large scale compute vision tasks and matches SOTA
CNNs’ performance. Specifically, ViTs first split an input image into a series
of patches which are embedded into tokens before passing into the ViT blocks;
Each ViT block consists of a stacked MSA and multi-layer perceptron (MLP)
module to extract the global relationship among input tokens; Multiple heads
further enable ViT blocks to extract different features via different heads,
improving the expressiveness of ViTs. The success of ViTs on large scale image
recognition tasks (e.g., ImageNet dataset (Deng et al. 2009)) has inspired a
series of following works to further exploit the expressive power of ViTs from
different perspectives: (Touvron et al. 2021) performs an exhaustive search
for the optimal training recipe for training ViTs; (Chen, Fan, and Panda 2021)
exploits the impact of patch sizes on ViTs’ achievable performance and
proposes a parallel ViT architecture to simultaneously utilize information
from both small and large patches to achieve better performance; and (Liu et
al. 2021) proposes hierarchical structures for ViTs like ResNet (He et al.
2016), and (Dong et al. 2021) modifies the shape of patches from square to
cross to balance the global and local attentions. Among them, LeViT (Graham et
al. 2021) achieves impressive performance by exploring the potential of
applying convolutional layers before ViTs. Specifically, LeViT is the first
ViT network that achieves a better accuracy-efficiency trade-off than
EfficientNet (Tan and Le 2019) under certain FLOPs ranges.

Input-adaptive Inference. Adaptively activating different components of a deep
neural network (DNN) in an input-dependent manner has been proved to be an
effective way to reduce the inference cost, which has been widely explored for
CNNs. Existing techniques in this regard can be roughly summarized into two
categories: (1) making early predictions by introducing multiple side branch
classifiers and dynamically exiting from one branch by analyzing the
confidence of intermediate feature maps (Kaya, Hong, and Dumitras 2019;
Teerapittayanon, McDanel, and Kung 2016) and (2) adaptively skipping certain
components of the model, such as blocks, channels and even bit-width (Wu et
al. 2018; Wang et al. 2020; Shen et al. 2020; Fu et al. 2020). Nevertheless,
the opportunities of input-adaptive inference have not yet been extensively
explored in the scope of ViTs as existing works merely focus on dynamically
adjust the input tokens, either by pruning out certain tokens (Rao et al.
2021) or by adjusting the input patch size (Wang et al. 2021b). Other design
dimensions in ViTs are still neglected, such as the number of attention heads
and model depth, which play important roles in the overparameterization of
ViTs.

Adversarial Robustness and Model Efficiency. DNNs’ robustness and efficiency
are two critical features required in real-world applications. There are some
pioneering works trying to simultaneously optimize both for CNNs. For example,
(Rakin et al. 2019; Ye et al. 2019; Sehwag et al. 2020; Fu et al. 2021b)
combine pruning techniques with adversarial training methods. (Fu et al.
2021a) leverages the poor adversarial transferability between different
precisions to win both robustness and efficiency, which can be accelerated by
customized accelerators (Fu et al. 2021c) for further boosted efficiency. And
(Rakin et al. 2018) uses dynamic quantization of activation functions to
defend against adversarial examples; (Hu et al. 2020) introduces input-
adaptive inference for simultaneously boosting robustness and efficiency.
However, existing works have shown that CNNs and ViTs behave differently under
adversarial attacks (Shao et al. 2021; Mao et al. 2021), and how to win both
adversarial robustness and model efficiency for ViTs is still an open
question.

##  3 The Proposed Techniques

In this section, we first present the MIA-Former framework which integrates a
MIA-Controller module to dynamically control the activated subparts of ViTs in
an input-dependent manner. After that, we introduce our proposed hybrid
supervised and reinforcement training method that can effectively train MIA-
Former based ViTs to be equipped with both improved robustness and efficiency.

Figure 1: Overview of the proposed MIA-Former framework integrating a MIA-
Controller and MIA-Block modules: (a) The MIA-Controller first decide whether
to skip the whole upcoming MIA-Block, and (b) if the upcoming MIA-Block is not
fully skipped, its MIA-Controller further dynamically masks out certain tokens
and heads to deactivate the corresponding modules and thus reduce
computational cost.

###  3.1 Overview

Fig. 1 illustrates an overview of the proposed MIA-Former framework. On top of
vanilla ViTs, MIA-Former adds a control module, dubbed MIA-Controller, to each
backbone ViT block, MIA-Blocks. Specifically, MIA-Former integrates a MIA-
Controller to first decide whether to mask out a given ViT block at a coarse
granularity, as shown in Fig. 1(a). In such cases, the outputs of a previous
block skip the current block and are directly fed into the next MIA-Controller
and MIA-Block. On the other hand, if the current MIA-Block is not skipped as a
whole, MIA-Former further uses the corresponding MIA-Controller to dynamically
mask out certain tokens and heads to deactivate the corresponding modules
during inference in an input-dependent manner, as shown in Fig. 1(b).

In vanilla ViT blocks, reducing the number of heads can only lead to linear
FLOPs reduction in MSA modules; moreover, the computational overhead in the
upcoming MLP module cannot be reduced, which is often the computational
bottleneck. Thus, we propose to further proportionally reduce both the input
and output size of the upcoming fully connected (FC) layers in the MIA-Block
module. In this way, MIA-Former can lead to linear FLOPs reduction in the MSA
module and nearly quadratic FLOPs reduction in the following MLP module when
skipping certain heads during inference, which can lead to large savings of
FLOPs in MIA-Former. Note that a skip connection is also added to pass the
information of masked heads and tokens to the outputs of the next block,
avoiding the permanent loss of information.

###  3.2 MIA-Controller

The key characteristic of the proposed MIA-Former is to dynamically adjust its
structure and thus complexity at three coarse-to-fine-grained granularities
(i.e., model depth, number of heads, and number of tokens). To achieve optimal
trade-offs between efficiency and accuracy, MIA-Former adopts a light-weight
controller to adaptively generate masks for skipping the corresponding
uninformative blocks/heads/tokens on top of its vanilla ViT backbone model.
However, it is non-trivial to derive such finer-grained learnable masks due to
the large skipping policy space. To tackle this problem, we propose a light-
weight controller, i.e., MIA-Controller, to generate the masks and thus
corresponding skipping policy. Specifically, MIA-Controller first examines
whether a ViT block shall be completely skipped; if not, it will skip certain
heads and tokens of this ViT block based on the input feature complexity. The
merit of introducing the above block-wise skipping is that it eliminates the
necessity of the additional efforts for computing head- and token-wise
skipping policy.

For each block $l$ in ViTs, we maintain three binary masks,
$D^{l}_{b}\in\\{0,1\\}$, $D^{l}_{h}\in\\{0,1\\}^{H}$, and
$D^{l}_{n}\in\\{0,1\\}^{N}$ for skipping blocks, heads, and tokens,
respectively, where $H$ and $N$ denote the number of heads and tokens,
respectively. In particular, an MIA-Controller is inserted ahead of each block
$l$ with the outputs of the previous MIA-Block
$I^{l}\in\mathcal{R}^{N_{h}\times N_{w}\times(HE)}$ as the inputs, where
$N=N_{h}N_{w}$ is the spatial dimension of the token array and $E$ is the
hidden dimension of each head in the ViT. The MIA-Controller first passes
$I^{l}$ through a two-layer CNN ($\text{CNN}_{b}$) with pooling to extract the
features $F^{l}_{b}$ which are then passed through an FC layer
($\text{FC}_{b}$) with Gumbel softmax to generate $D^{l}_{b}$, deciding
whether to skip the whole block. Such a skipping pipeline can be formulated
as:

$\displaystyle F^{l}_{b}$
$\displaystyle=\text{CNN}_{b}(I^{l})\in\mathcal{R}^{1\times 1\times
HE^{\prime}},$ (1) $\displaystyle G^{l}_{b}$
$\displaystyle=\text{FC}_{b}(F^{l}_{b})\in\mathcal{R},$ $\displaystyle
D^{l}_{b}$ $\displaystyle=\text{Round}(G^{l}_{b})\in\\{0,1\\},$

where $E^{\prime}<E$ is the hidden dimension of each head after the CNN with
pooling in the MIA-Controller, and we set $E^{\prime}=E/4$ in this work,
$Round()$ is rounding operation that rounds the value to the nearest integer.

If a ViT block is not skipped, the MIA-Controller further generates a token
mask $D^{l}_{n}$ and a head mask $D^{l}_{h}$ for determining the skipping
policy for the corresponding block’s heads and tokens, respectively. For
$D^{l}_{h}$, an FC layer ($\text{FC}_{h1}$) accepts $F^{l}_{b}$ as its inputs
to extract the head features $F^{l}_{h}$, and then another FC layer
($\text{FC}_{h2}$) with Gumbel softmax is adopted to generate the mask
$D^{l}_{h}$:

$\displaystyle F^{l}_{h}$
$\displaystyle=\text{FC}_{h1}(F^{l}_{b})\in\mathcal{R}^{H\times
E^{\prime\prime}}$ (2) $\displaystyle G^{l}_{h}$
$\displaystyle=\text{FC}_{h2}(F^{l}_{h})\in\mathcal{R}^{H}$ $\displaystyle
D^{l}_{h}$ $\displaystyle=\text{Round}(G^{l}_{h})\in\\{0,1\\}^{H}.$

For $D^{l}_{n}$, MIA-Former follows a similar procedure as (Rao et al. 2021).
Specifically, $I^{l}$ is reshaped to
$I^{l^{\prime}}\in\mathcal{R}^{N\times(HE)}$ and then applied to a two-layer
MLP ($\text{MLP}_{n}$) to extract the features $F^{l}_{n}$. After that, an FC
layer ($\text{FC}_{n}$) with Gumbel softmax is used to process $F^{l}_{n}$ and
generate the masks $D^{l}_{n}$.

$\displaystyle F^{L}_{n}$
$\displaystyle=\text{MLP}_{n}(I^{l^{\prime}})\in\mathcal{R}^{N\times
HE^{\prime}}$ (3) $\displaystyle G^{l}_{n}$
$\displaystyle=\text{FC}_{n}(F^{l}_{n})\in\mathcal{R}^{N}$ $\displaystyle
D^{l}_{n}$ $\displaystyle=\text{Round}(G^{l}_{n})\in\\{0,1\\}^{N}.$

The featuremap $I^{l}_{F}$ after the MIA-Controller to the MIA-Block $l$ are
computed as:

$\displaystyle I^{l}_{F}=$ $\displaystyle
I^{l}.\text{reshape}(N_{h},N_{w},H,E)\odot$ (4) $\displaystyle
D^{l}_{b}.\text{reshape}(1,1,1,1)\odot D^{l}_{h}.\text{reshape}(1,1,H,1)\odot$
$\displaystyle D^{l}_{n}.\text{reshape}(N_{h},N_{w},1,1),$

where $\odot$ represents the element-wise matrix multiplication (broadcast
will be performed if the matrix shapes do not match).

###  3.3 Hybrid Supervised and Reinforcement Training

Given the complexity of MIA-Former, directly training all its components leads
to unstable training and thus inferior performance. Thus, we propose a hybrid
supervised and reinforcement training pipeline, which consists of three steps:
(1) MIA-Controller pretraining, where we fix the parameters in the pretrained
MIA-Block and pretrain the MIA-Controller until they keep activating all
components of the MIA-Former without any skipping, (2) MIA-Former co-training,
where we co-train the MIA-Block and the MIA-Controller with a hybrid loss
function in a differentiable way, and (3) skipping policy finetuning with
hybrid RL, where we finetune the MIA-Block and the MIA-Controller with a
hybrid supervised and RL training method. We will illustrate the detail of the
above stages in the remaining part of this section.

MIA-Controller Pretraining. As a randomly initialized MIA-Controller randomly
skips different modules within the model, directly co-training the MIA-Block
and MIA-Controller lead to inferior performance. We conjecture this is due to
the fact that at the beginning of training, a large portion of the model is
randomly skipped, leading to a significant deviation from the original learned
distribution of pretrained ViTs. To tackle this, we propose to first pretrain
the MIA-Controller with the MIA-Block weights fixed until it does not skip any
components, and the whole MIA-Former should behave exactly the same as its
backbone ViT at the end of this stage. To achieve this, we train the MIA-
Controller with the pretrain loss defined as:

$\mathcal{L}_{pretrain}=\sum_{l=0}^{L}[(1-D^{l}_{b})+(1-D^{l}_{h})+(1-D^{l}_{n})],$
(5)

where $L$ is the total number of blocks in a MIA-Former. When
$\mathcal{L}_{pretrain}$ is minimized to zero, the MIA-Controller pretraining
stage is finished.

MIA-Former Co-training. The learning objective of MIA-Former is to reduce the
computational cost while preserving the model accuracy. To this end, we define
the training loss as follow,

$\mathcal{L}_{diff}=\mathcal{L}_{task}+\alpha\mathcal{L}_{cost},$ (6)

where $L_{task}$ is the task loss, $L_{cost}$ is the computational cost loss,
$\alpha$ is a weighted factor that trades off importance between accuracy and
computational budget. In this work, we define $L_{cost}$ as

$L_{cost}=\frac{\text{FLOPs}_{exec}}{\text{FLOPs}_{total}},$ (7)

where $\text{FLOPs}_{exec}$ and $\text{FLOPs}_{total}$ are the total FLOPs of
the executed parts in a MIA-Former and the total FLOPs when executing the
whole MIA-Former model.

To ensure that the MIA-Former can adaptively allocate computational budget
among all the input samples, instead of setting $\alpha$ to a fixed value, we
dynamically change the sign of $\alpha$ during training. Specifically, when
$\text{FLOPs}_{exec}$ is larger than the given target FLOPs
($\text{FLOPs}_{target}$), we set $\alpha>0$ to penalize MIA-Former for
adapting to a smaller computation budget. On the other hand, when
$\text{FLOPs}_{exec}$ is smaller than $\text{FLOPs}_{target}$, we set
$\alpha<0$ to encourage the model to execute more components in the MIA-
Former.

Figure 2: Benchmarking model performance (top-1 accuracy) and inference
efficiency (FLOPs) trade-off, where we compare our proposed MIA-Former with
SOTA dynamic ViTs (DynamicViT) and other SOTA image classification ViTs/CNNs
on the ImageNet dataset.  Figure 3: Comparison of MIA-Former with DynamicViT
and its backbone transformer models on the achieved ImageNet top-1 accuracy
and robustness accuracy under both PGD and FGSM attacks.

#### Skipping Policy Finetuning with Hybrid RL.

We further finetune the MIA-Former with a hybrid RL method to achieve better
performance. Specifically, we use an A2C-based (Mnih et al. 2016) RL learning
method to train the MIA-Controller and MIA-Block modules in a differentiable
manner. To make the MIA-Controller compatible with the above A2C algorithm, we
replace the last FC layer in each of the skipping dimensions (i.e., depth,
heads, and tokens) in an MIA-Controller with two parallel FC layers as the
actor and critic network in the A2C algorithm, respectively, constructing a
simple RL agent. To guarantee stable training, we further inherit all trained
weights from the MIA-Block and 75% of weights in the remaining parts of the
MIA-Controller. In this way, the newly introduced RL agent can start from the
learned features from the trained MIA-Controller, while the MIA-Block is also
adapted to the dynamically skipping mechanism. The reward function of RL is
defined as:

$\mathcal{R}=Y+\beta(\text{FLOPs}_{target}-\text{FLOPs}_{exec})$ (8)

where $Y$ is a binary value indicating the task accuracy, and $\beta$ is a
weighted factor that trades off ViTs’ model accuracy and efficiency.

To make the MIA-Block compatible with RL-based MIA-Controller, we also co-
train the MIA-Block alone with MIA-Controller in this stage with the
optimization objective defined as:

$\min_{\omega}\mathcal{L}_{hybrid}=\mathcal{L}_{task}-\mathcal{R},$ (9)

where $\omega$ is the weights of the MIA-Block.

##  4 Experiment Results

In this section, we first introduce the experiment setup, including models,
datasets, training hyperparams, baselines, and evaluation metrics. Then, we
validate the superiority of MIA-Former in boosting both efficiency and
robustness. In particular, we final that (1) MIA-Former can boost the
efficiency, i.e., achieving better accuracy vs. efficiency trade-off than both
vanilla models and models with SOTA dynamic ViT techniques, and (2) MIA-Former
can boost the robustness, i.e., achieving higher adversarial robustness while
preserving the accuracy compared with both vanilla ViTs and dynamic ViT
variants. Furthermore, we perform ablation study on (1) the comparison of the
redundancy between different dimensions of MIA-Former based on the relative
accuracy gap when running MIA-Former with different subsets of the input-
adaptive granularities, and (2) which training strategy for learning the
skipping policy can win better adversarial robustness in MIA-Former. Finally,
we visualize the skipping policies on a subset of input samples as well as the
skipping ratio distributions of different modules in the proposed MIA-Former.

###  4.1 Settings

#### Models, Datasets, and Baselines.

We evaluate our proposed MIA-Former over three ViT models (i.e., DeiT-Small
(Touvron et al. 2021), LeViT-192 and LeViT-256 (Graham et al. 2021)) on
ImageNet-1K dataset (Deng et al. 2009). We benchmark our method with vanilla
ViTs, SOTA dynamic ViT method DynamicViT (Rao et al. 2021) and other SOTA
ViTs/CNNs designs. For adversarial robustness, we evaluate the proposed MIA-
Former using Projected Gradient Descent (PGD) (Madry et al. 2017) attack under
$L_{Inf}$ constraint with a perturbation strength of 0.002 and Fast Signed
Gradient Matching (FGSM) attack under $L_{2}$ constraint with a perturbation
strength of 0.03 (Athalye, Carlini, and Wagner 2018), following the
adversarial attack setting in (Chen et al. 2020).

#### Training Recipe.

We adopt a three-stage training strategy: Stage 1: MIA-Controller pretraining:
we use an Adam (Kingma and Ba 2014) optimizer with a learning rate of 1e-4 to
train the MIA-Controller with fixed MIA-Block until $\mathcal{L}_{pretrain}$
is decreased to 0. Stage 2: MIA-Former co-training: we use an AdamW
(Loshchilov and Hutter 2017) optimizer with a batch size of 1024 and a
learning rate of 1e-5/1e-3 to train the MIA-Block/MIA-Controller,
respectively, for 200 epochs. We set the $\alpha$ to
$0.1\times\frac{L_{cls}}{L_{cost}}$. Stage 3: Skipping policy finetuning with
hybrid RL: after inserting the RL agents, we first train the RL agent for 20
epochs with all other parameter fixed and then unfreeze other parameters and
co-train the MIA-Former for a total of 50 epochs.

Table 1: Comparing MIA-Former with SOTA ViTs/CNNs on ImageNet. We compare our
MIA-Former with SOTA models under comparable FLOPs. We refer to MIA-Former
with different backbones as MIA-Former-BACKBONE. We also include the results
of vanilla models as references.

Model | GFLOPs | Accuracy (%)
---|---|---
MobileNet (Howard et al. 2017) | 0.58 | 70.6
PVT-v2 (Wang et al. 2021a) | 0.61 | 76.9
PiT-Ti (Heo et al. 2021) | 0.70 | 74.6
EfficientNet-B1 (Tan and Le 2019) | 0.70 | 79.1
LeViT-192 (Graham et al. 2021) | 0.66 | 80.0
IPE (Chen et al. 2021b) | 0.88 | 78.6
MIA-Former-LeViT-192 | 0.61 | 79.9
DeiT-Tiny (Touvron et al. 2021) | 1.30 | 72.2
PVT-v2 (Wang et al. 2021a) Swin (Liu et al. 2021) | 0.98 | 74.5
PiT-XS (Heo et al. 2021) | 1.40 | 79.1
LocalViT (Li et al. 2021) | 1.05 | 74.9
CoaT (Xu et al. 2021) | 0.95 | 73.8
EfficientNet-B2 (Tan and Le 2019) | 1.00 | 80.1
LeViT-256 (Graham et al. 2021) | 1.12 | 81.5
MIA-Former-LeViT-256 | 1.06 | 81.5

###  4.2 Benchmark with SOTA Designs

#### Enhanced accuracy-efficiency trade-off.

We apply the MIA-Former framework on top of three ViT models, including DeiT-
Small (one of the most widely used ViTs), LeViT-192, and LeViT-256, which are
SOTA ViT models with the optimal accuracy-efficiency trade-off. As shown in
Fig. 2, we observe that the MIA-Former can achieve a 20.1% reduction in FLOPs
on top of vanilla DeiT-Small with comparable accuracy. We further benchmark
the proposed MIA-Former with other SOTA CNNs and ViTs in Tab. 1 and show that
the proposed MIA-Former pushes forward the frontier of the achievable
accuracy-efficiency trade-off with a 5.4% reduction in FLOPs on top of
LeViT-256 and a comparable accuracy.

Figure 4: Visualizing the blockwise skipping ratios along different dynamic
dimensions on MIA-Former-DeiT-Small. Figure 5: Visualizing MIA-Former’s input-
adaptive skipping policy on different input samples. Each square in the figure
represents a block in MIA-Former, the size of blue square indicates the
executed part in the corresponding block, dashed square without any color
represents fully skipped block. The height and width of the blue square
indicate the executed number of heads and number of tokens, respectively.

#### Robustness improvement.

We compare the adversarial robustness of MIA-Former with vanilla DeiT-Small
and DynamicViT (Rao et al. 2021) models. As shown in Fig. 3, DynamicViT
suffers from the reduction in robust accuracy in both cases of attacks, while
our proposed method can achieve even higher model robustness compared with the
original model, leading to a win-win in robustness and efficiency.
Specifically, MIA-Former-DeiT-Small achieves up to a 2.4%/3.0% higher robust
accuracy and a 26.1% less FLOPs with comparable natural accuracy compared with
the original DeiT-Small model under PGD/FGSM attacks, respectively.

Table 2: Ablation study on different combination of dynamic dimensions.

Dynamic Dimension | GFLOPs | Accuracy (%)
---|---|---
Head | Depth | Token
|  |  | 4.6 | 79.9
✓ |  |  | 3.9 | 78.7
| ✓ |  | 4.1 | 76.2
|  | ✓ | 3.8 | 79.7
✓ | ✓ |  | 4.0 | 78.6
✓ |  | ✓ | 3.7 | 79.8
| ✓ | ✓ | 3.9 | 79.3
✓ | ✓ | ✓ | 3.9 | 79.9

Table 3: Ablation study on the effect of hybrid training on model’s
robustness.

| Target GFLOPs | Inherit Weight (%)
---|---|---
50 | 75 | 100 | Pretrain
Clean Accuracy | 3.4 | 79.2 | 79.9 | 79.9 | 79.9
Robust Accuracy | 21.3 | 19.8 | 15.7 | 14.2
Clean Accuracy | 3.9 | 78.3 | 79.6 | 79.6 | 79.5
Robust Accuracy | 22.8 | 20.2 | 17.5 | 16

###  4.3 Ablation Study

#### Analysis of the redundancy along each dimension.

To better understand the contribution of each dimension (i.e., head-wise,
depth-wise, and token-wise) to the finally achieved accuracy-efficiency trade-
off by MIA-Former, we conduct an ablation study by only enabling input-
adaptive skipping at certain dimensions of MIA-Former on top of DeiT-Small. In
Tab. 2, the first row is the vanilla DeiT-Small model without any input-
adaptive mechanism and the last row is the proposed MIA-Former with input-
adaptive skipping enabled along all dimensions. As shown in Tab. 2, by
activating different dimensions, the performance of the MIA-Former varies
significantly under similar inference FLOPs. Specifically, only activating
dynamic depth suffers from $3.7\%$ lower accuracy with $0.2$ GFLOPs higher
than the fully activated MIA-Former, indicating that the redundancy of the
model in the depth-wise dimension is the lowest. On the other hand, even when
only activating token-wise skipping, the accuracy only drops for $0.2\%$ with
comparable FLOPs compared to MIA-Former, indicating the high redundancy in
this dimension. This observation aligns with our intuition that depth-wise
skipping is of the most coarse granularity while token-wise skipping has the
finest granularity in ViTs.

#### Effectiveness of finetuning the skipping policy with hybrid RL.

How to inherit weights from the pretrained MIA-Controller when introducing the
RL agent is critical to the finally achieved natural and robust accuracy. We
study the impact of the portion of the inherited MIA-Controller weight at the
beginning of the skipping policy finetuning stage on top of DeiT-Small under
different FLOPs constraints. As shown in Tab. 3, there exists a trade-off
between the natural and robust accuracy when using different weight inheriting
strategies. In particular, inheriting a sufficient portion of pretrained MIA-
Controller weights leads to significantly better natural accuracy. On the
other hand, starting with a fully pretrained MIA-Controller without any
reinitialization will degrade the robust accuracy. We suspect that this is
because the inherited MIA-Controller makes the RL agent observe a similar
feature as the differentiablly co-trained MIA-Controller in the previous
training stage. The RL agent may try to make the decision in the same way as
the differentiablly co-trained MIA-Controller. Thus, we pick a sweet spot from
the trade-off and inherit 75% weights throughout the paper.

###  4.4 Skipping Policy Visualization

We first summarize the statistical characteristic of the generated skipping
policy on the validation set of ImageNet-1k (Deng et al. 2009). As in Fig. 4,
we have the following observations: (1) deeper blocks have significantly
higher redundancy than the shallower layers among all skipping dimensions in
MIA-Former, and (2) the token-wise dimension has the highest skipping
probability compared with the depth-wise and head-wise dimension, which is
consistent with the ablation study on different combinations of skipping
dimensions.

To better understand the behavior of MIA-Former, we visualize the generated
skipping policy of MIA-Former on different input samples. As shown in Fig. 5,
MIA-Former generates different skipping policies based on different input
samples. We can observe that MIA-Former can adaptively generate different
policy based on the difficulty of input samples. For example, according to the
comparison between Fig. 5(a) and (b), MIA-Former skips more from the token-
wise dimension when processing (b) while skipping more from the head-wise
dimension when processing (a) since (a) is full of corn in the image while (b)
has a clean background with a small object in the front.

##  5 Conclusion

In this paper, we propose, develop, and validate MIA-Former, a multi-grained
input-adaptive ViT framework, that is compatible with most of SOTA ViTs to
achieve a higher accuracy-efficiency trade-off by dynamically skipping ViTs’
blocks, heads, and tokens at coarse-to-fine granularities in an input-adaptive
manner. The proposed hybrid supervised and reinforcement training method for
effectively training the MIA-Former not only improves the achievable accuracy-
efficiency trade-off of MIA-Former, but also boosts the robustness accuracy
under different adversarial attacks, leading to a triple-win benefits.
Extensive experiments and ablation studies show that the proposed MIA-Former
achieves both efficiency and robustness improvement when being applied on top
of various SOTA ViTs, pushing forward the frontier of ViTs’ accuracy-
efficiency trade-offs. To the best of our knowledge, this work is the first to
simultaneously target multi-grained dynamic inference for ViTs.

## References

  * Athalye, Carlini, and Wagner (2018) Athalye, A.; Carlini, N.; and Wagner, D. 2018.  Obfuscated gradients give a false sense of security: Circumventing defenses to adversarial examples.  In _International conference on machine learning_ , 274–283. PMLR.
  * Caron et al. (2021) Caron, M.; Touvron, H.; Misra, I.; Jégou, H.; Mairal, J.; Bojanowski, P.; and Joulin, A. 2021.  Emerging properties in self-supervised vision transformers.  _arXiv preprint arXiv:2104.14294_.
  * Chen, Fan, and Panda (2021) Chen, C.-F.; Fan, Q.; and Panda, R. 2021.  Crossvit: Cross-attention multi-scale vision transformer for image classification.  _arXiv preprint arXiv:2103.14899_.
  * Chen et al. (2021a) Chen, H.; Wang, Y.; Guo, T.; Xu, C.; Deng, Y.; Liu, Z.; Ma, S.; Xu, C.; Xu, C.; and Gao, W. 2021a.  Pre-trained image processing transformer.  In _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition_ , 12299–12310.
  * Chen et al. (2021b) Chen, P.; Chen, Y.; Liu, S.; Yang, M.; and Jia, J. 2021b.  Exploring and Improving Mobile Level Vision Transformers.  _arXiv preprint arXiv:2108.13015_.
  * Chen et al. (2020) Chen, T.; Zhang, Z.; Liu, S.; Chang, S.; and Wang, Z. 2020.  Robust overfitting may be mitigated by properly learned smoothening.  In _International Conference on Learning Representations_.
  * Deng et al. (2009) Deng, J.; Dong, W.; Socher, R.; Li, L.-J.; Li, K.; and Fei-Fei, L. 2009.  Imagenet: A large-scale hierarchical image database.  In _2009 IEEE conference on computer vision and pattern recognition_ , 248–255. Ieee.
  * Dong et al. (2021) Dong, X.; Bao, J.; Chen, D.; Zhang, W.; Yu, N.; Yuan, L.; Chen, D.; and Guo, B. 2021\.  CSWin Transformer: A General Vision Transformer Backbone with Cross-Shaped Windows.  _arXiv preprint arXiv:2107.00652_.
  * Dosovitskiy et al. (2020) Dosovitskiy, A.; Beyer, L.; Kolesnikov, A.; Weissenborn, D.; Zhai, X.; Unterthiner, T.; Dehghani, M.; Minderer, M.; Heigold, G.; Gelly, S.; et al. 2020\.  An image is worth 16x16 words: Transformers for image recognition at scale.  _arXiv preprint arXiv:2010.11929_.
  * Fu et al. (2020) Fu, Y.; You, H.; Zhao, Y.; Wang, Y.; Li, C.; Gopalakrishnan, K.; Wang, Z.; and Lin, Y. 2020.  Fractrain: Fractionally squeezing bit savings both temporally and spatially for efficient dnn training.  _arXiv preprint arXiv:2012.13113_.
  * Fu et al. (2021a) Fu, Y.; Yu, Q.; Li, M.; Chandra, V.; and Lin, Y. 2021a.  Double-Win Quant: Aggressively Winning Robustness of Quantized Deep Neural Networks via Random Precision Training and Inference.  In _International Conference on Machine Learning_ , 3492–3504. PMLR.
  * Fu et al. (2021b) Fu, Y.; Yu, Q.; Zhang, Y.; Wu, S.; Ouyang, X.; Cox, D. D.; and Lin, Y. 2021b.  Drawing Robust Scratch Tickets: Subnetworks with Inborn Robustness Are Found within Randomly Initialized Networks.  In _Thirty-Fifth Conference on Neural Information Processing Systems_.
  * Fu et al. (2021c) Fu, Y.; Zhao, Y.; Yu, Q.; Li, C.; and Lin, Y. 2021c.  2-in-1 Accelerator: Enabling Random Precision Switch for Winning Both Adversarial Robustness and Efficiency.  In _MICRO-54: 54th Annual IEEE/ACM International Symposium on Microarchitecture_ , 225–237.
  * Graham et al. (2021) Graham, B.; El-Nouby, A.; Touvron, H.; Stock, P.; Joulin, A.; Jégou, H.; and Douze, M. 2021.  LeViT: a Vision Transformer in ConvNet’s Clothing for Faster Inference.  _arXiv preprint arXiv:2104.01136_.
  * He et al. (2016) He, K.; Zhang, X.; Ren, S.; and Sun, J. 2016.  Deep residual learning for image recognition.  In _Proceedings of the IEEE conference on computer vision and pattern recognition_ , 770–778.
  * Heo et al. (2021) Heo, B.; Yun, S.; Han, D.; Chun, S.; Choe, J.; and Oh, S. J. 2021.  Rethinking spatial dimensions of vision transformers.  _arXiv preprint arXiv:2103.16302_.
  * Howard et al. (2017) Howard, A. G.; Zhu, M.; Chen, B.; Kalenichenko, D.; Wang, W.; Weyand, T.; Andreetto, M.; and Adam, H. 2017.  Mobilenets: Efficient convolutional neural networks for mobile vision applications.  _arXiv preprint arXiv:1704.04861_.
  * Hu et al. (2020) Hu, T.-K.; Chen, T.; Wang, H.; and Wang, Z. 2020.  Triple wins: Boosting accuracy, robustness and efficiency together by enabling input-adaptive inference.  _arXiv preprint arXiv:2002.10025_.
  * Kaya, Hong, and Dumitras (2019) Kaya, Y.; Hong, S.; and Dumitras, T. 2019.  Shallow-deep networks: Understanding and mitigating network overthinking.  In _International Conference on Machine Learning_ , 3301–3310. PMLR.
  * Kingma and Ba (2014) Kingma, D. P.; and Ba, J. 2014.  Adam: A method for stochastic optimization.  _arXiv preprint arXiv:1412.6980_.
  * Kurakin et al. (2016) Kurakin, A.; Goodfellow, I.; Bengio, S.; et al. 2016.  Adversarial examples in the physical world.
  * Li et al. (2021) Li, Y.; Zhang, K.; Cao, J.; Timofte, R.; and Van Gool, L. 2021.  Localvit: Bringing locality to vision transformers.  _arXiv preprint arXiv:2104.05707_.
  * Liu et al. (2021) Liu, Z.; Lin, Y.; Cao, Y.; Hu, H.; Wei, Y.; Zhang, Z.; Lin, S.; and Guo, B. 2021\.  Swin transformer: Hierarchical vision transformer using shifted windows.  _arXiv preprint arXiv:2103.14030_.
  * Loshchilov and Hutter (2017) Loshchilov, I.; and Hutter, F. 2017.  Decoupled weight decay regularization.  _arXiv preprint arXiv:1711.05101_.
  * Madry et al. (2017) Madry, A.; Makelov, A.; Schmidt, L.; Tsipras, D.; and Vladu, A. 2017.  Towards deep learning models resistant to adversarial attacks.  _arXiv preprint arXiv:1706.06083_.
  * Mao et al. (2021) Mao, X.; Qi, G.; Chen, Y.; Li, X.; Ye, S.; He, Y.; and Xue, H. 2021.  Rethinking the Design Principles of Robust Vision Transformer.  _arXiv preprint arXiv:2105.07926_.
  * Mnih et al. (2016) Mnih, V.; Badia, A. P.; Mirza, M.; Graves, A.; Lillicrap, T.; Harley, T.; Silver, D.; and Kavukcuoglu, K. 2016.  Asynchronous methods for deep reinforcement learning.  In _International conference on machine learning_ , 1928–1937. PMLR.
  * Rakin et al. (2019) Rakin, A. S.; He, Z.; Yang, L.; Wang, Y.; Wang, L.; and Fan, D. 2019.  Robust sparse regularization: Simultaneously optimizing neural network robustness and compactness.  _arXiv preprint arXiv:1905.13074_.
  * Rakin et al. (2018) Rakin, A. S.; Yi, J.; Gong, B.; and Fan, D. 2018.  Defend deep neural networks against adversarial examples via fixed and dynamic quantized activation functions.  _arXiv preprint arXiv:1807.06714_.
  * Rao et al. (2021) Rao, Y.; Zhao, W.; Liu, B.; Lu, J.; Zhou, J.; and Hsieh, C.-J. 2021.  DynamicViT: Efficient Vision Transformers with Dynamic Token Sparsification.  _arXiv preprint arXiv:2106.02034_.
  * Ridnik et al. (2021) Ridnik, T.; Ben-Baruch, E.; Noy, A.; and Zelnik-Manor, L. 2021.  Imagenet-21k pretraining for the masses.  _arXiv preprint arXiv:2104.10972_.
  * Sehwag et al. (2020) Sehwag, V.; Wang, S.; Mittal, P.; and Jana, S. 2020.  Hydra: Pruning adversarially robust neural networks.  _arXiv preprint arXiv:2002.10509_.
  * Shao et al. (2021) Shao, R.; Shi, Z.; Yi, J.; Chen, P.-Y.; and Hsieh, C.-J. 2021.  On the adversarial robustness of visual transformers.  _arXiv preprint arXiv:2103.15670_.
  * Shen et al. (2020) Shen, J.; Wang, Y.; Xu, P.; Fu, Y.; Wang, Z.; and Lin, Y. 2020.  Fractional skipping: Towards finer-grained dynamic cnn inference.  In _Proceedings of the AAAI Conference on Artificial Intelligence_ , volume 34, 5700–5708.
  * Steiner et al. (2021) Steiner, A.; Kolesnikov, A.; Zhai, X.; Wightman, R.; Uszkoreit, J.; and Beyer, L. 2021.  How to train your ViT? Data, Augmentation, and Regularization in Vision Transformers.  _arXiv preprint arXiv:2106.10270_.
  * Strudel et al. (2021) Strudel, R.; Garcia, R.; Laptev, I.; and Schmid, C. 2021.  Segmenter: Transformer for Semantic Segmentation.  _arXiv preprint arXiv:2105.05633_.
  * Tan and Le (2019) Tan, M.; and Le, Q. 2019.  Efficientnet: Rethinking model scaling for convolutional neural networks.  In _International Conference on Machine Learning_ , 6105–6114. PMLR.
  * Teerapittayanon, McDanel, and Kung (2016) Teerapittayanon, S.; McDanel, B.; and Kung, H.-T. 2016.  Branchynet: Fast inference via early exiting from deep neural networks.  In _2016 23rd International Conference on Pattern Recognition (ICPR)_ , 2464–2469. IEEE.
  * Touvron et al. (2021) Touvron, H.; Cord, M.; Douze, M.; Massa, F.; Sablayrolles, A.; and Jégou, H. 2021.  Training data-efficient image transformers & distillation through attention.  In _International Conference on Machine Learning_ , 10347–10357. PMLR.
  * Vaswani et al. (2017) Vaswani, A.; Shazeer, N.; Parmar, N.; Uszkoreit, J.; Jones, L.; Gomez, A. N.; Kaiser, Ł.; and Polosukhin, I. 2017.  Attention is all you need.  In _Advances in neural information processing systems_ , 5998–6008.
  * Wang et al. (2021a) Wang, W.; Xie, E.; Li, X.; Fan, D.-P.; Song, K.; Liang, D.; Lu, T.; Luo, P.; and Shao, L. 2021a.  Pvtv2: Improved baselines with pyramid vision transformer.  _arXiv preprint arXiv:2106.13797_.
  * Wang et al. (2018) Wang, X.; Yu, F.; Dou, Z.-Y.; Darrell, T.; and Gonzalez, J. E. 2018.  Skipnet: Learning dynamic routing in convolutional networks.  In _Proceedings of the European Conference on Computer Vision (ECCV)_ , 409–424.
  * Wang et al. (2021b) Wang, Y.; Huang, R.; Song, S.; Huang, Z.; and Huang, G. 2021b.  Not All Images are Worth 16x16 Words: Dynamic Vision Transformers with Adaptive Sequence Length.  _arXiv preprint arXiv:2105.15075_.
  * Wang et al. (2020) Wang, Y.; Shen, J.; Hu, T.-K.; Xu, P.; Nguyen, T.; Baraniuk, R.; Wang, Z.; and Lin, Y. 2020.  Dual dynamic inference: Enabling more efficient, adaptive, and controllable deep inference.  _IEEE Journal of Selected Topics in Signal Processing_ , 14(4): 623–633.
  * Wu et al. (2018) Wu, Z.; Nagarajan, T.; Kumar, A.; Rennie, S.; Davis, L. S.; Grauman, K.; and Feris, R. 2018.  Blockdrop: Dynamic inference paths in residual networks.  In _Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition_ , 8817–8826.
  * Xu et al. (2021) Xu, W.; Xu, Y.; Chang, T.; and Tu, Z. 2021.  Co-scale conv-attentional image transformers.  _arXiv preprint arXiv:2104.06399_.
  * Ye et al. (2019) Ye, S.; Xu, K.; Liu, S.; Cheng, H.; Lambrechts, J.-H.; Zhang, H.; Zhou, A.; Ma, K.; Wang, Y.; and Lin, X. 2019.  Adversarial robustness vs. model compression, or both?  In _Proceedings of the IEEE/CVF International Conference on Computer Vision_ , 111–120.
  * Zhou et al. (2021) Zhou, D.; Kang, B.; Jin, X.; Yang, L.; Lian, X.; Jiang, Z.; Hou, Q.; and Feng, J. 2021.  Deepvit: Towards deeper vision transformer.  _arXiv preprint arXiv:2103.11886_.