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

2026-02-26

AI Summary

This work addresses the limited generalization of existing unified image restoration models under out-of-distribution degradation scenarios. To this end, we propose BaryIR, a novel framework that, for the first time, introduces the Wasserstein barycenter space into image restoration. By aligning multi-source degraded features in this space, BaryIR explicitly decouples degradation-agnostic shared content from degradation-specific information and constructs orthogonal subspaces to enable adaptive restoration. The method integrates Wasserstein barycenter modeling, residual subspace contrastive learning, and feature orthogonality-based disentanglement, significantly enhancing generalization to unseen and mixed degradations. Extensive experiments demonstrate that BaryIR achieves state-of-the-art performance across diverse degradation types and exhibits remarkable robustness in real-world complex scenarios.

Technology Category

• Computer Vision: Multi-modal Vision
• Machine Learning: Multimodal Learning
• Intelligent Robots: Multimodal Perception & Sensor Fusion

Application Category

• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Graph Algorithms and Modeling for the Web: Representation, reconstruction, and subgraph or motif discovery in Web-related graphs
• Systems and Infrastructure for Web, Mobile and WoT: Applied ML and AI for Web-based mobile applications

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.

Problem

• image restoration
• out-of-distribution generalization
• degradation-agnostic
• Wasserstein barycenter
• all-in-one

Innovation

• Wasserstein Barycenter
• Degradation-Agnostic Representation
• Feature Disentanglement
• All-in-One Image Restoration
• Out-of-Distribution Generalization