Research map / Image generation / Image restoration
Debiased All-in-one Image Restoration with Task Uncertainty Regularization
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes Task Uncertainty Regularization (TUR), a plug-and-play homoscedastic-uncertainty loss reweighting that retrains existing all-in-one image restoration models (AirNet, PromptIR, TransWeather, SRResNet, Uformer) to balance degradation tasks during training.
Problem
Existing all-in-one image restoration methods use uniform multi-task learning that mixes degradation tasks with equal treatment, ignoring heterogeneous optimization across degradations and yielding biased suboptimal results.
Approach
Model per-task homoscedastic uncertainty with a learnable auxiliary head (UEM) and derive a Bayesian multi-task loss that adaptively reweights each task's reconstruction loss plus a log-uncertainty regularizer; the uncertainty module is discarded at inference.
Key results
- AirNet retrained with TUR improves average PSNR from 31.20 to 32.36 dB on the three-task (deraining, dehazing, denoising) setting, surpassing PromptIR.
- TransWeather with TUR improves average PSNR from 25.22 to 28.56 dB on five-task setting, and AirNet from 25.49 to 27.30 dB.
- Uformer with TUR reaches 32.79 dB average PSNR on seven-degradation setting, +2.09 dB over Uformer baseline.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Add Task Uncertainty Regularization to Uformer | Uformer · 7 degradation tasks | average PSNR: 30.70 -> 32.79 | Helped |
| Add Task Uncertainty Regularization to SRResNet | SRResNet · 7 degradation tasks | average PSNR: 29.52 -> 30.45 | Helped |
| Add Task Uncertainty Regularization | AirNet · Rain-Haze-Noise | average PSNR: 31.20 -> 32.36 | Helped |
| Add Task Uncertainty Regularization | PromptIR · Rain-Haze-Noise | average PSNR: 32.06 -> 32.67 | Helped |
| Add Task Uncertainty Regularization | Transweather · Rain-Haze-Noise-Blur-Dark | average PSNR: 25.22 -> 28.56 | Helped |
| Add Task Uncertainty Regularization | AirNet · Rain-Haze-Noise-Blur-Dark | average PSNR: 25.49 -> 27.30 | Helped |
| Add Task Uncertainty Regularization | TransWeather · Synthetic deweathering average | PSNR: 29.44 -> 30.66 | Helped |
| Add Task Uncertainty Regularization | TransWeather · Real-world deweathering average | PSNR: 26.71 -> 29.96 | Helped |
| Add task-grouped regularization to uncertainty estimation (ablation) | SRResNet with uncertainty estimation (UE) · 7 degradation tasks | average PSNR: 30.09 -> 30.45 | Helped |
| Use uncertainty estimation without task-grouped regularization instead of vanilla SRResNet (ablation) | SRResNet · 7 degradation tasks | average PSNR: 29.52 -> 30.09 | Helped |
Surprising findings
- Retraining an older AirNet with TUR outperforms the newer PromptIR on the three-task setting.
- On LOL low-light enhancement, PSNR decreased slightly for retrained AirNet while SSIM improved.
Limitations
- Focuses on intra-task uncertainty and ignores inter-task relationships, as the authors note.
- Slight PSNR drop on LOL low-light enhancement for retrained AirNet.
Experimental setup
Datasets: DIV2K Flickr2K BSD400 WED Rain100L SOTS GoPro LOL Snow100K Raindrop Outdoor-Rain SPA
Benchmarks: Set of 7 degradation tasks (Kong et al. 2024) BSD68 Rain100L SOTS GoPro LOL Outdoor-Rain Snow100K RainDrop SPA+ RealSnow REVIDE
Metrics: PSNR SSIM
Baselines: AirNet PromptIR TransWeather SRResNet-S Uformer-S MPRNet NAFNet IDR WGWSNet WeatherDiff128
Builds on: Kendall and Gal multi-task uncertainty weighting AirNet PromptIR TransWeather
Closest papers by meaning
- Parameter Efficient Adaptation for Image Restoration with Heterogeneous Mixture-of-Experts (NeurIPS 2024 · similarity 0.88)
- Learning from History: Task-agnostic Model Contrastive Learning for Image Restoration (AAAI 2024 · similarity 0.88)
- UP-Restorer: When Unrolling Meets Prompts for Unified Image Restoration (AAAI 2025 · similarity 0.88)
- Continuous Expert Assembly: Instance-Conditioned Low-Rank Residuals for All-in-One Image Restoration (NeurIPS 2026 · similarity 0.87)
- UARE: A Unified Vision-Language Model for Image Quality Assessment, Restoration, and Enhancement (CVPR 2026 · similarity 0.87)
- Complexity Experts are Task-Discriminative Learners for Any Image Restoration (CVPR 2025 · similarity 0.87)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.