FLASHSWIN: Unlocking Large Windows and Dense Tokens in Swin Vision Transformers with Memory Efficient Attention

Tushar Kataria; Gerald Sabin; Ponnuswamy Sadayappan; Shireen Y. Elhabian

cs.CVOct 3, 2026

FLASHSWIN: Unlocking Large Windows and Dense Tokens in Swin Vision Transformers with Memory Efficient Attention

Authors: Tushar Kataria, Gerald Sabin, Ponnuswamy Sadayappan, Shireen Y. Elhabian

Organizations: Scientific Computing and Imaging Institute, University of Utah, Salt Lake City, UT, USA · Kahlert School of Computing, University of Utah, Salt Lake City, UT, USA · RNET Technologies, Dayton, OH, USA

Abstract

High-resolution vision backbones have long been forced to trade away local token density to afford larger receptive fields. Hierarchical Swin transformers impose this compromise because standard windowed attention materializes an $M^2\times M^2$ score matrix per window, incurring $O(M^4)$ memory as windows or token grids grow. Furthermore, Swin adds a learned relative-position bias elementwise to attention scores, requiring full materialization of the score matrix and its gradient. This keeps Swin and SwinV2 trapped in a small-window( $M=8,16$ ), coarse-token regime with patch size $4\times4$ ( $p=4$ ), limiting performance for fine-grained tasks. We introduce FLASHSWIN, which replaces standard windowed attention with a FlashAttention implementation that computes exact softmax attention without materializing the score matrix, reducing per-window memory from $O(M^4)$ to $O(M^2)$ . This enables higher token density and larger receptive fields without inflating memory overhead. Training memory is flat across window sizes: at a $32\times32$ window, FLASHSWIN-T requires only $12.4$ ,GB, unchanged from $8\times8$ , compared to $70/90$ ,GB for SwinV2/V1-T. However, applying FlashAttention directly to Swin creates a trade-off: bypassing the score matrix precludes Swin's additive relative-position bias, forfeiting spatial information in exchange for memory efficiency. FLASHSWIN restores position information as window-local learnable 2D RoPE, making large windows and dense token grids both affordable and accurate. At matched scale, FLASHSWIN-T outperforms Swin variants. With dense tokens and wide windows ( $p=2,M=32$ ), the same Tiny model reaches $84.1\%$ ImageNet-1K, $44.1$ COCO box AP, and $47.28$ ADE20K mIoU---gains of $+1.3$ , $+5.1$ , and $+1.82$ over SwinV2-T at $M=16$ , respectively. At fixed $M=32$ , halving the patch size yields roughly $3\times$ larger gains in boundary quality than in mIoU.

Figures & tables

Figure 1: Token density, architecture, and memory scaling comparison between Swin and FlashSwin . (A) An overlapping stem ( $P=2$ ) creates dense token grids that preserve spatial boundaries better than standard coarse patches ( $P=4$ ). (B) FlashSwin replaces Swin’s materialized score matrix and additive relative bias ( $O(M^{4})$ memory) with Window-Local 2D RoPE and a fused FlashAttention [ 12 ] kernel. (C) Peak training memory vs. window size $M$ on an H200 GPU ( $256\times 256$ resolution, batch size $128$ ). FlashSwin maintains a flat memory footprint as $M$ increases, whereas Swin-V1/V2 grow rapidly and run out of memory at dense settings.

Method	Params	$M$	$p$	Res	GFLOPs	IN-1K $\uparrow$	IN-V2 $\uparrow$
Matched compute
Swin-T	28M	8	4	256	5.9	81.9	70.2
SwinV2-T	28M	8	4	256	5.9	81.9	70.3
FlashSwin -T	28M	8	4	256	5.9	82.3	70.8
Large-window
SwinV2-T	28M	16	4	256	6.6	82.8	71.4

Table 1: ImageNet-1K classification.

M

is window size,

p

