LocalMamba: Visual State Space Model with Windowed Selective Scan

In plain words

Newer image models borrow a sequence model called Mamba. To use it on a picture, the picture's grid of small patches has to be read one at a time, like a line of text. A plain row-by-row reading order puts patches that touch vertically far apart in that line, so the model loses some of the neighbourhood structure.

LocalMamba changes the reading order. It cuts the grid into small square windows and reads each window in full before moving to the next one, so patches that sit next to each other stay close in the sequence. The model still reads the whole image in other orders too, so it keeps the big-picture view. Different layers seem to prefer different reading orders, so the method also picks which orders each layer should use from a set of options, rather than using one order everywhere.

The authors report higher image-classification accuracy than the Vim baseline at about the same compute, and results on object detection and segmentation.

TL;DR

• Vision Mamba models read image patches in a flat order, which separates patches that are neighbours in 2D §1 fig. 1.
• LocalMamba adds a windowed ("local") scan that keeps each small window contiguous, and runs four scan directions per block §4.1.
• A DARTS-style differentiable search selects four of eight candidate scan directions for each layer §4.2.
• Author-reported: LocalVim-T 76.2% top-1 on ImageNet-1K at 1.5G FLOPs vs Vim-Ti 73.1% tab. 2; LocalVMamba-T 82.7% vs VMamba-T 82.2% tab. 2.
• The official code is Apache-2.0 and implements the search and the fixed models, but its channel attention and search normalisation differ from the paper (see Paper vs code).

Problem

Mamba-style selective state space models (SSMs) process a 1-D sequence causally, so an image must first be flattened into a sequence §3.2 §4.1. Vim flattens rows, which pushes vertically adjacent tokens far apart in the sequence, and VMamba adds a second, column-wise scan but still does not keep every local neighbourhood contiguous §4.1 fig. 1. The authors argue that this weakens local 2-D dependencies, and that the best scan order differs across layers §1 §4.2.

The paper assumes that a global view is still needed. Its local scan therefore sits next to global scans instead of replacing them §4.1 (remark).

Method

Overview. Each LocalMamba block scans the token grid in several directions, runs one selective SSM per direction, undoes each scan, reweights channels per branch, and sums the branches §4.1 fig. 3. The directions are fixed per layer after a search over eight candidates §4.2. Two backbones are built: LocalVim, a plain Vim-style model, and LocalVMamba, a hierarchical VMamba-style model §4.3 tab. 1.

Step 1: SSM background. A continuous linear state space is discretised with a zero-order hold, which gives a recurrence over the sequence eq. 1 eq. 2 eq. 3. The recurrence can also be computed as a convolution with a kernel built from powers of the discrete state matrix eq. 4 §3.1. Mamba (selective SSM, S6) makes the input-dependent parameters B, C and Δ depend on the input §3.2.

Step 2: local scan §4.1. Divide the H × W token grid into non-overlapping windows of size w × w, scan each window, then move to the next window. The paper uses w = 2 and w = 7 in its search space §4.2. The flipped variant reads the same sequence from tail to head §4.1.

# Local window scan  [§4.1] [fig. 1c]
# Input:  grid of tokens (B, C, H, W); window size w
# Output: sequence (B, C, L) with L = H * W after padding

def local_scan(grid, window, flip=False):
    # Pad bottom and right so H and W are multiples of w; the paper does
    # not say this, the code pads with zeros
    # [code: classification/lib/models/mamba/local_scan.py#L253-L256]
    grid = pad_zeros_to_multiple(grid, window)
    # Cut the grid into w x w windows, visited row by row  [§4.1]
    windows = split_into_windows(grid, window)  # (B, C, Hg, Wg, w, w)
    # Read the tokens of each window row by row, window after window  [§4.1]
    # [code: classification/lib/models/mamba/local_scan.py#L260]
    sequence = flatten_window_major(windows)    # (B, C, L)
    # Flipped direction reads the sequence from tail to head  [§4.1]
    if flip:
        sequence = reverse_last_axis(sequence)
    return sequence

