Progressive Split Mamba: Effective State Space Modelling for Image Restoration

The paper proposes Progressive Split-Mamba (PS-Mamba), a topology-aware hierarchical state-space framework that addresses the spatial distortion and long-range decay limitations of standard Mamba models in image restoration by employing geometry-consistent partitioning and symmetric cross-scale shortcuts to effectively balance local structural preservation with global coherence.

The Gist

Here is an explanation of the paper "Progressive Split Mamba" using simple language and creative analogies.

The Big Problem: Fixing a Broken Photo

Imagine you have a beautiful, high-resolution photo that has been ruined by noise, blur, or compression artifacts (like a blurry JPEG). Your goal is to fix it. This is called Image Restoration.

To do this, computers need to understand two things simultaneously:

1. The Tiny Details: The texture of a cat's fur or the grain of wood (Local details).
2. The Big Picture: The fact that the cat's ear belongs to the cat's head, not the background (Global context).

The Old Ways (and Why They Failed)

Before this new method, computers tried two main approaches, both of which had flaws:

• The "Zoom Lens" Approach (CNNs): These models look at the image through a small window, moving pixel by pixel.
  • The Flaw: It's like trying to understand a whole novel by reading only one word at a time. You get the local details, but you lose the story. You can't see how the beginning connects to the end.
• The "Telepathy" Approach (Transformers): These models look at every pixel and instantly "talk" to every other pixel to understand the whole picture.
  • The Flaw: It's like trying to have a conversation with 1 million people at once. It's incredibly slow and expensive. Also, because they are so focused on the "big picture," they sometimes miss the tiny, fine details (like the texture of a leaf).

The New Contender: Mamba

Recently, a new architecture called Mamba arrived. It's like a super-efficient reader that can scan a book very quickly (linear time) while remembering the whole story. It's fast and handles long-range connections well.

However, Mamba has a fatal flaw when applied to 2D images:
To read an image, Mamba has to flatten it into a long, single line of pixels (like unrolling a carpet).

