# HRSAM: Efficiently Segment Anything in High-Resolution Images

You Huang, Wenbin Lai, Jiayi Ji, Liujuan Cao, Shengchuan Zhang, Rongrong Ji
Key Laboratory of Multimedia Trusted Perception and Efficient Computing,
Ministry of Education of China, Xiamen University
Corresponding author

###### Abstract

The Segment Anything Model (SAM) has significantly advanced interactive
segmentation but struggles with high-resolution images crucial for high-
precision segmentation. This is primarily due to the quadratic space
complexity of SAM-implemented attention and the length extrapolation issue in
common global attention. This study proposes HRSAM that integrates Flash
Attention and incorporates Plain, Shifted and newly proposed Cycle-scan Window
(PSCWin) attention to address these issues. The shifted window attention is
redesigned with padding to maintain consistent window sizes, enabling
effective length extrapolation. The cycle-scan window attention adopts the
recently developed State Space Models (SSMs) to ensure global information
exchange with minimal computational overhead. Such window-based attention
allows HRSAM to perform effective attention computations on scaled input
images while maintaining low latency. Moreover, we further propose HRSAM++
that additionally employs a multi-scale strategy to enhance HRSAM’s
performance. The experiments on the high-precision segmentation datasets
HQSeg44K and DAVIS show that high-resolution inputs enable the SAM-distilled
HRSAM models to outperform the teacher model while maintaining lower latency.
Compared to the SOTAs, HRSAM achieves a $1.56$ improvement in interactive
segmentation’s NoC95 metric with only $31\%$ of the latency. HRSAM++ further
enhances the performance, achieving a $1.63$ improvement in NoC95 with just
$38\%$ of the latency. Code is available at
https://github.com/YouHuang67/High-Resolution-Segment-Anything.git.

##  1 Introduction

The Segment Anything Model (SAM) [36] stands as a cornerstone in interactive
segmentation [3, 105, 37, 8, 51, 57] and extends its utility to various
downstream tasks in computer vision [63, 65, 39, 89, 102, 85]. As an
interactive segmentation model, SAM takes simple manual inputs like clicks
[31], bounding boxes [36] and coarse masks [37] to predict precise
segmentation results, thus reducing the manual labeling costs associated with
image segmentation. SAM [36] enhances the overall efficiency of interactive
segmentation, setting new benchmarks for speed and performance. Conversely,
interactive segmentation serves as a standard for evaluating foundational
models [36, 111, 46] on fine-grained visual tasks. This study focuses on
mainstream click-based interactive segmentation and further enhances SAM’s
performance.

SAM utilizes large Vision Transformers (ViTs) [16, 47] to encode images into
embeddings, which are then decoded into segmentation results by a lightweight
decoder. Despite SAM’s considerable success, the current implementation of SAM
is limited to handling fixed input resolutions of $1024\times 1024$.
Nevertheless, real-world applications frequently require visual models to
handle high-resolution images, such as those exceeding $4096\times 4096$,
which offer richer details and potentially improved segmentation results.
Directly scaling SAM’s input resolution introduces two substantial challenges,
_i.e.,_ memory inefficiency and lack of scalability to input sizes, which
significantly affect performance and efficiency. Firstly, memory inefficiency
arises because SAM’s ViTs utilize global attention mechanisms, leading to a
quadratic increase in space complexity. Even with $2048^{2}$111In this paper,
any resolution expressed as $N^{2}$ (e.g., $1024^{2}$, $2048^{2}$, $3072^{2}$,
$4096^{2}$) refers to an image resolution of $N\times N$. resolution inputs,
SAM fails to perform inference on commonly used 3090 GPUs due to insufficient
GPU memory. Secondly, the scalability issue becomes apparent. Previous studies
indicate that merely increasing the image sizes for SAM cannot proportionally
improve performance [78]. This issue originates from the length extrapolation
problem [41, 78], where the model’s attention mechanisms struggle to adapt to
varying lengths of image token sequences during training and testing,
resulting in suboptimal performance given the extensive computational costs
for the scaled input image sizes.

Figure 1: Scalability analysis of global attention and the proposed PSCWin
attention. (a) illustrates the global attention, and (b) depicts PSCWin
attention. SAM-ViT-Huge is used as the teacher model for distilling into ViTs
equipped with either global or PSCWin attention. Latency metrics are
determined by the time required to process an image, and segmentation
performance is measured by 5-mIoU with five clicks in interactive segmentation
over HQSeg44K dataset. Remarkably, the scaled $2048^{2}$ resolution enables
PSCWin to achieve a 5-mIoU of $87.31\%$ with just $192$ms in latency,
surpassing the teacher’s $87.21\%$ at $424$ms in both performance and speed.
Moreover, PSCWin achieves comparable performance to Global Attention with just
$14\%$ of the latency in the scaled $2048^{2}$ resolution, demonstrating its
scalability to large resolutions.

To address the above challenges, we introduce HRSAM to harness the performance
gains from high-resolution images efficiently. To combat memory inefficiency,
we pioneer the integration of Flash Attention [81, 18] within the SAM
framework. Flash Attention has been proven across various domains to reduce
space complexity from quadratic to linear while maintaining high performance
[17, 81]. Notably, the original SAM’s attention mechanism employs relative
positional encoding, which is incompatible with Flash Attention. To resolve
this, we substitute it with Rotary Position Embedding (RoPE) [80]. Regarding
the scalability issue, we incorporate the innovative Plain, Shifted and Cycle-
scan Window (PSCWin) attention mechanisms. Plain window attention divides
image embeddings into non-overlapping windows, facilitating efficient self-
attention. We have modified the vanilla shifted window attention [60, 59] by
integrating padding to maintain consistent window sizes. This modification
enhances performance during length extrapolation while requiring only a
marginal increase in computational resources. Furthermore, our novel cycle-
scan window attention marries the newly proposed State Space Models (SSMs)
[21] with window attention. SSMs [58, 109] treat image embeddings as
sequences, scanning them with approximately linear computational complexity
and significantly boosting the efficiency of global information exchange. All
the window-based attention mechanisms ensure consistent attention computations
throughout both the training and testing phases, thereby addressing length
extrapolation problems.

HRSAM further integrates a multi-scale strategy, evolving into HRSAM++,
inspired by the successful applications of this strategy [76]. For each image,
in addition to resizing to standard resolutions such as $1024^{2}$ or
$2048^{2}$, we also resize images to a smaller resolution of $512^{2}$. These
multi-scale inputs are simultaneously processed by HRSAM++ to analyze image
features across different scales with high parallelism. This simultaneous
processing is enabled by an indexing operator that efficiently organizes the
image embeddings into the proposed PSC windows. In specific blocks, the SSMs
[21] are employed for effective multi-scale feature fusion.

We evaluate HRSAM and HRSAM++ on the high-precision image segmentation
datasets HQSeg44K [34] and DAVIS [69]. The results indicate that HRSAM
outperforms the previous SOTA by achieving a $1.59\%$ improvement in 5-mIoU
while only requiring $31\%$ of the latency. HRSAM++ further sets a new
benchmark for high-precision interactive segmentation, delivering SOTA
performance with just $38\%$ of the previous SOTA’s latency. Additionally,
when distilled solely from SAM-ViT-Huge with low-resolution inputs, without
any additional labeled data training, HRSAM models exhibit superior
performance with high-resolution inputs while maintaining lower latency than
their teacher model. This showcases a cost-effective training strategy for
developing high-performance visual foundation models.

Our main contributions are as follows:

  * •

We explore efficient SAM for interactive segmentation to handle high-
resolution images. Specifically, we propose HRSAM, utilizing Flash Attention
and PSCWin to address memory inefficiency and length extrapolation issues.

  * •

We further propose HRSAM++ that additionally adopts a multi-scale strategy to
boost HRSAM’s performance on large-sized input images with slight extra
computational costs.

  * •

Our experiments show that HRSAM and HRSAM++ outperform the previous SOTA
method in high-precision interactive segmentation with only $31\%$ and $38\%$
of the latency respectively. When merely distilled from SAM with low-
resolution inputs, the HRSAMs surpass the teacher in segmentation performance
with scaled input sizes while maintaining lower latency.

##  2 Related Work

Interactive Segmentation. The integration of deep networks into interactive
segmentation begins with DIOS [93], leading to advancements in click-based
methods [64, 48, 52, 33, 77]. Subsequent methods focus on enhancing various
aspects of interactive segmentation [3, 105, 37, 8, 51, 57, 56, 31, 92, 106,
43, 95, 108, 72, 42]. SAM [36] improves the inference latency by reusing image
features, achieving robust zero-shot capabilities and leading to various
downstream applications [63, 65, 39, 89, 102, 85, 94, 9]. However, SAM
struggles with the high-precision segmentation on high-resolution images,
limiting its broader applications. This paper introduces HRSAM to address
these issues.

Efficient Attention. Attention mechanisms [83] have made significant strides
in computer vision [16, 22, 35, 79, 90, 103, 18, 86]. Meanwhile, the high
computational complexity of attention mechanisms leads to extensive research
on efficient attention [60, 110, 23, 90, 60, 82, 47, 25, 28, 32, 11, 87, 97,
15, 73]. Flash Attention [13, 12] re-implements the vanilla attention
mechanism efficiently at the CUDA level, reducing the quadratic space
complexity to linear and accelerating computations. This study explores
window-based attention to facilitate the attention computations on high-
resolution images.

Visual Length Extrapolation. Length extrapolation refers to a model’s ability
to generalize to longer inputs than those it is trained on [78], which has
been explored in NLP [70, 10]. In vision, addressing length extrapolation
involves slight modifications of ViTs, e.g., adjustments to positional
embeddings and patch sizes [40, 4, 5, 20, 101, 49, 26, 38, 100], and
sophisticated training methods, e.g., sequence packing [14] and masking most
of image tokens [41]. Recent efforts adapt post-training attention
computations to handle scaled input images [78]. Our PSCWin attention achieves
scalability to large image size with simple distillation training.

Visual State Space Model. Several efforts extend the recently developed State-
Space Models (SSMs) [21] to vision, including generation [19, 29, 75], multi-
modal tasks [98, 45, 71, 84, 107], medical image analysis [62, 88, 99, 54, 74]
and remote sensing [7, 6, 53]. Various studies [109, 58, 44, 96, 30, 61, 68,
67] explore 2D scanning strategies to enhance SSM applications. The proposed
cycle-scan in this study simplifies scanning by replicating and concatenating
image token sequences.

##  3 Method

This study introduces HRSAM to tackle SAM’s issues of memory inefficiency and
length extrapolation on high-resolution images. Section 3.1 provides an
overview of HRSAM. Section 3.2 presents both the plain window and the improved
shifted window attention. Section 3.3 proposes the cycle-scan window
attention. Section 3.4 discusses HRSAM++’s multi-scale strategy.

