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

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 M2× M2 score matrix per window, incurring O(M4) 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×4 (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(M4) to O(M2). This enables higher token density and larger receptive fields without inflating memory overhead. Training memory is flat across window sizes: at a 32×32 window, FLASHSWIN-T requires only 12.4,GB, unchanged from 8×8, 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× larger gains in boundary quality than in mIoU.