Step 3: the four-branch block §4.1 fig. 3a. The fixed-direction block uses a horizontal global scan, its flip, the local scan and its flip. Each branch is scanned, passed through its own selective SSM, un-scanned, reweighted by channel attention, and summed.

# LocalMamba block with fixed directions  [§4.1] [fig. 3a]
# Input:  tokens (B, L, D) on the grid; directions: four scan names
# Output: tokens (B, L, D)

def local_block(tokens, directions, ssms, attentions):
    # Scan the grid once per direction  [§4.1]
    scanned = [scan(tokens, d) for d in directions]
    # One selective SSM per branch, with its own parameters  [§3.2] [§4.1]
    processed = [ssm(s) for ssm, s in zip(ssms, scanned)]
    # Undo each scan so every branch is back in grid order  [§4.1]
    restored = [unscan(p, d) for p, d in zip(processed, directions)]
    # Channel reweighting per branch (SCAttn, see Step 4)  [§4.1] [fig. 3b]
    weighted = [attn(r) for attn, r in zip(attentions, restored)]
    # Merge the branches by summing them  [§4.1]
    # [code: classification/lib/models/mamba/multi_mamba.py#L340]
    return sum_over_branches(weighted)

Step 4: SCAttn §4.1 fig. 3b §4.3. The paper describes a channel branch and a token (spatial) branch. Channel weights come from the token-average of the branch features, passed through a bottleneck that reduces channels by a factor r = 8. Both branches are gated by sigmoid.

# SCAttn, channel branch  [§4.1] [fig. 3b] [§4.3]
# Input:  features (B, L, D) of one branch; reduction r = 8
# Output: features (B, L, D) with channel weights applied

def channel_attention(features):
    # Average over tokens to get one global vector per channel  [§4.1]
    # The code normalises first
    # [code: classification/lib/models/mamba/multi_mamba.py#L149-L152]
    global_vector = mean_over_tokens(layer_norm(features))  # (B, 1, D)
    # Reduce by r = 8, then activation  [§4.1] [§4.3]
    hidden = gelu(linear_reduce(global_vector))  # (B, 1, D / r)
    # Expand back and gate each channel with a sigmoid  [§4.1]
    gate = sigmoid(linear_expand(hidden))        # (B, 1, D)
    return features * gate
    # Token-wise (spatial) gating is in §4.1 but commented out in the code
    # [code: classification/lib/models/mamba/multi_mamba.py#L156-L161]

Step 5: direction search §4.2 eq. 5. The search space holds eight directions per layer: horizontal and vertical scans, each standard and flipped, plus local scans with w = 2 and w = 7, each standard and flipped §4.2. A learnable vector α per layer gives softmax weights over the eight branches, and the output is their weighted sum eq. 5. The supernet is trained with all eight branches, and afterwards the four highest-probability directions are kept for each layer §4.2.

# Differentiable direction search  [§4.2] [eq. 5]
# Input:  one learnable vector alpha per layer (8 values each)
# Output: four directions per layer

def search_branch(x, alpha, directions):
    # Softmax over the 8 candidate directions of this layer  [eq. 5]
    weights = softmax(alpha)                    # (8,)
    # One SSM output per candidate direction  [§4.2]
    outputs = [ssm_d(scan(x, d)) for d, ssm_d in zip(directions, ssms)]
    # Weighted sum of the candidate outputs  [eq. 5]
    # The code also applies a LayerNorm to each output before weighting
    # [code: classification/lib/models/mamba/multi_mamba.py#L43-L47]
    return sum(w * o for w, o in zip(weights, outputs))

def pick_directions(alpha):
    # After the supernet has been trained, keep the four largest probabilities
    # per layer  [§4.2] [code: classification/tools/vis_search_prob.py#L13-L16]
    probs = softmax(alpha)                      # (8,)
    return top_k_indices(probs, 4)              # sorted by candidate index

Step 6: architectures §4.3 tab. 1.

• LocalVim replaces each Vim block with the four-branch block. The original Vim block has two directions, so the paper reduces the depth from 24 to 20 blocks to keep the compute similar §4.3.
• LocalVMamba keeps the VMamba layout and swaps in the four-branch block, with no change to the structure §4.3.

