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

이 논문은 RAG 기반의 데이터 증류 파이프라인과 다목적 재가중 최적화 전략을 결합한 'UniRain'을 제안하여, 다양한 야간 및 주간 조건에서 비줄무늬와 빗방울로 인한 왜곡을 모두 효과적으로 제거하는 범용 이미지 비 제거 프레임워크를 구축했습니다.

핵심

1. 문제: "너무 많은 비, 너무 다른 비"

지금까지 비 제거 기술은 마치 비 오는 날마다 다른 우산을 써야 하는 상황과 비슷했습니다.

• 낮에 내리는 빗줄기 (스트릭) 를 지우려면 A 우산이 필요하고,
• 밤에 내리는 빗방울 (드롭) 을 지우려면 B 우산이 필요했습니다.
• 게다가 훈련용 데이터 (비 오는 사진) 들의 품질도 제각각이었습니다. 어떤 건 선명하고, 어떤 건 흐릿해서 AI 가 혼란을 겪었습니다.

2. 해결책 1: "명품 사냥꾼 (RAG 기반 데이터 증류)"

UniRain 은 먼저 **가장 좋은 훈련용 데이터만 골라내는 '명품 사냥꾼'**을 도입했습니다.

• 비유: 전 세계에 흩어진 비 오는 사진 200 만 장을 모두 모았다고 칩시다. 그중에는 화질이 나쁜 사진도 섞여 있습니다.
• 작동 방식: AI 가 "이 사진은 비가 진짜처럼 보이나요?"라고 **지식인 (VLM, 시각 - 언어 모델)**들에게 물어봅니다. 마치 **검색 엔진 (RAG)**이 가장 관련성 높은 정보만 찾아오듯, AI 는 화질이 좋고 비가 선명하게 찍힌 '명품' 사진들만 골라내서 훈련시킵니다.
• 효과: 쓰레기 같은 데이터는 버리고, 진짜 비를 잘 보여주는 데이터만 모아서 AI 를 가르칩니다.

3. 해결책 2: "유연한 요리사 (다목적 가중치 최적화)"

AI 가 여러 종류의 비를 한 번에 지우려 할 때, 어떤 비를 먼저 지울지 고민이 생깁니다.

• 쉬운 비 (예: 밤의 빗줄기) 는 금방 지워지지만, 어려운 비 (예: 낮의 빗방울) 는 시간이 오래 걸립니다.

• 문제: AI 가 쉬운 비만 잘 지우려고 노력하다 보니, 어려운 비는 제대로 지우지 못하게 됩니다. (학습 불균형)

• 해결책: UniRain 은 현명한 요리사처럼 상황을 조절합니다.

  • "오늘은 빗방울 지우기가 너무 느리네? 빗줄기 지우기보다 빗방울 지우기에 더 집중하자!"라고 **점수 (가중치)**를 실시간으로 조절합니다.
  • 어려운 작업일수록 더 많은 에너지를 쏟게 만들어, 모든 비를 골고루 잘 지우도록 돕습니다.

4. 해결책 3: "두 가지 모자를 쓴 전문가 (비대칭 MoE)"

마지막으로, AI 의 뇌 구조를 두 가지 모자를 쓴 전문가 팀처럼 설계했습니다.

• 엔코더 (입력 단계) = '부드러운 모자 (Soft-MoE)':
  • 모든 비의 특징을 모두 살펴봅니다. "아, 이건 빗줄기야, 저건 빗방울이야, 밤이야, 낮이야"라고 다양한 전문가들이 의견을 나누며 정보를 수집합니다.
• 디코더 (출력 단계) = '단단한 모자 (Hard-MoE)':
  • 이제 실제 사진을 복구할 때는 가장 적합한 전문가만 딱 1~2 명 뽑아 집중합니다. "이 부분은 빗방울을 지우는 전문가가 맡아!"라고 명확하게 지시하여 선명한 이미지를 만듭니다.

---

요약: 왜 UniRain 이 특별한가요?