Figure 2: Overview of our proposed HRSAM. HRSAM employs a four-stage encoder
with Plain, Shifted and Cycle-scan Window attention blocks. The cycle-scan
window block comprises a cycle-scan module and a plain/shifted attention
block. These blocks are distributed evenly across the stages. Outputs from all
stages are fused through summation and refined via a convolutional block to
create final image embeddings fed into the SAM decoder.

###  3.1 Overview of HRSAM

SAM pipeline. SAM integrates a heavy ViT-based image encoder, a compact prompt
encoder and a light decoder to efficiently generate segmentation results.
Prior to interactive segmentation, images are resized to $H\times W$
(typically $1024\times 1024$) resolution and undergo offline preprocessing by
the image encoder. Such preprocessing produces image embeddings that are
repeatedly used in the following steps. During online interaction, the prompt
encoder dynamically transforms user prompts such as clicks, bounding boxes and
previous segmentation masks into prompt embeddings. The decoder then combines
these prompt embeddings with the preprocessed image embeddings to produce the
final segmentation results through a series of cross-attention computations.
The SAM encoder constitutes $>95\%$ of the overall parameters and has a
critical impact on the overall performance. Thus, our HRSAM focuses on
optimizing the SAM encoder without modifying the other modules.

HRSAM encoder. The SAM encoder [36] utilizes ViTs [16] with minor
modifications, employing both window-based and global attention in specific
blocks [47] and incorporating relative positional encoding during attention
computations. Specifically, the SAM-ViT-Base comprises 12 standard transformer
blocks, organized into four stages with each stage containing two window-
attention blocks and one conventional global-attention block. As depicted in
Figure 2, HRSAM also consists of four stages but integrates Plain, Shifted and
Cycle-scan Window attention blocks, distributed evenly across each stage.
Moreover, HRSAM enhances the encoder by fusing outputs from all four stages
through summation, followed by a convolutional block to produce the final
$C$-dimensional image embeddings
$\mathbf{F}\in\mathbb{R}^{B\times\frac{H}{16}\times\frac{W}{16}\times C}$ ($B$
is the batch size). Additionally, HRSAM replaces SAM’s relative encoding with
RoPE [80] to facilitate the use of Flash Attention [13, 12] since the relative
encoding introduces biases in the attention maps that are incompatible with
Flash Attention [81, 18].

Flash Attention. Flash Attention [13, 12] addresses the memory inefficiency of
the conventional attention mechanism [83] by minimizing high-bandwidth memory
(HBM) read/write operations. It achieves this through tiling, where attention
inputs are processed in smaller, independent blocks, which recalculates the
intermedia results when needed instead of accessing HBM. Additionally, Flash
Attention utilizes the fast cache, SRAM, with low latency. These collectively
reduce memory overhead and transform the attention’s space complexity from
quadratic to linear, leading to significant memory efficiency for large-scale
attention computations. We provide the details in Appendix A.

Positional encoding. SAM’s attention employs the sophisticated relative
positional encoding, which hinders the application of Flash Attention [13,
12]. Thus, we replace the encoding with Rotary Position Embedding (RoPE) [80].

###  3.2 Plain and Shifted Window Attention

Conventional attention. Given a batch of $B$ sequentially arranged embedding
sets $\mathbf{X}\in\mathbb{R}^{B\times L\times C}$ with length $L$,
conventional attention first performs linear projections on $\mathbf{X}$ to
produce the $Q,K,V\in\mathbb{R}^{B\times L\times C}$. Then, the attention
outputs are formulated as

$\begin{split}\text{Attention}(\mathbf{X})=\text{Linear}\left(\dfrac{\text{Softmax}\left(\text{Linear}_{Q}(\mathbf{X})\text{Linear}_{K}(\mathbf{X})^{\top}\right)\text{Linear}_{V}(\mathbf{X})}{\sqrt{d}}\right),\end{split}$
(1)

where the matrix multiplication and transpose operators are performed
independently for each sample. The discussion of multi-head implementation is
omitted since it could be easily extended.

Plain window attention. Given the image embeddings
$\mathbf{F}\in\mathbb{R}^{B\times\frac{H}{16}\times\frac{W}{16}\times C}$ and
window size $w$, the plain window attention [47] firstly rearranges
$\mathbf{F}$ into $\mathbf{F}_{w}\in\mathbb{R}^{\frac{BHW}{16^{2}w^{2}}\times
w^{2}\times C}$, and then performs the conventional attention, i.e.,
$\text{Attention}(\mathbf{F}_{w})$, which is finally rearranged into the
original shape.

Vanilla shifted window attention. As depicted in Figure 3 (a), the vanilla
shifted window attention introduced in the Swin Transformer [60, 59] shifts
the embedding windows to enable interactions across neighboring windows. This
process involves moving each embedding within $\mathbf{F}$ along the x-y axes
by a predetermined number of units, with boundary embeddings wrapping around
to the opposite side to maintain the dimensions of $\mathbf{F}$. This design
facilitates more dynamic receptive fields and connectivity between adjacent
windows, enhancing the window attention’s ability to integrate contextual
information across the entire image.

Padding shifted window attention. However, the vanilla shifted window leads to
inconsistent window sizes across the embedding windows. To address this, we
introduce a padding strategy as depicted in Figure 3 (b). This padding
strategy utilizes a learnable embedding $\mathbf{p}\in\mathbb{R}^{C}$ to add
padding around $\mathbf{F}$, ensuring uniform window sizes throughout.
Specifically, when applying the vanilla shifting strategy, i.e., shifting
$\mathbf{F}$ along the x-axis by $S_{x}$ units and along the y-axis by $S_{y}$
units, the proposed padding shifted window instead performs padding of
$w-S_{x}$ units on the left side along the x-axis and $w-S_{y}$ units on the
top side along the y-axis. Additional padding is applied to the opposite sides
as necessary to ensure a complete number of windows. Although this padding
strategy increases the size of $\mathbf{F}$, we reduce computational overhead
through the replication operator. Before performing the padding, $\mathbf{F}$
and the learnable padding embedding $\mathbf{p}$ are both projected through
the attention’s QKV-linear layer. Then, the $(Q,K,V)$ of $\mathbf{p}$ are
replicated and concatenated to the $(Q,K,V)$ of $\mathbf{F}$. After the
attention computation, the paddings are discarded, retaining only the region
corresponding to $\mathbf{F}$. The efficiency and performance impacts of such
a padding strategy are detailed in Figure 3 (c).

###  3.3 Cycle-scan Window Attention

Figure 3: Illustration of proposed padding shifted window and Cycle-scan
module.

State Space Models. Structured State Space Models (SSMs) [21] are a class of
sequence models designed for sequence-to-sequence transformation, adept at
capturing long dependencies among sequential tokens. Specifically, SSMs map a
one-dimensional sequence $x(t)\in\mathbb{R}$ to $y(t)\in\mathbb{R}$ through an
intermediate latent state $h(t)\in\mathbb{R}^{N}$, formulated as:

$\begin{split}h^{\prime}(t)&=\mathbf{A}h(t)+\mathbf{B}x(t),\\\
y(t)&=\mathbf{C}h(t),\\\ \end{split}$ (2)

where the matrices $\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 predefined. Actually, the SSMs are discretized through a zero-order
hold rule with a given sample timescale $\Delta\in\mathbb{R}$, formulated as

$\begin{split}\bar{\mathbf{A}}=e^{\Delta\mathbf{A}},\quad\bar{\mathbf{B}}=(\Delta\mathbf{A})^{-1}\left(e^{\Delta\mathbf{A}}-\mathbf{I}\right)\Delta\mathbf{B},\quad\bar{\mathbf{C}}=\mathbf{C},\end{split}$
(3)

which implies the following iterative process:

$\begin{split}h_{t}&=\bar{\mathbf{A}}h_{t-1}+\bar{\mathbf{B}}x_{t},\\\
y_{t}&=\bar{\mathbf{C}}h_{t},\end{split}$ (4)

where $1\leq t\leq L$. To maintain the efficiency of parallel computation,
these iterative steps are transformed into a global convolutional operation:

$\begin{split}\mathbf{y}&=\mathbf{x}\ast\bar{\mathbf{K}}\text{, with
}\bar{\mathbf{K}}=\left(\bar{\mathbf{C}}\bar{\mathbf{B}},\bar{\mathbf{C}}\bar{\mathbf{A}}\bar{\mathbf{B}},\ldots,\bar{\mathbf{C}}\bar{\mathbf{A}}^{L-1}\bar{\mathbf{B}}\right),\end{split}$
(5)

where $\ast$ is the convolution operation. In summary, the SSM operator takes
the input sequence $\mathbf{x}=\\{x_{t}\\}_{t=1}^{L}$ as inputs, and produces
outputs of the same dimension, $\mathbf{y}=\\{y_{t}\\}_{t=1}^{L}$. SSMs are
capable of modeling global information with linear computational complexity.
Additionally, SSMs can be readily extended to handle $d$-dimensional sequences
$\mathbf{X}\in\mathbb{R}^{d\times L}$, processing each channel independently.
In this study, we leverage the recent advancements in selective SSMs,
specifically the Mamba [21]. For a comprehensive understanding of Mamba, we
refer readers to the detailed descriptions available in [21].

Cycle-scan module. Mamba’s SSM operator has been employed to scan images from
multiple directions to establish spatial dependencies among image tokens [58].
Instead of using sophisticated scanning strategies, we introduce the cycle-
scan strategy that simplifies the scanning process by expanding the image
token sequence through replication. As illustrated in Figure 3 (d), the
proposed cycle-scan strategy repeats the image token sequence three times and
connects them sequentially. The Mamba SSM operator then scans the tokens in
this order. After scanning, the results are split into the corresponding three
sequences and are subsequently merged through summation.

Cycle-scan window attention. The cycle-scan module is uniformly integrated
into all four stages of the HRSAM, as depicted in Figure 2. Specifically, the
scanning is conducted prior to the final block of each stage to facilitate
global information exchange among image tokens. The combination of the cycle-
scan with either plain or shifted window attention forms the cycle-scan window
attention.

Figure 4: Overview of HRSAM++. HRSAM++ is a multi-scale version of HRSAM that
takes both the original and downsampled images as inputs. Multi-scale features
are fused at the end of each stage through the multi-scale cycle-scan module.
The outputs of these modules are then split by the input scale, resized to a
fixed size, and concatenated to form the final output.

###  3.4 Multi-scale Fusion

