PaperlayerView exact revision on arXiv

Source-linked AI summary

MambaIRv2: Attentive State Space Restoration

Hang Guo, Yong Guo, Yaohua Zha, Yulun Zhang, Wenbo Li, Tao Dai, Shu-Tao Xia, Yawei Li

arXiv:2411.15269v2eess.IVcs.CVcs.LG

TL;DR

Mamba-based restoration is limited by causal scanning, which prevents tokens from fully using image-wide information. MambaIRv2 adds attentive state-space modeling and semantic-guided neighboring to provide non-causal interaction and single-scan processing, with reported gains over prior restoration backbones. Its broader interpretation and task coverage remain areas for future work.

  • Problem

    Causal Mamba modeling restricts each token to preceding scanned pixels, limiting image-wide pixel utilization and complicating efficient non-causal restoration.

  • Method

    MambaIRv2 combines an Attentive State-space Equation that queries beyond the scanned sequence with Semantic Guided Neighboring that places similar pixels closer together.

  • Results

    0.34dB PSNR: MambaIRv2 suppresses standard MambaIR on 2× Urban100 while reducing parameters by 43% and computational burden by 50%.

  • Takeaways & Limitations

    MambaIRv2 provides an efficient image-restoration backbone that integrates ViT-like non-causal abilities into Mamba-based models.

  • Takeaways & Limitations

    The paper calls for deeper interpretability analysis and notes that image-restoration tasks such as deblurring, dehazing, and deraining remain to be explored.

Abstract

from arXiv · show

The Mamba-based image restoration backbones have recently demonstrated significant potential in balancing global reception and computational efficiency. However, the inherent causal modeling limitation of Mamba, where each token depends solely on its predecessors in the scanned sequence, restricts the full utilization of pixels across the image and thus presents new challenges in image restoration. In this work, we propose MambaIRv2, which equips Mamba with the non-causal modeling ability similar to ViTs to reach the attentive state space restoration model. Specifically, the proposed attentive state-space equation allows to attend beyond the scanned sequence and facilitate image unfolding with just one single scan. Moreover, we further introduce a semantic-guided neighboring mechanism to encourage interaction between distant but similar pixels. Extensive experiments show our MambaIRv2 outperforms SRFormer by even 0.35dB PSNR for lightweight SR even with 9.3\% less parameters and suppresses HAT on classic SR by up to 0.29dB. Code is available at https://github.com/csguoh/MambaIR.

1. Introduction

MambaIRv2 addresses the causal limitations of existing Mamba-based restoration by adding ViT-like non-causal modeling and semantic interactions between distant similar pixels.

  • Motivation: Existing Mamba-based restoration methods model each pixel only from preceding tokens in a scanned sequence, underusing helpful pixels elsewhere in the image.This causal design also creates challenges for non-causal image restoration.
  • Proposed approach: MambaIRv2 uses the state-space output matrix to query relevant pixels beyond the scanned sequence, removing the need for multi-directional scanning.The approach integrates ViT-like non-causal modeling into Mamba.
  • Proposed approach: Semantic Guided Neighboring restructures the image so semantically similar pixels become spatially closer in the 1D sequence, mitigating long-range decay.This encourages interaction between distant yet similar pixels.
  • Proposed approach: The proposed Attentive State-space Equation uses prompt learning to represent semantically similar pixels across the image and enables single-pass scanning.Prompts are incorporated into the state-space output matrix through residual addition.

2. Related Work

Image restoration has progressed from CNNs to ViTs and Mamba, with Mamba offering a route toward global reception and efficient computation.

  • CNN-based restoration: CNNs were early image-restoration backbones for super-resolution, denoising, and JPEG artifact reduction.Representative methods include SRCNN, DnCNN, and ARCNN.
  • Transformer-based restoration: ViTs improved restoration performance but vanilla self-attention incurs quadratic computational complexity with input size.Windowing and patch-based methods were developed to improve attention efficiency.
  • Mamba-based restoration: Mamba has been explored as a backbone balancing efficient computation with global receptive fields in image restoration.MambaIR introduced Mamba for restoration, while later methods adapted state-space modeling to tasks such as deraining and low-light enhancement.
  • Mamba-based restoration: MambaIR analyzes scanned features across four directions and all layers, motivating examination of redundancy in multi-directional scanning.The analysis is described in Figure 2(a).

3. Motivation