1. 한 번에 해결: 낮, 밤, 빗줄기, 빗방울을 구분하지 않고 한 번에 처리합니다.
2. 질 좋은 학습: 쓰레기 같은 데이터는 버리고, 진짜 같은 고화질 데이터만 골라 학습시켰습니다.
3. 균형 잡힌 학습: 쉬운 비와 어려운 비를 공평하게 다룰 수 있도록 학습 방식을 조절했습니다.

결론적으로, UniRain 은 비가 온 날 찍은 모든 사진과 영상을 맑고 선명하게 만들어주는 만능 비 제거 전문가입니다. 자율주행차나 드론이 비 오는 날에도 길을 잘 볼 수 있도록 도와주는 기술입니다.

기술 요약

1. 문제 정의 (Problem Statement)

기존의 이미지 빗물 제거 (Deraining) 연구는 특정 유형의 빗물 손상 (예: 낮의 빗줄기, 밤의 빗줄기, 빗방울 등) 에만 특화되어 개발되는 경향이 있습니다. 이로 인해 다음과 같은 한계가 존재합니다:

• 일반화 부족: 특정 시나리오에 최적화된 모델은 다른 유형의 빗물 (예: 낮의 빗방울) 이나 복잡한 실제 환경 (야간, 다양한 조명) 에 적용 시 성능이 급격히 저하됩니다.
• 데이터의 불균형과 품질 문제: 모든 공개된 합성 및 실제 빗물 데이터를 단순히 합쳐서 학습시키면, 데이터 품질의 편차 (Synthetic vs. Real, 해상도 차이 등) 로 인해 모델 학습이 불안정해지고 수렴이 저해됩니다.
• 학습 불균형: 다양한 빗물 유형을 단일 모델로 학습할 때, 손실 함수가 수렴하기 쉬운 유형 (예: 밤의 빗줄기) 에 편향되어 복잡한 유형 (예: 낮의 빗방울) 의 복원 품질이 낮아지는 문제가 발생합니다.

2. 제안 방법론 (Methodology)

저자들은 UniRain이라는 통합 이미지 빗물 제거 프레임워크를 제안하며, 이는 크게 세 가지 핵심 기술로 구성됩니다.

2.1. RAG 기반 데이터 증류 (RAG-based Dataset Distillation)

대규모 공개 데이터셋 (>200 만 쌍) 에서 고품질 학습 샘플을 선별하기 위해 검색 증강 생성 (RAG) 기술을 도입했습니다.

• 검색 단계 (Retrieval Stage): 쿼리 이미지에 대해 텍스트 설명 (Caption), 시각적 특징, 구조적 특징을 추출하여 데이터베이스와 계층적 유사도 매칭 (의미론적, 시각적, 구조적 유사도) 을 수행합니다. 이를 통해 실제와 가장 유사한 참조 이미지를 검색합니다.
• 생성 단계 (Generation Stage): 검색된 참조 이미지와 쿼리 이미지를 비전 - 언어 모델 (VLM, Vision-Language Model) 에 입력하여 데이터의 신뢰성을 평가합니다. 여러 VLM 의 투표 (Ensemble Voting) 를 통해 '신뢰할 수 있는' 고품질 데이터만 선별하여 학습 데이터셋 (RainRAG) 을 증류합니다.

2.2. 비대칭 혼합 전문가 구조 (Asymmetric MoE Architecture)

다양한 빗물 손상을 효과적으로 처리하기 위해 인코더와 디코더에 서로 다른 MoE (Mixture-of-Experts) 방식을 적용했습니다.

• Soft-MoE 인코더: 연속적인 라우팅 가중치를 사용하여 여러 전문가 (Expert) 들의 특징을 유연하게 결합합니다. 이는 다양한 손상 패턴을 포괄적으로 포착하는 데 도움을 줍니다.
• Hard-MoE 디코더: Top-k 라우팅을 통해 가장 관련성 높은 전문가만 선택적으로 활성화합니다. 이는 미세한 질감과 구조적 세부 사항을 정밀하게 복원하는 데 집중합니다.

2.3. 다목적 재가중 최적화 전략 (Multi-objective Reweighted Optimization)

