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

📄 arXiv: 2603.03967v1 📥 PDF

作者: Qianfeng Yang, Qiyuan Guan, Xiang Chen, Jiyu Jin, Guiyue Jin, Jiangxin Dong

分类: cs.CV

发布日期: 2026-03-04

备注: Accepted by CVPR 2026; Project Page: https://github.com/QianfengY/UniRain


💡 一句话要点

UniRain:提出基于RAG的数据集蒸馏和多目标重加权优化的统一图像去雨框架

🎯 匹配领域: 支柱二:RL算法与架构 (RL & Architecture)

关键词: 图像去雨 数据集蒸馏 检索增强生成 多目标优化 混合专家模型

📋 核心要点

  1. 现有去雨方法通常只针对特定类型的雨天退化,难以泛化到不同的真实雨天场景,通用性不足。
  2. UniRain通过RAG的数据集蒸馏pipeline选择高质量训练样本,并结合多目标重加权优化策略,提升模型泛化能力。
  3. 实验结果表明,UniRain在多个公开数据集和自建基准测试中,性能优于当前最先进的去雨模型。

📝 摘要(中文)

本文提出UniRain,一个有效的统一图像去雨框架,能够恢复白天和夜晚条件下受雨条纹和雨滴退化的图像。为了更好地增强统一模型的泛化能力,我们构建了一个基于检索增强生成(RAG)的智能数据集蒸馏流程,该流程从所有公共去雨数据集中选择高质量的训练样本,以实现更好的混合训练。此外,我们将一个简单而有效的多目标重加权优化策略融入到非对称混合专家(MoE)架构中,以促进一致的性能并提高各种场景中的鲁棒性。大量实验表明,我们的框架在我们提出的基准和多个公共数据集上优于最先进的模型。

🔬 方法详解

问题定义:现有图像去雨方法通常针对特定类型的雨水退化设计,例如雨条纹或雨滴,并且在白天或夜晚条件下表现各异。这导致模型在面对真实世界中复杂多样的雨天场景时泛化能力不足。因此,如何构建一个能够处理各种雨水退化情况的通用去雨模型是一个关键问题。

核心思路:UniRain的核心思路是利用检索增强生成(RAG)的数据集蒸馏方法,从多个公开数据集中选择高质量的训练样本,从而构建一个更具代表性的混合训练集。此外,采用多目标重加权优化策略,平衡不同类型雨水退化场景下的性能,提高模型的鲁棒性和泛化能力。

技术框架:UniRain框架主要包含三个部分:1) 基于RAG的数据集蒸馏pipeline,用于选择高质量的训练样本;2) 非对称混合专家(MoE)架构,用于建模不同类型的雨水退化;3) 多目标重加权优化策略,用于平衡不同场景下的性能。整个流程首先使用RAG pipeline构建训练集,然后使用该训练集训练MoE模型,并在训练过程中使用多目标重加权优化策略。

关键创新:UniRain的关键创新在于:1) 提出了基于RAG的数据集蒸馏方法,能够有效地从多个数据集中选择高质量的训练样本,从而提高模型的泛化能力;2) 引入了多目标重加权优化策略,能够平衡不同类型雨水退化场景下的性能,提高模型的鲁棒性。

关键设计:在RAG pipeline中,使用了预训练的视觉语言模型来评估图像质量和相关性,并根据评估结果选择样本。MoE架构采用了非对称设计,不同的专家负责处理不同类型的雨水退化。多目标重加权优化策略根据不同场景下的损失函数值动态调整权重,从而平衡不同场景下的性能。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

UniRain在多个公开数据集和自建基准测试中取得了显著的性能提升。例如,在Rain13数据集上,UniRain的PSNR指标比当前最先进的方法提高了约0.5dB。在自建的包含多种雨水退化类型的测试集上,UniRain也表现出更强的鲁棒性和泛化能力。实验结果表明,UniRain能够有效地处理各种雨水退化情况,并且在不同场景下都能够取得良好的去雨效果。

🎯 应用场景

UniRain具有广泛的应用前景,例如智能监控、自动驾驶、图像编辑等领域。在智能监控中,可以提高雨天场景下的目标检测和识别精度。在自动驾驶中,可以提高雨天环境下的感知能力,保障行车安全。在图像编辑中,可以方便地去除图像中的雨水,提高图像质量。未来,该研究可以进一步扩展到视频去雨领域,并与其他图像增强技术相结合,实现更强大的图像处理功能。

📄 摘要(原文)

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.