Learning Continuous Wasserstein Barycenter Space for Generalized All-in-One Image Restoration

Xiaole Tang; Xiaoyi He; Jiayi Xu; Xiang Gu; Jian Sun

cs.CVFeb 26, 2026

Learning Continuous Wasserstein Barycenter Space for Generalized All-in-One Image Restoration

Authors: Xiaole Tang, Xiaoyi He, Jiayi Xu, Xiang Gu, Jian Sun

Organizations: School of Mathematics and Statistics, Xi’an Jiaotong University, Shaanxi, P.R. China.

Abstract

Despite substantial advances in all-in-one image restoration for addressing diverse degradations within a unified model, existing methods remain vulnerable to out-of-distribution degradations, thereby limiting their generalization in real-world scenarios. To tackle the challenge, this work is motivated by the intuition that multisource degraded feature distributions are induced by different degradation-specific shifts from an underlying degradation-agnostic distribution, and recovering such a shared distribution is thus crucial for achieving generalization across degradations. With this insight, we propose BaryIR, a representation learning framework that aligns multisource degraded features in the Wasserstein barycenter (WB) space, which models a degradation-agnostic distribution by minimizing the average of Wasserstein distances to multisource degraded distributions. We further introduce residual subspaces, whose embeddings are mutually contrasted while remaining orthogonal to the WB embeddings. Consequently, BaryIR explicitly decouples two orthogonal spaces: a WB space that encodes the degradation-agnostic invariant contents shared across degradations, and residual subspaces that adaptively preserve the degradation-specific knowledge. This disentanglement mitigates overfitting to in-distribution degradations and enables adaptive restoration grounded on the degradation-agnostic shared invariance. Extensive experiments demonstrate that BaryIR performs competitively against state-of-the-art all-in-one methods. Notably, BaryIR generalizes well to unseen degradations (e.g., types and levels) and shows remarkable robustness in learning generalized features, even when trained on limited degradation types and evaluated on real-world data with mixed degradations.

Figures & tables

Fig. 1 : (a) BaryIR decouples the latent space of multisource degraded images into a barycenter space that captures the degradation-agnostic invariance, and residual subspaces that retain degradation-specific knowledge, enhancing generalization to unseen data. (b) A zero-shot example: BaryIR restores an underwater image without being trained on such degradation.

Fig. 2 : Overview of BaryIR. BaryIR learns a barycenter map to reshape the multisource latent space of degraded images into an inherent degradation-agnostic barycenter space beyond training domain. Residual subspaces are constructed to retain the degradation-specific knowledge, with the residual embeddings contrasted with each other while maintaining orthogonal to the WB embeddings. By integrating the WB and residual embeddings in the decoding layers, BaryIR captures the degradation-agnostic semantics while retaining adaptive degradation-specific knowledge, thereby enabling generalized restoration.

Fig. 3 : t-SNE visualization of WB and residual embeddings with BaryIR trained with three degradations ( i.e., rain, haze and noise). The WB embeddings exhibit degradation-agnostic distribution while the residual embeddings are clearly separated according to the specific degradations. Notably, the WB and residual embeddings remain robust on unseen degradations ( i.e., blur and low light) for capturing degradation-agnostic and degradation-specific semantics.

Panel of Table 4 — TABLE I : The All-in-One comparison of our BaryIR with the state-of-the-art methods on three degradations.

Panel of Figure 5 — Fig. 4 : Visual comparison of generalization results on unseen degradation types , i.e., JPEG artifact correction (row 1) on BSD500 (QF=10) and underwater image enhancement on UIEB (row 2). BaryIR restores high-quality images with more faithful structural contents such as textures and colors.

Fig. 5 : Numerical comparison of robustness of generalization to training degradation numbers. As the degradation number decreases, BaryIR remain superior generalization to the unseen degradation types while achieving the best quantitative performance in comparison with other methods.

