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

arXiv: https://arxiv.org/abs/2602.23169 · HTML: https://mlsanity.com/paper/2602.23169/

## Key facts

- Cited by 15 papers in the map (data through 14 Oct 2026).
- Builds on 56 of its 75 references that are in the map.
- First posted on arXiv 26 Feb 2026.

## Builds on (most cited first, 30 of 56)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [Learning Transferable Visual Models From Natural Language Supervision](https://mlsanity.com/paper/2103.00020/) (2021)
- [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://mlsanity.com/paper/2010.11929/) (2020)
- [Swin Transformer: Hierarchical Vision Transformer using Shifted Windows](https://mlsanity.com/paper/2103.14030/) (2021)
- [The Unreasonable Effectiveness of Deep Features as a Perceptual Metric](https://mlsanity.com/paper/1801.03924/) (2018)
- [SwinIR: Image Restoration Using Swin Transformer](https://mlsanity.com/paper/2108.10257/) (2021)
- [Restormer: Efficient Transformer for High-Resolution Image Restoration](https://mlsanity.com/paper/2111.09881/) (2021)
- [Multi-Stage Progressive Image Restoration](https://mlsanity.com/paper/2102.02808/) (2021)
- [Deep Multi-scale Convolutional Neural Network for Dynamic Scene Deblurring](https://mlsanity.com/paper/1612.02177/) (2016)
- [Deep Retinex Decomposition for Low-Light Enhancement](https://mlsanity.com/paper/1808.04560/) (2018)
- [Simple Baselines for Image Restoration](https://mlsanity.com/paper/2204.04676/) (2022)
- [FFDNet: Toward a Fast and Flexible Solution for CNN based Image Denoising](https://mlsanity.com/paper/1710.04026/) (2017)
- [Deep Laplacian Pyramid Networks for Fast and Accurate Super-Resolution](https://mlsanity.com/paper/1704.03915/) (2017)
- [Benchmarking Single Image Dehazing and Beyond](https://mlsanity.com/paper/1712.04143/) (2017)
- [Unified-IO: A Unified Model for Vision, Language, and Multi-Modal Tasks](https://mlsanity.com/paper/2206.08916/) (2022)
- [Multi-Modal Fusion Transformer for End-to-End Autonomous Driving](https://mlsanity.com/paper/2104.09224/) (2021)
- [Deep Joint Rain Detection and Removal from a Single Image](https://mlsanity.com/paper/1609.07769/) (2016)
- [ULIP: Learning a Unified Representation of Language, Images, and Point Clouds for 3D Understanding](https://mlsanity.com/paper/2212.05171/) (2022)
- [Deep Continuous Fusion for Multi-Sensor 3D Object Detection](https://mlsanity.com/paper/2012.10992/) (2020)
- [An Underwater Image Enhancement Benchmark Dataset and Beyond](https://mlsanity.com/paper/1901.05495/) (2019)
- [TransWeather: Transformer-based Restoration of Images Degraded by Adverse Weather Conditions](https://mlsanity.com/paper/2111.14813/) (2021)
- [Spatial Attentive Single-Image Deraining with a High Quality Real Rain Dataset](https://mlsanity.com/paper/1904.01538/) (2019)
- [Image Restoration with Mean-Reverting Stochastic Differential Equations](https://mlsanity.com/paper/2301.11699/) (2023)
- [SpeechT5: Unified-Modal Encoder-Decoder Pre-Training for Spoken Language Processing](https://mlsanity.com/paper/2110.07205/) (2021)
- [InstructIR: High-Quality Image Restoration Following Human Instructions](https://mlsanity.com/paper/2401.16468/) (2024)
- [Fast and Accurate Image Super-Resolution with Deep Laplacian Pyramid Networks](https://mlsanity.com/paper/1710.01992/) (2017)
- [PromptIR: Prompting for All-in-One Blind Image Restoration](https://mlsanity.com/paper/2306.13090/) (2023)
- [AdaIR: Adaptive All-in-One Image Restoration via Frequency Mining and Modulation](https://mlsanity.com/paper/2403.14614/) (2024)
- [Controlling Vision-Language Models for Multi-Task Image Restoration](https://mlsanity.com/paper/2310.01018/) (2023)

## Built on by (most cited first, 15 of 15)

- [Learning Adaptive Dynamical Features via Multi-τ Liquid-Mamba for All-in-one Image Restoration](https://mlsanity.com/paper/2606.22801/) (2026)
- [CoRE-UIR: Prior-guided common and residual experts for efficient all-in-one remote sensing image restoration](https://mlsanity.com/paper/2607.27898/) (2026)
- [TeX-1500: A Paired Real-World LWIR Hyperspectral Dataset and Benchmark for Temperature-Emissivity-Texture Decomposition](https://mlsanity.com/paper/2606.03806/) (2026)
- [Causal-AgentIR: Self-Evolving Causal Memory for Adaptive Image Restoration Agents](https://mlsanity.com/paper/2607.21125/) (2026)
- [Binding Multiple Modalities via Multimodal Wasserstein Barycenter](https://mlsanity.com/paper/2609.33800/) (2026)
- [M2IR: Proactive All-in-One Image Restoration via Mamba-style Modulation and Mixture-of-Experts](https://mlsanity.com/paper/2603.14816/) (2026)
- [EvoIR-Agent: Self-Evolving Image Restoration Agentic System via Experience-Driven Learning](https://mlsanity.com/paper/2605.22208/) (2026)
- [DVANet: Degradation-aware Visual-prior Alignment Network for Image Restoration](https://mlsanity.com/paper/2606.19097/) (2026)
- [What to Remove, What to Preserve: Dual-Ambiguity Rectification for All-in-One Image Restoration](https://mlsanity.com/paper/2607.28526/) (2026)
- [Learning to Restore More: Continual Capability Expansion for Pretrained Image Restoration Models](https://mlsanity.com/paper/2608.30305/) (2026)
- [GraLoD: Graphics-Inspired Continuous Level-of-Detail Learning for Image Restoration](https://mlsanity.com/paper/2609.16578/) (2026)
- [Diff-RF: Mutually Reinforced Image Registration and Fusion via Degradation-Aware Learning](https://mlsanity.com/paper/2609.28235/) (2026)
- [TaskIR: Task-Driven Image Restoration via Degradation Adaptation and Task Feedback](https://mlsanity.com/paper/2609.31170/) (2026)
- [Simultaneous Neural Optimal Transport](https://mlsanity.com/paper/2609.37424/) (2026)
- [UniDeblur: Unified Motion Deblurring via Fourier Convolution Learning](https://mlsanity.com/paper/2610.12941/) (2026)

## Abstract (arXiv)

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