MambaIR: A Simple Baseline for Image Restoration with State-Space Model

In plain words

Image restoration means taking a damaged picture (too small, noisy, blurry) and producing a clean one. A good restorer looks at many pixels when it rebuilds each output pixel. Convolutional networks only look at a small neighbourhood. Transformers could look at the whole image, but comparing every pixel with every other pixel gets too expensive, so restoration transformers such as SwinIR only attend inside small windows.

MambaIR swaps the attention for Mamba, a "state-space model". Think of it as a reader that walks along a sequence once and keeps a running memory, so the cost grows linearly with the number of pixels while information can still travel across the whole image. Because an image is 2D, the feature map is unrolled into four 1D sequences (row by row, column by column, and both backwards), each one is scanned, and the results are added.

Unrolling an image puts some neighbouring pixels far apart in the sequence, and Mamba's large memory leaves many channels unused. The authors patch both problems with plain tools: a small convolution after the scan to bring back local detail, and a channel-attention gate that reweights channels. On super-resolution and denoising benchmarks the result matches or beats SwinIR at similar compute, by the authors' own numbers.

TL;DR

MambaIR is a SwinIR-shaped restoration network (shallow 3×3 conv → stacked residual groups → reconstruction head) whose core block, the RSSB, is LN → VSSM (4-direction selective scan) + scaled skip followed by LN → conv bottleneck → channel attention + scaled skip. It targets classic, lightweight and real-world SR, Gaussian color denoising, real denoising and JPEG artifact removal. Author-reported gains over SwinIR are up to 0.45 dB on SR at similar cost. The paper leaves most architecture sizes unstated; they are only in the official code (Apache-2.0), which also differs from the paper in a few training details (see Paper vs code).

Problem

Restoration quality tends to improve with receptive field, but CNN backbones have a small effective receptive field, and Transformers either pay quadratic cost in the number of pixels or restrict attention to shifted windows, which loses the global receptive field §1 fig. 1. Mamba offers linear-cost long-range modelling, but it was built for causal 1D sequences: flattening an image separates spatially adjacent pixels ("local pixel forgetting"), and the large hidden state leads to many inactive channels ("channel redundancy") §1 §3.3 fig. 3.

Method

Three stages §3.2 fig. 2:

1. Shallow features: a 3×3 conv maps I_LQ ∈ R^{H×W×3} to F_S ∈ R^{H×W×C} §3.2.
2. Deep features: several Residual State-Space Groups (RSSGs), each a stack of RSSBs followed by a conv; then F_R = F_D^L + F_S §3.2. In code each RSSG also has its own residual connection, and a LayerNorm plus a 3×3 conv follow the last group mambair_arch.py mambair_arch.py.
3. Reconstruction: produces I_HQ from F_R §3.2. In code: conv + LeakyReLU + pixel-shuffle upsampler + conv for classic SR; a single conv + pixel shuffle for lightweight SR; and for denoising a conv whose output is added to the input image mambair_arch.py mambair_arch.py.

SSM background §3.1. Continuous system eq. 1:

h'(t) = \mathbf{A}h(t) + \mathbf{B}x(t), \quad y(t) = \mathbf{C}h(t) + \mathbf{D}x(t)

with state size N, A ∈ R^{N×N}, B ∈ R^{N×1}, C ∈ R^{1×N}, D ∈ R. Zero-order-hold discretization as printed in the paper eq. 2:

\overline{\mathbf{A}} = \exp(\Delta\mathbf{A}), \quad \overline{\mathbf{B}} = (\Delta\mathbf{A})^{-1}(\exp(\mathbf{A}) - \mathbf{I})\cdot\Delta\mathbf{B}

Recurrent form eq. 3:

h_k = \overline{\mathbf{A}}h_{k-1} + \overline{\mathbf{B}}x_k, \quad y_k = \mathbf{C}h_k + \mathbf{D}x_k

Equivalent convolutional form over a length-L sequence eq. 4:

\overline{\mathbf{K}} \triangleq (\mathbf{C}\overline{\mathbf{B}}, \mathbf{C}\overline{\mathbf{A}}\overline{\mathbf{B}}, \dots, \mathbf{C}\overline{\mathbf{A}}^{L-1}\overline{\mathbf{B}}), \quad \mathbf{y} = \mathbf{x} \circledast \overline{\mathbf{K}}

Mamba makes B̄, C and Δ input-dependent §3.1.

Residual State-Space Block (RSSB) §3.3 fig. 2a, with learnable per-channel scales s, s' ∈ R^C:

Z^l = \mathrm{VSSM}(\mathrm{LN}(F_D^l)) + s \cdot F_D^l eq. 5

F_D^{l+1} = \mathrm{CA}(\mathrm{Conv}(\mathrm{LN}(Z^l))) + s' \cdot Z^l eq. 6

Conv is a bottleneck that compresses channels by a factor γ to C/γ and expands back §3.3. CA is squeeze-and-excitation channel attention §3.3.

Vision State-Space Module (VSSM) §3.4 fig. 2b, following VMamba, with channel expansion factor λ eq. 7:

X_1 = \mathrm{LN}(\text{2D-SSM}(\mathrm{SiLU}(\mathrm{DWConv}(\mathrm{Linear}(X))))) X_2 = \mathrm{SiLU}(\mathrm{Linear}(X)), \quad X_{out} = \mathrm{Linear}(X_1 \odot X_2)

2D Selective Scan (2D-SSM) §3.5 fig. 2c: flatten the feature map into four sequences, run the selective scan on each, sum the four outputs and reshape back to 2D.

# MambaIR forward pass (classic SR head), shapes for one image
# Input: low_quality (B, 3, H, W); output: (B, 3, s*H, s*W)

def mambair(low_quality):
    # subtract the fixed DIV2K RGB mean
    # [code: basicsr/archs/mambair_arch.py#L541-L545]
    x = low_quality - rgb_mean
    # shallow 3x3 conv to C channels  [§3.2]
    shallow = conv3x3(x, out=C)                      # (B, C, H, W)
    # tokens = pixels, LayerNorm after embedding
    # [code: basicsr/archs/mambair_arch.py#L560-L565]
    tokens = layer_norm(flatten_hw(shallow))         # (B, H*W, C)
    for group in groups:
        group_input = tokens
        for block in group.blocks:
            tokens = rssb(block, tokens, H, W)
        # per-group conv and residual
        # [§3.2] [code: basicsr/archs/mambair_arch.py#L765-L766]
        grid = unflatten(tokens)
        tokens = flatten_hw(group.conv3x3(grid)) + group_input
    # final norm and conv, then the long skip F_R = F_D + F_S  [§3.2]
    deep = conv3x3(unflatten(layer_norm(tokens)))
    features = deep + shallow
    # conv + LeakyReLU, pixel-shuffle upsampling, conv to RGB
    # [code: basicsr/archs/mambair_arch.py#L615-L620]
    out = conv_last(upsample(leaky_relu(conv3x3(features, out=64))))
    return out + rgb_mean

def rssb(block, tokens, H, W):
    x = reshape(tokens, (B, H, W, C))
    # global branch with scaled skip  [eq. 5]
    z = vssm(block, layer_norm(x)) + block.s * x
    # local branch: conv bottleneck then channel attention  [eq. 6]
    local = channel_attention(conv_bottleneck(layer_norm(z)))
    return flatten(local + block.s_prime * z)