Multi-scale inputs. Multi-scale strategies have been extensively explored in
visual tasks due to the proven effectiveness [76]. In this study, we propose a
specific multi-scale strategy to address the issue of length extrapolation.
This strategy involves resizing any image to a fixed size (e.g., $512^{2}$),
which provides an overview of the entire scene. Such overview enhances the
model’s global understanding and alleviates the issues posed by variable
receptive fields across different input sizes. Thus, we further propose an
enhanced version of our model, the HRSAM++, which incorporates the additional
downsampled input image as depicted in Figure 4. Both the original and
additional images are processed concurrently. Multi-scale features from these
images are integrated at the end of each stage in the HRSAM++. Further details
are discussed subsequently.

Multi-scale attention. In HRSAM++, each attention operator is replaced by a
multi-scale version, which is facilitated through the block-diagonal mask from
Flash Attention [13, 12], enabling parallel processing of multi-scaled inputs.
Specifically, both the original $H\times W$ image and the downsampled
$H_{s}\times W_{s}$ image are patchified and concatenated to form a sequence
of image tokens with a length of $(HW+H_{s}W_{s})/16^{2}$, as illustrated in
Figure 4. During the window attention computations, this token sequence is
reorganized by indexing to ensure tokens within the same window are sequential
and tokens from the same scale remain contiguous. This reorganization enables
the use of Flash Attention for attention computation across the entire token
sequence, utilizing the block-diagonal mask to skip unnecessary cross-window
attention. This implementation is more efficient and computationally
equivalent to performing attention on isolated windows. Additionally, for
padding shifted windows, the learnable padding embedding is firstly appended
at the sequence’s end. Then, this padded sequence is similarly reorganized
with appropriate token duplication, to maintain size consistency of the
windows and undergo the same attention processing. We provide the details in
Appendix B C.

Multi-scale cycle-scan module. For multi-scale inputs, the cycle-scan module
operates in two modes with subtle differences. As shown in Figure 4, HRSAM++
extends the four stages of HRSAM by incorporating an additional multi-scale
cycle-scan module, while retaining the original cycle-scan as a single-scale
operation within each stage. Both the single and multi-scale cycle-scan
modules follow the same process, differing only in the SSMs. The single-scale
cycle-scan first splits the tokens by scale, scan each scale’s tokens
separately by the SSM, and then concatenates them for post-computation. In
contrast, the multi-scale cycle-scan directly performs the SSM across the
tokens from all the scales, facilitating the multi-scale fusion.

##  4 Experiments

Section 4.1 details the experimental setup. Section 4.2 discusses the main
results. Section 4.3 evaluate HRSAM’s scalability to input sizes. Section 4.4
provides the ablation study of HRSAM’s modules. Section 4.5 presents the
qualitative results.

###  4.1 Experimental Setting

We provide the details in Appendix D and the code in the supplementary
materials.

Datasets. We adopt COCO [50], LVIS [24] and HQSeg44K [34] datasets for our
two-stage training process. In the first stage, we use only the images from
COCO and LVIS without labels, which differs from previous interactive
segmentation methods [8, 57, 56, 31]. In the second stage, we finetune the
weights pretrained at the first stage using the HQSeg44K training dataset. For
testing, we evaluate exclusively on the HQSeg44K validation [34] and DAVIS
[69] datasets, both of which have high-precision annotations, following recent
advances in high-precision interactive segmentation [55].

Training strategy. HRSAM and HRSAM++ undergo the same two-stage training
process. In the first stage, we use images from the COCO and LVIS datasets
without labels. We employ the vanilla MSE as the loss in the SAM distillation,
and minimize the MSE between HRSAM and the teacher’s image embeddings. In the
second stage, we further finetune the models using HQSeg44K training samples
with segmentation annotations, employing the commonly used normalized focal
loss [8, 57, 56, 31].

Evaluation. We evaluate HRSAMs against previous methods [37, 8, 56, 31, 55,
36, 104, 91]. In testing, click simulation places clicks at the centers of
erroneously predicted regions, aligning with the previous methods [8, 56, 31].
We assess all models’ segmentation performance and speed. Segmentation
performance is measured using 5-mIoU and NoC metrics. The 5-mIoU represents
the average IoU after the fifth click. The NoC metric indicates the average
minimum clicks required to reach a specified IoU. We focus on NoC@90 and
NoC@95 within 20 clicks. Speed is quantified as Seconds Per Click (SPC) on
GPUs, measuring the average inference latency per click, averaged over 20
clicks. For SAM series models, this includes the preprocessing time by the
backbones.

Model | Training Data | Latency $\downarrow$ |  HQSeg44K ${}_{\textbf{Max H/W > 4000}}$ |  DAVIS ${}_{\textbf{Max H/W < 1000}}$
---|---|---|---|---
20-SPC (ms) | 5-mIoU $\uparrow$ | NoC90 $\downarrow$ | NoC95 $\downarrow$ | 5-mIoU $\uparrow$ | NoC90 $\downarrow$ | NoC95 $\downarrow$
RITM-HRNet32 400 [37] | COCO+LVIS | 46 | 77.72 | 10.01 | 14.58 | 89.75 | 5.34 | 11.45
FocalClick-SegF-B3-S2 256 [8] | COCO+LVIS | 24 | 84.63 | 8.12 | 12.63 | 90.82 | 5.17 | 11.42
FocalClick-SegF-B3-S2 384 [8] | COCO+LVIS | 26 | 85.45 | 7.03 | 10.74 | 91.22 | 4.90 | 10.40
SimpleClick-ViT-B 448 [56] | COCO+LVIS | 31 | 85.11 | 7.47 | 12.39 | 90.73 | 5.06 | 10.37
InterFormer-ViT-B 1024 [31] | COCO+LVIS | 13 | 82.62 | 7.17 | 10.77 | 87.79 | 5.45 | 11.88
SAM-ViT-B 1024 [36] | SA-1B | 12 | 86.16 | 7.46 | 12.42 | 90.95 | 5.14 | 10.74
MobileSAM-ViT-T 1024 [104] | SA-1B | 8 | 81.98 | 8.70 | 13.83 | 89.18 | 5.83 | 12.74
EfficientSAM-ViT-T 1024 [91] | ImageNet+SA-1B | 6 | 77.90 | 10.11 | 14.60 | 85.26 | 7.37 | 14.28
EfficientSAM-ViT-T 2048 [91] | ImageNet+SA-1B | 32 | 74.20 | 9.47 | 13.13 | 84.10 | 8.00 | 14.37
EfficientSAM-ViT-S 1024 [91] | ImageNet+SA-1B | 7 | 79.01 | 8.84 | 13.18 | 87.55 | 6.37 | 12.26
EfficientSAM-ViT-S 2048 [91] | ImageNet+SA-1B | 60 | 74.91 | 8.27 | 11.97 | 85.17 | 6.86 | 12.49
SegNext (SA$\times$1) ViT-B 1024 [55] | COCO+LVIS | 42 | 85.41 | 7.47 | 11.94 | 90.13 | 5.46 | 13.31
SegNext (SA$\times$2) ViT-B 1024 [55] | COCO+LVIS | 58 | 85.71 | 7.18 | 11.52 | 89.85 | 5.34 | 12.80
SAM-ViT-H 1024 (Teacher) [36] | SA-1B | 29 | 87.21 | 6.85 | 11.57 | 90.82 | 5.20 | 10.04
HRSAM-ViT-B 1024 (Ours) |  †COCO+LVIS ${}_{\textbf{w/o labs}}$ | 10 | 85.86 | 7.66 | 12.54 | 89.54 | 5.48 | 11.61
HRSAM-ViT-B 2048 (Ours) |  †COCO+LVIS ${}_{\textbf{w/o labs}}$ | 18 | 87.31 | 6.76 | 11.17 | 90.36 | 5.39 | 10.15
HRSAM++-ViT-B 1024 (Ours) |  †COCO+LVIS ${}_{\textbf{w/o labs}}$ | 12 | 86.25 | 7.47 | 12.33 | 90.06 | 5.41 | 11.35
HRSAM++-ViT-B 2048 (Ours) |  †COCO+LVIS ${}_{\textbf{w/o labs}}$ | 22 | 88.66 | 6.11 | 10.48 | 90.78 | 5.08 | 9.73
Trained on HQSeg44K (HQ) |  |  |  |  |  |  |  |
HQ-SAM-ViT-B 1024 [63] | SA-1B+HQ | 9 | 89.85 | 6.49 | 10.79 | 91.77 | 5.26 | 10.00
SegNext (SA$\times$2) ViT-B 1024 [55] | COCO+LVIS+HQ | 58 | 91.75 | 5.32 | 9.42 | 91.87 | 4.43 | 10.73
HRSAM-ViT-B 1024 (Ours) |  †COCO+LVIS ${}_{\textbf{w/o labs}}$+HQ | 10 | 91.81 | 5.42 | 9.27 | 91.34 | 4.82 | 11.86
HRSAM-ViT-B 2048 (Ours) |  †COCO+LVIS ${}_{\textbf{w/o labs}}$+HQ | 18 | 93.34 | 4.37 | 7.86 | 92.63 | 4.22 | 8.83
HRSAM++-ViT-B 1024 (Ours) |  †COCO+LVIS ${}_{\textbf{w/o labs}}$+HQ | 12 | 91.84 | 5.32 | 9.18 | 91.25 | 5.02 | 11.64
HRSAM++-ViT-B 2048 (Ours) |  †COCO+LVIS ${}_{\textbf{w/o labs}}$+HQ | 22 | 93.32 | 4.20 | 7.79 | 92.73 | 4.12 | 8.72

Table 1: Quantitative results. $\dagger$ indicates the model is trained by
distillation without using the samples’ labels. Scaling inputs to $2048^{2}$
enables SAM-distilled HRSAM models to surpass their teacher model SAM-ViT-Huge
in both performance and speed. The further HQ-finetuned HRSAM++ achieves
state-of-the-art performance on the HQSeg44K and DAVIS datasets, demonstrating
superior segmentation performance and faster SPC latency compared to previous
SOTA SegNext.

###  4.2 Main Results

As shown in Table 1, when the HRSAM and HRSAM++ are trained by SAM-
distillation on $1024^{2}$ inputs, scaling the testing input to $2048^{2}$
allows them to surpass their teacher SAM-ViT-Huge in performance while
retaining a significant speed advantage. These results highlight the
scalability of our HRSAM models with respect to input size. Moreover, the
further HQ-finetuned HRSAM++ achieves state-of-the-art performance on the
high-precision segmentation datasets HQSeg44K and DAVIS, while maintaining a
faster SPC latency compared to the previous state-of-the-art SegNext.

###  4.3 Scale Analysis

