Your Pre-trained Diffusion Model Secretly Knows Restoration

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

• 배경: 사전 훈련된 대규모 텍스트 - 이미지/비디오 확산 모델 (Diffusion Models, 예: Stable Diffusion, FLUX, WAN) 은 자연 이미지 분포에 대한 강력한 사전 지식 (Prior) 을 가지고 있어 고품질 콘텐츠 생성에 탁월합니다. 최근 이러한 모델들을 'All-in-One Restoration (AiOR, 모든 종류의 열화 제거)' 작업에 활용하려는 시도가 증가하고 있습니다.
• 기존 방법의 한계:
  • 기존 AiOR 방법론들은 주로 모델의 파인튜닝 (Fine-tuning) 이나 ControlNet 스타일의 조건부 모듈을 추가하여 사전 훈련된 지식을 활용합니다. 이는 막대한 계산 비용이 들고, 훈련 데이터 분포에 과도하게 의존하여 실제 복잡한 열화 상황에서의 강건성이 떨어질 수 있습니다.
  • 플러그 - 앤 - 플레이 (Plug-and-play) 방식은 학습이 필요 없으나, 명시적인 열화 연산자 (Degradation operators) 가 필요하여 복잡한 실제 환경에서는 적용이 어렵습니다.
• 핵심 질문: "대규모 확산 모델의 사전 지식에 복원 (Restoration) 과 관련된 지식이 이미 내재되어 있는가? 만약 그렇다면, 이를 어떻게 효과적으로 꺼낼 수 있는가?"
• 발견된 문제점:
  1. 자연어 프롬프트 및 토큰 최적화의 실패: 자연어 프롬프트 (" haze 제거" 등) 나 텍스트 토큰 공간 (Token-space) 에서 프롬프트를 최적화하더라도, 모델은 열화를 제거하지 않고 단순히 노이즈만 제거하여 원래 열화된 입력으로 수렴하는 경향이 있습니다.
  2. 학습 - 추론 궤적 불일치 (Train-Test Trajectory Mismatch): 열화된 이미지에 노이즈를 추가하는 전방 과정 (Forward process) 으로 프롬프트를 학습하면, 학습 시와 추론 시 (Reverse sampling) 에 방문하는 중간 상태 (Intermediate states) 의 분포가 달라져 프롬프트의 성능이 불안정해지고 아티팩트가 발생합니다.

2. 제안된 방법론 (Methodology)

저자들은 사전 훈련된 확산 모델의 백본 (Backbone) 을 동결 (Freeze) 시키고, 오직 조건부 프롬프트 (Conditioning Prompt) 만 학습하는 경량화된 프레임워크를 제안합니다.

A. 임베딩 공간 프롬프트 최적화 (Embedding-Space Prompt Optimization)

• 핵심 통찰: 텍스트 토큰 공간이 아닌, 텍스트 인코더의 출력 공간 (Context Embedding Space) 에서 직접 조건부 벡터를 학습해야 합니다.
• 구현: 텍스트 인코더의 출력에서 직접 학습 가능한 임베딩 벡터 $p$를 최적화하여, 동결된 확산 모델이 열화된 잠재 공간 (Latent) 에서 깨끗한 이미지 ($z_{clean}$) 로 매핑되도록 유도합니다.
• 효과: 이 방식은 모델이 내재하고 있던 복원 능력을 활성화하여, 노이즈가 섞인 열화된 입력을 명확한 이미지로 복원하게 합니다.

B. 확산 브릿지 (Diffusion Bridge) 를 통한 궤적 정렬

• 문제 해결: 학습 시 열화된 상태 ($z_{deg}$) 에 고정된 노이즈를 추가하는 방식과 추론 시 열화에서 깨끗한 이미지로 점진적으로 이동하는 방식 사이의 불일치를 해결하기 위해 확산 브릿지 (Diffusion Bridge) 개념을 도입합니다.
• EBR (Energy-Oriented Diffusion Bridge) 적용:
  • 학습과 추론 모두에서 열화된 상태 ($z_{deg}$) 에서 깨끗한 상태 ($z_{clean}$) 로 부드럽게 이동하는 중간 상태들을 생성합니다.
  • 특히, 초기 단계에서는 열화된 입력에 노이즈가 포함된 상태를, 후기 단계에서는 점점 더 깨끗한 상태로 이동하는 단조 감소 (Monotone) 노이즈 스케줄을 사용합니다.
  • 이를 통해 학습된 프롬프트가 추론 시에도 동일한 중간 상태 분포를 경험하도록 하여 궤적 불일치를 제거하고 안정적인 복원을 가능하게 합니다.

