MiM-DiT: MoE in MoE with Diffusion Transformers for All-in-One Image Restoration

이 논문은 다양한 종류의 이미지 열화를 효과적으로 복원하기 위해 사전 학습된 확산 모델에 이중 계층 혼합 전문가 (MoE) 구조를 통합하여, 열화 유형별 coarse-grained 적응과 세부 변형에 대한 fine-grained 조절을 동시에 수행하는 통합 이미지 복원 프레임워크 'MiM-DiT'를 제안합니다.

핵심

이 논문은 **"MiM-DiT"**이라는 새로운 이미지 복원 기술을 소개합니다. 이를 쉽게 이해하기 위해 **'만능 사진 수리소'**와 **'전문가 팀'**의 비유를 들어 설명해 드릴게요.

1. 문제 상황: "한 명의 장인이 모든 일을 하기엔 너무 어렵다"

우리가 사진을 복구할 때 겪는 문제는 다양합니다.

• 흐린 사진 (Blur): 카메라가 흔들려서 흐릿합니다.
• 안개 낀 사진 (Haze): 안개 때문에 뿌옇습니다.
• 어두운 사진 (Low-light): 너무 어둡습니다.
• 노이즈 (Noise): 쌀알 같은 점들이 가득합니다.

기존의 AI 모델들은 보통 "한 가지 방식"으로 모든 사진을 고치려고 했습니다. 마치 한 명의 장인이 안개 낀 사진, 흐린 사진, 어두운 사진을 모두 똑같은 도구로 고치려다 보니, 안개는 잘 뚫었는데 흐린 사진은 오히려 더 흐릿해지거나, 어두운 사진은 색이 변하는 등 모든 문제를 완벽하게 해결하지 못하는 모순이 발생했습니다.

2. 해결책: "MiM-DiT" - 전문가들이 모여 만든 '초능력 팀'

이 논문은 이 문제를 해결하기 위해 **두 단계로 나누어진 '전문가 팀 (MoE, Mixture-of-Experts)'**을 도입했습니다. 이를 **'MiM (MoE in MoE)'**이라고 부릅니다.

🏢 1 단계: '팀장'이 맡은 역할 (Inter-MoE)

먼저, 들어온 사진을 보고 **"이건 어떤 종류의 문제인가?"**를 판단하는 **팀장 (라우터)**이 있습니다.

• 이 팀장은 네 가지 서로 다른 전문가 그룹을 관리합니다.
  1. 공간 전문가: 흐릿한 선을 바로잡는 데 특화됨.
  2. 채널 전문가: 색감이나 밝기 조절에 특화됨.
  3. 스윈 (Swin) 전문가: 전체적인 구조와 국소적인 디테일을 동시에 봄.
  4. SE 전문가: 안개나 어두운 빛을 처리하는 데 특화됨.

비유: 팀장은 "아, 이 사진은 안개 낀 거야!"라고 판단하면, 채널 전문가와 SE 전문가 그룹에게 "너희가 이걸 맡아!"라고 지시합니다. 반대로 "흐린 사진이네?"라면 공간 전문가 그룹을 불러모읍니다. 이렇게 문제의 종류에 따라 가장 적합한 팀을 선택하는 것입니다.

🧩 2 단계: '팀원'들이 맡은 역할 (Intra-MoE)

하지만 같은 '안개' 문제라도 가벼운 안개일 수도 있고 짙은 안개일 수도 있습니다.

• 이때 선택된 전문가 그룹 안에는 다시 **여러 명의 하위 전문가 (서브-전문가)**들이 있습니다.
• 팀장은 다시 한번 "이건 가벼운 안개니까 A 팀원, 짙은 안개니까 B 팀원"처럼 세부적인 상황에 맞춰 가장 적합한 전문가 하나를 골라 작업을 시킵니다.

비유: 안개 팀장 밑에 "가벼운 안개 처리반"과 "짙은 안개 처리반"이 있는데, 상황에 따라 가장 잘하는 한 명만 골라 작업을 시키는 것입니다.

3. 핵심 기술: "이미지 생성의 달인 (Diffusion Transformer)"

이 모든 전문가 팀은 **'생성형 AI (Diffusion Model)'**라는 거대한 엔진 위에 얹혀 있습니다.

• 기존 AI는 단순히 흐린 부분을 지우거나 채우는 방식이라서, 결과물이 너무 매끄럽고 (Over-smoothed) 디테일이 사라지는 경우가 많았습니다.
• 하지만 이 모델은 미리 훈련된 '생성형 AI'를 베이스로 사용합니다. 이는 마치 화가가 빈 캔버스에 새로운 그림을 그리는 능력을 가지고 있는 것과 같습니다.
• 전문가 팀이 "여기에 나무가 있어야 해", "여기는 빛이 반사되어야 해"라고 지시하면, 생성형 AI 는 그 지시를 받아 자연스럽고 선명한 디테일을 새로 만들어냅니다.

4. 왜 이 기술이 특별한가요?