The motivation centers on causal state-space modeling: scanned tokens see only their predecessors, while long-range decay weakens interactions with distant pixels.

  • Causal modeling: Mamba’s discrete state-space equation makes each token depend on its previous tokens, giving the model causal properties.The control, input, and output matrices define the discrete formulation.
  • Causal modeling: Because a scanned pixel sees only a constrained portion of the image, existing methods use multi-directional scans to broaden the receptive field at higher computational cost.The resulting scanned sequences can also be highly correlated, with cosine similarity above 0.7 on all testing datasets.
  • Long-range decay: The learned control matrix A is statistically less than 1, so the interaction A^k weakens as pairwise distance k increases.This produces long-range decay in the state-space interaction.
  • Long-range decay: Long-range decay causes current methods to underuse distant scanned pixels that may be useful for restoration.The weakness becomes more pronounced when two pixels are far apart.

4. Attentive State Space Restoration

MambaIRv2 bridges attention and state-space modeling to give Mamba non-causal image restoration. Its ASSM combines ASE-based global querying, semantic-guided sequence rearrangement, and local window attention.

  • Bridging Attention and State-Space: Causal Mamba models restrict each token to preceding scanned pixels, creating challenges for non-causal image restoration.The paper motivates replacing this restricted dependency pattern with non-causal modeling.
  • Bridging Attention and State-Space: Attention and state space share analogous roles, with C corresponding to the attention query and motivating global information injection into C.The analysis identifies hi ∼ Si, B ∼ K^⊤, and C ∼ Q.
  • Attentive State Space Module: ASSM enables non-causal Mamba modeling by applying positional encoding, semantic-guided unfolding, ASE modeling, inverse folding, and linear projection.ASSM operates on feature maps x ∈ R^(H×W×C), with L = HW flattened tokens in ASE.
  • Attentive State Space Module: ASE routes instance-specific prompts from a learnable prompt pool and adds them residually to C so tokens can query information from unscanned pixels.Prompt routing uses a projection, LogSoftmax, Gumbel-softmax, and a routing matrix R.
  • Semantic Guided Neighboring: Semantic Guided Neighboring groups semantically similar pixels into nearby positions in the unfolded sequence, addressing long-range decay during Mamba scanning.It redefines token neighborhoods using prompt categories learned by ASE.
  • Overall Network Architecture: ASSM’s single-direction scan leaves parameter budget for window MHSA, which is added to strengthen local interactions alongside global dependency modeling.The paper describes window attention and ASSM as complementary basic elements of MambaIRv2.

5. Experiments

Experiments evaluate MambaIRv2 across super-resolution, JPEG artifact reduction, and denoising, showing that ASE and SGN improve restoration performance while preserving efficiency.

  • Experimental Setup: The experiments cover classic and lightweight image super-resolution, JPEG compression artifact reduction, and Gaussian color image denoising.Training uses 64 × 64 patches for SR and 128 × 128 patches for denoising.
  • Ablation Study: ASE improves Urban100/Manga109 performance by 0.05 dB/0.09 dB, while adding SGN yields 32.97/39.24 dB PSNR.The ablation uses a MambaIRv2-light 2× SR model trained on DIV2K for 250K iterations.
  • Ablation Study: Adding prompts to C performs slightly better than adding them to B, whereas inserting them into ∆ or y gives unsatisfactory results.C is closer to the output end of the state-space equation and therefore has greater impact on the final result.
  • Lightweight Image Super-Resolution: 0.35 dB: MambaIRv2-light outperforms SRFormer-light on 2× Urban100 with 79K fewer parameters.The comparison uses parameter count and MACs as efficiency metrics.
  • Classic Image Super-Resolution: MambaIRv2 achieves the best performance across most classic SR datasets and scales, including a 0.16 dB gain over HAT on 2× Manga109.MambaIRv2-L adds 0.18 dB PSNR on 4× Urban100 over its base counterpart.
  • Discussion: A single semantic-space scan reduces parameters by 43% and computation by 50% versus standard MambaIR while improving 2× Urban100 by 0.34 dB PSNR.The attentive state space lets query pixels attend to corresponding semantic parts beyond the scanned sequence.

6. Conclusion

MambaIRv2 addresses Mamba’s causal modeling limitation by combining attentive state-space modeling with semantic-guided neighboring and single-scan processing.

  • 6. Conclusion: ASE uses prompt learning to enlarge token perception, while SGN positions similar pixels closer to handle Mamba’s long-range decay.Together, these components integrate ViT-like non-causal abilities into Mamba-based restoration.