def vssm(block, x):                                  # x (B, H, W, C)
    # one linear gives both branches, each lambda*C wide
    # [eq. 7] [code: basicsr/archs/mambair_arch.py#L374-L375]
    branch_ssm, gate = split(linear(x, out=2 * lam * C))
    # depthwise 3x3 conv + SiLU
    # [eq. 7] [code: basicsr/archs/mambair_arch.py#L377-L378]
    branch_ssm = silu(depthwise_conv3x3(branch_ssm))
    # four-direction selective scan, outputs summed  [§3.5]
    y = sum(selective_scan_2d(branch_ssm))
    # LayerNorm, gate with SiLU, project back to C  [eq. 7]
    return linear(layer_norm(y) * silu(gate), out=C)

def selective_scan_2d(feature):                      # (B, D, H, W)
    # row-major, column-major, and both reversed
    # [code: basicsr/archs/mambair_arch.py#L343-L344]
    rows = flatten_row_major(feature)
    cols = flatten_column_major(feature)
    sequences = [rows, cols, reverse(rows), reverse(cols)]
    # separate dt/B/C projections, A and D per direction
    # [code: basicsr/archs/mambair_arch.py#L346-L362]
    outputs = []
    for direction, sequence in enumerate(sequences):
        outputs.append(mamba_scan(sequence, params[direction]))
    # undo reversal and transpose so all four align row-major
    # [code: basicsr/archs/mambair_arch.py#L365-L369]
    return realign(outputs)

The bottleneck and channel attention in the classic and denoising models are conv3x3(C→C/3) → GELU → conv3x3(C/3→C) → CA, with CA = avgpool → 1×1 conv (C→C/30) → ReLU → 1×1 conv → sigmoid → multiply mambair_arch.py mambair_arch.py.

Loss §3.6: L1 for SR eq. 8; Charbonnier with ε = 1e-3 for denoising eq. 9:

\mathcal{L} = \|I_{HQ} - I_{LQ}\|_1 \qquad \mathcal{L} = \sqrt{\|I_{HQ} - I_{LQ}\|^2 + \epsilon^2}

(As printed; the target should be the ground truth, see Pitfalls.)

Hyperparameters

Paper-stated values:

From the official code (not stated in the paper):

Results

Author-reported. PSNR/SSIM on the Y channel of YCbCr §4.1. Classic SR trained on DIV2K + Flickr2K; lightweight SR on DIV2K only; Gaussian color denoising on DIV2K + Flickr2K + BSD500 + WED; real denoising on 320 SIDD images §4.1.

Ablations (setting not stated; the "baseline" row equals the ×2 classic SR model in tab. 3) tab. 1 tab. 2:

Classic SR (PSNR / SSIM) tab. 3:

Lightweight SR (MACs for a 1280×720 output) tab. 4:

Gaussian color denoising (PSNR) tab. 5:

Real denoising (PSNR / SSIM) tab. 6:

Other results: real-world ×4 SR is shown only as visual comparisons, since there is no ground truth §4.3 fig. 6. GPU memory as the input grows from 48×48 to 84×84 scales roughly like SwinIR and far below full attention §4.3 fig. 5; no numbers are given in the text. JPEG artifact-reduction results and further ablations (e.g. the number of RSSBs) are in the supplementary material, which is not part of the arXiv source §4.1 §4.2.

Paper vs code

