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

TL;DR

This paper introduces BaryIR, a novel framework that aligns degraded image features in Wasserstein barycenter space to improve generalization and robustness in all-in-one image restoration across various degradations.

Contribution

BaryIR uniquely models a degradation-agnostic distribution using Wasserstein barycenters and disentangles shared content from degradation-specific features for better generalization.

Findings

BaryIR outperforms state-of-the-art methods on multiple benchmarks.

It generalizes well to unseen degradations and real-world data.

It maintains robustness even with limited training degradation types.

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