To analyze the scalability of HRSAM with increased input sizes, we evaluate
both HRSAM and HRSAM++ using input sizes of $1024^{2}$, $2048^{2}$, $3072^{2}$
and $4096^{2}$. The performance of the models are tested on the HQSeg44K
dataset, focusing on NoC@90 and NoC@95 metrics. As shown in Table 2, both
HRSAM and HRSAM++ scale from the $1024^{2}$ training resolution to $2048^{2}$
testing size effectively. However, for larger input sizes, such as $3072^{2}$
and $4096^{2}$, HRSAM’s performance deteriorates. Even when the images larger
than $2048^{2}$ are resized to $3072^{2}$ or $4096^{2}$, the performance is
inferior to resizing to $2048^{2}$. This decline can be attributed to two main
reasons. First, resizing to larger images requires a more robust resizing
strategy. The linear interpolation used in this study cannot adequately
preserve image details. Second, the models are trained on $1024^{2}$ inputs,
but handling inputs larger than $3072^{2}$ involves processing image token
sequences more than $9\times$ longer, exceeding the models’ length
extrapolation capabilities. On the other hand, HRSAM++ shows that for larger
images (greater than $2048^{2}$), resizing to $3072^{2}$ can further improve
performance, confirming the effectiveness of HRSAM++’s multi-scale strategy.

Model | Size (0, 1024] | Size (1024, 2048] | Size (2048, 3072] | Size (3072, 4096]
---|---|---|---|---
NoC90 $\downarrow$ |  NoC95 $\downarrow$ |  NoC90 $\downarrow$ |  NoC95 $\downarrow$ |  NoC90 $\downarrow$ |  NoC95 $\downarrow$ |  NoC90 $\downarrow$ |  NoC95 $\downarrow$
SAM-ViT-H 1024 [36] | 5.23 | 10.27 | 7.94 | 13.00 | 7.67 | 11.87 | 7.79 | 12.45
HRSAM-ViT-B 1024 | 5.80 | 10.99 | 8.60 | 13.82 | 8.19 | 12.69 | 9.47 | 13.25
HRSAM-ViT-B 2048 | 5.67 | 10.16 | 7.87 | 12.82 | 7.08 | 11.00 | 7.28 | 11.56
HRSAM-ViT-B 3072 | 6.91 | 11.24 | 8.76 | 13.96 | 7.75 | 11.59 | 7.76 | 12.24
HRSAM-ViT-B 4096 | 8.54 | 13.25 | 10.69 | 15.69 | 8.99 | 13.20 | 8.76 | 13.56
HRSAM++-ViT-B 1024 | 5.58 | 10.78 | 8.53 | 13.65 | 7.97 | 12.30 | 9.18 | 13.19
HRSAM++-ViT-B 2048 | 4.95 | 9.42 | 7.18 | 12.15 | 6.46 | 10.33 | 6.65 | 11.14
HRSAM++-ViT-B 3072 | 5.10 | 9.87 | 7.28 | 12.41 | 6.27 | 9.98 | 6.49 | 10.74
HRSAM++-ViT-B 4096 | 5.57 | 10.82 | 8.28 | 13.20 | 6.36 | 10.68 | 6.79 | 11.33

Table 2: Input size scalability analysis. The metrics NoC@90 and NoC@95 are
measured over the HQSeg44K dataset, with differently sized inputs.

Global | Plain | Van | Pad | SS | MS | 1024$\times$1024 img | 2048$\times$2048 img
---|---|---|---|---|---|---|---
Win | Swin | Swin | CScan | CScan | Lat (ms) $\downarrow$ | NoC90 $\downarrow$ | NoC95 $\downarrow$ | Lat (ms) $\downarrow$ | NoC90 $\downarrow$ | NoC95 $\downarrow$
✓ |  |  |  |  |  | 104 | 7.47 | 12.30 | 1338 | 6.89 | 11.61
✓ | ✓ |  |  |  |  | 61 | 7.56 | 12.39 | 528 | 6.97 | 11.57
| ✓ | ✓ |  |  |  | 37 | 7.68 | 12.71 | 122 | 6.82 | 11.44
| ✓ |  | ✓ |  |  | 44 | 7.68 | 12.69 | 136 | 6.81 | 11.36
| ✓ | ✓ |  | ✓ |  | 52 | 7.59 | 12.59 | 178 | 6.87 | 11.34
| ✓ |  | ✓ | ✓ |  | 58 | 7.66 | 12.54 | 192 | 6.76 | 11.17
| ✓ |  | ✓ | ✓ | ✓ | 98 | 7.47 | 12.33 | 277 | 6.11 | 10.48

Table 3: Ablation study of HRSAM’s modules. The modules include Global
attention (Global), Plain Window attention (Plain Win), Vanilla Shifted Window
attention (Van Swin), Padding Shifted Window attention (Pad Swin), Single-
Scale Cycle-scan (SS CScan), and Multi-Scale Cycle-scan (MS CScan). Latency
and performance metrics (NoC90 and NoC95) are evaluated under input sizes of
$1024^{2}$ and $2048^{2}$ on HQSeg44K.

###  4.4 Ablation Study

In the ablation study, we train HRSAM and its various ablation variants using
the same distillation protocol in the main experiments. We analyze the impact
of different modules, including Global attention (Global), Plain Window
attention (Plain Win), Vanilla Shifted Window attention (Van Swin), Padding
Shifted Window attention (Padding Swin), Single-Scale Cycle-scan (SS CScan)
and Multi-Scale Cycle-scan (MS CScan). We compare the latency and performance
of these variant backbones under $1024^{2}$ and $2048^{2}$ input sizes. As
shown in Table 3, the Global attention achieves better performance with
$1024^{2}$ inputs due to its strong representation capability, which allows it
to effectively learn from the teacher model. However, when scaled to
$2048^{2}$ inputs, its performance deteriorated significantly compared to
other HRSAM variants, with extremely high latency. Conversely, while some
HRSAM variants show only moderate performance over $1024^{2}$ inputs, they
generally perform well when scaled to $2048^{2}$. Notably, HRSAM++ with the
Multi-Scale Cycle-scan module demonstrates significant improvements over
$2048^{2}$ inputs, highlighting the effectiveness of our proposed HRSAM and
HRSAM++ in enhancing model scalability to input image sizes.

Figure 5: Qualitative comparison of our HRSAM++ and previous SOTA SegNext.

###  4.5 Qualitative Results

In Figure 5, we present a qualitative comparison between our HRSAM++ model and
the previous SOTA SegNext [55], using a set of challenging samples with thin
edges. We visualize the segmentation results across four different clicks,
clearly demonstrating the superior performance of HRSAM++ in high-precision
segmentation tasks.

##  5 Limitation

HRSAM cannot scale effectively to excessively large images, as shown in our
multi-scale analysis. Over-enlarging input images or processing particularly
large images at extremely high resolutions does not improve performance. This
presents a challenge in determining the optimal input image size, which HRSAM
currently cannot address. Future work will focus on developing adaptive
methods to set input sizes more effectively.

##  6 Conclusion

In this paper, we propose HRSAM to address the limitations of SAM in handling
high-resolution images critical for high-precision segmentation tasks. HRSAM
leverages Flash Attention to resolve the fundamental issue of memory
insufficiency on large input images. Then, HRSAM incorporates PSCWin attention
ensuring computational consistency during both training and testing. The
enhanced version, HRSAM++ additionally adopts a multi-scale strategy. HRSAM
models set new benchmarks in high-precision interactive segmentation.
Furthermore, our findings reveal a cost-effective training strategy for
building high-performance visual foundation models through visual length
extrapolation.