1. Locality Distortion: When you unroll a carpet, the pixels that were right next to each other (like a cat's nose and whisker) might end up at opposite ends of the line. The model gets confused about who is a neighbor.
2. Long-Range Decay: Mamba works like a game of "Telephone." As the message passes from pixel 1 to pixel 10,000, the information gets weaker and weaker. By the time it reaches the end, the important details have faded away.

The Solution: Progressive Split-Mamba (PS-Mamba)

The authors of this paper invented PS-Mamba. Think of it as a smart way to organize a messy room before cleaning it.

Instead of unrolling the whole carpet into one long line, PS-Mamba uses a "Progressive Split" strategy:

1. The "Pizza Slicing" Analogy (Topology-Aware Splitting)

Imagine the image is a giant pizza.

• Old Mamba: Cuts the pizza into a single, long strip of crust and toppings, then tries to eat it from one end to the other. The pepperoni at the start is far from the pepperoni at the end.
• PS-Mamba: Instead of one long strip, it cuts the pizza into halves, then quarters, then eighths.
  • It processes these small, manageable slices independently.
  • Why this helps: Neighbors stay neighbors! The pepperoni is still next to the cheese because they are in the same small slice. This preserves the "local structure" perfectly.

2. The "High-Speed Elevator" Analogy (Symmetric Shortcuts)

Even with slices, if the pizza is huge, information might still get lost as it travels through the layers of the network.

• The Fix: PS-Mamba installs symmetric cross-scale shortcuts.
• The Analogy: Imagine a building with many floors. If you have to walk down every single step to get to the basement, you get tired (information decays). PS-Mamba adds elevators that connect the top floor directly to the bottom floor.
• Result: The "global context" (the big picture) is sent directly to the deep layers without fading away. It ensures the model remembers the whole image while fixing the tiny details.

How It Works in Practice

The system works in three steps:

1. Split: It breaks the image into geometric chunks (halves, quarters, octants) so neighbors stay together.
2. Process: It uses the efficient Mamba engine to fix the details inside each chunk.
3. Merge & Refine: It puts the chunks back together, but uses those "elevator shortcuts" to make sure the global structure is consistent and the colors match up perfectly.

The Results: Why It Matters

The paper tested this on three difficult tasks:

• Super-Resolution: Making a small, blurry photo huge and sharp.
• Denoising: Removing static/noise from an old photo.
• JPEG Artifact Reduction: Fixing the blocky squares that appear in low-quality JPEGs.

The Outcome:
PS-Mamba beat the previous best models (like MambaIR and SwinIR) in almost every test.

• It produced sharper edges.
• It kept textures (like hair or fabric) looking natural.
• It did all this faster and with fewer computer resources than the heavy Transformers.

The Takeaway

Progressive Split-Mamba is like a master restorer who doesn't try to fix the whole painting at once (too slow) or just one tiny dot at a time (too confused). Instead, they break the painting into logical sections, fix the details in each section while keeping the neighbors together, and use a special "magic wire" to ensure the whole painting stays connected and consistent.

It solves the "Telephone game" problem of AI image restoration, giving us clearer, sharper, and more realistic images.

Technical Summary

Here is a detailed technical summary of the paper "Progressive Split Mamba: Effective State Space Modelling for Image Restoration."

1. Problem Statement

Image restoration (e.g., super-resolution, denoising, JPEG artifact reduction) requires balancing two conflicting requirements: preserving fine-grained local structures and maintaining long-range spatial coherence.

• CNNs struggle with long-range dependencies due to limited receptive fields.
• Transformers capture global context well but suffer from quadratic computational complexity and often lose local structural details due to window partitioning or global attention mechanisms.
• State Space Models (SSMs/Mamba) offer linear-time complexity for long-range modeling but face two intrinsic shortcomings when naively applied to 2D images:
  1. Locality Distortion: Flattening 2D feature maps into 1D sequences (rasterization) disrupts spatial topology. Neighboring pixels in 2D space become distant in the 1D sequence, forcing the model to learn local textures through inefficient long-range transitions.
  2. Long-Range Decay: The stability constraints of SSMs (requiring spectral radius < 1) cause information to decay exponentially as it propagates through the sequence. In high-resolution images with long sequences, global and low-frequency signals weaken before reaching distant positions, leading to inconsistent global structures.

2. Methodology: Progressive Split-Mamba (PS-Mamba)

The authors propose PS-Mamba, a topology-aware hierarchical framework designed to reconcile locality preservation with efficient global propagation without token reordering or multi-directional scanning.

Core Components:

1. Progressive Split Hierarchy:

   • Instead of flattening the entire image, PS-Mamba performs geometry-consistent partitioning. It splits feature maps hierarchically into halves, quadrants, and octants.
   • This maintains neighbourhood integrity (4-connected adjacency) within each patch before state-space processing.
   • By processing shorter sequences within patches ($L_j \ll HW$), the model avoids the adjacency errors of full rasterization and reduces the recurrent depth, mitigating information decay.

2. Symmetric Cross-Scale Shortcut Pathways:

   • To counteract the long-range decay inherent in linear SSMs, the authors introduce symmetric skip connections.
   • These pathways directly transmit low-frequency global context across hierarchical levels, effectively collapsing long dependency distances and stabilizing information flow.

3. Hybrid Architecture:

   • Convolutional Preprocessing: A lightweight residual convolution branch extracts short-range textures and edge-aware features before feeding them into the Mamba module.
   • Patch-Level Mamba Core: Each partitioned patch is processed independently through the Mamba sequence, preserving local structure.
   • Attention-Based Fusion: A content-adaptive gating mechanism fuses convolutional features with Mamba features. This is followed by a dual attention refinement (Channel Attention and Spatial Attention) to emphasize discriminative structures and reinforce global consistency.

4. Training Objective:

   • Unlike progressive networks that use stage-wise losses, PS-Mamba employs a single unified L1 loss (or Charbonnier loss for denoising) on the final output. This simplifies optimization and encourages consistent feature refinement across all stages.

3. Key Contributions

1. Topology-Preserving State-Space Module: Introduces a progressive split strategy (halves $\to$ quadrants $\to$ octants) that processes 2D images as geometry-aligned patches. This eliminates raster-induced adjacency errors while retaining the linear complexity of SSMs.
2. Symmetric Cross-Scale Skip Links: A novel mechanism that bypasses long Mamba chains to deliver global structural information directly, effectively mitigating the exponential decay of information in linear SSMs.
3. Unified Restoration Architecture: Integrates adaptive Mamba-Conv blocks and dual attention refinement, creating a framework that unifies the strengths of state-space modeling (global efficiency) and locality consideration (structural fidelity).
4. State-of-the-Art Performance: Demonstrates that PS-Mamba achieves superior results in detail fidelity and global coherence compared to existing Mamba-based and Transformer-based models, often with fewer parameters.

4. Experimental Results

The method was evaluated on Super-Resolution (SR), Denoising, and JPEG Compression Artifact Reduction (CAR) across standard benchmarks (Set5, Set14, BSDS100, Urban100, Manga109, CBSD68, etc.).

• Lightweight SR: PS-Mamba-Light consistently outperformed lightweight baselines (e.g., SwinIR-light, MambaIRv2-light).
  • Example: On Set5 ($\times2$), it achieved 38.31 dB PSNR, surpassing MambaIRv2-light (38.26 dB) with significantly lower MACs (60.1G vs. 126.7G).
• Classic SR: The large variant (PS-Mamba-Large, 21.2M params) outperformed the much larger MambaIRv2-L (34.2M params).
  • Example: On Manga109 ($\times2$), it reached 40.90 dB, a clear margin over MambaIRv2-L (40.55 dB), despite using ~13M fewer parameters.
• Denoising & CAR: The model set new benchmarks in Gaussian color image denoising (e.g., 34.55 dB on CBSD68) and JPEG artifact reduction, outperforming U-Net based models like Restormer and other Mamba variants.
• Ablation Studies:
  • Split Size: Octant-level splitting yielded the best performance, balancing local continuity and global aggregation.
  • Channel Dimension: A fixed channel increment of 48 provided the optimal trade-off between reconstruction power and computational cost.

5. Significance

• Solving the 2D-SSM Mismatch: PS-Mamba addresses the fundamental limitation of applying 1D SSMs to 2D data without resorting to complex token reordering or multi-directional scanning strategies used in previous works (like MambaIRv2).
• Efficiency vs. Performance: It proves that high-fidelity image restoration does not require the massive parameter counts of Transformers or the complex scanning mechanisms of other SSM adaptations. It achieves a better balance between model capacity and computational efficiency.
• Generalizability: The framework is shown to be a robust backbone for diverse restoration tasks, suggesting that preserving geometric locality during state propagation is a critical design principle for future vision models based on State Space Models.