Fast Vision Mamba: Pooling Spatial Dimensions for Accelerated Processing
FastVim halves Mamba scan depth through alternating spatial pooling, achieving up to 72.5% faster inference on 2048² images.
Key Findings
Methodology
FastVim extends Vision Mamba (Vim) with parameter-free mean pooling after 1D convolution. A two-dimensional h×w token grid is compressed along rows or columns before Mamba’s selective scan, reducing the sequence to h or w tokens. The scan still uses input-dependent B, C, and Δ; outputs are repeated to the original grid. Transposition between blocks alternates pooling directions, enabling information exchange across both spatial axes.
Key Results
- On supervised ImageNet-1K training, FastVim-T/S/B achieves 75.4%, 81.1%, and 82.6% Top-1 accuracy, compared with Vim’s 76.1%, 80.5%, and LayerNorm-stabilized 82.6%. FastVim-S uses 4.43G FLOPs versus 5.9G for Vim-S while preserving or improving accuracy.
- On H100 GPUs with batch size 128 and float32 inference, FastVim reaches up to 72.5% end-to-end speedup at 2048×2048 resolution. Its SSM component is 324% faster; FastVim-T reduces FLOPs by 35% at 224 resolution and 38.5% at 2048.
- MAE-pretrained FastMaskVim obtains 83.0%/84.9%/86.1% for Base/Large/Huge and 86.7% at 448×448, a Mamba-based high-resolution result. On JUMP-CP, FastChannelVim reaches 73.6% at patch/8; ChannelVim reaches 83.0% and FastChannelVim 83.1%.
Significance
The work addresses a practical bottleneck in high-resolution vision: although Mamba avoids Transformer attention’s quadratic token interactions, its scan length still grows with the image area. FastVim shows that structured, sparse interaction can preserve recognition accuracy while substantially improving throughput. The approach is relevant to pathology, microscopy, remote sensing, and industrial imaging, where resolution and channel count are costly. It also clarifies that Mamba’s recurrent contextualization can exploit compression differently from Transformer attention.
Technical Contribution
The principal contribution is to reduce the scan length from L=h² to h for a square token grid, changing parallel depth from log(h²)=2log(h) to log(h). Algorithm 1 specifies reshape, pooling, selective parameter projection, discretization, SSM scan, repetition, and residual processing. FastMaskVim adapts transpose and pooling to sparse masked grids; FastChannelVim handles per-channel tokenization, hierarchical channel sampling, channel ordering, and alternative scan paths.
Novelty
Unlike token pruning or sparse attention in ViTs, FastVim keeps the number of inter-block tokens unchanged and introduces no learnable pooling parameters. It compresses only inside the Mamba contextualization module and reconstructs the grid afterward. The essential novelty is alternating spatial axes: pooling only columns or only rows permanently blocks one interaction direction, whereas alternation restores two-dimensional communication across depth.
Limitations
- Repeating a pooled representation broadcasts one contextual vector across several positions, potentially erasing fine spatial distinctions. The method therefore relies on multiple alternating blocks; shallow networks or boundary-sensitive segmentation may suffer.
- Pooling does not reduce every block component: MLP and gating layers retain original token scaling, while pooling and repetition add memory-access overhead. Evidence is concentrated on Vim and its extensions rather than a comprehensive evaluation across all Mamba backbones.
Future Work
Future research should test learned or content-adaptive pooling, residual detail paths, and hardware-fused pool-repeat kernels. Important directions include irregular non-square grids, dense segmentation, pathology multiple-instance learning, and broader VMamba or MambaVision integration. More systematic studies of mixed precision, scan ordering, channel compression, and accuracy-versus-resolution trade-offs are also needed.
AI Executive Summary
High-resolution vision is exposing a weakness in otherwise efficient sequence models. Vision Transformers suffer quadratic self-attention, while Vision Mamba (Vim) replaces it with linear-complexity selective state-space scanning. Yet an image still produces an area-proportional number of tokens, so the parallel scan depth remains log(L). At 2048×2048 resolution, this recurrent contextualization becomes a major throughput bottleneck.
FastVim attacks that bottleneck with a deliberately simple operation. After a 1D convolution, it mean-pools the two-dimensional token grid along one spatial axis, runs Mamba’s input-dependent selective scan on the compressed sequence, and repeats the result back to the original grid. Blocks transpose the grid and alternate axes, allowing row-wise and column-wise information to propagate over depth. For an h×h grid, scan depth falls from log₂(h²)=2log₂h to log₂h.
The accuracy cost is small. On ImageNet-1K, FastVim-T/S/B reaches 75.4%, 81.1%, and 82.6% Top-1, with FastVim-T using 35–38.5% fewer FLOPs than Vim-T across 224–2048 resolutions. On H100 GPUs, 2048² inference is up to 72.5% faster overall, while the SSM layer alone gains 324%. MAE-pretrained FastMaskVim reaches 86.7% at 448×448, and FastChannelVim supports multichannel cell imaging. The work suggests that Mamba can tolerate structured sparsity, although detail loss, memory overhead, and transfer beyond Vim remain open challenges.
Deep Analysis
Background
S4 introduced structured state-space parameterizations for efficient sequence modeling; Mamba added input-dependent B, C, and Δ through selective scan. Vim adapted Mamba to images with forward and backward SSM branches and causal 1D convolutions. Compared with ViT, Vim avoids quadratic attention, but its token count L=H×W/P² still grows with image area. Parallel scan reduces recurrence to log(L), not to a resolution-independent cost.
Core Problem
The target is the recurrent contextualization bottleneck rather than parameter count. On an h×h grid, Vim scans h² tokens and requires log(h²)=2log(h) parallel stages. Simply pruning tokens can damage spatial reasoning, while parallel scan still incurs synchronization, memory traffic, and long-sequence costs at high resolution.
Innovation
FastVim compresses one spatial dimension before scanning and restores the grid afterward. Alternating transpose and pooling directions prevents permanent loss of row or column interactions. FastMaskVim uses sparse indexing for masked or irregular grids and normalizes row sums by the nominal width. FastChannelVim extends the design to ChannelViT-style per-channel tokens, while preserving channel order and hierarchical channel sampling.
Methodology
- �� Reshape x∈R^(B×L×D) into (B,h,w,D).
- �� Apply normalization, expansion, and Conv1D; mean-pool one axis to obtain (B,h,1,D).
- �� Generate input-dependent B and C with Linear_N; generate Δ with softplus(Parameter+sΔ(x)); discretize A and B using zero-order hold.
- �� Run forward and backward Mamba selective scans only on compressed tokens.
- �� Repeat outputs along the pooled axis, then apply the residual/skip path and LayerNorm.
- �� Transpose the grid between blocks to alternate axes.
- �� For FastMaskVim, use sparse indexing and mask-aware aggregation; for FastChannelVim, scan spatial-first or channel-first sequences with sorted hierarchical channel sampling.
Experiments
ImageNet-1K contains 1.28M training and 50K validation images. Supervised models use 300 epochs, AdamW, batch size 1,024, learning rate 1e-3, five-epoch warmup, cosine decay, and weight decay 0.05. MAE uses 1,600 epochs and a 0.75 masking ratio. Efficiency is measured on H100 GPUs with batch size 128 and float32, comparing FastVim, Vim, and ViT. JUMP-CP uses the BR00116991 plate: 127K/45K/45K train/validation/test images, eight channels, and 160-way perturbation prediction.
Results
FastVim-T/S/B reaches 75.4/81.1/82.6% Top-1 with 1.17/4.43/17.23G FLOPs. FastVim-T saves 35% FLOPs at 224 and 38.5% at 2048 relative to Vim-T. At 2048, SSM speedup is 324% and overall speedup 72.5%. FastMaskVim-Huge reaches 86.1% at 224 and 86.7% at 448. On JUMP-CP patch/8, ChannelVim and FastChannelVim score 83.0% and 83.1%, far above ChannelViT’s 74.8%.
Applications
Potential uses include high-resolution classification, detection, segmentation, MAE or DINOv2 representation learning, pathology multiple-instance learning, microscopy perturbation prediction, and multispectral satellite analysis. Regular grids can use FastVim; masked or sparse layouts require FastMaskVim; complementary imaging channels benefit from FastChannelVim. Deployment requires an efficient Mamba scan implementation and careful validation of local-boundary accuracy.
Limitations & Outlook
Repeating pooled outputs may blur object boundaries and fine textures, making performance dependent on depth and alternating propagation. MLP and gating operations retain original token scaling, and pool-repeat introduces runtime overhead despite low FLOP cost. The strategy reportedly fails in ViT, indicating architecture-specific behavior. The paper does not provide broad evidence across every Mamba backbone, task, resolution, or accelerator, so real-world gains may vary.
Plain Language Accessible to non-experts
Imagine a huge warehouse whose floor is divided into thousands of squares. Each square has a worker who sees local information, and workers must pass a summary through the warehouse. In the original system, every worker joins a very long communication line. Parallel organization helps, but the line still becomes enormous when the warehouse grows.
FastVim first asks workers in each column to hold a short meeting and produce one combined report. Only these reports travel through the long communication process. The result is then sent back to all squares in that column. In the next round, the warehouse is turned sideways and workers hold row meetings instead. Repeating this alternation lets information travel in both directions without forcing every square into the same giant line.
The floor is not permanently shrunk; compression happens only during communication. That is why accuracy remains close to the original model. But averaging can hide differences: two squares containing a cat and a fire extinguisher might produce an unhelpful shared summary. FastVim is therefore especially attractive when broad context matters, while detail-sensitive tasks may need an additional local-detail path. The reported 72.5% speedup at 2048² shows the practical value of this trade-off.
ELI14 Explained like you're 14
Picture a video game with an enormous map. Every tiny tile has a little scout that reports what it sees. On a small map, scouts can share information quickly. On a giant 2048×2048 map, there are so many tiles that communication becomes the laggy part—even if the game engine is clever.
FastVim gives nearby scouts a quick team meeting. Instead of sending every tiny report, the team sends one short summary. After the summary travels, it is copied back to the tiles. Then the next game round changes the meeting direction: first columns, then rows. Why? If scouts only talk vertically, they miss important horizontal clues. Alternating directions lets information spread across the whole map.
This is different from deleting tiles. The original map positions remain, but communication is temporarily compressed. On ImageNet-1K, FastVim-Base scores 82.6% Top-1 accuracy, and at 2048×2048 it can make the whole model 72.5% faster. That is like keeping the same giant game world while making its messaging system much quicker!
There is a catch: a summary can miss tiny details, such as a hidden enemy or a sharp object boundary. Also, not every part of the model becomes cheaper. FastVim is strongest for large images, long-range context, and many imaging channels. Future versions could keep a special “detail camera” alongside the fast summaries.
Glossary
State-Space Model (状态空间模型)
A sequence model that stores history in a hidden state and updates it over time. Technically, it uses equations such as h_t=Ah_{t-1}+Bx_t and y_t=Ch_t+Dx_t.
FastVim uses Mamba’s state-space module for visual contextualization.
Selective Scan (选择性扫描)
A Mamba mechanism whose state parameters depend on the current input, allowing content-dependent information retention. It uses token-dependent B, C, and Δ.
FastVim applies selective scan after spatial compression.
Parallel Scan (并行扫描)
A parallel algorithm that reorganizes recurrence from L sequential operations into approximately log(L) stages. It improves hardware utilization but does not remove sequence-length dependence.
FastVim changes the depth from 2log(h) to log(h) on square grids.
Mean Pooling (平均池化)
A parameter-free operation that averages representations along one dimension. It creates a shorter summary sequence at the cost of some local detail.
It is FastVim’s default pre-scan compression operation.
Masked Autoencoder, MAE (掩码自编码器)
A self-supervised method that masks image patches and trains a model to reconstruct missing content. It learns visual representations without class labels.
FastMaskVim uses MAE with a 0.75 masking ratio on ImageNet-1K.
Per-Channel Tokenization (逐通道token化)
A representation in which image channels produce separate tokens rather than being merged immediately. It preserves complementary information in microscopy or satellite data.
FastChannelVim applies it to eight-channel JUMP-CP images.
Open Questions Unanswered questions from this research
- 1 How much fine detail is lost by pooling and repetition, and which object sizes or textures are most vulnerable, remains unclear. Controlled frequency and boundary analyses are needed.
- 2 Why the strategy works in Vim but fails in ViT is unresolved. A theory should compare recurrent state aggregation with permutation-invariant self-attention.
- 3 End-to-end gains depend on memory traffic, kernel fusion, precision, and accelerator design. More reproducible hardware benchmarks are required.
Applications
Immediate Applications
High-resolution medical imaging
Pathology slides and microscopy images can use FastMaskVim to process masked, sparse, or gapped token layouts with shorter scans. Teams need sparse indexing, task-specific fine-tuning, and boundary-accuracy checks before deployment.
Multichannel cell profiling
Drug-screening platforms can use FastChannelVim to preserve information across eight or more imaging channels. JUMP-CP results show strong perturbation prediction, but channel ordering and hierarchical channel sampling must be implemented consistently.
Long-term Vision
Real-time remote sensing and industrial vision
FastVim could reduce latency and memory for large satellite scenes, inspection images, and robotic perception. Adoption depends on fused pooling kernels, local-detail preservation, and validation across diverse accelerators and dense prediction tasks.
Abstract
State Space Models (SSMs) with selective scan (Mamba) have been adapted into efficient vision models. Mamba, unlike Vision Transformers, achieves linear complexity for token interactions through a recurrent hidden state process. This sequential processing is enhanced by a parallel scan algorithm, which reduces the computational time of recurrent steps from $L$ sequential steps to $log(L)$ parallel steps with respect to the number of input tokens ($L$). In this work, we propose Fast Vision Mamba (FastVim), that further reduces the computational time of the SSM block by reducing the number of recurrent steps in Vision Mamba models while still retaining model performance. By alternately pooling tokens along image dimensions across Mamba blocks, we obtain a 2$\times$ reduction in the number of parallel steps in SSM block. Our model offers up to $72.5\%$ speedup in inference speed compared to baseline Vision Mamba models on high resolution (2048$\times$2048) images. Our experiments demonstrate state-of-the-art performance with dramatically improved throughput in a range of tasks such as image classification, cell perturbation prediction, segmentation, and object detection. Code is made available at https://github.com/insitro/FastVim