C. 설계 선택 사항

• 잔여 프롬프트 (Residual Prompt): 텍스트 인코더 출력을 직접 최적화하는 대신, 'Null-text' 컨텍스트에 잔여 프롬프트를 추가하고 주입하는 방식을 사용하여 파라미터 효율성을 높이고 성능을 개선했습니다.
• 혼합 열화 처리: 각 열화 유형 (안개, 비, 눈 등) 마다 별도의 프롬프트를 학습하며, 혼합된 열화의 경우 샘플링 시 각 프롬프트의 예측 속도를 평균화하여 처리합니다.

3. 주요 기여 (Key Contributions)

1. 복원 지식의 존재 증명 및 활성화: 자연어 프롬프트나 토큰 최적화로는 접근 불가능하지만, 텍스트 인코더 출력 공간의 임베딩을 직접 학습함으로써 사전 훈련된 확산 모델이 강력한 복원 능력을 내재하고 있음을 증명하고 이를 활성화했습니다.
2. 궤적 불일치 해결: 학습과 추론 간의 궤적 불일치 문제를 식별하고, 확산 브릿지 (EBR 스타일) 를 도입하여 학습 및 추론 동역학을 정렬함으로써 프롬프트 학습의 불안정성을 해결했습니다.
3. 경량화 및 범용성: 파인튜닝이나 복잡한 제어 모듈 없이, WAN (비디오 모델, 13 억 파라미터) 과 FLUX (이미지 모델, 120 억 파라미터) 에 경량 프롬프트만 적용하여 SOTA(State-of-the-Art) 수준의 성능을 달성했습니다.

4. 실험 결과 (Results)

• 데이터셋: 다양한 열화 (안개, 비, 눈, 모션 블러, 저조도) 에 대한 이미지 및 비디오 데이터셋 (RESIDE, Snow100K, LOLv1, GoPro 등) 과 OOD(Out-of-Distribution) 데이터셋을 사용했습니다.
• 성능:
  • 이미지 복원 (FLUX): 기존 AiOR 방법 (AutoDIR, PixWizard, FoundIR 등) 과 비교하여 다양한 열화 조건에서 경쟁력 있는 성능을 보였습니다. 특히 PSNR/SSIM 같은 왜곡 지표뿐만 아니라 CLIPIQA, MUSIQ 같은 지각적 지표에서도 우수한 성능을 기록했습니다.
  • 비디오 복원 (WAN): ViWS-Net, AverNet 등 기존 비디오 복원 모델보다 OOD 데이터셋에서 더 높은 DOVER(비디오 품질 평가 지표) 점수를 기록하며, 시간적 일관성 (Temporal consistency) 이 뛰어난 것을 입증했습니다.
  • 일반화 능력: 훈련되지 않은 열화 (Unseen degradations) 나 혼합된 열화 조건에서도 강력한 일반화 능력을 보여주었습니다.
• 정성적 평가: 시각적 결과에서 기존 방법들이 열화를 제거하지 못하거나 아티팩트를 생성하는 반면, 제안된 방법은 디테일을 유지하면서 열화를 효과적으로 제거하는 것을 확인했습니다.

5. 의의 및 결론 (Significance)

• 패러다임 전환: 대규모 확산 모델을 AiOR 작업에 적용할 때, 무거운 파인튜닝이나 추가 모듈 대신 경량 프롬프트 학습만으로도 모델이 가진 잠재된 복원 능력을 unlock 할 수 있음을 보여주었습니다.
• 효율성: 모델 백본을 동결하고 오직 작은 프롬프트만 학습하므로, 메모리 및 계산 비용이 크게 절감됩니다.
• 실용성: 복잡한 실제 환경의 열화에도 강건하게 작동하여, 다양한 저수준 비전 (Low-level vision) 작업에 적용 가능한 범용적인 솔루션을 제시합니다.

이 논문은 사전 훈련된 생성 모델이 단순히 생성뿐만 아니라 복원 (Restoration) 과 같은 저수준 작업에도 깊은 지식을 내재하고 있으며, 올바른 조건부 메커니즘 (임베딩 공간 프롬프트 + 브릿지 정렬) 을 통해 이를 효과적으로 활용할 수 있음을 규명한 중요한 연구입니다.