💻 computer science

DPMambaIR: All-in-One Image Restoration via Degradation-Aware Prompt State Space Model

The paper proposes DPMambaIR, an All-in-One image restoration framework that integrates a fine-grained degradation extractor with a Degradation-Aware Prompt State Space Model and a High-Frequency Enhancement Block to achieve superior performance across multiple degradation types by dynamically adapting to fine-grained degradation features.

Original authors: Zhanwen Liu, Sai Zhou, Yuchao Dai, Yang Wang, Yisheng An, Xiangmo Zhao

Published 2026-02-06

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Zhanwen Liu, Sai Zhou, Yuchao Dai, Yang Wang, Yisheng An, Xiangmo Zhao

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you have a collection of old, damaged photographs. Some are blurry, some are covered in rain streaks, some are too dark, and others are grainy with noise.

In the past, fixing these photos was like hiring a different specialist for every problem. You needed one expert to fix blur, another to remove rain, and a third to brighten dark rooms. If you didn't know exactly what was wrong with a photo, you were stuck.

This paper introduces DPMambaIR, a "super-restorer" that acts as a single, all-in-one expert capable of fixing all these different problems at once. Here is how it works, using simple analogies:

1. The "Detective" (The Degradation Extractor)

Before trying to fix a photo, the system first acts like a detective. It looks at the damaged image and asks, "What exactly is wrong here? Is it heavy rain or just a light mist? Is the blur from a moving car or a shaky hand?"

Instead of just guessing a category (like "Rain" or "Blur"), this detective creates a detailed, continuous report (called a "degradation embedding"). Think of it like a doctor not just saying "You have a fever," but measuring the exact temperature and identifying the specific virus. This allows the system to understand the severity and type of damage with great precision.

2. The "Smart Brain" (The Degradation-Aware Prompt State Space Model)

Once the detective has the report, it hands it to the "Smart Brain" (the core of the system).

The Old Way: Imagine a factory machine with fixed gears. If you feed it a blurry photo, it tries to sharpen it using the same gears it uses for a noisy photo. It's rigid and often struggles to do both well at the same time.
The New Way (DPMambaIR): This system is like a shapeshifting machine. When it receives the detective's report, it instantly reconfigures its own internal gears and settings to match the specific problem.
- If the photo is noisy, it tightens its settings to smooth out the grain (like a low-pass filter).
- If the photo is too dark, it loosens its settings to amplify weak signals and bring out the light.
- If the photo is blurry, it adjusts to focus on edges.

This "reconfiguration" happens dynamically for every part of the image, allowing the system to adapt its physics to the specific damage it sees.

3. The "Detail Polisher" (High-Frequency Enhancement Block)

Sometimes, when a machine tries to fix many different things at once, it gets a bit lazy with the tiny details (like the texture of fabric or the sharpness of an eye). It focuses on the big picture but misses the fine lines.

To fix this, the authors added a small, lightweight "Detail Polisher." Think of this as a final touch-up artist who specifically looks for the tiny, high-frequency details that the main brain might have smoothed over too much. It ensures the restored photo looks crisp and sharp, not just "okay."

The Results

The researchers tested this "super-restorer" on a mix of seven different types of damage (blur, rain, snow, fog, noise, low light, and JPEG compression).

Performance: It beat all other current "all-in-one" methods, achieving the highest scores in image quality.
Versatility: It didn't just do well on the mix; it was also excellent at fixing specific problems individually, often outperforming machines designed only for that single task.
Surprise Test: They even tested it on damage types it had never seen before (like pixelation and raindrops). Because its "detective" understands the nature of damage rather than just memorizing labels, it could still fix these new problems surprisingly well.

The Limitation

The paper admits one boundary: If a photo has a massive, opaque object blocking the view (like a thick, heavy raindrop completely covering a face), the system cannot "invent" the missing face behind it. It can fix the damage around it, but it can't magically hallucinate content that is completely hidden.

In summary: DPMambaIR is a flexible, smart system that first diagnoses the exact nature of an image's damage and then instantly reshapes its own internal logic to fix that specific problem, all while keeping the tiny details sharp. It's a single tool that replaces the need for a whole toolbox.

Technical Summary: DPMambaIR

Problem Definition

All-in-One image restoration aims to recover high-quality images from various degraded inputs using a single unified model, addressing the practical limitations of designing dedicated models for each specific degradation type (e.g., noise, blur, rain, low-light). Existing approaches generally fall into two categories: Mixture-of-Experts (MoE) and Prompt-based methods. However, these suffer from significant drawbacks:

MoE-based methods often incur high computational costs and lack explicit physical modeling of degradations, limiting adaptability to unseen scenarios.
Prompt-based methods typically rely on coarse-grained, static priors (discrete class labels or fixed semantic tokens) that fail to capture the continuous nature and spatial variability of real-world degradations. Furthermore, they often use an additive fusion paradigm where degradation cues are merely appended to visual features, leaving the underlying network architecture governed by fixed parameters that cannot dynamically adapt to heterogeneous restoration demands.

Additionally, joint optimization of heterogeneous degradations often induces an implicit bias toward low-frequency structures, making the recovery of fine high-frequency details challenging.

Methodology