patch size; accuracy is top-1 (%) on ImageNet-1K and on the held-out ImageNet-V2 set. Only the first block is compute-matched; later blocks trade GFLOPs for window size and token density at matched parameter count, the exception being FlashBias-T, whose low-rank factors scale with

M^{2}

. Bold : best in block.

Backbone	M/p/Res	IN-1K $\uparrow$	IN-A $\uparrow$	IN-R $\uparrow$	Sketch $\uparrow$	IN-C $\downarrow$
Swin-T	8/4/256	81.9	25.1	41.7	28.0	58.5
SwinV2-T	8/4/256	81.9	24.7	41.3	27.9	58.3
FlashSwin -T	8/4/256	82.3	28.1	43.6	30.4	56.9
SwinV2-T	16/4/256	82.8	31.2	44.6	31.4	54.1
FlashSwin -T	16/4/256	82.9	33.2	45.6	32.9	53.7
FlashBias-T	32/4/256	82.7	28.4	43.0	29.4	54.7

Table 2: Out-of-distribution robustness. Top-1 (%) on ImageNet-A, -R, and -Sketch; mean corruption error (mCE) on ImageNet-C. IN-1K repeats clean accuracy for reference and is excluded from bolding; bold : best in each window block. ImageNet-A and Sketch peak at

M{=}16

while mCE keeps falling; dense tokenization is strongest on all four. FlashBias-T fuses the kernel but keeps a rank-

16

additive bias: matched to FlashSwin -T on clean accuracy, it scores at the

M{=}8

level on natural shift despite

M{=}32

, while remaining competitive on synthetic corruption.

Backbone	M/p/Res	mIoU $\uparrow$	Boundary-IoU $\uparrow$	BF $\uparrow$
Swin-T	8/4/256	43.20	31.88	45.63
SwinV2-T	8/4/256	43.47	32.26	46.17
FlashSwin -T	8/4/256	44.39	33.11	47.36
SwinV2-T	16/4/256	45.46	34.14	48.22
FlashSwin -T	16/4/256	46.03	34.73	48.98
FlashBias-T	32/4/256	45.15	33.77	48.30

Table 3: ADE20K segmentation and boundary quality. Single-scale mIoU, Boundary-IoU, and F1. FlashSwin -T leads across all matched windows. Larger windows improve region and boundary metrics equally, while token density (

p{=}2

) yields

2\text{--}3\times

larger gains on boundaries than mIoU. FlashBias-T, matched to FlashSwin -T on ImageNet, trails it by

1.7

mIoU and by a similar relative margin on both boundary metrics, performing at SwinV2-T’s

M{=}16

level despite

M{=}32

. Bold : best in block.

		Box $\uparrow$			Mask $\uparrow$
Backbone	M/p/Res	AP	AP 50	AP 75	AP	AP 50	AP 75
Swin-T	8/4/256	37.1	58.2	40.1	34.2	54.9	36.2
SwinV2-T	8/4/256	37.3	58.8	40.3	34.4	55.3	36.4
FlashSwin -T	8/4/256	38.5	59.8	41.4	35.3	56.8	37.6
SwinV2-T	16/4/256	39.0	60.3	42.2	35.7	57.0	37.9
FlashSwin -T	16/4/256	39.0	60.2	42.2	36.0	57.3	38.1

Table 4: COCO detection and instance segmentation. FlashSwin -T improves with window size and matches or beats the Swin baselines at each matched window; halving the patch at

M{=}32

gives the largest gain. FlashBias-T fuses the kernel but keeps an additive bias as a learned rank-16 factorization: it ties FlashSwin on ImageNet (

82.65

82.94

) yet loses

3.3

box AP, performing at Swin-T’s

M{=}8

level despite

M{=}32

Backbone	$p$	$M$	PT Res.	IN-1K	IN-V2	mIoU	B-IoU	BF
Swin-L	4	7	$224^{2}$	86.51	76.81	50.13	38.73	54.31
SwinV2-L	4	12	$192^{2}$	86.93	77.48	48.35	37.54	52.72
FlashSwin -L	2	32	$192^{2}$	86.91	78.02	52.58	41.75	57.29