• 유연함: 사진의 상태 (흐림, 안개, 어둠 등) 에 따라 AI 가 사용하는 '두뇌 회로'를 실시간으로 바꿉니다.
• 정교함: 큰 문제 (안개 vs 흐림) 를 구분하는 것뿐만 아니라, 세부적인 문제 (가벼운 안개 vs 짙은 안개) 까지 세밀하게 처리합니다.
• 자연스러움: 단순히 픽셀을 고치는 게 아니라, AI 가 상상력을 발휘해 원래 사진에 없던 디테일까지 자연스럽게 복원합니다.

📝 요약

이 논문은 **"하나의 AI 모델이 모든 사진 문제를 해결할 수 있을까?"**라는 질문에 대해, **"단순히 많은 전문가를 모으는 게 아니라, '팀장 (Inter-MoE)'이 문제를 분류하고, '팀원 (Intra-MoE)'이 세부 상황에 맞춰 일하는 계층 구조를 만들고, 여기에 '창의적인 화가 (Diffusion)'의 능력을 더했다"**고 답합니다.

그 결과, 흐릿한 사진, 안개 낀 사진, 어두운 사진 등 어떤 상황에서도 가장 선명하고 자연스러운 사진을 만들어내는 **'만능 사진 수리소'**가 탄생한 것입니다.

기술 요약

1. 연구 배경 및 문제 정의 (Problem)

**All-in-One Image Restoration (올인원 이미지 복원)**은 안개 (Haze), 흐림 (Blur), 노이즈 (Noise), 저조도 (Low-light) 등 다양한 열화 유형을 하나의 모델로 처리하여 고품질 이미지를 복원하는 것을 목표로 합니다. 그러나 기존 방법론들은 다음과 같은 근본적인 한계에 직면해 있습니다.

• 상충되는 모델링 요구사항: 열화 유형마다 복원 전략이 다릅니다. 예를 들어, 안개 제거는 전역적 (Global) 통계에 의존하는 반면, 흐림 제거는 날카로운 국소적 (Local) 구조 재구성이 필요합니다. 단일 모델이 이러한 상충되는 요구사항을 동시에 만족시키기 어렵습니다.
• 기존 MoE (Mixture-of-Experts) 의 한계: 기존 MoE 기반 방법들은 결정론적 (Deterministic) 손실 함수를 사용하여 과부드러짐 (Over-smoothing) 이 발생하고 질감 생성이 제한적입니다.
• 기존 Diffusion Model 의 한계: 확산 모델 (Diffusion Model) 은 생성 능력이 뛰어나지만, 모든 열화 유형에 대해 균일한 (Uniform) 확산 과정을 적용하여 열화 특성을 무시함으로써 구조적 왜곡이나 세부 정보 손실을 초래합니다.
• 단순 통합의 비효율성: 기존 MoE 와 Diffusion 모델을 단순히 결합하더라도, 모든 열화에 동일한 아키텍처를 사용하는 '동질적 (Homogeneous)' 설계로 인해 각 작업별 고유한 귀납적 편향 (Inductive Bias) 을 반영하지 못합니다.

2. 제안된 방법론 (Methodology)

저자들은 **MiM-DiT (MoE in MoE with Diffusion Transformers)**라는 새로운 프레임워크를 제안합니다. 이는 사전 학습된 Diffusion Transformer (DiT) 백본에 이중 계층 MoE (Hierarchical MoE-in-MoE) 아키텍처를 통합한 것입니다.

A. 계층적 MoE-in-MoE (Hierarchical MiM) 구조

모델은 두 가지 수준의 라우팅 메커니즘을 통해 작동합니다.

1. Inter-MoE (상위 계층 - 구조적 적응):

   • 목적: 서로 다른 열화 유형 (Inter-class) 에 대응하기 위해 이질적인 (Heterogeneous) 전문가 그룹들을 통합합니다.
   • 구현: 4 가지 서로 다른 어텐션 메커니즘 기반의 전문가 그룹을 구성합니다.
     • Spatial Self-Attention (장거리 의존성 모델링)
     • Channel Self-Attention (채널별 특징 재조정)
     • Swin Attention (효율적인 국소 및 전역 컨텍스트)
     • SE Attention (전역 조명 모델링, 안개/저조도 처리에 유리)
   • 라우팅: **Dense Router(밀집 라우터)**를 사용하여 모든 전문가 그룹의 출력을 입력 열화 특성에 따라 적응적으로 가중치 합성 (Fusion) 합니다. 이는 특정 열화에 가장 적합한 구조적 사전 지식을 동적으로 선택하게 합니다.

2. Intra-MoE (하위 계층 - 세부 적응):

   • 목적: 각 전문가 그룹 내부의 세밀한 열화 변형 (Intra-class, 예: 가벼운 안개 vs 짙은 안개) 에 대응합니다.
   • 구현: 각 전문가 그룹 내에서 동일한 어텐션 구조를 공유하지만 파라미터가 다른 여러 서브-전문가 (Sub-experts) 를 배치합니다.
   • 라우팅: **Sparse Router(희소 라우터, Top-k)**를 사용하여 각 입력에 가장 관련성이 높은 서브-전문가들만 활성화합니다. 이는 계산 효율성을 유지하면서 세부적인 적응을 가능하게 합니다.