## References

  * [1] Xformers - flash attention.  https://facebookresearch.github.io/xformers/components/ops.html#xformers.ops.fmha.attn_bias.BlockDiagonalMask.  Accessed: 2024-05-18.
  * [2] Deconstructing LLM-Flash Attention: A Comprehensive Guide Starting from Softmax.  https://zhuanlan.zhihu.com/p/663932651, 2023.  Accessed: May 21, 2024.
  * [3] David Acuna, Huan Ling, Amlan Kar, and Sanja Fidler.  Efficient interactive annotation of segmentation datasets with polygon-rnn++.  In 2018 IEEE Conference on Computer Vision and Pattern Recognition, CVPR 2018, Salt Lake City, UT, USA, June 18-22, 2018, pages 859–868. IEEE Computer Society, 2018.
  * [4] Lucas Beyer, Pavel Izmailov, Alexander Kolesnikov, Mathilde Caron, Simon Kornblith, Xiaohua Zhai, Matthias Minderer, Michael Tschannen, Ibrahim Alabdulmohsin, and Filip Pavetic.  Flexivit: One model for all patch sizes.  In CVPR, 2023.
  * [5] Han Cai, Chuang Gan, Tianzhe Wang, Zhekai Zhang, and Song Han.  Once for all: Train one network and specialize it for efficient deployment.  ICLR, Apr 2020.
  * [6] Hongruixuan Chen, Jian Song, Chengxi Han, Junshi Xia, and Naoto Yokoya.  Changemamba: Remote sensing change detection with spatio-temporal state space model, 2024.
  * [7] Keyan Chen, Bowen Chen, Chenyang Liu, Wenyuan Li, Zhengxia Zou, and Zhenwei Shi.  Rsmamba: Remote sensing image classification with state space model, 2024\.
  * [8] Xi Chen, Zhiyan Zhao, Yilei Zhang, Manni Duan, Donglian Qi, and Hengshuang Zhao.  Focalclick: towards practical interactive image segmentation.  pages 1300–1309, 2022.
  * [9] Junlong Cheng, Jin Ye, Zhongying Deng, Jianpin Chen, Tianbin Li, Haoyu Wang, Yanzhou Su, Ziyan Huang, Jilong Chen, Lei Jiang, et al.  Sam-med2d.  arXiv preprint arXiv:2308.16184, 2023.
  * [10] Ta-Chung Chi, Ting-Han Fan, Peter J Ramadge, and Alexander Rudnicky.  Kerple: Kernelized relative positional embedding for length extrapolation.  NeurIPS, 35, 2022.
  * [11] Xiangxiang Chu, Zhi Tian, Yuqing Wang, Bo Zhang, Haibing Ren, Xiaolin Wei, Huaxia Xia, and Chunhua Shen.  Twins: Revisiting the design of spatial attention in vision transformers.  Neural Information Processing Systems,Neural Information Processing Systems, Dec 2021.
  * [12] Tri Dao.  FlashAttention-2: Faster attention with better parallelism and work partitioning.  2023\.
  * [13] Tri Dao, Daniel Y. Fu, Stefano Ermon, Atri Rudra, and Christopher Ré.  FlashAttention: Fast and memory-efficient exact attention with IO-awareness.  In Advances in Neural Information Processing Systems, 2022.
  * [14] Mostafa Dehghani, Basil Mustafa, Josip Djolonga, Jonathan Heek, Matthias Minderer, Mathilde Caron, Andreas Steiner, Joan Puigcerver, Robert Geirhos, Ibrahim Alabdulmohsin, et al.  Patch n’pack: Navit, a vision transformer for any aspect ratio and resolution.  arXiv preprint arXiv:2307.06304, 2023.
  * [15] 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 2022 IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), Jun 2022.
  * [16] Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszkoreit, and Neil Houlsby.  An image is worth 16x16 words: Transformers for image recognition at scale.  In 9th International Conference on Learning Representations, ICLR 2021, Virtual Event, Austria, May 3-7, 2021. OpenReview.net, 2021.
  * [17] Yuxin Fang, Quan Sun, Xinggang Wang, Tiejun Huang, Xinlong Wang, and Yue Cao.  Eva-02: A visual representation for neon genesis.  arXiv preprint arXiv:2303.11331, 2023.
  * [18] Yuxin Fang, Wen Wang, Binhui Xie, Quan Sun, Ledell Wu, Xinggang Wang, Tiejun Huang, Xinlong Wang, and Yue Cao.  Eva: Exploring the limits of masked visual representation learning at scale.  In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 19358–19369, 2023.
  * [19] Linjie Fu, Xia Li, Xiuding Cai, Yingkai Wang, Xueyao Wang, Yali Shen, and Yu Yao.  Md-dose: A diffusion model based on the mamba for radiotherapy dose prediction, 2024.
  * [20] Chengyue Gong, Dilin Wang, Meng Li, Xinlei Chen, Zhicheng Yan, Yuandong Tian, Vikas Chandra, et al.  Nasvit: Neural architecture search for efficient vision transformers with gradient conflict aware supernet training.  In ICLR, 2021.
  * [21] Albert Gu and Tri Dao.  Mamba: Linear-time sequence modeling with selective state spaces.  arXiv preprint arXiv:2312.00752, 2023.
  * [22] Jiaqi Gu, Hyoukjun Kwon, Dilin Wang, Wei Ye, Meng Li, Yu-Hsin Chen, Liangzhen Lai, Vikas Chandra, and David Z Pan.  Multi-scale high-resolution vision transformer for semantic segmentation.  In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 12094–12103, 2022.
  * [23] Qipeng Guo, Xipeng Qiu, Pengfei Liu, Yunfan Shao, Xiangyang Xue, and Zheng Zhang.  Star-transformer.  In Proceedings of the 2019 Conference of the North, Jan 2019.
  * [24] Agrim Gupta, Piotr Dollár, and Ross B. Girshick.  LVIS: A dataset for large vocabulary instance segmentation.  In IEEE Conference on Computer Vision and Pattern Recognition, CVPR 2019, Long Beach, CA, USA, June 16-20, 2019, pages 5356–5364. Computer Vision Foundation / IEEE, 2019.
  * [25] Dongchen Han, Xuran Pan, Yizeng Han, Shiji Song, and Gao Huang.  Flatten transformer: Vision transformer using focused linear attention.  In Proceedings of the IEEE/CVF International Conference on Computer Vision, pages 5961–5971, 2023.
  * [26] Kaiming He, Xinlei Chen, Saining Xie, Yanghao Li, Piotr Dollar, and Ross Girshick.  Masked autoencoders are scalable vision learners.  In CVPR, Jun 2022.
  * [27] Kaiming He, Georgia Gkioxari, Piotr Dollár, and Ross Girshick.  Mask r-cnn.  In Proceedings of the IEEE international conference on computer vision, pages 2961–2969, 2017.
  * [28] Jonathan Ho, Nal Kalchbrenner, Dirk Weissenborn, and Tim Salimans.  Axial attention in multidimensional transformers.  Cornell University - arXiv,Cornell University - arXiv, Sep 2019\.
  * [29] Vincent Tao Hu, Stefan Andreas Baumann, Ming Gui, Olga Grebenkova, Pingchuan Ma, Johannes Fischer, and Björn Ommer.  Zigma: A dit-style zigzag mamba diffusion model, 2024.
  * [30] Tao Huang, Xiaohuan Pei, Shan You, Fei Wang, Chen Qian, and Chang Xu.  Localmamba: Visual state space model with windowed selective scan, 2024\.
  * [31] You Huang, Hao Yang, Ke Sun, Shengchuan Zhang, Liujuan Cao, Guannan Jiang, and Rongrong Ji.  Interformer: Real-time interactive image segmentation.  In Proceedings of the IEEE/CVF International Conference on Computer Vision, pages 22301–22311, 2023.
  * [32] Zilong Huang, Xinggang Wang, Yunchao Wei, Lichao Huang, Humphrey Shi, Wenyu Liu, and Thomas S. Huang.  Ccnet: Criss-cross attention for semantic segmentation.  IEEE Transactions on Pattern Analysis and Machine Intelligence, page 1–1, Jan 2020.
  * [33] Won-Dong Jang and Chang-Su Kim.  Interactive image segmentation via backpropagating refinement scheme.  In IEEE Conference on Computer Vision and Pattern Recognition, CVPR 2019, Long Beach, CA, USA, June 16-20, 2019, pages 5297–5306. Computer Vision Foundation / IEEE, 2019.
  * [34] Lei Ke, Mingqiao Ye, Martin Danelljan, Yifan Liu, Yu-Wing Tai, Chi-Keung Tang, and Fisher Yu.  Segment anything in high quality.  In NeurIPS, 2023.
  * [35] Salman Khan, Muzammal Naseer, Munawar Hayat, Syed Waqas Zamir, Fahad Shahbaz Khan, and Mubarak Shah.  Transformers in vision: A survey.  ACM computing surveys (CSUR), 54(10s):1–41, 2022.
  * [36] Alexander Kirillov, Eric Mintun, Nikhila Ravi, Hanzi Mao, Chloe Rolland, Laura Gustafson, Tete Xiao, Spencer Whitehead, Alexander C. Berg, Wan-Yen Lo, Piotr Dollár, and Ross Girshick.  Segment anything.  arXiv:2304.02643, 2023.
  * [37] Anton Konushin Konstantin Sofiiuk, Ilia A. Petrov.  Reviving iterative training with mask guidance for interactive segmentation.  arXiv: Computer Vision and Pattern Recognition, 2021.
  * [38] Aditya Kusupati, Gantavya Bhatt, Aniket Rege, Matthew Wallingford, Aditya Sinha, Vivek Ramanujan, William Howard-Snyder, Kaifeng Chen, Sham Kakade, Prateek Jain, and Ali Farhadi.  Matryoshka representations for adaptive deployment.  arXiv preprint arXiv:2205.13147, May 2022.
  * [39] Xin Lai, Zhuotao Tian, Yukang Chen, Yanwei Li, Yuhui Yuan, Shu Liu, and Jiaya Jia.  Lisa: Reasoning segmentation via large language model.  arXiv preprint arXiv:2308.00692, 2023.
  * [40] Kenton Lee, Mandar Joshi, Iulia Raluca Turc, Hexiang Hu, Fangyu Liu, Julian Martin Eisenschlos, Urvashi Khandelwal, Peter Shaw, Ming-Wei Chang, and Kristina Toutanova.  Pix2struct: Screenshot parsing as pretraining for visual language understanding.  In ICML, 2023.
  * [41] Vincent Leroy, Jerome Revaud, Thomas Lucas, and Philippe Weinzaepfel.  Win-win: Training high-resolution vision transformers from two windows.  arXiv preprint arXiv:2310.00632, 2023.
  * [42] Kehan Li, Yian Zhao, Zhennan Wang, Zesen Cheng, Peng Jin, Xiangyang Ji, Li Yuan, Chang Liu, and Jie Chen.  Multi-granularity interaction simulation for unsupervised interactive segmentation.  In Proceedings of the IEEE/CVF International Conference on Computer Vision, pages 666–676, 2023.
  * [43] Kun Li, George Vosselman, and Michael Ying Yang.  Interactive image segmentation with cross-modality vision transformers.  In Proceedings of the IEEE/CVF International Conference on Computer Vision, pages 762–772, 2023.
  * [44] Shufan Li, Harkanwar Singh, and Aditya Grover.  Mamba-nd: Selective state space modeling for multi-dimensional data, 2024\.
  * [45] Wenrui Li, Xiaopeng Hong, and Xiaopeng Fan.  Spikemba: Multi-modal spiking saliency mamba for temporal video grounding, 2024.
  * [46] Xiangtai Li, Haobo Yuan, Wei Li, Henghui Ding, Size Wu, Wenwei Zhang, Yining Li, Kai Chen, and Chen Change Loy.  Omg-seg: Is one model good enough for all segmentation?  In CVPR, 2024.
  * [47] Yanghao Li, Hanzi Mao, Ross Girshick, and Kaiming He.  Exploring plain vision transformer backbones for object detection.
  * [48] Zhuwen Li, Qifeng Chen, and Vladlen Koltun.  Interactive image segmentation with latent diversity.  In 2018 IEEE Conference on Computer Vision and Pattern Recognition, CVPR 2018, Salt Lake City, UT, USA, June 18-22, 2018, pages 577–585. IEEE Computer Society, 2018.
  * [49] Mingbao Lin, Mengzhao Chen, Yuxin Zhang, Ke Li, Yunhang Shen, Chunhua Shen, and Rongrong Ji.  Super vision transformer.  IJCV, May 2022.
  * [50] 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.  Lecture Notes in Computer Science, 2014.
  * [51] Zheng Lin, Zheng-Peng Duan, Zhao Zhang, Chun-Le Guo, and Ming-Ming Cheng.  Focuscut: Diving into a focus view in interactive segmentation.  pages 2637–2646, 2022.
  * [52] Zheng Lin, Zhao Zhang, Lin-Zhuo Chen, Ming-Ming Cheng, and Shao-Ping Lu.  Interactive image segmentation with first click attention.  In 2020 IEEE/CVF Conference on Computer Vision and Pattern Recognition, CVPR 2020, Seattle, WA, USA, June 13-19, 2020, pages 13336–13345. IEEE, 2020.
  * [53] Chenyang Liu, Keyan Chen, Bowen Chen, Haotian Zhang, Zhengxia Zou, and Zhenwei Shi.  Rscama: Remote sensing image change captioning with state space model, 2024.
  * [54] Jiarun Liu, Hao Yang, Hong-Yu Zhou, Yan Xi, Lequan Yu, Yizhou Yu, Yong Liang, Guangming Shi, Shaoting Zhang, Hairong Zheng, and Shanshan Wang.  Swin-umamba: Mamba-based unet with imagenet-based pretraining, 2024.
  * [55] Qin Liu, Jaemin Cho, Mohit Bansal, and Marc Niethammer.  Rethinking interactive image segmentation with low latency, high quality, and diverse prompts.  arXiv preprint arXiv:2404.00741, 2024.
  * [56] Qin Liu, Zhenlin Xu, Gedas Bertasius, and Marc Niethammer.  Simpleclick: Interactive image segmentation with simple vision transformers.  arXiv preprint arXiv:2210.11006, 2022.
  * [57] Qin Liu, Meng Zheng, Benjamin Planche, Srikrishna Karanam, Terrence Chen, Marc Niethammer, and Ziyan Wu.  Pseudoclick: Interactive image segmentation with click imitation.  pages 728–745, 2022.
  * [58] 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.
  * [59] 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 Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pages 12009–12019, 2022.
  * [60] Ze Liu, Yutong Lin, Yue Cao, Han Hu, Yixuan Wei, Zheng Zhang, Stephen Lin, and Baining Guo.  Swin transformer: Hierarchical vision transformer using shifted windows.  In Proceedings of the IEEE/CVF international conference on computer vision, pages 10012–10022, 2021.
  * [61] Shaocong Long, Qianyu Zhou, Xiangtai Li, Xuequan Lu, Chenhao Ying, Yuan Luo, Lizhuang Ma, and Shuicheng Yan.  Dgmamba: Domain generalization via generalized state space model, 2024\.
  * [62] Jun Ma, Feifei Li, and Bo Wang.  U-mamba: Enhancing long-range dependency for biomedical image segmentation, 2024.
  * [63] Jun Ma and Bo Wang.  Segment anything in medical images.  arXiv preprint arXiv:2304.12306, 2023.
  * [64] Kevis-Kokitsi Maninis, Sergi Caelles, Jordi Pont-Tuset, and Luc Van Gool.  Deep extreme cut: From extreme points to object segmentation.  In 2018 IEEE Conference on Computer Vision and Pattern Recognition, CVPR 2018, Salt Lake City, UT, USA, June 18-22, 2018, pages 616–625. IEEE Computer Society, 2018.
  * [65] Maciej A Mazurowski, Haoyu Dong, Hanxue Gu, Jichen Yang, Nicholas Konz, and Yixin Zhang.  Segment anything model for medical image analysis: an experimental study.  Medical Image Analysis, 89:102918, 2023.
  * [66] Maxim Milakov and Natalia Gimelshein.  Online normalizer calculation for softmax.  arXiv preprint arXiv:1805.02867, 2018.
  * [67] Badri N. Patro and Vijay S. Agneeswaran.  Simba: Simplified mamba-based architecture for vision and multivariate time series, 2024.
  * [68] Xiaohuan Pei, Tao Huang, and Chang Xu.  Efficientvmamba: Atrous selective scan for light weight visual mamba, 2024\.
  * [69] Federico Perazzi, Jordi Pont-Tuset, Brian McWilliams, Luc Van Gool, Markus H. Gross, and Alexander Sorkine-Hornung.  A benchmark dataset and evaluation methodology for video object segmentation.  In 2016 IEEE Conference on Computer Vision and Pattern Recognition, CVPR 2016, Las Vegas, NV, USA, June 27-30, 2016, pages 724–732. IEEE Computer Society, 2016.
  * [70] Ofir Press, Noah A Smith, Mike Lewis, et al.  Train short, test long: Attention with linear biases enables input length extrapolation.  arXiv preprint arXiv:2108.12409, 2021.
  * [71] Yanyuan Qiao, Zheng Yu, Longteng Guo, Sihan Chen, Zijia Zhao, Mingzhen Sun, Qi Wu, and Jing Liu.  Vl-mamba: Exploring state space models for multimodal learning, 2024.
  * [72] Amit Kumar Rana, Sabarinath Mahadevan, Alexander Hermans, and Bastian Leibe.  Dynamite: Dynamic query bootstrapping for multi-object interactive segmentation transformer.  In Proceedings of the IEEE/CVF International Conference on Computer Vision, pages 1043–1052, 2023.
  * [73] Pengzhen Ren, Changlin Li, Guangrun Wang, Yun Xiao, and QingDuXiaodanLiangXiaojun Chang.  Beyond fixation: Dynamic window visual transformer.
  * [74] Jiacheng Ruan and Suncheng Xiang.  Vm-unet: Vision mamba unet for medical image segmentation, 2024.
  * [75] Qiuhong Shen, Xuanyu Yi, Zike Wu, Pan Zhou, Hanwang Zhang, Shuicheng Yan, and Xinchao Wang.  Gamba: Marry gaussian splatting with mamba for single view 3d reconstruction, 2024.
  * [76] Baifeng Shi, Ziyang Wu, Maolin Mao, Xin Wang, and Trevor Darrell.  When do we not need larger vision models?  arXiv preprint arXiv:2403.13043, 2024.
  * [77] Konstantin Sofiiuk, Ilia A. Petrov, Olga Barinova, and Anton Konushin.  F-BRS: rethinking backpropagating refinement for interactive segmentation.  In 2020 IEEE/CVF Conference on Computer Vision and Pattern Recognition, CVPR 2020, Seattle, WA, USA, June 13-19, 2020, pages 8620–8629. IEEE, 2020.
  * [78] Yiran Song, Qianyu Zhou, Xiangtai Li, Deng-Ping Fan, Xuequan Lu, and Lizhuang Ma.  Ba-sam: Scalable bias-mode attention mask for segment anything model.  arXiv preprint arXiv:2401.02317, 2024.
  * [79] Robin Strudel, Ricardo Garcia, Ivan Laptev, and Cordelia Schmid.  Segmenter: Transformer for semantic segmentation.  pages 7262–7272, 2021.
  * [80] Jianlin Su, Murtadha Ahmed, Yu Lu, Shengfeng Pan, Wen Bo, and Yunfeng Liu.  Roformer: Enhanced transformer with rotary position embedding.  Neurocomputing, 568:127063, 2024.
  * [81] Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, et al.  Llama: Open and efficient foundation language models.  arXiv preprint arXiv:2302.13971, 2023.
  * [82] Ashish Vaswani, Prajit Ramachandran, Aravind Srinivas, Niki Parmar, Blake Hechtman, and Jonathon Shlens.  Scaling local self-attention for parameter efficient visual backbones.  In 2021 IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), Jun 2021.
  * [83] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin.  Attention is all you need.  In Isabelle Guyon, Ulrike von Luxburg, Samy Bengio, Hanna M. Wallach, Rob Fergus, S. V. N. Vishwanathan, and Roman Garnett, editors, Advances in Neural Information Processing Systems 30: Annual Conference on Neural Information Processing Systems 2017, December 4-9, 2017, Long Beach, CA, USA, pages 5998–6008, 2017.
  * [84] Zifu Wan, Yuhao Wang, Silong Yong, Pingping Zhang, Simon Stepputtis, Katia Sycara, and Yaqi Xie.  Sigma: Siamese mamba network for multi-modal semantic segmentation, 2024\.
  * [85] Jiaqi Wang, Zhengliang Liu, Lin Zhao, Zihao Wu, Chong Ma, Sigang Yu, Haixing Dai, Qiushi Yang, Yiheng Liu, Songyao Zhang, et al.  Review of large vision models and visual prompt engineering.  arXiv preprint arXiv:2307.00855, 2023.
  * [86] 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.  In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 14408–14419, 2023.
  * [87] Wenxiao Wang, Lu Yao, Long Chen, Deng Cai, Xiaofei He, and Wei Liu.  Crossformer: A versatile vision transformer based on cross-scale attention.  arXiv: Computer Vision and Pattern Recognition,arXiv: Computer Vision and Pattern Recognition, Jul 2021.
  * [88] Ziyang Wang, Jian-Qing Zheng, Yichi Zhang, Ge Cui, and Lei Li.  Mamba-unet: Unet-like pure visual mamba for medical image segmentation, 2024.
  * [89] Junde Wu, Rao Fu, Huihui Fang, Yuanpei Liu, Zhaowei Wang, Yanwu Xu, Yueming Jin, and Tal Arbel.  Medical sam adapter: Adapting segment anything model for medical image segmentation.  arXiv preprint arXiv:2304.12620, 2023.
  * [90] Enze Xie, Wenhai Wang, Zhiding Yu, Animashree Anandkumar, JoseM. Alvarez, and Ping Luo.  Segformer: Simple and efficient design for semantic segmentation with transformers.  Cornell University - arXiv,Cornell University - arXiv, Dec 2021\.
  * [91] Yunyang Xiong, Bala Varadarajan, Lemeng Wu, Xiaoyu Xiang, Fanyi Xiao, Chenchen Zhu, Xiaoliang Dai, Dilin Wang, Fei Sun, Forrest Iandola, et al.  Efficientsam: Leveraged masked image pretraining for efficient segment anything.  arXiv preprint arXiv:2312.00863, 2023.
  * [92] Long Xu, Yongquan Chen, Rui Huang, Feng Wu, and Shiwu Lai.  Structured click control in transformer-based interactive segmentation.  arXiv preprint arXiv:2405.04009, 2024.
  * [93] Ning Xu, Brian L. Price, Scott Cohen, Jimei Yang, and Thomas S. Huang.  Deep interactive object selection.  In 2016 IEEE Conference on Computer Vision and Pattern Recognition, CVPR 2016, Las Vegas, NV, USA, June 27-30, 2016, pages 373–381. IEEE Computer Society, 2016.
  * [94] Shilin Xu, Haobo Yuan, Qingyu Shi, Lu Qi, Jingbo Wang, Yibo Yang, Yining Li, Kai Chen, Yunhai Tong, Bernard Ghanem, et al.  Rap-sam: Towards real-time all-purpose segment anything.  arXiv preprint arXiv:2401.10228, 2024.
  * [95] Cilin Yan, Haochen Wang, Jie Liu, Xiaolong Jiang, Yao Hu, Xu Tang, Guoliang Kang, and Efstratios Gavves.  Piclick: Picking the desired mask in click-based interactive segmentation.  arXiv preprint arXiv:2304.11609, 2023.
  * [96] Chenhongyi Yang, Zehui Chen, Miguel Espinosa, Linus Ericsson, Zhenyu Wang, Jiaming Liu, and Elliot J. Crowley.  Plainmamba: Improving non-hierarchical mamba in visual recognition, 2024\.
  * [97] Jianwei Yang, Chunyuan Li, Pengchuan Zhang, Xiyang Dai, Bin Xiao, Lu Yuan, Jianfeng Gao, MicrosoftResearchAt Redmond, Microsoft Cloud, and + Ai.  Focal self-attention for local-global interactions in vision transformers.
  * [98] Yuhuan Yang, Chaofan Ma, Jiangchao Yao, Zhun Zhong, Ya Zhang, and Yanfeng Wang.  Remamber: Referring image segmentation with mamba twister, 2024.
  * [99] Zi Ye, Tianxiang Chen, Fangyijie Wang, Hanwei Zhang, Guanxi Li, and Lijun Zhang.  P-mamba: Marrying perona malik diffusion with mamba for efficient pediatric echocardiographic left ventricular segmentation, 2024.
  * [100] Hongxu Yin, Arash Vahdat, Jose Alvarez, Arun Mallya, Jan Kautz, and Pavlo Molchanov.  Adavit: Adaptive tokens for efficient vision transformer.  arXiv preprint arXiv:2112.07658, 2021.
  * [101] Jiahui Yu, Pengchong Jin, Hanxiao Liu, Gabriel Bender, Pieter-Jan Kindermans, Mingxing Tan, Thomas Huang, Xiaodan Song, Ruoming Pang, and Quoc Le.  Bignas: Scaling up neural architecture search with big single-stage models.  In ECCV, Jan 2020.
  * [102] Tao Yu, Runseng Feng, Ruoyu Feng, Jinming Liu, Xin Jin, Wenjun Zeng, and Zhibo Chen.  Inpaint anything: Segment anything meets image inpainting.  arXiv preprint arXiv:2304.06790, 2023.
  * [103] Yuhui Yuan, Rao Fu, Lang Huang, Weihong Lin, Chao Zhang, Xilin Chen, and Jingdong Wang.  Hrformer: High-resolution vision transformer for dense predict.  Advances in Neural Information Processing Systems, 34:7281–7293, 2021.
  * [104] Chaoning Zhang, Dongshen Han, Yu Qiao, Jung Uk Kim, Sung-Ho Bae, Seungkyu Lee, and Choong Seon Hong.  Faster segment anything: Towards lightweight sam for mobile applications.  arXiv preprint arXiv:2306.14289, 2023.
  * [105] Shiyin Zhang, Jun Hao Liew, Yunchao Wei, Shikui Wei, and Yao Zhao.  Interactive object segmentation with inside-outside guidance.  In 2020 IEEE/CVF Conference on Computer Vision and Pattern Recognition, CVPR 2020, Seattle, WA, USA, June 13-19, 2020, pages 12231–12241. IEEE, 2020.
  * [106] Tiezheng Zhang, Xiaoxi Chen, Chongyu Qu, Alan Yuille, and Zongwei Zhou.  Leveraging ai predicted and expert revised annotations in interactive segmentation: Continual tuning or full training?  arXiv preprint arXiv:2402.19423, 2024.
  * [107] Han Zhao, Min Zhang, Wei Zhao, Pengxiang Ding, Siteng Huang, and Donglin Wang.  Cobra: Extending mamba to multi-modal large language model for efficient inference, 2024.
  * [108] Minghao Zhou, Hong Wang, Qian Zhao, Yuexiang Li, Yawen Huang, Deyu Meng, and Yefeng Zheng.  Interactive segmentation as gaussion process classification.  In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 19488–19497, 2023.
  * [109] 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.
  * [110] Shen Zhuoran, Zhang Mingyuan, Zhao Haiyu, Yi Shuai, and Li Hongsheng.  Efficient attention: Attention with linear complexities.  In 2021 IEEE Winter Conference on Applications of Computer Vision (WACV), Jan 2021.
  * [111] Xueyan Zou, Jianwei Yang, Hao Zhang, Feng Li, Linjie Li, Jianfeng Wang, Lijuan Wang, Jianfeng Gao, and Yong Jae Lee.  Segment everything everywhere all at once.  Advances in Neural Information Processing Systems, 36, 2024.

