Gang Wu (吴刚), Junjun Jiang (江俊君), Yijun Wang (王奕钧), Kui Jiang (江奎), and Xianming Liu (刘贤明)
AIIA Lab, Faculty of Computing, Harbin Institute of Technology, Harbin 150001, China.
 Framework |  Task Uncertainty |
All-in-one image restoration is a fundamental low-level vision task with significant real-world applications. The primary challenge lies in addressing diverse degradations within a single model. While current methods primarily exploit task prior information to guide the restoration models, they typically employ uniform multi-task learning, overlooking the heterogeneity in model optimization across different degradation tasks. To eliminate the bias, we propose a task-aware optimization strategy, that introduces adaptive task-specific regularization for multi-task image restoration learning. Specifically, our method dynamically weights and balances losses for different restoration tasks during training, encouraging the implementation of the most reasonable optimization route. In this way, we can achieve more robust and effective model training. Notably, our approach can serve as a plug-and-play strategy to enhance existing models without requiring modifications during inference. Extensive experiments across diverse all-in-one restoration settings demonstrate the superiority and generalization of our approach. For instance, AirNet retrained with TUR achieves average improvements of 1.3 dB on three distinct tasks and 1.81 dB on five distinct all-in-one tasks. These results underscore TUR's effectiveness in advancing the SOTAs in all-in-one image restoration, paving the way for more robust and versatile image restoration.
 |  |
 |  |
Detail experiment settings and datasets can be found at MioIR.
Our retrained Mixed Uformer-TUR.
Detail experiment settings and datasets can be found at AirNet and PromptIR project pages.
Our retrained AirNet-TUR .
Our retrained Transweather-TUR.
Detail experiment settings and datasets can be found at Transweather and WGWS-Net.
Our retrained Transweather-TUR for AllWeather and Real-World deweathering.
If our project helps your research or work, please cite our paper or star this repo. Thank you!
@inproceedings{wu2025debiased,
title={Debiased All-in-one Image Restoration with Task Uncertainty Regularization},
author={Gang Wu, Junjun Jiang, Yijun Wang, Kui Jiang, and Xianming Liu},
booktitle={Proceedings of the AAAI Conference on Artificial Intelligence},
year={2025}
}
This project is based on AirNet, PromptIR, MioIR, Transweather, and WGWS-Net, thanks for their nice sharing.