Learning Continuous Wasserstein Barycenter Space for Generalized All-in-One Image Restoration
This work aims to improve the generalization of all-in-one image restoration models for real-world applications by making them more robust to unseen degradations, which is an incremental improvement for practitioners in computer vision.
This paper addresses the challenge of out-of-distribution degradations in all-in-one image restoration by proposing BaryIR, a framework that aligns multisource degraded features in a Wasserstein barycenter (WB) space. This approach disentangles degradation-agnostic invariant content from degradation-specific knowledge, leading to competitive performance against state-of-the-art methods and improved generalization to unseen degradation types and levels.
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.