• Adam β2: paper 0.999 §4.1; code 0.99 train_MambaIR_SR_x2.yml.
• Charbonnier ε: paper sqrt(d² + ε²) with ε = 1e-3 eq. 9; code computes sqrt(d² + eps) per element with eps = 1e-3, i.e. an effective ε of about 0.032 losses.py train_MambaIR_ColorDN_level15.yml.
• SR patch size: paper says 64×64 SR patches §4.1. Code crops HR at 128 for ×2 and 192 for ×3/×4, so LR patches are 64 (×2, ×3) but 48 for ×4 train_MambaIR_SR_x2.yml train_MambaIR_SR_x4.yml.
• Scan order: the paper lists four corner-to-corner directions (TL→BR, BR→TL, TR→BL, BL→TR) §3.5. The code scans row-major, column-major, and the reverse of each mambair_arch.py.
• ×3/×4 halving: paper says to halve LR and iterations when initialising from ×2 §4.1. The classic configs do this train_MambaIR_SR_x4.yml; the lightweight ×3/×4 configs load ×2 weights but keep 2e-4 and 500k iterations train_MambaIR_lightSR_x4.yml.
• Lightweight block: not described in the paper. Code swaps the bottleneck to 1×1 (C→C/2) → depthwise 3×3 → GELU → 1×1 → depthwise 3×3 with dilation 2 → CA and uses a one-step pixel-shuffle head mambair_arch.py mambair_arch.py.
• Real denoising: paper describes one architecture and Charbonnier loss for denoising §3.6. The real-denoising config uses a different U-Net model (MambaIRUNet, dim 48, blocks [4, 6, 6, 8], 4 refinement blocks, λ = 1.5), L1 loss, AdamW (3e-4, weight decay 1e-4, β2 0.999), gradient clipping, mixup, and progressive patches 128→384 over 300k iterations train_MambaIR_RealDN.yml train_MambaIR_RealDN.yml.
• Color denoising metric: paper reports Y-channel PSNR for everything §4.1; the color-denoising test configs compute RGB PSNR with test_y_channel: false test_MambaIR_ColorDN_level15.yml.
• σ25/σ50 denoising: initialised from the σ15 model with LR 1e-4 and 800k iterations; not in the paper train_MambaIR_ColorDN_level50.yml.
• Extra pieces not in the paper: subtraction of the RGB mean, a LayerNorm after patch embedding, a residual around each RSSG, a final LayerNorm, and a global input skip for denoising mambair_arch.py mambair_arch.py.

Pitfalls

• Most architecture sizes are only in code. The paper gives no values for C, the group/block counts, state size N, λ or γ; use the code tables above.
• Eq. 2 as printed is inconsistent: B̄ uses exp(A) while Ā uses exp(ΔA). The code never implements eq. 2 itself; it calls mamba_ssm's selective_scan_fn with softplus on Δ plus a bias mambair_arch.py. Reuse that kernel rather than coding eq. 2 from the paper.
• Loss equations compare output with input. Eqs. 8 and 9 are written as I_HQ − I_LQ; the code compares the prediction with the ground truth losses.py.
• Four scans means four sets of SSM parameters. Each direction has its own Δ/B/C projection, Δ projection, A and D mambair_arch.py. The scan itself runs in float32 even under mixed precision mambair_arch.py.
• Hard CUDA dependency. The model imports mamba_ssm's fused selective scan mambair_arch.py. A pure-PyTorch port must reproduce its Δ softplus, bias and D-skip behaviour.
• Narrow channel attention. Squeeze factor 30 leaves 6 hidden channels at C = 180 and 2 at C = 60 mambair_arch.py.
• Noise synthesis. Gaussian noise of σ/255 is added on the fly to float images in [0, 1] with no clipping; validation fixes seed 0 paired_image_dataset.py.
• Ablation setup unstated. tab. 1 and tab. 2 give no training setting, and they leave out the parameter and MAC counts that the variants change.
• Long training. Classic ×2 SR runs 500k iterations on 8 GPUs, and σ15 denoising 1.6M iterations train_MambaIR_SR_x2.yml train_MambaIR_ColorDN_level15.yml.
• Repository also holds MambaIRv2. mambairv2_arch.py and the mambairv2 configs belong to a later model; the files for this paper are mambair_arch.py and options/*/mambair/.

Lineage

• Builds on: Mamba (arXiv:2312.00752) for the selective scan; VMamba (arXiv:2401.10166) for the VSSM and the 2D four-direction scan; SwinIR's three-stage pipeline and experimental setup; squeeze-and-excitation channel attention §3.3 §3.4 §3.5 §4.1.
• Compared against: SwinIR, SRFormer, IPT, ELAN, RCAN, EDSR, HAT (receptive-field plots), Restormer, DRUNet, Uformer and lightweight SR models tab. 3-6.