Selective Hourglass Mapping for Universal Image Restoration Based on Diffusion Model
DiffUIR employs a selective hourglass mapping strategy with strong condition guidance and shared distribution terms, achieving state-of-the-art performance in multi-task image restoration with only 0.89M parameters.
Key Findings
Methodology
This paper introduces DiffUIR, a diffusion-based selective hourglass mapping strategy that integrates strong conditional guidance with a flexible shared distribution term (SDT). During the forward diffusion process, the model gradually reduces the influence of the condition (e.g., degraded images), causing different degradation distributions to approach a shared, impure Gaussian distribution. In the reverse process, the model leverages the combined effect of SDT and strong condition guidance to iteratively recover the task-specific distribution from the shared one. This approach enables the model to learn a unified shared distribution across multiple degradation types, avoiding the parameter-heavy multi-partite mappings typical of existing methods. The training objective maximizes the likelihood by predicting residuals, and the architecture is based on a lightweight U-Net, making the entire system highly parameter-efficient. Experimental results across five image restoration tasks and 22 benchmarks demonstrate that this strategy outperforms existing universal models, achieving SOTA performance with only 0.89M parameters.
Key Results
- On five core image restoration tasks, DiffUIR surpasses previous multi-task methods, with PSNR improvements ranging from 0.36dB to 1.85dB. For instance, in deraining, PSNR increased by 0.36dB; in low-light enhancement, by 0.95dB; in desnowing, dehazing, and deblurring, improvements were 1.85dB, 1.5dB, and 1.3dB respectively. Despite using only a fraction of the parameters of large-scale models, it achieves comparable or superior performance, demonstrating remarkable parameter efficiency.
- In zero-shot generalization tests, lightweight versions like DiffUIR-T (0.89M parameters) maintained high restoration quality on unseen tasks and scenes, confirming strong generalization capabilities. The ablation studies confirmed that combining strong condition guidance with shared distribution mapping yields the best results, with performance degrading when either component is removed.
- The method's simplicity—modifying only the mapping strategy—allows it to outperform existing universal methods significantly. It achieves these results without complex training pipelines or large pre-trained models, making it highly practical for real-world applications.
Significance
This work addresses a long-standing challenge in universal image restoration: how to efficiently learn a single model capable of handling diverse degradation types with high quality and generalization. By introducing a shared distribution framework within the diffusion process, the authors bridge the gap between task-specific and universal models, reducing parameter complexity while enhancing performance. The approach's ability to generalize zero-shot to unseen degradations and scenes makes it particularly valuable for real-world applications such as autonomous driving, robotic perception, and mobile imaging. Furthermore, the lightweight design opens avenues for deployment on resource-constrained devices, broadening the impact of multi-task learning in computer vision.
Technical Contribution
The core technical innovation lies in integrating a shared distribution term (SDT) into the diffusion process, enabling the gradual convergence of different degradation distributions into a common, impure Gaussian distribution during the forward process. This is coupled with a strong condition guidance mechanism that ensures high-fidelity task-specific recovery during the reverse process. The authors develop a novel mapping strategy—selective hourglass mapping—that dynamically adjusts the influence of conditions and shared distribution, effectively balancing shared knowledge and task-specific details. The entire framework is implemented with a lightweight U-Net architecture, trained with a simplified objective that predicts residuals, resulting in a model with only 0.89M parameters yet achieving state-of-the-art results. This approach fundamentally redefines how multi-task image restoration can be formulated within the diffusion paradigm, emphasizing parameter efficiency, robustness, and generalization.
Novelty
This research is the first to incorporate a shared distribution term directly into the diffusion model's mapping process, enabling the gradual unification of multiple degradation distributions into a single shared distribution. Unlike prior methods relying on multi-branch architectures or prompt-based large models, DiffUIR achieves multi-task learning through a simple yet effective mapping strategy that modulates the influence of conditions and shared distribution coefficients over diffusion steps. The combination of strong condition guidance with a flexible shared distribution approach introduces a new paradigm for universal image restoration, significantly reducing model complexity while maintaining or surpassing SOTA performance. This innovation opens new avenues for efficient multi-task generative modeling.
Limitations
- While the lightweight model performs well across multiple tasks, its performance may still degrade under extremely severe degradations such as very high noise levels or complex combined degradations, which require further robustness enhancements.
- The current framework is primarily validated on 2D images; extending it to videos or 3D data involves additional challenges related to temporal consistency and spatial complexity.
- Although the model is parameter-efficient, training still depends on large annotated datasets, and its performance in data-scarce environments remains to be explored. Incorporating self-supervised or unsupervised learning strategies could address this limitation.
Future Work
Future directions include extending DiffUIR to handle video and 3D data, integrating self-supervised learning to reduce reliance on annotated datasets, and exploring hardware-friendly model compression techniques for deployment on edge devices. Additionally, incorporating multi-modal inputs, such as depth or semantic information, could further improve restoration quality and robustness. The authors also plan to investigate adaptive coefficient scheduling strategies to dynamically balance shared and task-specific influences during diffusion, aiming for even better generalization and efficiency.
AI Executive Summary
In recent years, the field of image restoration has seen rapid advancements, driven by deep learning techniques that target specific degradation types such as denoising, deblurring, and super-resolution. However, real-world scenarios often involve multiple simultaneous degradations—images may be blurry, dark, snowy, or foggy at the same time—posing a significant challenge for traditional single-task models. Existing universal approaches attempt to unify multiple tasks within a single framework, but they often rely on complex multi-branch architectures or large-scale pre-trained models, which are computationally expensive and lack flexibility.
This paper introduces DiffUIR, a novel diffusion-based model that leverages a selective hourglass mapping strategy to address these issues. The core idea is to unify multiple degradation distributions into a shared, impure Gaussian distribution during the forward diffusion process, and then gradually recover task-specific details during the reverse process. This is achieved through a combination of strong condition guidance—ensuring accurate task-specific generation—and a shared distribution term (SDT), which modulates the influence of the condition over diffusion steps. The innovative mapping strategy dynamically balances the shared knowledge and task-specific nuances, enabling the model to learn a universal representation that can adapt to various degradations.
The authors demonstrate that this approach significantly outperforms existing multi-task models across five key image restoration tasks, including deraining, low-light enhancement, desnowing, dehazing, and deblurring. Remarkably, the lightweight model with only 0.89 million parameters achieves state-of-the-art results on 22 benchmarks, surpassing larger models by substantial margins. The zero-shot generalization experiments further validate the robustness and versatility of DiffUIR, showing its ability to handle unseen degradations and scenes without additional training.
This work marks a major step forward in the development of efficient, scalable, and generalizable multi-task image restoration models. By rethinking the role of distribution mapping within diffusion processes, the authors open new avenues for deploying powerful visual enhancement tools in real-world applications such as autonomous vehicles, robotics, and mobile imaging. The simplicity and effectiveness of DiffUIR suggest that future research can build upon this foundation to explore broader multi-modal and multi-task learning scenarios, ultimately bringing smarter and more adaptable visual systems closer to reality.
Deep Analysis
Background
Image restoration一直是计算机视觉中的基础任务，经历了从传统滤波、非局部均值到深度学习的快速发展。早期方法如BM3D、非局部均值等，依赖手工设计的滤波器，效果有限。随着深度神经网络的兴起，U-Net、ResNet等架构在去雨、去雪、去雾等子任务中表现出色，但模型庞大、训练复杂，难以满足实际多场景、多任务的需求。近年来，预训练大模型如SwinIR、MIRNet-v2在单任务中表现优异，但在多任务场景中参数膨胀严重，泛化能力不足。为此，研究者开始探索多任务统一模型，采用多编码器、多提示等策略，但多分支结构导致信息隔离，难以充分利用任务间的潜在关联。扩散模型作为近年来崛起的生成模型，以其强大的生成能力被引入图像修复领域，出现RainDiffusion、RDDM等变体，提升了恢复质量，但仍未解决多任务共享分布的问题。本研究在此基础上，提出基于扩散模型的选择性钟形映射策略，旨在实现多任务的高效统一。
Core Problem
多任务图像恢复的核心难题在于如何在一个模型中同时学习多种退化分布，并保证恢复质量。传统方法多采用多分支或条件提示，导致模型参数庞大，训练复杂，且难以捕获任务间的潜在联系。扩散模型虽能生成高质量图像，但在多任务场景中，端点的高斯噪声缺乏任务信息，限制了跨任务的共享能力。强条件引导虽然能提升生成质量，但会限制模型的泛化能力，导致多分支映射，忽略任务间的潜在联系。这些问题限制了多任务模型在实际复杂环境中的应用效果。
Innovation
本研究的创新点在于：第一，将强条件引导与共享分布项（SDT）结合，设计出一种新的扩散映射策略，实现不同退化分布的逐步融合与任务特异性恢复。第二，提出选择性钟形映射机制，通过逐步减弱条件影响，将多样的退化分布映射到一个非纯高斯的共享分布，增强模型的泛化能力。第三，采用极简架构的U-Net，参数仅0.89M，却在五个任务和22个基准测试中实现SOTA，极大地提升了参数效率和实用性。这些创新突破了多任务图像恢复的瓶颈，为扩散模型在多任务中的应用提供了新思路。
Methodology
- �� 采样正向过程：在扩散模型中引入条件机制，利用残差映射Ires，将退化图像与干净图像的差异作为条件输入。
- �� 共享分布映射：在正向扩散中加入逐步减弱的共享分布系数δt，使不同退化分布逐渐融合到一个非纯高斯的共享分布。
- �� 反向生成：结合SDT和强条件引导，逐步将共享分布引导回任务特定的分布，利用Reparameterization和DDIM采样策略实现高效逆向采样。
- �� 训练目标：最大化条件下的似然估计，最小化残差预测误差，采用L1损失，确保模型在不同任务中的泛化能力。
- �� 模型架构：基于U-Net设计，参数量极少，支持多任务训练，支持零样本泛化。
- �� 训练细节：采用Adam优化器，训练76小时，批次大小10，数据增强包括翻转和直方图均衡，输入裁剪为256×256。
Experiments
- �� 数据集：涵盖去雨（Merged Rain Dataset）、低光（LOL）、除雪（Snow100K）、去雾（RESIDE）和去模糊（GoPro）等五个任务，均为公开标准数据集。• 评估指标：采用PSNR、SSIM、NIQE、LPIPS等多指标，全面衡量恢复质量。• 实验设置：模型在不同参数版本（DiffUIR-T、S、B、L）下训练，参数从0.89M到64M不等，测试在全分辨率下进行，采用3个时间步长。• 比较基线：包括任务专用模型（SwinIR、MIRNet-v2）和现有多任务模型（Restorer、AirNet、Prompt-IR等），确保公平性。• 消融研究：验证强条件引导、共享分布项的贡献，分析不同参数配置的性能差异。
Results
- �� 定量指标：在五个任务中，DiffUIR在PSNR和SSIM上均优于现有多任务方法，例如在去雨任务中提升0.36dB PSNR，在去雪和去雾任务中分别提升1.85和1.3dB，参数仅为其他大模型的五分之一。• 零样本泛化：在未见过的场景和任务中，轻量版本依然表现出色，验证了模型的泛化能力。• 消融分析：去除SDT后，性能下降明显，验证其在共享分布中的作用；去除强条件引导，则生成质量降低，说明两者协同作用的重要性。
Applications
- �� 实时图像增强：可应用于自动驾驶、无人机视觉等场景中的多任务图像恢复，减少模型部署复杂度。• 机器人感知：在复杂环境下实现多任务感知，如同时进行去雾和去雪，提升环境理解能力。• 智能监控：结合多任务恢复能力，增强监控系统在不同天气和光照条件下的表现，提升安全性。
Limitations & Outlook
- �� 在极端退化条件下，模型仍可能出现恢复不足的问题，尤其是在噪声极大或模糊严重的场景。• 当前模型主要在二维图像任务中验证，扩展到视频或三维场景涉及时间一致性和空间复杂性。• 模型训练依赖大量标注数据，尽管参数少，但在数据匮乏环境中表现有限，未来需结合自监督或无监督技术提升鲁棒性。
Plain Language Accessible to non-experts
想象你有一个神奇的修理工，他可以帮你修复各种受损的照片，比如模糊、暗淡、下雨或下雪的照片。以前，这个修理工每次只专注一种损坏类型，比如只修模糊的照片，或者只修雨天的照片。这样的话，如果照片同时受多种损坏，比如既模糊又下雨，修理工就得用不同的工具，效率很低，也不够聪明。
现在，这个新修理工变得特别聪明，它学会了一个秘密：它可以先把所有不同的损坏类型都变成一种“共享的损坏状态”，就像把各种不同的污渍都变成一种特殊的油漆，然后再用一种特别的方法，把这个“共享油漆”逐渐变回原本的清晰照片。在这个过程中，它会用一些“条件”来指导自己，比如告诉自己“这是模糊的照片”或“这是下雪的场景”，但同时它也会逐渐减少对这些条件的依赖，让自己变得更聪明，能处理更多不同的损坏情况。
这个方法的厉害之处在于，它只需要很少的“修理工具”，参数非常少（只有0.89M），就能在多种任务中表现出色。它还能在没有看到新照片的情况下，直接修复未见过的损坏类型，表现出极强的泛化能力。这个新修理工的出现，意味着未来我们可以用一个“万能修理包”，轻松应对各种照片损坏问题，不再需要为每一种损坏准备不同的工具箱。
Abstract
Universal image restoration is a practical and potential computer vision task for real-world applications. The main challenge of this task is handling the different degradation distributions at once. Existing methods mainly utilize task-specific conditions (e.g., prompt) to guide the model to learn different distributions separately, named multi-partite mapping. However, it is not suitable for universal model learning as it ignores the shared information between different tasks. In this work, we propose an advanced selective hourglass mapping strategy based on diffusion model, termed DiffUIR. Two novel considerations make our DiffUIR non-trivial. Firstly, we equip the model with strong condition guidance to obtain accurate generation direction of diffusion model (selective). More importantly, DiffUIR integrates a flexible shared distribution term (SDT) into the diffusion algorithm elegantly and naturally, which gradually maps different distributions into a shared one. In the reverse process, combined with SDT and strong condition guidance, DiffUIR iteratively guides the shared distribution to the task-specific distribution with high image quality (hourglass). Without bells and whistles, by only modifying the mapping strategy, we achieve state-of-the-art performance on five image restoration tasks, 22 benchmarks in the universal setting and zero-shot generalization setting. Surprisingly, by only using a lightweight model (only 0.89M), we could achieve outstanding performance. The source code and pre-trained models are available at https://github.com/iSEE-Laboratory/DiffUIR
References (20)
Residual Denoising Diffusion Models
Jiawei Liu, Qiang Wang, Huijie Fan et al.
Denoising Diffusion Implicit Models
Jiaming Song, Chenlin Meng, S. Ermon
Denoising Diffusion Probabilistic Models
Jonathan Ho, Ajay Jain, P. Abbeel
RestoreFormer: High-Quality Blind Face Restoration from Undegraded Key-Value Pairs
Zhouxia Wang, Jiawei Zhang, Runjian Chen et al.
Contrast enhancement based on layered difference representation
Chulwoo Lee, Chulwoo Lee, Chang-Su Kim
Score-Based Generative Modeling through Stochastic Differential Equations
Yang Song, Jascha Narain Sohl-Dickstein, Diederik P. Kingma et al.
Benchmarking Single-Image Dehazing and Beyond
Boyi Li, Wenqi Ren, Dengpan Fu et al.
The Unreasonable Effectiveness of Deep Features as a Perceptual Metric
Richard Zhang, Phillip Isola, Alexei A. Efros et al.
Deep Retinex Decomposition for Low-Light Enhancement
Chen Wei, Wenjing Wang, Wenhan Yang et al.
An Introduction to Variational Autoencoders
Diederik P. Kingma, M. Welling
A General Decoupled Learning Framework for Parameterized Image Operators
Qingnan Fan, Dongdong Chen, Lu Yuan et al.
Generative Modeling by Estimating Gradients of the Data Distribution
Yang Song, S. Ermon
Human-Aware Motion Deblurring
Ziyi Shen, Wenguan Wang, Xiankai Lu et al.
PyTorch: An Imperative Style, High-Performance Deep Learning Library
Adam Paszke, Sam Gross, Francisco Massa et al.
Decision-Making with Auto-Encoding Variational Bayes
Romain Lopez, Pierre Boyeau, N. Yosef et al.
Image Restoration for Under-Display Camera
Yuqian Zhou, D. Ren, N. Emerton et al.
Multi-Scale Progressive Fusion Network for Single Image Deraining
Kui Jiang, Zhongyuan Wang, Peng Yi et al.
All in One Bad Weather Removal Using Architectural Search
Ruoteng Li, R. Tan, L. Cheong
Real-World Blur Dataset for Learning and Benchmarking Deblurring Algorithms
Jaesung Rim, H. Chwa, Sunghyun Cho
Visualizing Data using t-SNE
L. Maaten, Geoffrey E. Hinton
Cited By (20)
UARE: A Unified Vision-Language Model for Image Quality Assessment, Restoration, and Enhancement
Expandable, Compressible, Mineable: Open-World Thermal Image Restoration
Decoupled Residual Denoising Diffusion Models for Unified and Data Efficient Image-to-Image Translation
Breaking Degradation Coupling: A Structural Entropy Guided Decoupled Framework and Benchmark for Infrared Enhancement
Bridging Degradation Discrimination and Generation for Universal Image Restoration
Unifying Heterogeneous Degradations: Uncertainty-Aware Diffusion Bridge Model for All-in-One Image Restoration
Overlapped Wavelet Diffusion for Low-Light Image Enhancement
P2DNet: A Physics-Constrained Dual-Branch Network for Underwater Polarimetric Image Restoration
FoundIR-v2: Optimizing Pre-Training Data Mixtures for Image Restoration Foundation Model
Virtual Consistency Model for All-in-One Image Restoration
PixRestore: Unified Image Restoration via Pixel Diffusion Transformer
Bilevel Layer-Positioning LoRA for Real Image Dehazing
StarIR: Convolutional Image Restoration With Spatial-Frequency Fusion
M2IR: Proactive All-in-One Image Restoration via Mamba-style Modulation and Mixture-of-Experts
UnSCAR: Universal, Scalable, Controllable, and Adaptable Image Restoration
Energy-oriented Diffusion Bridge for Image Restoration with Foundational Diffusion Models
Visual-in-Visual: A Unified and Efficient Baseline for Image Restoration
Time-frequency image enhancement for low SNR communication signals using generative diffusion models
Restore, Assess, Repeat: A Unified Framework for Iterative Image Restoration
A survey on image restoration methods based on denoising diffusion probabilistic models series models