PaperlayerView exact revision on arXiv

Source-linked AI summary

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

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

arXiv:2602.23169v1cs.CV

TL;DR

All-in-one restoration methods remain vulnerable to out-of-distribution degradations, motivating a shared representation that generalizes beyond limited training types. BaryIR learns a Wasserstein barycenter space for degradation-agnostic content alongside orthogonal residual spaces for degradation-specific knowledge, and experiments report strong performance and generalization to unseen degradations. The framework also provides theoretical error bounds, while the authors identify empirical barycenter-weight selection and failure on intense localized rain streaks as limitations.

  • Problem

    Existing all-in-one restoration methods remain vulnerable to out-of-distribution degradations, limiting generalization in real-world scenarios.

  • Method

    BaryIR aligns multisource degraded features in a continuous Wasserstein barycenter space and uses orthogonal residual subspaces to preserve degradation-specific knowledge.

  • Results

    Experiments on synthetic and real-world data report state-of-the-art AIR performance, superior generalization to unseen degradations, and robust generalized features with limited degradation types.

  • Takeaways & Limitations

    The reported disentanglement supports adaptive restoration grounded in shared degradation-agnostic invariance while retaining degradation-specific information.

  • Takeaways & Limitations

    Barycenter weights are currently chosen empirically from each source’s training-sample count, and intense localized rain streaks can remain due to residual-space insensitivity.

Abstract