Method	BSD500	UIEB
Restormer [ 9 ]	25.60/0.742/0.177/71.24	17.34/0.770/0.300/47.12
PromptIR [ 22 ]	25.71/0.748/0.172/65.33	17.56/0.778/0.285/40.26
DiffUIR [ 32 ]	26.10/0.762/0.175/75.22	17.86/0.784/0.295/36.23
InstructIR [ 29 ]	25.54/0.746/0.185/76.68	17.51/0.780/0.288/43.26
DA-RCOT [ 28 ]	26.02/0.765/ 0.148 / 45.71	18.34/ 0.802 / 0.248 / 30.53
MoCE-IR [ 33 ]	26.42 / 0.768 /0.162/58.96	18.66 /0.800/0.256/33.25

TABLE IV : Generalization performance on unseen degradation types, i.e., JPEG artifact correction on BSD500 (QF=10) and underwater image enhancement on UIEB. The metrics are reported as PSNR(

\uparrow

)/SSIM(

\uparrow

)/LPIPS(

\downarrow

)/FID(

\downarrow

)

Method	Rain100L	Rain100H
Restormer [ 9 ]	28.76/0.901/0.140/63.21	14.50/0.464/0.484/250.2
IR-SDE [ 12 ]	28.49/0.897/0.123/55.21	13.55/0.422/0.465/234.5
PromptIR [ 22 ]	31.82/0.931/0.078/38.41	14.28/0.444/0.472/242.7
DA-CLIP [ 23 ]	32.87/0.944/0.066/35.12	14.40/0.435/0.438/228.6
DiffUIR [ 32 ]	33.20/0.942/0.036/34.64	14.78/0.487/0.442/235.5
InstructIR [ 29 ]	33.65/0.951/0.030/28.24	14.67/0.468/0.460/238.3