Table 5: Large-scale results. PT Res. denotes IN-22K pretraining resolution; all IN-1K/IN-V2 results use

256^{2}

fine-tuning. FlashSwin -L uses

4\times

more stage-1 tokens than SwinV2-L at the same

192^{2}

pretraining resolution. Published Swin-L/SwinV2-L results at

384^{2}

fine-tuning are not directly comparable. Higher is better for all metrics.

$p$	Model	$M{=}8$	$M{=}16$	$M{=}32$	$M{=}8$	$M{=}16$	$M{=}32$	Params (M)
		Memory (GB) $\downarrow$			Train throughput (img/s) $\uparrow$
4	Swin-T	15.7/6.4	31.0/11.6	90.1/37.7	1173/893	828/612	442/314	28.5
	SwinV2-T	15.0/5.4	29.8/12.0	70.4/34.9	871/670	643/484	371/272	28.4
	FlashBias-T	13.2/5.2	13.2/5.1	13.2/5.2	934/741	924/733	826/657	28.5–29.5
	FlashSwin -T	12.4/4.2	12.4/4.8	12.4/5.0	1058/815	1103/850	1043/808	28.3
2	Swin-T	64.4/23.2	123.3/45.6	OOM/OOM	317/239	210/153	OOM/OOM	28.5

Table S1: Peak training memory and throughput without/with gradient checkpointing, measured at batch size

128

with

256^{2}

inputs on one H200. FlashBias-T uses a trainable rank-

16

bias factorization. Both FlashBias-T and FlashSwin -T remove the memory dependence on window size, while FlashSwin -T uses less memory and higher throughput at

M\!\geq\!16

. FlashBias parameters increase with

M

, whereas FlashSwin -T remains fixed at

28.3

\dagger

denotes a configuration that FlashBias cannot launch at batch

128

due to its CUDA grid-size limit.

Variant	RoPE	Overlap	IN-1K $\uparrow$
Naive Flash	–	–	80.0
+ learnable RoPE	✓	–	81.8
+ overlap stem	✓	✓	82.3
SwinV2-T (ref.)	relative-position bias		81.9

Table S2: Component ablation at

M{=}8

p{=}4

256^{2}

. Dropping Swin’s relative-position bias by applying the fused FlashAttention kernel costs

1.9

points against SwinV2-T.

Largest gains				Largest losses
Class	#im	B-IoU	BF	Class	#im	B-IoU	BF
pot	141	$+7.21$	$+9.79$	mountain	161	$-1.78$	$-1.89$
sconce	108	$+6.64$	$+6.54$	earth	190	$-0.80$	$-0.76$
vase	122	$+5.73$	$+5.97$	railing	100	$-0.71$	$-0.38$
flower	138	$+4.11$	$+4.78$	grass	227	$-0.66$	$-0.29$
bottle	114	$+4.09$	$+5.98$

Table S3: Per-class boundary change from halving the patch size at

M{=}32

, FlashSwin -L

p{=}4\!\rightarrow\!p{=}2

on ADE20K, restricted to the

45

classes appearing in at least

100

validation images.

41/45

