UniRain: Unified Image Deraining with RAG-based Dataset Distillation and Multi-objective Reweighted Optimization

TL;DR

UniRain is proposed, an effective unified image deraining framework capable of restoring images degraded by rain streak and raindrop under both daytime and nighttime conditions and is incorporated into the asymmetric mixture-of-experts (MoE) architecture to facilitate consistent performance and improve robustness across diverse scenes.

Abstract

Despite significant progress has been made in image deraining, we note that most existing methods are often developed for only specific types of rain degradation and fail to generalize across diverse real-world rainy scenes. How to effectively model different rain degradations within a universal framework is important for real-world image deraining. In this paper, we propose UniRain, an effective unified image deraining framework capable of restoring images degraded by rain streak and raindrop under both daytime and nighttime conditions. To better enhance unified model generalization, we construct an intelligent retrieval augmented generation (RAG)-based dataset distillation pipeline that selects high-quality training samples from all public deraining datasets for better mixed training. Furthermore, we incorporate a simple yet effective multi-objective reweighted optimization strategy into the asymmetric mixture-of-experts (MoE) architecture to facilitate consistent performance and improve robustness across diverse scenes. Extensive experiments show that our framework performs favorably against the state-of-the-art models on our proposed benchmarks and multiple public datasets.

Figures (8)

• Figure 1: Overview of motivation. (a) Rainy image samples from public datasets, illustrating the noticeable differences in data quality. (b) Directly merging existing synthetic and real datasets enlarges data volume, yet quality disparity hinders performance, as shown by PSNR results. (c) The loss curves of DRS, DRD, NRS, and NRD (denoting daytime/nighttime rain streaks and raindrops) show different convergence rates, leading to imbalance in unified training. (d) The PSNR curves indicate that the model tends to favor simpler degradations but struggles with complex ones.
• Figure 2: Overall framework of UniRain. (Left) The RAG-based dataset distillation pipeline retrieves real rainy references consistent with the query image via multi-level similarity search and employs vision language models to evaluate its quality, thereby distilling reliable samples from public datasets. (Right) The asymmetric MoE architecture consists of soft-MoE encoder and hard-MoE decoder, optimized via the multi-objective reweighted strategy to achieve balanced learning and robust performance across multiple rain degradation types.
• Figure 3: Visual comparison of image deraining results on the RainRAG-NRS and RainDS-real-RD datasets. Zoom in for a better view.
• Figure 4: Visual comparison of image restoration results on the real-world benchmark (i.e., WeatherBench guan2025weatherbench). Compared to the state-of-the-art methods, our UniRain restores a high-quality image with clear details, even outperforming the GT by removing residual raindrops.
• Figure 5: Generalization results across multiple scenarios. Our method achieves the best restoration results in driving scenes (first row), UAV scenes (second row), and maritime scenes (third row).