[논문 리뷰] Learning A Coarse-to-Fine Diffusion Transformer for Image Restoration

Liyan Wang, Qinyu Yang외 4명|arXiv (Cornell University)|2023. 08. 17.

MRI in cancer diagnosis인용 수 5

한 줄 요약

이 논문은 이미지 복원을 위한 코arse-to-fine 확산 트랜스포머(C2F-DFT)를 제안하며, 노이즈 제약 외에 샘플링된 출력을 진수 이미지와 대조하여 개선하는 새로운 훈련 방식을 도입함으로써 복원 품질을 향상시킨다. 이 방법은 시간 단계 조건부 확산 자기주도 및 피드포워드 네트워크를 활용하여 강도 렌즈, 빗방울 제거 및 노이즈 제거 작업에서 최신 기술 수준의 성능을 달성한다.

ABSTRACT

Recent years have witnessed the remarkable performance of diffusion models in various vision tasks. However, for image restoration that aims to recover clear images with sharper details from given degraded observations, diffusion-based methods may fail to recover promising results due to inaccurate noise estimation. Moreover, simple constraining noises cannot effectively learn complex degradation information, which subsequently hinders the model capacity. To solve the above problems, we propose a coarse-to-fine diffusion Transformer (C2F-DFT) for image restoration. Specifically, our C2F-DFT contains diffusion self-attention (DFSA) and diffusion feed-forward network (DFN) within a new coarse-to-fine training scheme. The DFSA and DFN respectively capture the long-range diffusion dependencies and learn hierarchy diffusion representation to facilitate better restoration. In the coarse training stage, our C2F-DFT estimates noises and then generates the final clean image by a sampling algorithm. To further improve the restoration quality, we propose a simple yet effective fine training scheme. It first exploits the coarse-trained diffusion model with fixed steps to generate restoration results, which then would be constrained with corresponding ground-truth ones to optimize the models to remedy the unsatisfactory results affected by inaccurate noise estimation. Extensive experiments show that C2F-DFT significantly outperforms diffusion-based restoration method IR-SDE and achieves competitive performance compared with Transformer-based state-of-the-art methods on $3$ tasks, including image deraining, image deblurring, and real image denoising. Code is available at https://github.com/wlydlut/C2F-DFT.

연구 동기 및 목표

• 샘플링 중 정확하지 못한 노이즈 추정으로 인해 복원 품질이 열등해지는 확산 기반 이미지 복원 모델의 한계를 해결하기 위해.
• 단순한 노이즈 제약을 넘어서 복잡한 열화 패턴을 학습하는 이중 단계 훈련 방식을 통해 모델 용량을 향상시키기 위해.
• 확산 통합 트랜스포머 아키텍처를 사용하여 이미지 복원에서 장거리 의존성 모델링 및 계층적 특징 학습을 향상시키기 위해.
• 노이즈가 아닌 샘플링된 복원 결과를 사용하여 모델을 최적화하는 미세조정 전략을 개발하여 정확하지 못한 노이즈 추정에서 발생하는 오류를 수정하기 위해.
• 강도 렌즈, 빗방울 제거 및 실세계 노이즈 제거를 포함한 여러 이미지 복원 벤치마크에서 경쟁 가능한 성능을 달성하기 위해.

제안 방법

• 모델는 시간 단계 임bedding에 조건부로 지정된 확산 자기주도(DFSA) 및 확산 피드포워드 네트워크(DFN)를 갖춘 확산 트랜스포머 아키텍처를 사용하여 확산 과정의 시간 역동성을 모델링한다.
• 코어스-투-파인 훈련 방식이 도입되었으며, (a) 노이즈 예측을 제약하는 코어스 훈련과 (b) 샘플링 과정에서 유도된 출력을 진수 이미지와 대비하여 제약하는 파인 훈련으로 구성된다.
• 파인 훈련 단계에서는 노이즈가 아닌 코어스 트레이닝된 모델에서 유도된 샘플링된 복원 결과를 사용하여 정확하지 못한 노이즈 추정에서 발생하는 오류를 수정한다.
• 일관성과 효율성을 확보하기 위해 고정된 수의 샘플링 단계(4단계)를 사용하여 미세조정 과정을 수행하며, 시각적 품질과 계산 비용 간의 균형을 이룬다.
• 코어스 훈련 동안 패치 사이클 훈련 전략을 사용하여 일반화 및 효율성을 향상시켰으며, 유사한 훈련 시간 동안 고정 패치 기반 모델보다 뛰어난 성능을 보였다.
• 시간 단계 조건부는 DFSA 및 DFN에 모두 통합되어 모델이 다양한 확산 단계에서 행동을 적응시킬 수 있도록 하여 복원 성능을 크게 향상시켰다.

실험 결과

연구 질문

• RQ1 노이즈가 아닌 샘플링된 출력을 개선하는 코어스-투-파인 훈련 방식이 확산 모델의 이미지 복원 품질 향상에 기여하는가?
• RQ2 확산 트랜스포머에 통합된 시간 단계 임베딩이 장거리 의존성 및 복원 정확도를 파악하는 데 모델의 능력에 어떤 영향을 미치는가?
• RQ3 파인 튜닝 단계에서 생성된 출력에 제약을 가하는 것이 표준 노이즈 제약 훈련보다 더 뛰어난 성능을 내는가?
• RQ4 파인 튜닝 파이프라인에서 복원 품질과 추론 속도를 균형 잡는 데 최적의 샘플링 단계 수는 얼마인가?
• RQ5 패치 사이클 훈련 전략은 고정 패치 훈련 대비 모델의 일반화 및 효율성 향상에 얼마나 효과적인가?

주요 결과

• 제안된 C2F-DFT는 강도 렌즈 벤치마크 Rain100H에서 PSNR 31.38 dB, SSIM 0.921을 기록하여 기준 코어스 훈련(30.93 dB PSNR) 및 IR-SDE를 능가한다.
• Test2800 데블러링 세트에서 모델은 34.03 dB PSNR 및 0.944 SSIM을 달성하여 코어스 오직 훈련 대비 0.29 dB 향상되었다.
• 파인 튜닝 단계는 시각적으로 도식된 그림 9 및 10에서 볼 수 있듯이 코어스 훈련 대비 잡음과 잔상 감소로 시각적 품질을 크게 향상시켰다.
• 시간 단계 임bedding는 PSNR를 4.08 dB 향상시켰다(26.85에서 30.93 dB로), 이는 성능 향상에 있어 핵심적인 역할을 한다는 것을 입증한다.
• 파인 튜닝에서 4단계의 샘플링 단계를 사용하면 시각적 품질(LPIPS)과 추론 시간 간의 유리한 균형을 이룰 수 있었으며, 더 적거나 더 많은 단계보다 뛰어난 성능을 보였다.
• 패치 사이클 훈련 전략은 유사한 훈련 시간 동안 고정 패치 훈련 대비 더 뛰어난 성능을 달성하여 그 효과성을 확인했다.