Progressive Split Mamba: Effective State Space Modelling for Image Restoration

TL;DR

Progressive Split-Mamba (PS-Mamba), a topology-aware hierarchical state-space framework designed to reconcile locality preservation with efficient global propagation, is proposed, a topology-aware hierarchical state-space framework designed to reconcile locality preservation with efficient global propagation.

Abstract

Image restoration requires simultaneously preserving fine-grained local structures and maintaining long-range spatial coherence. While convolutional networks struggle with limited receptive fields, and Transformers incur quadratic complexity for global attention, recent State Space Models (SSMs), such as Mamba, provide an appealing linear-time alternative for long-range dependency modelling. However, naively extending Mamba to 2D images exposes two intrinsic shortcomings. First, flattening 2D feature maps into 1D sequences disrupts spatial topology, leading to locality distortion that hampers precise structural recovery. Second, the stability-driven recurrent dynamics of SSMs induce long-range decay, progressively attenuating information across distant spatial positions and weakening global consistency. Together, these effects limit the effectiveness of state-space modelling in high-fidelity restoration. We propose Progressive Split-Mamba (PS-Mamba), a topology-aware hierarchical state-space framework designed to reconcile locality preservation with efficient global propagation. Instead of sequentially flattening entire feature maps, PS-Mamba performs geometry-consistent partitioning, maintaining neighbourhood integrity prior to state-space processing. A progressive split hierarchy (halves, quadrants, octants) enables structured multi-scale modelling while retaining linear complexity. To counteract long-range decay, we introduce symmetric cross-scale shortcut pathways that directly transmit low-frequency global context across hierarchical levels, stabilising information flow over large spatial extents. Extensive experiments on super-resolution, denoising, and JPEG artifact reduction show consistent improvements over recent Mamba-based and attention-based models with a clear margin.

Figures (3)

• Figure 1: Motivation of using PS-Mamba compared to basic Mamba and recent methods in image restoration. (a) The input image is conceptually split into pixel groups (shown enlarged for illustration). (b) PS-Mamba processes these geometry-aligned patches independently, preserving neighbourhood proximity and retaining local structure. This is contrasted with (c) the basic Mamba formulation, which flattens the entire image into a single long sequence and disrupts spatial adjacency, and (d) re-ordering strategies that require an additional learning stage to embed local textures into a closer sequence space, as in Guo2025MambaIRv2. The stars represent the query (red) and reference (blue) pixels for the sequence modelling.
• Figure 2: Pipeline of the proposed PS-Mamba framework: The input image is progressively split into multi-scale, geometry-aligned regions(e.g. halves, quadrants, and octants) while maintaining the original spatial resolution. These patches enable patch-level processing that preserves locality and mitigates long-range decay. The channel dimension is increased by 48 to enhance representation capacity, and it is symmetrically reduced by 48 during the merging stages to maintain balanced skip connections.
• Figure 3: Visual comparison of HR , EDSR, RCAN, IGNN, SwinIR, MambaIR, MambaIRv2, and PS-Mamba on four picture from Manga109 and Urban100. PS-Mamba shows the closest to HR