Supplementary Material

MambaIRv2 is reported to retain global pixel utilization with efficiency comparable to HAT across large input resolutions.

  • A. Efficiency Comparison on Large Inputs: MambaIRv2-B uses fewer MACs than HAT from 64 × 64 through 1024 × 1024, including 30% savings at 256 × 256.Its parameter count is roughly similar to HAT, and the model still outperforms HAT by a noticeable margin.

B. More Ablation Results

These ablations and visualizations examine prompt and neighboring mechanisms alongside computational cost, showing how semantic restructuring supports interactions among similar pixels.

  • Prompt Learning: The prompt-learning ablation varies the prompt-pool size T and semantic-decoupling rank r to assess their effects on performance.The experiments explicitly target these two hyperparameters.
  • Semantic Guided Neighboring: Semantic Guided Neighboring restructures the image so semantically similar pixels become spatially close in the unfolded 1D sequence.This design targets interactions among distant but similar pixels.
  • Computational Cost: The computational-cost comparison evaluates MambaIRv2-B and HAT on 4× classical super-resolution across output sizes from 64× to 1024×1024.The table reports MACs at multiple image resolutions.

C. Comparison on Receptive Field

The receptive-field comparisons visualize pixel activation and effective receptive fields, while SGN places semantically similar pixels closer to strengthen their sequence interactions.

  • LAM Attribution: MambaIRv2 activates more pixels than HAT, exhibiting a wider LAM attribution and higher diffusion index.The diffusion index reflects the range of involved pixels, with higher values indicating wider utilization.
  • Effective Receptive Field: MambaIRv2 shows darker colors across the entire image in ERF visualization, indicating global perception.A larger ERF is represented by a more extensively distributed dark area.
  • Semantic Guided Neighboring: SGN moves semantically similar pixels from distant positions to spatially close positions in the unfolded sequence.This arrangement facilitates stronger interactions among those pixels.
  • Non-causal Modeling: MambaIR displays crisscrossing ERF patterns associated with causal scanning, whereas MambaIRv2 does not exhibit this unfavorable pattern.The comparison is presented as evidence for the proposed non-causal modeling.
  • Evaluation Setup: The evaluation uses Set5, Set14, B100, Urban100, and Manga109 for super-resolution testing.These datasets are used to evaluate different super-resolution methods.
  • Comparison to ATD: MambaIRv2 achieves performance comparable to ATD on 2× classic super-resolution.ATD is described as a highly optimized Transformer-based method.

F. Limitation and Future Works

The authors identify interpretability, broader task coverage, and the early stage of Mamba-based restoration as directions for future work.

  • Interpretability: Future work should analyze what Mamba and ViT learn during image restoration to improve interpretability and network design.The paper frames Mamba as a third backbone option alongside CNNs and ViTs.
  • Task Coverage: The study covers multiple restoration tasks but leaves deblurring, dehazing, and deraining for future exploration.The authors also suggest implementing a U-shaped MambaIRv2 backbone for these tasks.
  • Field Maturity: Mamba-based image restoration networks remain in an early stage despite the promising results reported here.The authors connect this boundary to increasing research interest in Mamba.

G. Proof for Long-range Decay

The proof explains long-range decay in causal Mamba: distant inputs receive increasingly small contributions to later outputs, weakening remote-pixel interactions.

  • Causal Interaction: Causal Mamba models each query token using only preceding tokens in the scanned sequence.The derivation concerns the interaction between query token x_k and previous tokens.
  • Derivation: The contribution of an earlier input to a later output is weighted by powers of A in the iterated state-space equation.The appendix derives the output recursively by setting i = 0, 1, 2, and continuing through k.
  • General Form: The appendix generalizes the k-th output as a function of inputs x_0 through x_k after continuous iteration of the causal recurrence.The displayed derivation summarizes the accumulated contributions across the scanned sequence.
  • Long-range Decay: When the mean value of A is statistically below 1, increasing sequence distance makes an earlier input’s contribution small.This is the stated long-range decay mechanism.
  • Restoration Consequence: Long-range decay can severely impair restoration when a distant earlier pixel is highly helpful to the query pixel.The limitation applies to non-causal image restoration interactions.
Loading 2411.15269v2…