TABLE V : The deraining results on unseen rain levels using the five-degradation models. The metrics are reported as PSNR(

\uparrow

)/SSIM(

\uparrow

)/LPIPS(

\downarrow

)/FID(

\downarrow

Method	O-HAZE	SPANet	LOL-v2-real
Restormer [ 9 ]	18.02/0.724/0.345/275.8	34.38/0.917/0.032/43.29	27.12/0.877/0.112/85.45
IR-SDE [ 12 ]	17.85/0.716/0.338/256.3	35.02/0.922/0.029/38.87	23.12/0.801/0.144/122.7
PromptIR [ 22 ]	18.38/0.730/0.336/260.1	35.34/0.938/0.026/33.12	27.65/0.870/0.107/80.1
DA-CLIP [ 23 ]	18.22/0.725/0.323/242.5	35.65/0.942/0.026/26.96	26.46/0.856/0.104/76.45
DiffUIR [ 32 ]	18.75/0.731/0.315/225.4	35.95/0.940/0.022/22.25	26.12/0.861/0.098/78.24
InstructIR [ 29 ]	18.85/0.738/0.308/236.5	36.42/0.946/0.028/30.54	28.02/0.901/0.110/81.85

TABLE VII : Generalization to unseen real-world O-HAZE [ 70 ] and SPANet [ 71 ] datasets with the five-degradation models. The metrics are reported as PSNR(

\uparrow

)/SSIM(

\uparrow

)/LPIPS(

\downarrow

)/FID(

\downarrow

Method	CDD-11 [ 73 ]	Real haze and rain	Real blur and noise
Method	PSNR/SSIM/LPIPS/FID	NIQE ( $\downarrow$ )/PIQE ( $\downarrow$ )	NIQE ( $\downarrow$ )/PIQE ( $\downarrow$ )
Restormer [ 9 ]	26.99/0.864/0.105/23.25	9.62/115.8	8.56/96.42
IR-SDE [ 12 ]	25.48/0.856/0.112/25.21	9.45/112.1	8.75/100.5
PromptIR [ 22 ]	25.90/0.850/0.105/28.54	8.05/102.4	7.22/78.44
DA-CLIP [ 23 ]	25.88/0.855/0.112/23.66	7.72/95.40	7.45/83.25
DiffUIR [ 32 ]	27.35/0.868/0.094/18.51	7.78/99.12	7.21/77.35

TABLE VIII : Quantitative comparison with state-of-the-art methods on mixed-degradation images from CDD-11 [ 73 ] , real-world SPANet [ 71 ] (haze and rain), and Lai [ 74 ] (blur and noise).

Panel of Figure 11 — Fig. 6 : Visual examples of generalization evaluation with five-degradation models on unseen real-world O-HAZE [ 70 ] and SPANet [ 71 ] .

Embedding or loss function components			In-distribution						Out-of-distribution
Original emb. $\bm{z}_{k}$	WB emb. $\bm{b}_{k}$	Residual emb. $\bm{r}_{k}$	SOTS	Rain100L	BSD68 $\sigma$ =25	GoPro	LOL	Average	O-HAZE	SPANet
✓	✗	✗	24.09/0.065	34.81/0.045	30.78/0.095	27.22/0.174	20.41/0.109	27.46/0.098	18.02/0.345	34.38/0.032
✗	✓	✗	30.27/0.015	37.23/0.025	31.05/0.088	28.05/0.155	22.86/0.096	29.89/0.076	22.04/0.278	38.53/0.022
✓	✗	✓	29.40/0.019	36.23/0.027	30.88/0.093	27.40/0.170	21.78/0.105	29.14/0.083	19.84/0.295	36.22/0.029
✗	✓	✓	31.20 / 0.009	38.10 / 0.011	31.43 / 0.086	29.51 / 0.132	23.37 / 0.090	30.72 / 0.066	22.98 / 0.248	39.24 / 0.012
MWB loss $\mathcal{L}_{\mathrm{MWB}}$	IRC loss $\mathcal{L}_{\mathrm{IRC}}$	BRO loss $\mathcal{L}_{\mathrm{BRO}}$	SOTS	Rain100L	BSD68 $\sigma$ =25	GoPro	LOL	Average	O-HAZE	SPANet

TABLE IX : Ablation studies on using different embeddings or loss functions for training . Metrics are reported as PSNR(

\uparrow

)/LPIPS(

\downarrow

Method	Dehazing	Deraining	Denoising			Average
Method	SOTS	Rain100L	BSD68 $\sigma$ =15	BSD68 $\sigma$ =25	BSD68 $\sigma$ =50	Average
uniform $\lambda_{1:K}$	31.27/0.980/0.007/4.655	38.55/0.982/0.012/5.854	34.02/0.932/0.040/24.21	31.42/0.888/0.082/52.56	28.10/0.800/0.166/89.24	32.67/0.916/0.062/35.31
portion-based $\lambda_{1:K}$	31.40 / 0.980 / 0.007 / 4.523	39.02 / 0.984 / 0.008 / 5.739	34.16 / 0.935 / 0.038 / 22.69	31.54 / 0.892 / 0.075 / 40.11	28.25 / 0.802 / 0.158 / 82.63	32.86 / 0.919 / 0.057 / 31.14

TABLE X : The All-in-One three-degradation results with different

\lambda_{1:K}

. The metrics are reported as PSNR(

\uparrow

)/SSIM(

\uparrow

)/LPIPS(

\downarrow

)/FID(

\downarrow

Panel of Figure 14 — Fig. 8 : Visual results to show the effect of embeddings during training . The model with WB embeddings generalize better to the OOD degradations, restoring sharper images with more faithful structural contents.

Fig. 9 : Illustration of the contributions of WB and residual embeddings at the image level during inference. The WB difference maps (R3-R1) primarily capture common image structures and contents. Conversely, the Res. difference maps (R3-R2) are spatially-adaptive, concentrating on the regions/edges which are heavily affected by the degradations, thereby complementing the WB features with degradation-specific detail refinements, such as high-frequency textures and local contrast adjustments.

Architecture for transport map $T_{\theta}$	In-distribution	Out-of-distribution
Architecture for transport map $T_{\theta}$	Average	O-HAZE	SPANet
w/o $T_{\theta}$ (Restormer)	27.46/0.098	18.02/0.345	34.38/0.032
MLP-based mapping	30.04/0.081	22.14/0.268	38.76/0.022
UNet-style mapping	30.40/0.072	22.55/0.264	39.09/0.019
ViT-style mapping	30.48/0.075	22.71/0.255	39.03/0.017
Ours	30.72/0.066	22.98/0.248	39.24/0.012

TABLE XI : Ablation studies on architectures for parameterizing the transport map

T_{\theta}

. Metrics are reported as PSNR(

\uparrow

)/LPIPS(

\downarrow

Fig. 10 : Impact of the training batch size. We report the average results on five-degradation test sets and on OOD O-HAZE+SPANet data.

Method	Restormer	PromptIR	DA-CLIP	DA-RCOT	MoCE-IR	BaryIR
#Param	26.1M	36.3M	174.1M	40.9M	25.4M	26.1M + 8.3M
Flops	118G	158G	118.5G	262G	142G	118G + 64G
Time	0.13s	0.15s	4.59s	0.21s	0.16s	0.13s + 0.03s

TABLE XII : Comparison of the number of parameters, model computational efficiency, and inference time. The flops and inference time are computed on rainy image of size 256

\times

256.

Method	Training cost		Inference cost		In-distribution	Out-of-distribution
Method	#Param	Time (1 epoch)	#Param	Time	Average	O-HAZE	SPANet
Restormer	26.1M	3.2 hours	26.1M	0.13s	27.46/0.098	18.02/0.345	34.38/0.032
BaryIR ${}_{\text{Restormer}}$	35.9M	3.8 hours	34.4M	0.16s	30.72/0.066	22.98/0.248	39.24/0.012
PromptIR	36.3M	3.5 hours	36.3M	0.15s	29.72/0.084	18.38/0.338	35.34/0.026
BaryIR ${}_{\text{PromptIR}}$	46.1M	4.0 hours	44.6M	0.18s	31.05/0.070	23.21/0.245	39.38/0.013

TABLE XIII : Cost-performance trade-off analysis. Metrics are reported as PSNR(

\uparrow

)/LPIPS(

\downarrow

Fig. 11 : The loss curves of the barycenter map and potentials during training. The loss of $T_{\theta}$ is scaled to $[0,1]$ . The loss of $f_{\omega_{1:K}}$ is scaled to $[0,1]$ and then takes the negative.

Fig. 12 : Effect of the trade-off hyperparameter $\alpha$ on validation performance.

Fig. 13 : Some failure cases on severe real-world mixed degradations.

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

May 2, 2026

Degradation-Aware Adaptive Context Gating for Unified Image Restoration

Unified image restoration using a single model often faces task interference due to diverse degradations. To address this, we propose DACG-IR (Degradation-Aware Adaptive Context Gating), which enables explicit perception of degradation characteristics to dynamically modulate feature representations. Our method constructs degradation-aware contextual representations from the input to modulate attention distribution, frequency-domain features, and feature aggregation. Specifically, a lightweight multi-scale degradation-aware module extracts coarse degradation information and generates layer-wise prompts. These prompts guide attention temperature and output gating in encoder and decoder blocks for adaptive feature extraction. Additionally, a spatial-channel dual-gated adaptive fusion mechanism refines encoder features, suppressing noise propagation from shallow to deep layers. This design effectively suppresses degradation-induced noise while preserving informative structures. Experiments show DACG-IR outperforms state-of-the-art methods in single-task, all-in-one, adverse weather removal, and composite degradation settings. Code: https://github.com/HlHomes/DACG-IR-code

Lei He, Jielei Chu, Fengmao Lv +4

the School of Computing and Artificial Intelligence, Southwest Jiaotong University, Chengdu 611756, China · the Engineering Research Center of Sustainable Urban Intelligent Transportation, Ministry of Education, Chengdu, Sichuan 611756, China · School of Computing and Artificial Intelligence, Jiangxi University of Finance and Economics, Nanchang, 330013, China +3

All-in-One Image Restoration Image Restoration Degradation-Aware Image Restoration

Ask PDF

Jul 16, 2026

QuReC: All-in-One Image Restoration with Query-Specific Guidance and Local-Global Response Calibration

All-in-one image restoration aims to recover clean images degraded by multiple corruption types using a single unified model. Existing methods typically rely on image-level prompts or shared guidance to handle diverse degradations. However, such a paradigm becomes inadequate when degradations are spatially heterogeneous or even coexist in mixed forms within a single image. Yet spatially adaptive guidance alone is not sufficient, since accurate restoration also requires each spatial query to reliably aggregate complementary information from local neighborhoods and global contexts. To this end, we propose QuReC, a unified framework for all-in-one image restoration. QuReC consists of a Degradation-Guided Query Reconstruction Module (DQRM) and a Local-Global Response Calibration Module (LGRCM). Specifically, DQRM matches each spatial query against a degradation prototype space to reconstruct a query-specific degradation-aware representation, thereby providing fine-grained spatially adaptive restoration guidance. To further stabilize this query-wise matching process, we introduce a weakly supervised prototype matching learning strategy to improve optimization stability and degradation semantic consistency. Meanwhile, LGRCM performs local-global dual-branch aggregation and calibrates the aggregated responses with learnable priors, improving the reliability of feature aggregation and the coordination between local detail modeling and global context modeling. Extensive experiments demonstrate that QuReC achieves superior performance on multiple all-in-one image restoration benchmarks. The code is released at https://github.com/zhoushen1/QuReC.

Shen Zhou, Jinghui Zhang, Wenbo Huang +7

Southeast University Nanjing, China · Hainan University Hainan, China

All-in-One Image Restoration Image Restoration Degradation-Aware Image Restoration

Ask PDF

Mar 16, 2026

M2IR: Proactive All-in-One Image Restoration via Mamba-style Modulation and Mixture-of-Experts

While Transformer-based architectures have dominated recent advances in all-in-one image restoration, they remain fundamentally reactive: propagating degradations rather than proactively suppressing them. In the absence of explicit suppression mechanisms, degraded signals interfere with feature learning, compelling the decoder to balance artifact removal and detail preservation, thereby increasing model complexity and limiting adaptability. To address these challenges, we propose M2IR, a novel restoration framework that proactively regulates degradation propagation during the encoding stage and efficiently eliminates residual degradations during decoding. Specifically, the Mamba-Style Transformer (MST) block performs pixel-wise selective state modulation to mitigate degradations while preserving structural integrity. In parallel, the Adaptive Degradation Expert Collaboration (ADEC) module utilizes degradation-specific experts guided by a DA-CLIP-driven router and complemented by a shared expert to eliminate residual degradations through targeted and cooperative restoration. By integrating the MST block and ADEC module, M2IR transitions from passive reaction to active degradation control, effectively harnessing learned representations to achieve superior generalization, enhanced adaptability, and refined recovery of fine-grained details across diverse all-in-one image restoration benchmarks. Our source codes are available at https://github.com/Im34v/M2IR.

Shiwei Wang, Yongzhen Wang, Bingwen Hu +3

School of Computer Science and Technology, Nanjing University of Aeronautics and Astronautics, Nanjing 210016, China · School of Computer Science and Technology, Anhui University of Technology, Ma’anshan 243032, China · Shenzhen Ubiquitous Data Enabling Key Laboratory, Tsinghua Shenzhen International Graduate School, Tsinghua University, Shenzhen 518055, China +1

Degradation-Aware Image Restoration All-in-One Image Restoration Image Restoration Mixture of Experts

Ask PDF

Abstract

All-in-One Image Restoration

Explore similar work

Degradation-Aware Adaptive Context Gating for Unified Image Restoration

QuReC: All-in-One Image Restoration with Query-Specific Guidance and Local-Global Response Calibration

M2IR: Proactive All-in-One Image Restoration via Mamba-style Modulation and Mixture-of-Experts