improve (mean

+2.30

Boundary-IoU, sd

1.83

Backbone	$w$	$p$	B-IoU	BF-score at $\tau/d$
				0.002	0.004	0.008	0.016	0.032
Swin-T	8	4	31.88	27.53	36.97	45.63	52.84	58.89
SwinV2-T	8	4	32.26	27.99	37.56	46.17	53.19	59.17
FlashSwin -T	8	4	33.11	28.81	38.61	47.36	54.48	60.47
SwinV2-T	16	4	34.14	29.59	39.50	48.22	55.13	60.74
FlashSwin -T	16	4	34.73	30.31	40.36	48.98	55.87	61.63

Table S4: BF-score under varying boundary tolerance

\tau

on ADE20K val . The ranking is unchanged at every tolerance.

		Box				Mask
Backbone	M/p/Res	AP	AP S	AP M	AP L	AP	AP S	AP M	AP L
Swin-T	8/4/256	37.1	16.8	40.9	55.9	34.2	11.9	36.7	56.5
SwinV2-T	8/4/256	37.3	17.1	41.3	55.5	34.4	12.1	37.0	56.5
FlashSwin -T	8/4/256	38.5	18.1	42.9	57.2	35.3	12.8	38.2	57.2
SwinV2-T	16/4/256	39.0	18.7	43.3	57.5	35.7	13.2	38.6	58.7
FlashSwin -T	16/4/256	39.0	19.3	43.3	57.7	36.0	13.5	38.9	58.5

Table S5: COCO detection with size-stratified AP. Setup as in Tab. 4 . Larger windows improve all object scales; the finer patch improves small objects almost exclusively.

Route	Patch	Res	GFLOPs	Mem (GB)	img/s $\uparrow$	IN-1K $\uparrow$
finer patch	2	256	37.4	24.4	287	84.1
higher res	4	512	37.4	23.8	284	84.0

Table S6: Two routes to identical token density. At matched FLOPs (

{\sim}37

G) and parameters (

{\approx}28

M) on a

128\!\times\!128

grid (H200, batch

64

, window

32

), both routes yield equivalent cost and accuracy.

	Top-1 $\uparrow$					mCE $\downarrow$
Backbone	IN-1K	IN-V2	IN-A	IN-R	Sketch	IN-C
Swin-L	86.51	76.81	62.79	65.02	49.24	39.74
SwinV2-L	86.93	77.48	66.39	66.28	50.48	37.70
FlashSwin -L	86.91	78.02	69.71	67.38	49.89	38.37

Table S7: Out-of-distribution generalization at Large scale. All models are evaluated without fine-tuning from their IN-22K

\rightarrow

IN-1K checkpoints, each at its own pretraining window. ImageNet-A and -R are scored over their

200

-class subsets; ImageNet-V2 uses the matched-frequency split; mCE is AlexNet-normalized over

15

corruptions at

5

severities. At essentially identical clean accuracy (

86.91

vs.

86.93

), FlashSwin -L leads on ImageNet-V2, -A and -R.

Figure S1 : Inference throughput versus window size at patch $4$ , $256^{2}$ input, batch $128$ , one H200. FlashSwin -T is flat across the sweep ( $3286\!\rightarrow\!3336$ images/s) while Swin-T falls $3.3\times$ ( $3498\!\rightarrow\!1068$ ) and SwinV2-T $2.8\times$ ( $2771\!\rightarrow\!1002$ ). At $M{=}8$ the fused kernel is slightly slower than the bias-table path, since a small window does not amortise its overhead; FlashSwin overtakes both baselines from $M{=}16$ onward and is $3.1\times$ faster than Swin-T at $M{=}32$ .

Figure S2 : Effective receptive field (ERF) of the final feature at the center location, accumulated over $200$ ImageNet-val images (log scale). FlashSwin -T’s ERF expands smoothly and coherently as the window grows ( $50\%$ -ERF radius $50\!\to\!61\!\to\!66\%$ of the input side for $M\!=\!8,16,32$ ), whereas SwinV2-T ( $M\!=\!8$ ) is markedly more local. Because memory is independent of window size, FlashSwin can realize the large windows that produce these wide receptive fields.

Figure S3 : Component-ablation training curves (ImageNet-1K val top-1, w8/p4/256, ${\sim}5.9$ GFLOPs). Naive memory-efficient attention (orange) plateaus below SwinV2-T (dashed, $81.9$ ); the learnable RoPE (blue) recovers parity and the overlapping stem (green) surpasses it ( $+2.3$ total). Right: zoom on the final $60$ epochs with end values.

Figure S4 : Learned RoPE frequencies depart from the geometric initialization. Inverse frequencies

\theta_{i}=\exp(\omega_{i})

per stage (averaged over the blocks in each stage) for FlashSwin -T, versus the standard geometric RoPE schedule used at initialization (dashed). The learned frequencies deviate markedly—

84\%

move by more than

10\%

—and each stage adopts a distinct profile, particularly in the mid-band indices, indicating that the optimal positional spectrum is stage-dependent rather than fixed.

Explore similar work

Gated-SwinRMT: Unifying Swin Windowed Attention with Retentive Manhattan Decay via Input-Dependent Gating
Apr 7, 2026Dipan Maity, Suman Mondal, Arindam RoyVision Transformer Gated Attention
Student Kolkata, West Bengal, India · Department of Computer Science Yogoda Satsanga Palpara Mahavidyalaya West Bengal, India · Department of Computer Science & Application Prabhat Kumar College Contai West Bengal, India
QFlash: Bridging Quantization and Memory Efficiency in Vision Transformer Attention
Apr 28, 2026Sehyeon Oh, Yongin Kwon, Jemin LeeSoftmax Attention GPU Acceleration
University of Science and Technology, Daejeon, Republic of Korea · Electronics and Telecommunications Research Institute, Daejeon, Republic of Korea · Pusan National University, Busan, Republic of Korea +1
Sliding-window beats linear attention
Aug 28, 2026Alexia Jolicoeur-Martineau, Rhea Sanjay Sukthanker, Pashmina Cameron +1KV Caching LLM Inference Acceleration
Microsoft Applied Sciences Group (ASG) · Independent

Apr 7, 2026J/K move · Enter open · S save

Gated-SwinRMT: Unifying Swin Windowed Attention with Retentive Manhattan Decay via Input-Dependent Gating

Dipan Maity, Suman Mondal, Arindam Roy

Student Kolkata, West Bengal, India · Department of Computer Science Yogoda Satsanga Palpara Mahavidyalaya West Bengal, India · Department of Computer Science & Application Prabhat Kumar College Contai West Bengal, India

Vision Transformer Gated Attention Image Classification Sliding-Window Attention

Open paper AskPDF arXiv

We introduce Gated-SwinRMT, a family of hybrid vision transformers that combine the shifted-window attention of the Swin Transformer with the Manhattan-distance spatial decay of Retentive Networks (RMT), augmented by input-dependent gating. Self-attention is decomposed into consecutive width-wise and height-wise retention passes within each shifted window, where per-head exponential decay masks provide a two-dimensional locality prior without learned positional biases. Two variants are proposed.Gated-SwinRMT-SWAT substitutes softmax with sigmoid activation, implements balanced ALiBi slopes with multiplicative post-activation spatial decay, and gates the value projection via SwiGLU; the Normalized output implicitly suppresses uninformative attention scores. Gated-SwinRMT-Retention retains softmax-normalized retention with an additive log-space decay bias and incorporates an explicit G1 sigmoid gate -- projected from the block input and applied after local context enhancement (LCE) but prior to the output projection~

W_O

-- to alleviate the low-rank

W_V \!\cdot\! W_O

bottleneck and enable input-dependent suppression of attended outputs. We assess both variants on Mini-ImageNet (

224{\times}224

, 100 classes) and CIFAR-10 (

32{\times}32

, 10 classes) under identical training protocols, utilizing a single GPU due to resource limitations. At

{\approx}77

79

,M parameters, Gated-SwinRMT-SWAT achieves

80.22\%

and Gated-SwinRMT-Retention

78.20\%

top-1 test accuracy on Mini-ImageNet, compared with

73.74\%

for the RMT baseline. On CIFAR-10 -- where small feature maps cause the adaptive windowing mechanism to collapse attention to global scope -- the accuracy advantage compresses from

+6.48

,pp to

+0.56

,pp.

Abstract

Explore similar work

Gated-SwinRMT: Unifying Swin Windowed Attention with Retentive Manhattan Decay via Input-Dependent Gating

QFlash: Bridging Quantization and Memory Efficiency in Vision Transformer Attention

Sliding-window beats linear attention