학습 중 발생하는 다양한 빗물 유형 간의 불균형을 해결하기 위해 동적 가중치 조정 전략을 도입했습니다.

• 수렴 기울기 추정 (Convergence Slope Estimation): 슬라이딩 윈도우 내에서 각 빗물 유형의 손실 (Loss) 변화율을 선형 회귀를 통해 추정합니다.
• 적응적 재가중 (Adaptive Reweighting):
  • TBS (Type Balance Score): 수렴이 느린 유형에 더 높은 가중치를 부여하여 모든 유형의 동시 수렴을 유도합니다.
  • TSS (Type Stability Score): 발산하거나 불안정한 유형의 가중치를 낮춥니다.
  • AF (Adaptivity Factor): 학습 단계에 따라 TBS 와 TSS 의 기여도를 동적으로 조절합니다.
  • 이를 통해 모델이 쉬운 유형에 편향되지 않고 모든 빗물 유형에서 균일한 성능을 내도록 합니다.

3. 주요 기여 (Key Contributions)

1. UniRain 프레임워크: 낮/밤 조건과 빗줄기/빗방울 유형을 모두 아우르는 최초의 통합 이미지 빗물 제거 모델 제안.
2. 지능형 데이터 증류 파이프라인: RAG 와 VLM 을 활용하여 대규모 데이터셋에서 고품질 샘플을 자동 선별하는 새로운 데이터 증류 기법 개발.
3. 동적 최적화 전략: 다목적 학습의 불균형을 해결하기 위한 적응적 재가중 최적화 알고리즘 제안.
4. 성능 입증: 제안된 벤치마크 (RainRAG) 와 여러 공개 데이터셋 (RealRain-1k, WeatherBench 등) 에서 SOTA(State-of-the-Art) 모델 대비 우수한 정량적/정성적 성능 입증.

4. 실험 결과 (Results)

• 정량적 성능: 제안된 RainRAG 데이터셋에서 UniRain 은 평균 PSNR 28.93 dB, SSIM 0.8515 를 기록하여 기존 SOTA 모델 (Restormer, NeRD-Rain, URIR 등) 보다 우세한 성능을 보였습니다. 특히 낮의 빗방울 (DRD) 과 밤의 빗줄기 (NRS) 등 어려운 유형에서 큰 성능 향상을 보였습니다.
• 실제 환경 적용: RealRain-1k, RainDS-real, WeatherBench 등 실제 빗물 데이터셋에서도 평균 PSNR 29.42 dB 를 기록하며, 기존 방법론보다 1.73 dB 이상 향상된 성능을 입증했습니다.
• 일반화 능력: 자율주행, 드론 (UAV), 해상 등 다양한 시나리오에서 뛰어난 일반화 능력을 보였으며, 기존 방법들이 실패하는 복잡한 빗방울 제거에서도 선명한 배경 복원을 달성했습니다.
• 모델 효율성: Restormer 등 기존 모델 대비 더 적은 파라미터 (24.388M) 와 FLOPs (126.541G) 로 동급 이상의 성능을 달성하여 효율성이 높음을 보였습니다.
• 확장성: 빗물 제거뿐만 아니라 눈, 안개 등 다양한 기상 조건을 제거하는 'All-in-one Weather Restoration' 작업에서도 우수한 성능을 발휘했습니다.

5. 의의 및 결론 (Significance)

이 논문은 단일 모델이 다양한 빗물 손상 유형과 환경 (낮/밤) 을 통합적으로 처리할 수 있는 새로운 패러다임을 제시합니다. 단순히 데이터를 늘리는 것이 아니라, RAG 기반의 지능형 데이터 증류를 통해 학습 데이터의 품질을 극대화하고, 다목적 재가중 최적화를 통해 학습 불균형을 해결함으로써 실제 세계의 복잡한 빗물 제거 문제에 대한 강력한 솔루션을 제공했습니다. 이는 저수준 컴퓨터 비전 (Low-level Vision) 분야에서 데이터 품질 관리와 최적화 전략의 중요성을 재조명하며, 향후 다양한 기상 조건을 통합적으로 처리하는 이미지 복원 연구의 방향성을 제시합니다.