B. DiT 백본과의 통합 (Conditional Generation)

• Zero-Linear 연결: MiM 모듈을 통해 추출된 열화 특이적 특징 (Feature) 은 'Zero-Linear' 계층을 거쳐 DiT 백본의 조건부 입력 (Conditional Input) 으로 주입됩니다.
• 동작 원리: 초기 학습 단계에서는 Zero-Linear 가 0 으로 초기화되어 DiT 의 사전 학습된 생성 능력을 보존하며, 학습이 진행됨에 따라 MiM 에서 추출된 특징이 확산 과정을 안내하여 고품질 복원을 수행합니다.

3. 주요 기여점 (Key Contributions)

1. 통합된 올인원 프레임워크: 사전 학습된 Diffusion 모델의 강력한 생성 능력과 MoE 패러다임을 결합하여 다양한 열화 유형을 단일 모델로 처리하면서도 생성 품질을 유지합니다.
2. 이중 계층 MoE-in-MoE 아키텍처:
   • Inter-MoE: 서로 다른 구조 (어텐션 메커니즘) 간 라우팅을 통해 전역적/국소적 의존성과 채널 간 관계를 포괄적으로 학습.
   • Intra-MoE: 동일 구조 내 세부 변형에 대한 희소 라우팅을 통해 세밀한 적응성 확보.
   • 이 설계는 열화 유형 간 (Inter-class) 과 유형 내 (Intra-class) 모두에 대한 동적이고 세분화된 적응을 가능하게 합니다.
3. 성능 입증: 다양한 이미지 복원 벤치마크에서 최신 기법 (SOTA) 대비 우수한 성능을 입증했습니다.

4. 실험 결과 (Results)

• 데이터셋: FoundIR (대규모 통합 데이터셋) 및 4KRD, RealRain-1K, HazeRD, UHD-LL 등 다양한 공개 벤치마크에서 평가.
• 정량적 평가:
  • 참조 기반 지표 (LPIPS, FID): 안개, 저조도, 노이즈 제거 등에서 가장 낮은 (우수한) 점수를 기록하여 높은 충실도 (Fidelity) 를 보임.
  • 비참조 지표 (NIQE, LIQE, MUSIQ, CLIP-IQA): 시각적 자연스러움과 지각적 품질 측면에서 기존 방법들을 크게 상회하는 최고 성능 달성.
• 정성적 평가:
  • 흐림 제거: 날카로운 에지와 미세한 질감을 보존하며 모션 블러를 효과적으로 제거.
  • 안개 제거: 색상 왜곡 없이 자연스러운 대비와 색상을 복원.
  • 저조도 향상: 기존 방법들이 겪는 과부드러짐이나 색상 편색 (Color cast) 없이 세부 구조와 정확한 색상을 복원.
• Ablation Study:
  • Intra-MoE 제거 시: 세부 질감 손실 및 과부드러짐 발생 (MUSIQ, CLIPIQA 감소).
  • Inter-MoE 동질화 시 (Single Attention): 다양한 열화 유형 처리 능력 저하.
  • 라우팅 설계: Dense Inter-MoE + Sparse Intra-MoE 조합이 가장 최적의 성능과 효율성을 보임.

5. 의의 및 결론 (Significance)

이 논문은 **구조적 다양성 (Structural Diversity)**과 **적응형 라우팅 (Adaptive Routing)**을 결합하여 올인원 이미지 복원의 난제를 해결했습니다.

• 기술적 혁신: 단일 모델이 서로 상충되는 복원 요구사항 (전역 vs 국소, 조명 vs 구조) 을 처리할 수 있도록, MoE 를 MoE 안에 중첩 (MoE-in-MoE) 하는 새로운 계층적 설계를 제시했습니다.
• 실용적 가치: 다양한 실제 환경 (Real-world scenarios) 에서 발생하는 복합적인 열화에 대해 강건한 (Robust) 복원 성능을 제공하며, 생성 모델의 잠재력을 저하 없이 활용합니다.
• 미래 방향: 이 프레임워크는 저수준 컴퓨터 비전 작업 (Low-level vision tasks) 에서 생성 모델과 MoE 를 결합하는 새로운 표준을 제시하며, 향후 더 복잡한 다중 태스크 학습 문제 해결에 중요한 시사점을 줍니다.

요약하자면, MiM-DiT는 다양한 열화 유형에 맞춰 동적으로 아키텍처와 세부 파라미터를 조정하는 이중 계층 MoE를 확산 모델에 접목함으로써, 기존 방법들의 한계를 극복하고 고충실도이면서 시각적으로 자연스러운 올인원 이미지 복원을 실현한 획기적인 연구입니다.