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

2026-03-04

AI Summary

This work addresses the limited generalization of existing image deraining methods, which struggle to handle diverse rain-induced degradations in real-world scenarios—such as rain streaks and droplets under varying lighting conditions like day and night. To overcome this challenge, the authors propose a unified deraining framework that innovatively integrates a retrieval-augmented generation (RAG)-driven cross-dataset distillation mechanism to select high-quality training samples. The framework further incorporates an asymmetric mixture-of-experts (MoE) architecture and a multi-objective reweighting optimization strategy to enhance model generalization. Extensive experiments demonstrate that the proposed method consistently outperforms state-of-the-art approaches across multiple public benchmarks and a newly introduced dataset, exhibiting superior robustness and cross-scenario adaptability.

Technology Category

• Computer Vision: Image and Video Retrieval
• Machine Learning: Mixture of Experts (MoE)
• Intelligent Robots: Multimodal Perception & Sensor Fusion

Application Category

• Search and Retrieval-Augmented AI: Retrieval-Augmented Generation (RAG) and multi-modal RAG
• User Modeling, Personalization and Recommendation: Fairness-aware retrieval and ranking
• Web Mining and Content Analysis: Large pretrained models with web data

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.

Problem

• image deraining
• rain degradation
• generalization
• unified framework
• real-world rainy scenes

Innovation

• Unified Image Deraining
• RAG-based Dataset Distillation
• Multi-objective Reweighted Optimization
• Mixture-of-Experts
• Rain Streak and Raindrop Removal