from arXiv · show

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 (\textit{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.

1 INTRODUCTION

All-in-one image restoration seeks one model for varied degradations, but existing methods struggle with degradation-agnostic features and unseen conditions. BaryIR addresses this by separating shared invariant content from degradation-specific knowledge in a continuous Wasserstein barycenter framework.

  • All-in-one restoration consolidates multiple degradation types in one model, addressing real-world settings where varied and unseen degradations occur.
  • Existing approaches often inject degradation-specific cues but struggle to capture degradation-agnostic features needed for generalization beyond training samples.
  • BaryIR models multisource degraded feature distributions as degradation-specific shifts from an underlying shared distribution whose recovery supports generalization across degradations.
  • BaryIR maps multisource degraded features into a Wasserstein barycenter space while residual subspaces preserve degradation-specific knowledge for adaptive restoration.
  • The framework explicitly constructs orthogonal WB and residual spaces, with the former encoding invariant content and the latter retaining degradation-specific knowledge.
  • An adversarial max-min algorithm learns a continuous barycenter map, while theoretical error bounds provide approximation guarantees for the recovered barycenter distribution.
  • Experiments on synthetic and real-world data report state-of-the-art AIR performance, stronger generalization to unseen degradations, and robustness with limited degradation types.

2 RELATED WORK

Related work uses degradation embeddings and unified representations to handle multiple sources, but existing unified representations often remain limited in source count or structural resolution. BaryIR instead learns a continuous barycenter space that scales to arbitrary source numbers.

  • AIR methods commonly guide unified restoration with degradation embeddings, including contrastive encoders, visual prompts, residual embeddings, and other degradation cues.
  • Unified representation methods align sources in shared latent spaces or use source-agnostic encoders, codebooks, and prototypes to express shared content.
  • Existing unified-representation approaches often target only two sources or use discrete spaces that hinder fine-grained modeling of multisource structures.
  • BaryIR learns a continuous barycenter space that is naturally scalable to an arbitrary number of sources through the optimal-transport barycenter formulation.

3 PRELIMINARIES

The preliminaries formulate optimal transport and Wasserstein barycenters for distributions from multiple sources. BaryIR uses a unified neural barycenter map to transform degraded features into a continuous shared embedding space.

  • 3.1 Optimal Transport: Optimal transport seeks a transport plan between distributions that minimizes a specified transport cost.
  • 3.1 Optimal Transport: The 1-Wasserstein distance results from using Euclidean transport cost, while squared Euclidean cost yields the squared 2-Wasserstein distance.
  • 3.2 Wasserstein Barycenter: A Wasserstein barycenter minimizes a weighted sum of optimal-transport distances from source distributions to a common distribution.
  • 3.2 Wasserstein Barycenter: Empirical source samples can define a continuous barycenter problem through a map that sends samples from each source toward the approximate barycenter space.
  • 3.2 Wasserstein Barycenter: In AIR, the barycenter models a distribution minimizing average Wasserstein distance to degraded feature distributions and encodes source-agnostic content.

4 METHOD

BaryIR learns a Wasserstein barycenter space to encode degradation-agnostic semantics, while residual subspaces preserve degradation-specific knowledge. Its training combines barycenter optimization with contrastive and orthogonality constraints, then integrates both representations for restoration.

  • Wasserstein Barycenter Space: BaryIR transforms multisource degraded features into a Wasserstein barycenter space intended to encode degradation-agnostic invariance.The barycenter map is learned from multisource latent distributions, with a dual reformulation enabling max-min optimization.
  • Wasserstein Barycenter Space: The continuous barycenter map is trained with an adversarial max-min objective and has an error bound relating recovered and true barycenter maps.The optimization alternates between neural-network potentials and the barycenter map; the bound applies under strong-convexity assumptions on transport costs.
  • Residual Feature Learning: Residual embeddings preserve degradation-specific information discarded by barycenter embeddings and are contrasted across degradation subspaces.The residual is defined as r_k = z_k − b_k, while the inter-residual contrastive loss attracts same-subspace embeddings and repels embeddings from other subspaces.
  • Residual Feature Learning: The barycenter-residual orthogonal loss separates degradation-agnostic WB contents from degradation-specific residual knowledge.It penalizes inner products between WB and residual embeddings, enforcing orthogonality for disentangled representation learning.
  • Representation Behavior: WB embeddings remain clustered across degradation types, whereas residual embeddings separate rain, haze, and noise semantics; WB clustering also persists for unseen blur and low-light degradations.These patterns are reported in the t-SNE visualization and support the intended separation of invariant and degradation-specific representations.

5 EXPERIMENTS

BaryIR is evaluated against state-of-the-art all-in-one restoration methods on three- and five-degradation benchmarks, using pixel-wise and perceptual metrics. It delivers consistent quantitative gains and balanced visual restoration across multiple degradations.

  • Experimental setup: The experiments compare BaryIR with state-of-the-art methods on synthetic and real-world all-in-one restoration data.Evaluation includes in-distribution benchmarks, out-of-distribution settings, and mixed degradations, using PSNR, SSIM, LPIPS, FID, NIQE, and PIQE.
  • Three degradations: 0.81 dB average PSNR gain over PromptIR and 0.26 dB over DA-RCOT are reported across three degradations.BaryIR also gains 0.66 dB on deraining over DA-RCOT.
  • Five degradations: 0.52 dB average PSNR gain over MoCE-IR is achieved across five degradations, including dehazing, deraining, denoising, deblurring, and low-light enhancement.The dehazing gain over MoCE-IR is 0.66 dB.
  • Five degradations: BaryIR produces balanced restoration with finer textures and colors under the five-degradation setting.The reported explanation is that barycenters capture common natural-image patterns and reduce overfitting to dominant training data.
  • Plug-in framework: Using PromptIR as backbone, BaryIR gains 1.08 dB and 1.33 dB over original PromptIR under three- and five-degradation settings.It also exceeds DA-RCOT by 0.53 dB and 0.65 dB in the respective settings.

5.2 Generalization to Unseen Degradations

BaryIR is tested on unseen degradation types, levels, and real-world datasets. It consistently outperforms competing methods, including degradation-agnostic approaches, across quantitative and visual evaluations.

  • Evaluation settings: BaryIR is evaluated on unseen JPEG artifact correction, underwater enhancement, severe rain levels, severe noise levels, and real-world haze, rain, and low-light data.The unseen-type tests use BSD500 at QF=10 and UIEB; unseen-level tests include σ = 60 and σ = 75.
  • Unseen degradation types: BaryIR consistently outperforms existing methods across all metrics on unseen degradation types.The reported comparison covers PSNR, SSIM, LPIPS, and FID for JPEG artifact correction and underwater enhancement.
  • Unseen degradation levels: 1.82 dB PSNR gain over MoCE-IR is achieved on Rain100L, while σ = 75 denoising surpasses the best competing method by 2.20 dB.These results concern degradation levels not included in training.
  • Real-world generalization: 2.09 dB, 1.68 dB, and 0.81 dB PSNR gains over second-best methods are reported on O-HAZE, SPANet, and LOL-v2-real, respectively.BaryIR achieves the best performance across all three unseen real-world datasets.
  • Visual results: Visual comparisons show more faithful structures, textures, and colors on unseen degradation types and real-world data.Competing methods are reported to leave degradations or distort structural details.

5.3 Robustness of Generalization Capability to the Number of Training Degradations

The robustness study examines BaryIR with progressively fewer training degradation types. BaryIR retains the strongest OOD performance and the smallest degradation as training coverage decreases.

  • Study design: The study evaluates generalized degradation-agnostic features under four settings with progressively fewer training degradation types.The full setting includes dehazing, deraining, denoising, deblurring, and low-light enhancement.
  • Quantitative results: Methods modeling degradation-agnostic representations outperform All-in-One models without agnostic modeling on OOD degradation types.The comparison identifies DiffUIR, MoCE-IR, and BaryIR as the agnostic-modeling methods.
  • Visual results: Visual examples evaluate five-degradation models on unseen real-world O-HAZE and SPANet data and on mixed haze-rain or blur-noise inputs.These examples complement the numerical robustness comparison.
  • Quantitative results: BaryIR achieves the best PSNR and LPIPS scores and the smallest performance drop as training degradation types decrease.The reported advantage is associated with WB embeddings capturing intrinsic invariance and residual embeddings retaining adaptive degradation-specific knowledge.

5.4 Handling Images with Mixed Degradations

BaryIR is evaluated on synthetic and real-world images containing mixed degradations. It consistently outperforms competing methods while removing multiple artifacts and preserving scene details.

  • Evaluation setup: The mixed-degradation evaluation covers synthetic CDD-11 and 49 real-world images with blur-noise or rain-haze combinations.The pretrained five-degradation BaryIR model is used for the real-world assessment, with NIQE and PIQE as no-reference metrics.
  • Results: BaryIR consistently outperforms other methods with quantitative and qualitative improvements on mixed-degradation images.Results are reported in Table 7 and Figure 8.
  • Visual results: BaryIR removes rain streaks and haze, eliminates blur without artifacts, and preserves scene details.The reported visual comparison emphasizes high-quality restoration relative to competing approaches.

5.5 Ablation Studies

Ablations show that WB embeddings provide the main generalization benefit, while residual embeddings and complementary losses refine degradation-specific restoration. The framework remains effective across transport-map architectures, weighting strategies, and batch sizes.

  • Embedding components: WB embeddings alone improve unseen-degradation generalization, while combining WB and residual embeddings yields optimal restoration performance.Residuals act as degradation-specific cues that refine WB-based common structures, supported by the barycenter-residual orthogonality constraint.
  • Loss functions: The multisource Wasserstein barycenter loss is the fundamental contributor to OOD generalization, while contrastive and orthogonal losses provide synergistic support.Using the latter two losses without the barycenter loss produces only marginal gains.
  • Transport-map architectures: All tested transport-map architectures improve over baseline, with transformer blocks using MDTA and GDFN achieving the best performance.The architectures preserve the same output dimensionality for alignment with the restoration backbone.
  • Barycenter weights: Proportion-based barycenter weights outperform uniform weights across all tasks, with higher PSNR/SSIM and lower LPIPS/FID.The weighting scheme follows the number of training samples per source.
  • Training batch size: PSNR stabilizes beyond batch size 8 while LPIPS decreases and levels off across both in-distribution and OOD scenarios.The results indicate that a moderate batch size is sufficient for robust barycenter optimization.

5.6 Discussion and Model Analysis

Model analysis finds moderate computational and training-cost overhead, stable adversarial optimization, and a validation-optimal trade-off coefficient of α = 0.05.

  • Computational efficiency: BaryIR adds 8.3M parameters and 64G FLOPs while retaining 0.16s inference time, remaining more efficient than DA-CLIP.The reported inference time is 0.13s backbone time plus 0.03s additional time.
  • Cost-performance trade-off: The OT-based framework increases training cost by approximately +9.8M parameters, while training and inference time increase only marginally.The added cost mainly comes from the transport-map and potential networks.
  • Optimization behavior: The barycenter-map and potential losses converge stably in an adversarial manner during training.The curves are shown for All-in-One settings with three and five degradations.
  • Hyperparameter sensitivity: α = 0.05 yields the best validation performance and is used in all settings.The sensitivity analysis trains on 90% of the five-degradation data and validates on the remaining 10%.

5.7 Analysis of Representative Failure Cases

Failure cases arise under severe mixed degradations: residual representations miss intense rain streaks, while barycenter representations can preserve artifacts and smooth genuine textures when scene details are obscured.

  • Rain and haze: Under rain plus haze, the model preserves high-intensity rain streaks because the residual space is insufficiently sensitive to localized intensity outliers.The residual decoupling mechanism treats these streaks as high-salience image structures.
  • Complex OOD degradations: Under complex OOD degradations such as underwater scattering and JPEG artifacts, BaryIR corrects color but fails to restore sharp textures or suppress artifacts.Reliable scene details are heavily obscured in these mixed degradations.
  • Representation trade-off: The barycenter representation emphasizes shared global structures in extreme mixed degradations, trading texture clarity and artifact suppression for color consistency.Systematic JPEG compression can become indistinguishable from common structures, causing artifacts to persist while genuine textures are smoothed.

6 CONCLUSION, LIMITATION, AND FUTURE WORK

BaryIR disentangles degradation-agnostic and degradation-specific representations through a continuous Wasserstein barycenter space and orthogonal residual subspaces. The paper reports strong OOD robustness but identifies empirical barycenter weighting as a limitation and future target for adaptation.

  • Conclusion: BaryIR learns a continuous WB space for invariant contents and orthogonal residual subspaces for dynamic degradation-specific knowledge.Adversarial max-min optimization produces a smooth, geometrically consistent barycenter mapping.
  • Conclusion: Experiments report state-of-the-art restoration performance, superior robustness to OOD degradations, and stable performance with limited training degradation types.The conclusion frames barycenter-based feature disentanglement as a principled approach for generalized restoration.
  • Limitations and future work: Barycenter weights are currently determined empirically from the number of training samples for each source.Future work proposes theoretical justification and more adaptive weight-selection strategies.
Loading 2602.23169v1…