Cost §4.3. Scanning only reorders tokens and adds no FLOPs. SCAttn reduces the token dimension by 1/r and adds a small amount of compute. The authors report LocalVMamba-T at 5.7G FLOPs against 5.6G for VMamba-T §4.3.

Hyperparameters

Results

Author-reported. Setup: ImageNet-1K classification at 224², 300 epochs §5.1; COCO val with Mask R-CNN and Cascade Mask R-CNN §5.2; ADE20K val with UperNet §5.3. Baselines are the models listed in each table.

ImageNet-1K top-1 tab. 2. Rows marked * are LocalVim without direction search (fixed horizontal and w2 scans) tab. 2 tab. 5.

Ablation on ImageNet-1K tab. 5. LocalVim-T* setting (no search, 2×2 local window) tab. 5.

COCO object detection and instance segmentation, val tab. 3. Mask R-CNN, 1× and 3× schedules §5.2.

COCO with the Vim detector setup, val tab. 6 app. A.1. LocalVim-T with Cascade Mask R-CNN and ViTDet, following Vim. Vim did not report params or FLOPs tab. 6.

ADE20K semantic segmentation, val, UperNet tab. 4. Single-scale (SS) and multi-scale (MS) mIoU; input 512 × 512; FLOPs at 512 × 2048. "-" means not reported by the source tab. 4.

Searched directions fig. 4 app. A.2. The paper reports that LocalVim favours local scans at the first and last layers and global horizontal and vertical scans in the middle. The 2×2 local scans cluster towards the end of the network and 7×7 towards the start. LocalVMamba leans more towards local scans and prefers 7×7 over 2×2 §5.5 fig. 4. The paper reports these as observations; no quantitative analysis is given §5.5.

Paper vs code

Code: https://github.com/hunto/LocalMamba at 75e5d1f007c303509c77d179df38f7fa19c0d05c (Apache-2.0). The repo contains classification/, detection/ and segmentation/ folders. The paper links this repo in its abstract.

• SCAttn token branch is missing. The paper describes a channel and a spatial (token-wise) attention branch §4.1 fig. 3b. The code implements only the channel branch; the spatial gate is commented out multi_mamba.py. LocalVMamba has the same channel-only BiAttn local_vmamba.py.
• Search adds a LayerNorm before the softmax mix. Eq. 5 is a plain softmax-weighted sum of SSM outputs eq. 5. The code applies a non-affine LayerNorm to each branch output before the weighted sum multi_mamba.py. The softmax logits start from 1e-3 * randn multi_mamba.py; the paper does not state an initialisation.
• Duplicate local_vim_small definition. The file defines local_vim_small twice, with @register_model on each. The first, with searched directions, is at local_vim.py. The second, with the fixed directions ('h', 'h_flip', 'w2', 'w2_flip') * 20, is at local_vim.py. Python binds the name to the later definition, so the name most likely resolves to the fixed-direction model. This is inferred from the source; it was not run.
• EMA differs. The paper uses EMA decay 0.9999 §5.1. The ImageNet config has model_ema: False and model_ema_decay: 0.999 config.yaml. The README lists --model-ema-decay 0.9999 as an option README.md.
• CutMix is enabled. The paper names MixUp (ratio 0.2) only §5.1. The config also sets cutmix: 1.0 and mixup_switch_prob: 0.5 config.yaml.
• Batch size. The paper states a base batch of 1024 §5.1. The README command uses -b 128 with 8 GPUs README.md. This matches 1024 only if -b is per GPU; the code does not show this.
• Drop path differs by model. The config and README use 0.1 config.yaml README.md. The LocalVMamba-T constructor defaults to 0.2 local_vmamba.py. The paper does not state a value.
• Segmentation LocalVim uses different norm settings. The classification LocalVim-T uses rms_norm=True and fused_add_norm=True local_vim.py. The ADE20K LocalVim-T config sets rms_norm=False upernet_local_vim-tiny_8xb2-160k_ade20k-512x512.py. The paper does not discuss this.
• Checkpoints are placeholders. The pretrained=True branches load from url="to.do" local_vim.py local_vim.py local_vim.py. The segmentation configs point to release checkpoints upernet_local_vim-tiny_8xb2-160k_ade20k-512x512.py upernet_local_vssm_tiny_8xb2-160k_ade20k-512x512.py. Load weights from those release files, not from the classifier's pretrained branches.
• Per-layer direction lists match the search space. Every fixed LocalVim and LocalVMamba list has four names from the eight candidates local_vim.py local_vmamba.py. The paper's counts (20 LocalVim blocks; 15 LocalVMamba-T blocks) match these lists.
• Searched lists are chosen by top-4 probability. vis_search_prob.py takes the four largest softmax values per layer and sorts them by candidate index vis_search_prob.py. This matches §4.2.