##  Appendix A Flash Attention

The following introduction to Flash Attention [13, 12] is based on [2].

###  A.1 Motivation

Flash Attention is motivated by the need to optimize memory bandwidth and
computation speed in GPUs, particularly leveraging High Bandwidth Memory (HBM)
and Static Random-Access Memory (SRAM). In GPUs like the A100-40GB, HBM
provides 40GB of memory with a bandwidth of 1.5TB/s. However, actual GPU
computations occur in SRAM, which has a much higher bandwidth of 19TB/s but is
limited to only 20MB of usable memory. The substantial speed difference (SRAM
being 12.67 times faster than HBM) and the limited SRAM capacity create a
bottleneck when memory exchanges between SRAM and HBM are required.

The primary goal of Flash Attention is to minimize memory exchanges between
SRAM and HBM during attention computation. Traditional attention mechanisms
require seven exchanges between these memory types, creating significant
overhead. Flash Attention aims to perform the entire attention computation
within the 20MB SRAM, thus avoiding these costly memory swaps and achieving up
to a $7.6\times$ speedup for $N\times N$ matrix computations.

###  A.2 SRAM Cache Overflow Example

To illustrate the efficiency of Flash Attention, consider a scenario where
SRAM can handle 10000 data points at a time. For a single vector, let’s say Q
and K are both of size $[100,100]$. Loading both vectors would require
$2\times 100\times 100=20000$ data points, which exceeds the 10000 data points
capacity of SRAM, necessitating write/read operations to HBM.

