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

Bingnan Wang; Bin Qin; Jiangmeng Li; Fanjiang Xu; Fuchun Sun; Hui Xiong

cs.CVMar 4, 2026

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

Authors: Bingnan Wang, Bin Qin, Jiangmeng Li, Fanjiang Xu, Fuchun Sun, Hui Xiong

Organizations: University of Chinese Academy of Sciences, Beijing 100190, China · Institute of Software, Chinese Academy of Science, Beijing 100190, China · Department of Computer Science and Technology, Tsinghua University, Beijing National Research Center for Information Science and Technology (BNRist), Beijing 100084, China · Artificial Intelligence Thrust, The Hong Kong University of Science and Technology (Guangzhou), Guangzhou 511458, China

Abstract

Image restoration represents a promising approach for addressing the inherent defects of image content distortion. Standard image restoration approaches suffer from high storage cost and the requirement towards the known degradation pattern, including type and degree, which can barely be satisfied in dynamic practical scenarios. In contrast, all-in-one image restoration (AiOIR) eliminates multiple degradations within a unified model to circumvent the aforementioned issues. However, according to our causal analysis, we disclose that two significant defects still exacerbate the effectiveness and generalization of AiOIR models: 1) the spurious correlation between non-degradation semantic features and degradation patterns; 2) the biased estimation of degradation patterns. To obtain the true causation between degraded images and restored images, we propose Causal-deconfounding Wavelet-disentangled Prompt Network (CWP-Net) to perform effective AiOIR. CWP-Net introduces two modules for decoupling, i.e., wavelet attention module of encoder and wavelet attention module of decoder. These modules explicitly disentangle the degradation and semantic features to tackle the issue of spurious correlation. To address the issue stemming from the biased estimation of degradation patterns, CWP-Net leverages a wavelet prompt block to generate the alternative variable for causal deconfounding. Extensive experiments on two all-in-one settings prove the effectiveness and superior performance of our proposed CWP-Net over the state-of-the-art AiOIR methods.

Figures & tables

Fig. 1: (a) The illustration of spurious correlation. We present the statistical analysis of four selected categories in the training data of five tasks. Detailed dataset information is provided in Sec. V-A . (b) The prediction accuracy of five degradation patterns on imbalanced test data and balanced test data. ResNet-34 [ 20 ] , VGG-16 [ 21 ] , and DA-CLIP [ 15 ] are utilized as degradation classifiers. The training and testing data used in conventional AiOIR models are considered imbalanced. Balanced means that the task data is not limited to the scenarios present in the training set, but rather it is rich and comprehensive. All classifiers are trained on imbalanced training data. The balanced test data is constructed with details outlined in Sec. V-D .

Fig. 2: The SCM graph of AiOIR. (a) A unified causal framework for existing ideal AiOIR methods. (b) Two challenges of practical AiOIR via causal lens. The dashed line $\dashleftarrow\dashrightarrow$ represents the spurious correlation, and the grey node denotes the unobservability of the variable. (c) Causal deconfounding through alternative variable $P$ . The red path and the symbol “ ${\color[rgb]{1,0,0}\times}$ ” indicate the use of the wavelet attention module to eliminate the spurious correlation between $C$ and $T$ , allowing the restoration network to focus on degraded regions rather than semantic features.

Fig. 3: The pipeline of our causal analysis and the corresponding causality-guided method, i.e., CWP-Net.

Fig. 4: (a) The distortion severity of five degradation patterns on each wavelet subband. The “PSNR” is calculated with a higher value indicating milder distortion. Here, we roughly consider items with a PSNR above 30 dB (green dashed line) as distortion-mild subbands, and those below 30 dB as distortion-severe subbands. (b) Experiment results of introducing prompt components to distortion-mild (e.g., “HH” and “HL” for derain task) and distortion-severe (e.g., “LH” and “LL” for derain task) subbands of features in AiOIR model.

Fig. 5: Architecture of our proposed CWP-Net. (a) provides an overview of the U-shaped network for AiOIR, in which WAE and WAD are symmetrically placed at each level of the encoder and decoder, and WPB is inserted into the skip connection. (b) WAE and WAD extract attention maps of different wavelet subbands in both spatial and channel dimensions. (c) WPB generates weighted prompted wavelet subbands for backdoor adjustment based on the low-frequency wavelet attention map. $Z_{*}$ and $P_{*}$ denote the wavelet subbands before and after prompt-based modulation, respectively.

Fig. 6: Visualization of the attention maps of the low-frequency subband in WAE. The first row depicts the input degraded image, and the second row illustrates the corresponding attention map. It provides evidence that the map focuses on the degraded regions and decouples the degradation features from the semantic features.

Method	Dehazing		Deraining		Denoising						Average
	Dehazing		Deraining		$\sigma=15$		$\sigma=25$		$\sigma=50$		Average
	PSNR	SSIM	PSNR	SSIM	PSNR	SSIM	PSNR	SSIM	PSNR	SSIM	PSNR	SSIM
