[논문 리뷰] Designing a Better Asymmetric VQGAN for StableDiffusion

Zixin Zhu, Xuelu Feng외 6명|arXiv (Cornell University)|2023. 06. 07.

Generative Adversarial Networks and Image Synthesis인용 수 5

한 줄 요약

이 논문은 StableDiffusion를 위한 비대칭 VQGAN을 제안하여, 이미지 복원 및 국소 편집 중 편집되지 않은 영역에서의 왜곡을 줄인다. 작업별 사전 지식(예: 복원에서의 마스크 해제된 영역)을 사용하는 조건부 디코더 브랜치와 인코더보다 훨씬 깊은 디코더를 도입함으로써, 높은 정밀도의 세부 정보를 유지하면서도 추론 비용을 낮추어, 복원 및 편집 벤치마크에서 최신 기술 수준의 성능을 달성한다.

ABSTRACT

StableDiffusion is a revolutionary text-to-image generator that is causing a stir in the world of image generation and editing. Unlike traditional methods that learn a diffusion model in pixel space, StableDiffusion learns a diffusion model in the latent space via a VQGAN, ensuring both efficiency and quality. It not only supports image generation tasks, but also enables image editing for real images, such as image inpainting and local editing. However, we have observed that the vanilla VQGAN used in StableDiffusion leads to significant information loss, causing distortion artifacts even in non-edited image regions. To this end, we propose a new asymmetric VQGAN with two simple designs. Firstly, in addition to the input from the encoder, the decoder contains a conditional branch that incorporates information from task-specific priors, such as the unmasked image region in inpainting. Secondly, the decoder is much heavier than the encoder, allowing for more detailed recovery while only slightly increasing the total inference cost. The training cost of our asymmetric VQGAN is cheap, and we only need to retrain a new asymmetric decoder while keeping the vanilla VQGAN encoder and StableDiffusion unchanged. Our asymmetric VQGAN can be widely used in StableDiffusion-based inpainting and local editing methods. Extensive experiments demonstrate that it can significantly improve the inpainting and editing performance, while maintaining the original text-to-image capability. The code is available at \url{https://github.com/buxiangzhiren/Asymmetric_VQGAN}.

연구 동기 및 목표

• StableDiffusion 기반 이미지 편집 중 편집되지 않은 영역에서 발생하는 왜곡 아티팩트의 근본 원인을 규명하고 이를 해결한다.
• 특히 텍스트와 같은 세밀한 구조를 가진 영역에서 편집된 이미지의 정밀도와 세부 정보 유지 능력을 향상시킨다.
• 기존 StableDiffusion 파ip라인과의 호환성을 유지하면서도 편집 작업 성능을 향상시킬 수 있는 VQGAN 변종을 설계한다.
• 새로운 아키텍처가 훈련 및 추론 모두에서 효율성을 유지하여 광범위한 구현이 가능하도록 보장한다.

제안 방법

• 디코더는 VQGAN 인코더 출력과 작업별 사전 지식(예: 복원에서의 마스크 해제된 영역)을 입력으로 받는 조건부 디코더로 재구성된다.
• 인코더보다 훨씬 깊고 넓은 디코더를 사용하여, 양자화된 잠재 표현에서 고정밀도 세부 정보를 더 잘 복원할 수 있도록 한다.
• 사전 훈련 중 일부 마스크를 전체 마스크로 교체함으로써, 디코더가 사전 지식 없이도 전체 이미지를 복원하도록 훈련시킨다.
• 디코더만 재훈련하며, 원본 VQGAN 인코더와 StableDiffusion 모델은 그대로 유지되어 플러그 앤 플레이 통합이 가능하다.
• 조건부 브랜치는 편집 작업 중에만 활성화되며, 텍스트 기반 이미지 생성 시에는 비활성화되어 생성 능력은 그대로 유지된다.
• 동적 입력 라우팅을 통해, 사전 지식이 있는 편집 작업과 사전 지식이 없는 순수 텍스트 기반 이미지 생성 작업을 모두 지원한다.

실험 결과

연구 질문

• RQ1 기존의 StableDiffusion 기반 편집 방법은 왜 편집되지 않은 영역에서 심각한 왜곡을 유발하는가?
• RQ2 작업별 사전 지식을 포함하는 조건부 디코더 브랜치가 편집되지 않은 영역의 재구성 아티팩트를 줄일 수 있는가?
• RQ3 인코더보다 디코더의 용량을 늘리면 세부 정보 복구 능력과 양자화로 인한 왜곡이 향상되는가?
• RQ4 제안된 비대칭 VQGAN은 메인 디퓨전 모델을 재훈련하지 않더라도 편집 및 텍스트 기반 이미지 생성 작업 모두에서 높은 성능을 유지할 수 있는가?
• RQ5 제안된 비대칭 설계에서 개선된 이미지 품질과 추론 비용 간의 상충 관계는 어떠한가?

주요 결과

• 제안된 비대칭 VQGAN은 이미지 복원 작업에서 Place 데이터셋에서 1.03 FID 점수를 기록하여 최신 기술 수준의 성능을 달성했다.
• COCOEE 데이터셋에서 모델은 페인팅-바이-예제 편집에서 86.35%의 CLIP 점수를 기록하여 이전 방법들보다 뚜렷이 뛰어난 성능을 보였다.
• 제거 분석 결과, 조건부 브랜치가 페인팅-바이-예제 작업에서 보존 오차를 588.86에서 0.76으로 줄였으며, 강력한 세부 정보 유지 능력을 입증했다.
• 조건부 브랜치 없이도 더 큰 디코더는 기준 모델보다 더 정교한 세부 정보를 복원할 수 있었으며, 이는 고정밀도 복원 능력이 내재되어 있음을 시사한다.
• 순수 텍스트 기반 이미지 생성 작업에서 모델은 원본 StableDiffusion와 유사하거나 더 뛰어난 성능을 유지하여 일반성의 타당성을 확인했다.
• 그림 5, 6, 7의 시각적 결과는 다양한 편집 및 생성 시나리오에서 조화와 세부 정보 유지 측면에서 일관된 향상을 확인한다.