Pitfalls

• Internal inconsistency in the text. The introduction and abstract give +3.1 for LocalVim-T over Vim-Ti; the results text says LocalVim-T exceeds Vim-Ti by 2.7 §5.1. Table 2 shows 76.2 − 73.1 = 3.1 for the full LocalVim-T. The 2.7 matches the no-search LocalVim-T* (75.8 − 73.1). Use the table values.
• Padding is not in the paper. Window scans need H and W divisible by w. The code pads the bottom and right with zeros and crops after reversing local_scan.py local_scan.py. At 224² input the patch grid is 14 × 14, which 2 and 7 divide, so classification needs no padding. Detection and segmentation inputs usually do.
• Within-window order is not in the paper. The code reads windows row-major and rows row-major inside each window local_scan.py. The column-first variant exists in the code but MultiScan does not call it local_scan.py.
• Search width differs from the final model. The supernet uses width 128 (LocalVim) and 32 (LocalVMamba) §5.1; the final models use 192 and 96. The searched directions therefore come from a different network. The paper does not justify transferring them.
• Search memory grows with the candidate set. The paper reports 28 GB for 20 blocks with 128 directions per block §4.2. The search supernet holds all eight branches per block multi_mamba.py.
• Hidden memory cost of the scans. The code's scan uses permute and reshape copies multi_mamba.py. A Triton local-scan kernel exists but is commented out in MultiScan multi_mamba.py. The paper's "no additional FLOPs" claim does not count this memory traffic.
• Selective-scan CUDA kernel required. The block calls mamba_ssm's selective-scan kernel; the import falls back to None if it is missing multi_mamba.py. The repo ships causal-conv1d and mamba-1p1p1 directories.
• Detection and segmentation are not comparable across papers. LocalVMamba uses Mask R-CNN as VMamba does; LocalVim uses Cascade Mask R-CNN with ViTDet, as Vim does §5.2. The authors implemented LocalVim segmentation themselves because Vim released no segmentation code §5.3.
• Table 3 typo. The LocalVMamba-S row shows "69N" for params tab. 3; the 3× row shows 69M.
• Unstated details. Seeds and variance, the optimiser's betas and eps, the search optimiser, α initialisation, the detection and segmentation epochs for LocalVim, and warmup for segmentation are not stated.
• The "no extra cost" claim covers FLOPs only. The paper reports FLOPs, not measured latency or throughput §4.3 tab. 1.

Lineage

• Builds on: Vim, a plain bidirectional Mamba vision backbone, [arXiv:2401.09417]; the LocalVim block changes its scan set and reduces depth from 24 to 20 §1 §4.3. VMamba, a hierarchical 2-D scan model, [arXiv:2401.10166]; LocalVMamba swaps in the new block without changing the layout §2 §4.3. Mamba (S6) [/papers/2312.00752] provides the selective SSM §3.2. DARTS, a differentiable architecture search, provides the relaxation §4.2.
• Compared against: RegNetY, ViT, DeiT, Swin tab. 2; Vim and VMamba tab. 2 tab. 3 tab. 4; ConvNeXt and ViT-Adapter tab. 3; ResNet-50 and DeiT with MLN on ADE20K tab. 4.
• Superseded by: not known from content already in this wiki.