To address this, Flash Attention splits the vectors into smaller chunks:

  * •

Split $Q[100,100]$ into $\\{Q1[50,100],Q2[50,100]\\}$

  * •

Split $K[100,100]$ into $\\{K1[50,100],K2[50,100]\\}$

This way, the computation of the attention score

$\text{score}_{ij}=Q_{i}K_{j}^{T}$

can be completed entirely within SRAM, avoiding the need to swap data with
HBM. By processing these smaller chunks within the high-speed SRAM, Flash
Attention efficiently computes the attention scores block-wise, thereby
maintaining high computational speed and minimizing memory bottlenecks.

In summary, Flash Attention leverages the speed of SRAM by minimizing memory
exchanges and efficiently handling data within the limited capacity of SRAM,
resulting in significant performance improvements in attention computation.
Next, we detail the algorithm of Flash Attention, beginning with the
reformulation of attention’s softmax operator.

###  A.3 Softmax

Numerically stable softmax. Before discussing further details of Flash
Attention, we reformulate the conventional softmax operator into a numerically
stable version as follows:

$\text{softmax}(x_{i})=\dfrac{\exp(x_{i}-\max_{k}x_{k})}{\sum_{j=1}^{n}\exp(x_{j}-\max_{k}x_{k})}.$

Online softmax. The softmax has a computationally equivalent online version
[66], implemented with three-pass as follows:

$\begin{split}m_{0}&\leftarrow-\infty\\\ d_{0}&\leftarrow 0.0\\\ \text{for
}i\leftarrow 1,N\\\ m_{i}&\leftarrow\max(m_{i-1},x_{i})\\\ \text{for
}i\leftarrow 1,N\\\ d_{i}&\leftarrow d_{i-1}+\exp(x_{i}-m_{N})\\\ \text{for
}i\leftarrow 1,N\\\ a_{i}&\leftarrow\exp(x_{i}-m_{N})/d_{N}\\\ \end{split}$

We then reformulate $d_{i}$ with the iterative version, $d_{i}^{\prime}$,
defined as

$\begin{split}d_{i}^{\prime}&=\sum_{j=1}^{i}\exp(x_{j}-m_{i})\\\
&=\left(\sum_{j=1}^{i-1}\exp(x_{j}-m_{i})+\exp(x_{i}-m_{i})\right)\\\
&=\exp(m_{i-1}-m_{i})\left(\sum_{j=1}^{i-1}\exp(x_{j}-m_{i-1})\right)+\exp(x_{i}-m_{i})\\\
&=\exp(m_{i-1}-m_{i})d_{i-1}^{\prime}+\exp(x_{i}-m_{i}),\\\ \end{split}$

which implies the following two-pass implementation of online softmax:

$\begin{split}m_{0}&\leftarrow-\infty\\\ d_{0}&\leftarrow 0.0\\\ \text{for
}i\leftarrow 1,N\\\ m_{i}&\leftarrow\max(m_{i-1},x_{i})\\\
d_{i}^{\prime}&\leftarrow\exp(m_{i-1}-m_{i})d_{i-1}^{\prime}+\exp(x_{i}-m_{i})\\\
\text{for }i\leftarrow 1,N\\\ a_{i}&\leftarrow\exp(x_{i}-m_{N})/d_{N}\\\
\end{split}$

###  A.4 Flash Attention by Online Softmax

Two-pass Attention. Based on the two-pass online softmax, we derive a two-pass
implementation of attention, with the pre-computated
$Q,K,V\in\mathbb{R}^{N\times d}$, as follows:

$\begin{split}\text{for }i\leftarrow 1,N\\\ x_{i}&\leftarrow
Q[k,:]K^{T}[:,i]\\\ m_{i}&\leftarrow\max(m_{i-1},x_{i})\\\
d^{\prime}_{i}&\leftarrow\exp(m_{i-1}-m_{i})d^{\prime}_{i-1}+\exp(x_{i}-m_{i})\\\
\text{for }i\leftarrow 1,N\\\
a_{i}&\leftarrow\frac{\exp(x_{i}-m_{N})}{d^{\prime}_{N}}\\\ o_{i}&\leftarrow
o_{i-1}+a_{i}V[i,:]\end{split}$

Vanilla Flash Attention. We reformulate the output $o_{i}$ as

$\begin{split}o_{i}&=o_{i-1}+\frac{\exp(x_{i}-m_{N})}{d^{\prime}_{N}}V[i,:]\\\
\implies
o_{i}&=\sum_{j=1}^{i}\frac{\exp(x_{j}-m_{{\color[rgb]{1,0,0}N}})}{d^{\prime}_{{\color[rgb]{1,0,0}N}}}V[j,:],\\\
\end{split}$

which inspires the following definition

$\begin{split}o^{\prime}_{i}&=\sum_{j=1}^{i}\frac{\exp(x_{j}-m_{{\color[rgb]{1,0,0}i}})}{d^{\prime}_{{\color[rgb]{1,0,0}i}}}V[j,:],\\\
\end{split}$

where $o_{i}=o^{\prime}_{i}$ holds if $i=N$. Then, we rewrite $o^{\prime}_{i}$
in an iterative form:

$\begin{split}o^{\prime}_{i}&=\sum_{j=1}^{i}\frac{\exp(x_{j}-m_{i})}{d^{\prime}_{i}}V[j,:]\\\
&=\left(\sum_{j=1}^{i-1}\frac{\exp(x_{j}-m_{i})}{d^{\prime}_{i}}V[j,:]\right)+\frac{\exp(x_{i}-m_{i})}{d^{\prime}_{i}}V[i,:]\\\
&=\left(\sum_{j=1}^{i-1}\frac{\exp(x_{j}-m_{i-1})}{d^{\prime}_{i-1}}\frac{\exp(x_{j}-m_{i})}{\exp(x_{j}-m_{i-1})}\frac{d^{\prime}_{i-1}}{d^{\prime}_{i}}V[j,:]\right)+\frac{\exp(x_{i}-m_{i})}{d^{\prime}_{i}}V[i,:]\\\
&=\left(\sum_{j=1}^{i-1}\frac{\exp(x_{j}-m_{i-1})}{d^{\prime}_{i-1}}V[j,:]\right)\frac{\exp(x_{j}-m_{i})}{\exp(x_{j}-m_{i-1})}\frac{d^{\prime}_{i-1}}{d^{\prime}_{i}}+\frac{\exp(x_{i}-m_{i})}{d^{\prime}_{i}}V[i,:]\\\
&=\left(\sum_{j=1}^{i-1}\frac{\exp(x_{j}-m_{i-1})}{d^{\prime}_{i-1}}V[j,:]\right)\frac{d^{\prime}_{i-1}\exp(m_{i-1}-m_{i})}{d^{\prime}_{i}}+\frac{\exp(x_{i}-m_{i})}{d^{\prime}_{i}}V[i,:]\\\
&=o^{\prime}_{i-1}\frac{d^{\prime}_{i-1}\exp(m_{i-1}-m_{i})}{d^{\prime}_{i}}+\frac{\exp(x_{i}-m_{i})}{d^{\prime}_{i}}V[i,:].\\\
\end{split}$