BRDNet [ 49 ]	23.23	0.895	27.42	0.895	32.26	0.898	29.76	0.836	26.34	0.836	27.80	0.843
LPNet [ 50 ]	20.84	0.828	24.88	0.784	26.47	0.778	24.77	0.748	21.26	0.552	23.64	0.738
FDGAN [ 51 ]	24.71	0.929	29.89	0.933	30.25	0.910	28.81	0.868	26.43	0.776	28.02	0.883
MPRNet [ 52 ]	25.28	0.955	33.57	0.954	33.54	0.927	30.89	0.880	27.56	0.779	30.17	0.899

TABLE I: Quantitative results of different methods under the five-pattern setting. PSNR(

\uparrow

) and SSIM(

\uparrow

) are reported. The best and second-best results are marked in bold and underlined , respectively.

Method	Deraining		Dehazing		Low-Light		Deblurrring		Denoising		Average
Method	PSNR	SSIM	PSNR	SSIM	PSNR	SSIM	PSNR	SSIM	PSNR	SSIM	PSNR	SSIM
NAFNet [ 61 ]	35.56	0.967	25.23	0.939	20.49	0.809	26.53	0.808	31.02	0.883	27.76	0.886
HINet [ 62 ]	35.67	0.969	24.74	0.937	19.47	0.800	26.12	0.788	31.00	0.881	27.40	0.875
MPRNet [ 52 ]	38.16	0.981	24.27	0.937	20.84	0.824	26.87	0.823	31.35	0.889	28.27	0.890
DGUNet [ 63 ]	36.62	0.971	24.78	0.940	21.87	0.823	27.25	0.837	31.10	0.883	28.32	0.891
MIRNetv2 [ 64 ]	33.89	0.954	24.03	0.927	21.52	0.815	26.30	0.799	30.97	0.881	27.34	0.875

TABLE II: Quantitative results of different methods under the seven-pattern setting. PSNR(

\uparrow

) and SSIM(

\uparrow

) are reported. The best and second-best results are marked in bold and underlined , respectively.

Fig. 7: Visual comparisons of all-in-one methods on image dehazing. PSNR of the input image and images recovered by different approaches are listed below the corresponding samples. Our method restores building colors and structural textures that are closest to the reference. Zoom in for the best view.

Fig. 8: Visual comparisons of all-in-one methods on image deraining. PSNR of the input image and images recovered by different approaches are listed below the corresponding samples. The restored results of our method exhibit fewer rain streaks and a cleaner background. Zoom in for the best view.

Fig. 9: Visual comparisons of all-in-one methods on image denoising. PSNR of the input image and images recovered by different approaches are listed below the corresponding samples. Our method is able to maintain fine detail while removing noise, rather than producing overly smooth results. Zoom in for the best view.

Method	Dehazing		Deraining		Denoising						Average
	Dehazing		Deraining		$\sigma=15$		$\sigma=25$		$\sigma=50$		Average
	PSNR	SSIM	PSNR	SSIM	PSNR	SSIM	PSNR	SSIM	PSNR	SSIM	PSNR	SSIM
AirNet	19.46	0.892	35.29	0.968	36.16	0.953	33.69	0.925	30.32	0.863	30.98	0.920
PromptIR	19.32	0.883	37.35	0.978	36.18	0.953	33.72	0.925	30.35	0.864	31.38	0.921
CWP-Net	21.34	0.916	38.47	0.984	36.34	0.955	33.91	0.928	30.58	0.869	32.13	0.931

TABLE III: Additional quantitative results on the balanced test set under the three-pattern setting. PSNR and SSIM are reported.

Model	Baseline	WAE	WAD	WPB	Derain	Dehaze	Deblur	Low-Light	Denoise	Average	Param(M)	FLOPs(G)
$(a)$	✓	✗	✗	✗	35.03	30.19	25.87	21.73	30.77	28.72	14.75	128.93
$(b)$	✓	✓	✗	✗	36.52	31.02	26.02	21.64	30.92	29.22	14.72	128.63
$(c)$	✓	✓	✓	✗	36.53	31.15	26.03	22.04	30.90	29.33	14.74	125.61
$(d)$	✓	✓	✗	✓	36.87	31.02	26.17	22.41	31.03	29.50	15.47	133.56
$(e)$	✓	✗	✓	✓	36.84	30.92	26.06	22.35	31.00	29.43	15.53	130.85
$(f)$	✓	✓	✓	✓	37.48	31.21	26.49	22.92	31.06	29.83	15.50	130.55

TABLE IV: Ablation studies of the proposed individual modules on overall performance. PSNR of different models is reported.

Method	Derain	Dehaze	Deblur	Low-light	Denoise	Average	Params (M)	FLOPs (G)
Concat $\&$ Self-Attention	37.22	31.13	26.51	22.24	31.06	29.63	16.25	135.52
Cross Attention	36.93	31.48	26.25	22.51	30.97	29.63	15.64	131.53
SFT (ours)	37.48	31.21	26.49	22.92	31.06	29.83	15.50	130.55

TABLE V: Various prompt-feature interaction modules for CWP-Net. PSNR of different models is reported. FLOPs are measured on the patch size of

256\times 256\times 3

Number of clusters	3	5	7
Average PSNR	29.77	29.83	29.80

