[논문 리뷰] Refusion: Enabling Large-Size Realistic Image Restoration with Latent-Space Diffusion Models

Ziwei Luo, Fredrik Gustafsson외 3명|arXiv (Cornell University)|2023. 04. 17.

Advanced Image Processing Techniques인용 수 6

한 줄 요약

Refusion는 U-Net 기반 잠재 확산 모델을 제안하여, 저해상도의 압축된 잠재 공간에서 노이즈 제거를 수행하면서 디코더에 부착된 잠재 벡터를 통해 고해상도 세부 정보를 유지함으로써 고정밀도, 대규모 이미지 복원을 가능하게 한다. 이는 실세계 이미지 복원 작업에서 최고 수준의 인지 성능을 달성하며, 2023년 NTIRE 그림자 제거 챌린지에서 2위를 차지했고, 적대적 훈련이나 데이터셋 정제 없이도 6000×4000 이미지를 성공적으로 처리한다.

ABSTRACT

This work aims to improve the applicability of diffusion models in realistic image restoration. Specifically, we enhance the diffusion model in several aspects such as network architecture, noise level, denoising steps, training image size, and optimizer/scheduler. We show that tuning these hyperparameters allows us to achieve better performance on both distortion and perceptual scores. We also propose a U-Net based latent diffusion model which performs diffusion in a low-resolution latent space while preserving high-resolution information from the original input for the decoding process. Compared to the previous latent-diffusion model which trains a VAE-GAN to compress the image, our proposed U-Net compression strategy is significantly more stable and can recover highly accurate images without relying on adversarial optimization. Importantly, these modifications allow us to apply diffusion models to various image restoration tasks, including real-world shadow removal, HR non-homogeneous dehazing, stereo super-resolution, and bokeh effect transformation. By simply replacing the datasets and slightly changing the noise network, our model, named Refusion, is able to deal with large-size images (e.g., 6000 x 4000 x 3 in HR dehazing) and produces good results on all the above restoration problems. Our Refusion achieves the best perceptual performance in the NTIRE 2023 Image Shadow Removal Challenge and wins 2nd place overall.

연구 동기 및 목표

• 복잡한 실세계의 열화를 가진 대규모 이미지 복원 작업에 확산 모델의 적용 가능성을 향상시키기 위해.
• 기존 확산 모델의 한계, 즉 높은 계산 비용, 알려진 열화 매개변수에 대한 의존성, 잠재 공간 압축에서의 불안정성 등을 극복하기 위해.
• 적대적 최적화 없이도 다양한 복원 작업에 일반화되고 안정적이며 효율적인 프레임워크를 개발하기 위해.
• 낮은 차원의 잠재 공간에서 작동함으로써 고해상도 이미지 복원(예: 6000×4000)을 가능하게 하되, 세부 정보를 유지하기 위해.

제안 방법

• 역 노이즈 제거 과정으로 평균 회귀 확률적 미분 방정식(SDE)을 사용하며, IR-SDE를 영감으로 받아 열화를 암묵적으로 모델링하고 작업 일반화를 가능하게 한다.
• 기존 잠재 확산 모델에서 사용하는 VAE-GAN 대신 U-Net 기반 압축 네트워크를 도입하여, 적대적 훈련 없이도 더 안정적이고 정확한 잠재 표현 학습을 가능하게 한다.
• 노이즈 예측 헤드로 비선형 활성화 없이 설계된 네트워크(NAFNet)를 사용하여 표준 U-Net 아키텍처보다 효율성과 성능을 향상시킨다.
• 잠재 공간은 노이즈 제거에 사용되며, 고해상도 특징은 디코더의 스킵 연결을 통해 유지되어 세부 정보 정확한 재구성 가능하다.
• 데이터셋과 노이즈 네트워크만 변경함으로써 다양한 이미지 복원 작업을 지원하여 유연성과 일반화 능력을 확보한다.
• 잠재 Refusion 변종은 잠재 공간에서 확산을 수행함으로써 약 15배의 FLOPs를 감소시켜, 6000×4000 이미지와 같은 대규모 이미지에서의 추론을 가능하게 한다.

실험 결과

연구 질문

• RQ1 적대적 훈련이나 VAE-GAN 압축에 의존하지 않고도 U-Net 기반 잠재 확산 모델이 안정적이고 고정밀도의 이미지 복원을 달성할 수 있는가?
• RQ2 U-Net 노이즈 네트워크를 NAFNet으로 교체할 경우 이미지 복원 작업의 성능과 효율성에 어떤 영향을 미치는가?
• RQ3 압축된 공간에서의 잠재 확산이 매우 큰 이미지(예: 6000×4000)의 효과적인 복원을 가능하게 하고 고해상도 세부 정보를 유지할 수 있는가?
• RQ4 제안된 방법이 그림자 제거, 안개 제거, 보케 효과 변환과 같은 다양한 실세계 이미지 복원 작업에 얼마나 일반화되는가?
• RQ5 대규모 데이터셋에서 기존의 확산 기반 방법과 비교해 복원의 인지 품질과 추론 속도 측면에서 모델 성능은 어떠한가?

주요 결과

• Refusion는 2023년 NTIRE 이미지 그림자 제거 챌린지에서 2위를 차지했으며, LPIPS(0.149)와 MOS(8.94)에서 가장 높은 인지 점수를 기록하여 모든 팀을 압도했다.
• 실세계 그림자 제거 작업에서 뛰어난 인지 성능을 달성했으며, LPIPS 0.149와 MOS 8.94로 DHAN과 같은 베이스라인 모델보다 뚜렷이 뛰어났다.
• 잠재 Refusion 변종은 전체 Refusion 모델 대비 약 15배의 FLOPs를 감소시켜, 전체 Refusion와 IR-SDE가 실패하는 6000×4000 이미지에서의 추론을 가능하게 했다.
• 보케 효과 변환 작업에서 잠재 Refusion는 Restormer보다 더 낮은 LPIPS 점수를 기록했고, 거의 동일한 속도로 작동하여 더 선명한 결과를 생성했다.
• U-Net 기반 잠재 압축 전략은 VAE-GAN 기반 방법보다 훨씬 더 안정적이었으며, 적대적 최적화 없이도 고정밀도 이미지 복원이 가능했다.
• 모델은 메모리 및 계산 자원 제약으로 인해 전체 해상도 확산 모델인 IR-SDE와 Refusion가 불가능한 6000×4000 이미지에서 비균일한 안개 제거도 성공적으로 수행했다.