Thus, we derive the one-pass iterative form of Flash Attention:

$\begin{split}\text{for }i\leftarrow 1,N\\\ x_{i}&\leftarrow
Q[k,:]K^{T}[:,i]\\\ m_{i}&\leftarrow\max(m_{i-1},x_{i})\\\
d^{\prime}_{i}&\leftarrow
d^{\prime}_{i-1}\exp(m_{i-1}-m_{i})+\exp(x_{i}-m_{i})\\\
o^{\prime}_{i}&\leftarrow
o^{\prime}_{i-1}\frac{d^{\prime}_{i-1}\exp(m_{i-1}-m_{i})}{d^{\prime}_{i}}+\frac{\exp(x_{i}-m_{i})}{d^{\prime}_{i}}V[i,:]\\\
\text{end}\\\ O[k,:]&\leftarrow o^{\prime}_{N}\end{split}$

###  A.5 Tiling Flash-Attention

When dealing with multiple data tiles, the Flash Attention algorithm can be
further refined as follows:

$\begin{split}\text{for }i\leftarrow 1\\#\text{tiles}\\\ x_{i}&\leftarrow
Q[k,:]K^{T}[:,(i-1)b:ib]\\\
m_{i}^{\text{local}}&\leftarrow\max_{j=1}^{b}(x_{i}[j])\\\
m_{i}&\leftarrow\max(m_{i-1},m_{i}^{\text{local}})\\\
d^{\prime}_{i}&\leftarrow
d^{\prime}_{i-1}\exp(m_{i-1}-m_{i})+\sum_{j=1}^{b}\exp(x_{i}[j]-m_{i})\\\
o^{\prime}_{i}&\leftarrow
o^{\prime}_{i-1}\frac{d^{\prime}_{i-1}}{d^{\prime}_{i}}\exp(m_{i-1}-m_{i})+\sum_{j=1}^{b}\frac{\exp(x_{i}[j]-m_{i})}{d^{\prime}_{i}}V[j+(i-1)b,:]\\\
\text{end}\\\ O[k,:]&\leftarrow o^{\prime}_{N}\end{split}$

###  A.6 Memory Efficiency of Flash Attention

Based on the above derivation, it can be observed that Flash Attention
requires only a fixed amount of memory to execute efficient tiling. Thus,
Flash Attention algorithm eliminates the need to store the entire attention
map, thereby reducing the quadratic space complexity to linear, primarily for
storing computational results.

###  A.7 Limitation

As outlined in the derivation process, the transition from online softmax to
two-pass softmax and one-pass attention computation heavily relies on
sophisticated mathematical derivations and strict formalism. Although Flash
Attention optimizes the conventional attention computation process, it is
incompatible with certain attention variants, such as the relative positional
encoding used in SAM [36]. This type of encoding requires adding an additional
attention bias to the attention map, i.e.,
$\text{softmax}(QK^{\top}+\text{bias})V$, which cannot be integrated into the
one-pass attention algorithm of Flash Attention. As a result, SAM cannot
directly utilize Flash Attention.

##  Appendix B Block-diagonal Attention

Figure 6: Illustration of block-diagonal attention [1].

Flash Attention supports block-diagonal mode [1] for computing attention. This
mode introduces a skip mechanism to the previously discussed one-pass
attention algorithm, which skips the ignored parts of the attention map.
Consequently, multiple attention computations can be merged using block-
diagonal attention.

As shown in Figure 6, multiple token sequences of different lengths can be
combined into a single token sequence. By constructing an appropriate mask to
block attention between different token sequences, the CUDA implementation of
Flash Attention can skip these irrelevant attentions. This takes full
advantage of the highly parallel nature of GPUs to process multiple token
sequences simultaneously.

Moreover, this strategy can be used for vanilla shifted window attention with
differently sized window attentions and is particularly effective for the
proposed padding shifted window. By providing the diagonal mask dividing the
image tokens into equal-sized windows, we can directly calculate attention
without reshaping the image token sequence to construct windows. This method
optimizes the computation process, ensuring efficient and simultaneous
processing of multiple token sequences. Additionally, the multi-scale image
tokens are similarly concatenated and processed by the block-diagonal
attention with the attention computed within the same-scale tokens.

##  Appendix C Indexing Operator

For instance, the multi-scale image token sequence is formulated as the
concatenated embeddings
$\mathbf{F}_{c}=\text{Concat}(\mathbf{F},\mathbf{F}_{s})\in\mathbb{R}^{(L+L_{s})\times
C}$, with $L=\frac{HW}{16^{2}}$ and $L_{s}=\frac{H_{s}W_{s}}{16^{2}}$. Assume
the window size is $S$, which is divisible by $\frac{H}{16}$ and
$\frac{W}{16}$ for simplicity. We then detail the plain and padding shifted
window operations over the multi-scale token sequences using the indexing
operator.

###  C.1 Indexing Multi-scale Plain Window.

The single-scale plain window is created by reshaping $\mathbf{F}$ to obtain
the windowed embeddings $\mathbf{F}_{w}\in\mathbb{R}^{\frac{L}{S^{2}}\times
S^{2}\times C}$, where each window
$\mathbf{F}_{w}[i]\in\mathbb{R}^{S^{2}\times C}$ is processed using
conventional attention. For the multi-scale plain window, $\mathbf{F}_{c}$ is
rearranged so that tokens within the same window are contiguous in the
sequence, as shown in Figure 7 (a). The block-diagonal attention is then
applied to the rearranged sequence with a mask that restricts attention to the
local window.

###  C.2 Indexing Multi-scale Padding Shifted Window.

Similarly, the multi-scale padding shifted window requires a rearranged token
sequence. The difference lies in the additional padding token, as illustrated
in Figure 7 (b). This padding token is appended at the end of the token
sequence and replicated according to the indices. The same block-diagonal
attention is then applied.

Figure 7: Illustration of the indexing operator. The red boxes in the sequence
indicate the regions where self-attention is computed.

##  Appendix D Experimental Details

###  D.1 Datasets.

We adopt COCO [50], LVIS [24] and HQSeg44K [34] datasets for our two-stage
training process. Specifically, the COCO dataset includes 118K training images
(1.2M instances). LVIS shares the same images as COCO with superior quality of
segmentation annotations. HQSeg44K comprises six existing image datasets,
including DIS, ThinObject-5K, FSS-1000, ECSSD, MSRA10K and DUT-OMRON,
containing extremely high-precision segmentation annotations.

In the first stage, we use only the images from COCO and LVIS without the
labels in our distillation training, which differs from previous interactive
segmentation methods [8, 57, 56, 31]. In the second stage, we finetune the
weights pretrained at the first stage using the HQSeg44K training dataset. For
testing, we evaluate exclusively on the HQSeg44K validation [34] and DAVIS
[69] datasets, both of which have high-precision annotations, following recent
advances in high-precision interactive segmentation [55].

###  D.2 Implementation details.

We utilize ViT-Base [16, 27] as the backbone for HRSAM, retaining SAM’s prompt
encoder and decoder. The ViT-Base consists of 12 blocks, equally divided into
four stages. The cycle-scan module is introduced before the last block of each
stage. For the enhanced version, HRSAM++, an additional multi-scale cycle-scan
module is inserted at the end of each stage. Similar to the Swin Transformer
[60], we alternate between plain window attention and the proposed padding
shifted window attention. Each block has an embedding dimension of 768, and
the output dimension of each stage is 256, aligning with SAM decoder’s inputs.
Besides, we initialize HRSAM-ViT-Base by the MAE-pretrained weights [27],
following the previous works [56, 31, 55]. In addition, the typical FFN has a
hidden layer dimension of $768\times 4$, with a uniform window size set to 16.
The hidden state of Mamba is configured as 32, and the number of attention’s
heads is set to 12.

###  D.3 Training strategy.

In the training process, we follow a two-stage approach. In the first stage,
images from the COCO and LVIS training datasets are used. Initially, the
teacher model SAM-ViT-Huge’s backbone is employed to perform inference on
these images, and the resulting image embeddings are stored. During HRSAM
training, these precomputed embeddings are loaded, and the training images
undergo simple color augmentation without other augmentation methods. HRSAM
computes its own image embeddings on these augmented images, and the
difference between HRSAM and the teacher embeddings is minimized using vanilla
MSE as the loss function. In the second stage, we further finetune the models
using HQSeg44K training samples with segmentation annotations. We employ the
commonly used normalized focal loss for this finetuning stage.

Both stages share most of the training settings. The optimizer is AdamW with
weight decay of 0.05 and betas of $(0.9,0.999)$. We use a learning rate
scheduler with a linear warm-up starting from $1\times 10^{-6}$ and a
polynomial decay starting from $1\times 10^{-4}$ over 160k iterations for the
first stage and 40k iterations for the second stage. Additionally, during the
second stage, we apply additional random flip augmentation to the images.
Training is conducted on 4 NVIDIA RTX 3090 GPUs with a batch size of 1 per
GPU.

###  D.4 Evaluation.

We evaluate HRSAMs against previous methods, i.e., RITM [37], FocalClick [8],
SimpleClick [56], InterFormer [31], SegNext [55], SAM [36], MobileSAM [104]
and EfficientSAM [91]. For non-SAM models, we use their default input sizes.
For SAM series models, we use both $1024^{2}$ and $2048^{2}$ inputs if the
model supports $2048^{2}$ inputs. In testing, click simulation places clicks
at the centers of erroneously predicted regions, with the binary label of each
click determined by the maximum distance to the boundaries of false negative
and false positive regions, aligning with the previous methods [8, 56, 31]. We
assess all models’ segmentation performance and speed. Segmentation
performance is measured using 5-mIoU and NoC metrics. The 5-mIoU represents
the average IoU after the fifth click. The NoC metric indicates the average
minimum clicks required to reach a specified IoU. We focus on NoC@90 and
NoC@95 within 20 clicks. Speed is quantified as Seconds Per Click (SPC) on
GPUs, measuring the average inference latency per click, averaged over 20
clicks. For SAM series models, this includes the preprocessing time by the
backbone.