TABLE VI: Ablation studies of numbers of clusters.

Fig. 10: The case study on the generalization performance of CWP-Net under distribution shift.

Method	Params(M)	FLOPs(G)	Inference time(ms)
BRDNet	1.11	278	62
LPNet	10.07	247.82	41.61
FDGAN	11.80	69.16	17.94
MPRNet	20.13	1707.36	84.21
DL	2.09	0.13	6.57
AirNet	5.77	301.27	109.83

TABLE VII: Efficiency comparisons of existing methods under the five-pattern setting. The inference time per image is calculated based on the average result of 100 forward propagations.

Method	Params(M)	FLOPs(G)	Inference time(ms)
NAFNet	17.11	15.96	18.90
HINet	88.67	170.49	18.20
MPRNet	15.74	850.04	60.52
DGUNet	17.33	866.17	105.10
MIRNetv2	5.86	140.92	46.41
SwinIR	0.91	58.72	122.46

TABLE VIII: Efficiency comparisons of existing methods under the seven-pattern setting. The inference time per image is calculated based on the average result of 100 forward propagations.

Fig. 11: SCM graph of the drug study.

Fig. 12: Subgraphs of G used in the derivation of causal effects.

Fig. 13: The output results of the FCI algorithm. Under the semantics of FCI output, a directed arrow

A\rightarrow B

indicates that

A

is a cause of

B

, a bidirected arrow

A\leftrightarrow B

indicates that there is a latent confounder of

A

and

B

, and a circle endpoint

A\ \text{o–o}\ B

indicates that no set separates

A

and

B

$e$	Deraining		Dehazing
$e$	Acc(%, $\uparrow$ )	FID( $\downarrow$ )	Acc(%, $\uparrow$ )	FID( $\downarrow$ )
1	43.25	147.62	52.38	133.54
2	27.84	172.57	30.86	154.28
3	29.59	156.78	38.74	148.97

TABLE X: The accuracy of classifiers and the Fréchet-type distribution distance of generative models under different environments

e

All-in-One Image Restoration

All-in-one image restoration trains unified models to remove diverse and co-occurring degradations while preserving scene content.

See all 88 papers

DPC-Net: Dual-Prior Collaborative Network for All-in-One Image Restoration
Aug 20, 2026Zhaokun He, Kangbiao Shi, Axi Niu +4Geometry-Guided Image Restoration All-in-One Image Restoration
Northwestern Polytechnical University · Singapore Management University · Xi’an University of Architecture and Technology
What to Remove, What to Preserve: Dual-Ambiguity Rectification for All-in-One Image Restoration
Jul 30, 2026Cencen Liu, Wen Yin, Dongyang Zhang +6Representation Learning Representation Disentanglement
University of Electronic Science and Technology of China · University of Electronic Science and Technology of China, Chengdu, Sichuan, China · Jiigan Technology +1
Universal Image Restoration via Internalized Chain-of-Thought Reasoning
Jun 16, 2026Yu Guo, Zhengru Fang, Shengfeng He +4All-in-One Image Restoration Image Restoration
Hong Kong JC Lab of Smart City and the Department of Computer Science, City University of Hong Kong · School of Computing and Information Systems, Singapore Management University · Computer Science and Information Engineering, National Taiwan University

Aug 20, 2026J/K move · Enter open · S save

DPC-Net: Dual-Prior Collaborative Network for All-in-One Image Restoration

Zhaokun He, Kangbiao Shi, Axi Niu, Jian Jin+3

Northwestern Polytechnical University · Singapore Management University · Xi’an University of Architecture and Technology

Geometry-Guided Image Restoration All-in-One Image Restoration Image Restoration

Open paper AskPDF arXiv

All-in-One Image Restoration (AiOIR) aims to handle diverse degradations within a unified model. However, existing methods often overlook image semantics in degradation modeling and lack low-level visual priors during reconstruction, leading to structural distortions and semantic inconsistencies. To address these issues, we propose a novel Dual-Prior Collaborative Network (DPC-Net), which achieves high-quality restoration by jointly exploiting degradation-semantic coupled priors and low-level visual priors. Specifically, degraded images are fed into a Degradation-Aware Network (DAN) to extract degradation-semantic coupled features. To this end, a Vision-Language Model (VLM) supervises DAN by constraining its features distribution, introducing image semantics into the encoding of degradation patterns. A Degradation-Semantic Modulation Module (DSMM) further translates this guidance into degradation-semantic coupling and propagates coupled representations to the decoder. During decoding, knowledge bases provide low-level visual priors, and the Dual-Prior Collaborative Reconstruction Module (DPCR) integrates dual-prior information to guide degradation removal while preserving structure and semantics, producing high-fidelity restored images. Extensive experiments on multiple restoration benchmarks demonstrate that DPC-Net achieves superior performance against state-of-the-art AiOIR methods.

Abstract

All-in-One Image Restoration

Explore similar work

DPC-Net: Dual-Prior Collaborative Network for All-in-One Image Restoration

What to Remove, What to Preserve: Dual-Ambiguity Rectification for All-in-One Image Restoration

Universal Image Restoration via Internalized Chain-of-Thought Reasoning