All-in-One Image Restoration via Causal-Deconfounding Wavelet-Disentangled Prompt Network

This paper proposes CWP-Net, a novel all-in-one image restoration framework that utilizes causal deconfounding and wavelet-disentangled prompts to eliminate spurious correlations and biased degradation estimation, thereby achieving superior performance over state-of-the-art methods.

The Gist

Imagine you have a photo album, but every picture is ruined in a different way. Some are blurry, some are covered in rain, some are foggy, and some are just grainy with noise.

The Old Way (The "Specialist" Problem):
Traditionally, to fix these photos, you needed a different "doctor" for each problem. You'd hire a "blur-remover," a "rain-eraser," and a "noise-cleaner."

• The Downside: This is expensive (you need to store all these different doctors) and slow. Worse, you have to tell the doctor exactly what's wrong with the photo before they start. If you hand a "rain-eraser" a foggy photo, they might get confused or make it worse.

The New Idea (The "All-in-One" Doctor):
Scientists wanted to build one "Super Doctor" (called All-in-One Image Restoration) that could fix any problem without being told what it is. They thought, "Let's just train one giant brain to recognize rain, blur, and fog all at once."

The Hidden Trap (The "Spurious Correlation"):
The researchers in this paper discovered that these "Super Doctors" were cheating. They weren't actually learning how to remove rain; they were learning bad habits.

• The Analogy: Imagine a doctor who notices that every time they see a photo of a dog, there is rain in the picture. But in photos of buildings, there is never rain.
• The Mistake: The doctor starts thinking, "Oh, if I see a dog, I must remove rain." But what if you show them a picture of a dog on a sunny day? The doctor gets confused and tries to "remove rain" from the dog's fur, ruining the picture.
• The Real Issue: The AI was linking the subject of the photo (the dog, the building) with the problem (rain, fog) because the training data was unbalanced. It learned a fake connection (spurious correlation) instead of the real cause.

The Solution: CWP-Net (The "Frequency Detective")
The authors built a new system called CWP-Net to fix this. They used a clever trick involving Wavelets (think of this as a special pair of glasses that lets you see a photo in layers of "frequency" rather than just a flat picture).

Here is how CWP-Net works, using simple metaphors:

1. The "Frequency Glasses" (Wavelet Attention)

Instead of looking at the whole picture at once, CWP-Net puts on "frequency glasses."