The authors propose DPMambaIR, a novel All-in-One framework that shifts from conventional feature fusion to explicit dynamic parameter modulation. The architecture is built upon a Mamba-based U-shaped network with an asymmetric encoder-decoder design and consists of three core components:

1. Degradation Extractor

To capture fine-grained, continuous degradation features, the authors design a pre-trained degradation extractor.

Mechanism: Instead of relying on discrete labels, the extractor uses a reconstruction-based objective to regress continuous degradation embeddings. It employs Central Difference Convolution (CDC) to identify gradient-based cues (edges, textures) and multi-scale branches to perceive degradations across different frequency bands.
Training: The extractor is pre-trained via a self-supervised reconstruction task where it maps a degraded image to a low-dimensional embedding ( $E_d$ ). A decoder then attempts to reconstruct the degraded image using $E_d$ and a content embedding derived from the clean ground truth. This ensures $E_d$ captures pure degradation information independent of image content.
Re-parameterization: During inference, a re-parameterization technique merges the multi-branch structure into a compact representation to balance accuracy and efficiency.

2. Degradation-Aware Prompt State Space Model (DP-SSM)

The core innovation is the DP-SSM, which integrates the extracted degradation embeddings directly into the State Space Model (SSM) dynamics.

Dynamic Modulation: Unlike standard SSMs where parameters are fixed or input-dependent, DP-SSM uses the degradation embedding to dynamically modulate the discretization step ( $\Delta$ $Δ$ ), the input matrix ( $B$ $B$ ), and the output matrix ( $C$ $C$ ).
- $\Delta$ Modulation: The model learns to assign smaller $\Delta$ values for high-frequency degradations (e.g., noise), creating a high-inertia state that acts as a low-pass filter to suppress noise. Conversely, larger $\Delta$ values are generated for low-frequency/structural degradations (e.g., blur, low-light), entering a high-gain regime to amplify weak signals.
- $B$ and $C$ Modulation: $B$ controls how strongly external observations drive state transitions, while $C$ injects global degradation priors into the feature readout.
Efficiency: This approach avoids the high computational cost of multi-directional scanning (used to mitigate causal limitations in standard SSMs) by instead using the modulated $C$ matrix to inject global context, maintaining linear complexity.

3. High-Frequency Enhancement Block (HEB)

To counteract the low-frequency bias inherent in joint optimization, a lightweight HEB is introduced.

Function: It explicitly extracts high-frequency components by subtracting globally averaged low-frequency features from the input features and adding them back. This complements the dynamic modulation of the DP-SSM to ensure the restoration of local textures and edges.

Key Contributions

Fine-Grained Degradation Extractor: A regression-based module that captures continuous, fine-grained degradation features (type and severity) rather than relying on discrete class labels.
Degradation-Aware Prompt State Space Model (DP-SSM): A novel mechanism that dynamically modulates SSM parameters ( $\Delta, B, C$ ) based on degradation priors. This allows the model to physically adapt its evolution dynamics (e.g., switching between low-pass filtering and signal amplification) to specific degradation characteristics.
High-Frequency Enhancement Block (HEB): A complementary module designed to alleviate the low-frequency bias in multi-task learning, ensuring the recovery of fine details.
State-of-the-Art Performance: Extensive experiments demonstrate that DPMambaIR achieves superior performance across seven degradation types in both All-in-One and degradation-specific settings.

Experimental Results

The method was evaluated on a mixed dataset containing seven degradation types (haze, rain, snow, low-light, blur, noise, JPEG) and specific benchmarks for individual tasks.

All-in-One Performance: On the mixed dataset, DPMambaIR achieved a PSNR of 27.69 dB and an SSIM of 0.893, outperforming existing All-in-One methods like AdaIR (27.26 dB) and OneRestore (26.64 dB). It achieved the best performance across all seven sub-tasks, effectively balancing multi-task conflicts.
Degradation-Specific Performance: The model also achieved competitive or superior results on specific tasks compared to dedicated SOTA methods:
- Deraining: 32.30 dB PSNR (surpassing AdaIR by 0.84 dB).
- Low-Light Enhancement: 24.20 dB PSNR.
- Deblurring: 31.42 dB PSNR.
- Dehazing: 30.89 dB PSNR.
Efficiency: DPMambaIR maintains comparable computational overhead (Params: 29.35M, GFLOPs: 153.33) to other SOTA methods while delivering significantly higher performance.
Generalization (OOD): The model demonstrated robust generalization on unseen degradation types (Raindrop and Pixelation) without fine-tuning, achieving top-tier performance. It successfully handled compound degradations (e.g., pixelation + low-light) where other methods failed.
Limitations: The authors note a failure boundary with heavy, opaque raindrops that completely occlude the background, attributing this to the need for generative hallucination (inpainting) rather than restoration cues.

Significance

The paper claims that DPMambaIR represents a significant advancement in All-in-One image restoration by moving beyond static, additive fusion paradigms. By introducing a degradation-aware prompt mechanism that dynamically re-parameterizes the State Space Model, the framework achieves a physical adaptability to diverse degradation dynamics that was previously unattainable. The work underscores the importance of fine-grained degradation modeling and dynamic feature modulation in creating unified, robust, and efficient restoration solutions. The results validate that explicitly modeling degradation characteristics and integrating them into the core system dynamics can resolve the inherent conflicts in multi-task learning, offering a promising direction for future unified image restoration frameworks.