• Low Frequency: This is the "skeleton" of the image (big shapes, colors, the dog's body).
• High Frequency: This is the "texture" and "noise" (the rain streaks, the grain, the blur).
• The Trick: The AI learns to look only at the high-frequency "noise" layers to figure out what's wrong. It ignores the "skeleton" (the dog or the building).
• Result: It stops guessing based on the subject. It sees the rain streaks regardless of whether they are on a dog or a building. This breaks the bad habit (spurious correlation).

2. The "Smart Hint System" (Wavelet Prompt Block)

Sometimes, the AI still gets confused about how bad the damage is.

• The Problem: Imagine trying to clean a window. If it's lightly misted, you wipe gently. If it's covered in mud, you scrub hard. The AI needs to know how to clean.
• The Solution: CWP-Net has a "Smart Hint System." It looks at the damage and generates a specific "recipe" (a prompt) for the cleaning process.
• How it works: It creates a custom "cleaning tool" for that specific photo. If the rain is heavy, the tool is a heavy-duty scrubber. If it's light, it's a soft cloth. This ensures the AI doesn't guess blindly; it adapts its strategy based on the actual damage, not the background.

3. The "Causal Detective"

The whole system is built on Causal Logic.

• Old AI: "I see a dog + rain streaks -> I remove rain." (Correlation)
• CWP-Net: "I see rain streaks (regardless of the dog) -> I remove rain." (Causation)

Why is this a big deal?

• One Model to Rule Them All: You only need to store one model, saving massive amounts of space.
• Works in the Real World: Real life is messy. You don't always know if a photo is blurry or foggy. This AI doesn't need to be told; it figures it out by looking at the "texture" of the damage, not the "story" of the picture.
• Better Results: Because it stopped cheating (relying on bad habits), it actually restores photos better, keeping the details sharp and the colors true.

In a Nutshell:
The paper teaches computers to stop looking at what is in the picture (the dog, the car) to figure out what's wrong, and instead teaches them to look strictly at how the picture is damaged (the streaks, the blur). By using "frequency glasses" and "smart hints," they built a universal photo fixer that actually works in the messy, unpredictable real world.

Technical Summary

1. Problem Statement

All-in-One Image Restoration (AiOIR) aims to remove multiple types of image degradations (e.g., noise, rain, haze, blur, low-light) using a single unified model, addressing the high storage costs and deployment complexity of task-specific models. However, existing AiOIR methods suffer from two critical defects that hinder their generalization and effectiveness:

1. Spurious Correlation: Degradation classifiers often learn false correlations between semantic features (e.g., "animals" or "furniture") and degradation patterns (e.g., rain or haze) due to dataset biases. For instance, if rain streaks frequently appear in animal scenes in the training data, the model may incorrectly associate "animals" with "rain," leading to failure when these correlations are absent in real-world scenarios.
2. Biased Estimation of Degradation Patterns: Existing methods rely on estimating the degradation type and degree to guide restoration. However, on imbalanced data, these estimators fail to predict degradation patterns accurately, especially for rare or complex cases, leading to suboptimal restoration.

The authors argue that these issues prevent the model from learning the true causal relationship between the degraded image ($X$) and the restored image ($Y$), instead learning a biased conditional probability $P(Y|X)$ rather than the causal effect $P(Y|do(X))$.

2. Methodology: CWP-Net

The proposed Causal-deconfounding Wavelet-disentangled Prompt Network (CWP-Net) utilizes a Structural Causal Model (SCM) framework to address these defects. The methodology consists of three core components:

A. Theoretical Framework (SCM)

The authors model the AiOIR process where:

• $X$: Degraded Image
• $Y$: Restored Image
• $C$: Semantic Features (Content)
• $T$: Degradation Pattern (Unobservable/Noisy)
• $D$: Degradation Features
• Goal: Block the confounding path $X \leftarrow D \leftarrow T \rightarrow Y$ and the spurious correlation $C \sim T$.

B. Wavelet Attention Modules (WAE & WAD)

To solve the Spurious Correlation challenge, the network explicitly disentangles degradation features from semantic features.

• Mechanism: The network employs Discrete Wavelet Transform (DWT) to decompose features into frequency subbands (LL, LH, HL, HH).
• Encoder (WAE) & Decoder (WAD): Attention maps are generated for each subband. Crucially, the low-frequency (LL) attention map is used as the degradation representation.
• Rationale: Degradation artifacts (like rain streaks or haze) are more prominent in specific frequency bands and are semantically independent. By focusing attention on these frequency-specific regions, the model suppresses semantic content ($C$) and isolates degradation features ($D$), effectively blocking the spurious correlation path.

C. Wavelet Prompt Block (WPB)

To solve the Biased Estimation challenge, the authors introduce an alternative variable for causal deconfounding via Backdoor Adjustment.

• Concept: Since the true degradation pattern $T$ is unobservable and hard to estimate, the model introduces a variable $P$ (Prompted Wavelet Subbands) such that $T \rightarrow P \rightarrow Y$.
• Components:
  1. Degradation-based Weight Estimator (DWE): Uses the clean degradation representation (from WAE) to estimate the probability distribution of different degradation patterns via K-Means clustering. It outputs weights ($\omega$) for specific wavelet subbands.
  2. Prompt-guided Weighted Spatial Feature Transform (PWSFT): Generates learnable prompts conditioned on the input. These prompts interact with wavelet subbands using Spatial Feature Transform (SFT) to modulate features.
• Operation: The WPB dynamically selects and weights specific wavelet subbands (e.g., emphasizing "LH" for rain) to serve as the adjustment variable $P$, allowing the model to approximate the causal effect $P(Y|do(X))$ without needing perfect knowledge of $T$.

D. Network Architecture

CWP-Net follows a U-shaped encoder-decoder architecture:

• Encoder: Contains CNN blocks and WAE modules.
• Decoder: Contains WAD modules and CNN blocks.
• Skip Connections: The WPB is inserted into the skip connections to modulate features during the restoration process.
• Loss Function: A combination of Reconstruction Loss ($L_1$) and Frequency Loss (L1 distance in Fourier domain) to ensure both structural and high-frequency detail recovery.

3. Key Contributions

1. Causal Analysis: The paper provides a rigorous causal analysis of AiOIR, identifying spurious correlations and biased estimation as the root causes of poor generalization.
2. Wavelet-Disentangled Design: It proposes the first integration of wavelet transforms into AiOIR specifically for causal deconfounding, using frequency-domain attention to separate semantics from degradations.
3. Backdoor Adjustment via Prompts: It introduces a novel Wavelet Prompt Block that uses an alternative variable (weighted wavelet subbands) to perform backdoor adjustment, bypassing the need for perfect degradation pattern estimation.
4. State-of-the-Art Performance: The method achieves superior results on multiple benchmarks without relying on large-scale pre-trained models (like CLIP or Stable Diffusion) used by competitors.

4. Experimental Results

The authors evaluated CWP-Net on two settings: Five-Pattern (Denoise $\sigma \in \{15,25,50\}$, Derain, Dehaze) and Seven-Pattern (adding Deblur and Low-light).

• Quantitative Performance:
  • Five-Pattern: Outperformed the previous best (Lin et al.) by +0.59 dB average PSNR. Specifically, it improved Dehazing by +1.58 dB and Deraining by +1.10 dB.
  • Seven-Pattern: Surpassed the second-best all-in-one method (IDR) by +2.22 dB average PSNR and outperformed general restoration models (like Restormer) significantly.
• Generalization (Balanced Test Set):
  • On a constructed balanced test set (where semantic-degradation correlations are removed), CWP-Net maintained high performance, whereas methods like AirNet and PromptIR suffered severe performance drops. This confirms CWP-Net successfully learns causal features rather than spurious correlations.
• Efficiency:
  • CWP-Net achieves these gains with 15.5M parameters and 130.55 G FLOPs, offering a faster inference time (34.32 ms) compared to heavy models like Lin et al. (128.51 ms) and PromptIR (82.93 ms).

5. Significance

This paper represents a significant shift in image restoration research by moving from purely data-driven pattern matching to causality-driven learning.

• Robustness: By explicitly addressing the "why" behind model failures (spurious correlations), CWP-Net offers a robust solution for real-world scenarios where training data distributions do not perfectly match test data.
• Methodological Innovation: The fusion of Wavelet Transforms with Prompt Learning and Causal Inference creates a new paradigm for disentangling complex image features.
• Practicality: It demonstrates that high-performance, general-purpose image restoration can be achieved without relying on massive external pre-trained models, making it more deployable in resource-constrained environments.