Pretraining is All You Need for Image-to-Image Translation (PITI) 논문 리뷰

사전 훈련된 확산 모델을 통한 image-to-image 변환

 

Project Page

PITI

 

 

 

Abstract

이전의 image-to-image 변환 방법은 전용 아키텍처 설계가 필요하고 처음부터 훈련하며 대규모의 쌍을 이루는 데이터셋이 필요하다.

본 논문에서는 image-to-image 변환을 다운스트림 작업으로 간주하고 다양한 종류의 변환을 수용하기 위해 사전 훈련된 확산 모델을 적용하는 간단하고 일반적인 프레임워크를 소개한다.

또한 생성 품질을 개선하기 위한 정규화된 가이드 샘플링과 함께 확산 모델 훈련에서 텍스처 합성을 향상하기 위한 적대적 훈련을 제안한다.

 

 

 

Introduction

이미지 변환은 사전 훈련된 신경망을 사용하여 자연스러운 이미지 매니폴드를 탐색하고 입력 의미와 관련된 실현 가능한 지점을 찾는 것과 같다. 이를 위해 합성 네트워크는 방대한 양의 이미지를 사용하여 사전 훈련되어야 하며 잠재 공간에서 샘플링할 때 그럴듯한 출력으로 이어지도록 하는 generative prior의 역할을 한다.

 

사전 훈련 Generative prior으로 GLIDE(Guided Diffusion + CLIP) 채택. 복잡한 장면을 모델링하고 전체 이미지 분포를 이상적으로 포착하기 위해 확산 모델이 적합하고, GLIDE가 의미론적 및 저차원 잠재성을 모델에 잘 반영하도록 설계되었기 때문이다.

 

다운스트림 작업을 수용하기 위해 사전 훈련된 모델의 잠재 공간에 변환 입력(e.g. 분할 마스크)을 투영하는 작업별 헤드를 훈련한다. 작업을 위한 네트워크는 인코더-디코더 구조이며, 인코더는 입력을 잠재 공간으로 변환하고 디코더는 확산 모델이다.

 

먼저 사전 훈련된 디코더를 고정하고 인코더만 업데이트한 다음 전체 네트워크를 공동으로 미세 조정한다. 이러한 단계별 훈련은 주어진 입력에 대한 충실도를 보장하면서 사전 훈련된 지식을 최대한 활용할 수 있다.

 

생성 품질을 위한 추가적인 기술

• 거친 이미지를 생성한 다음 초해상도를 수행하는 cascade 생성 전략을 채택. 하지만 확산 모델의 denoising 성질로 인해 지나치게 부드러운 결과를 생성하므로 denoising process 동안 적대적 훈련 도입.
• Classifier-free Guidance는 세부 정보가 씻겨나간 지나치게 포화된 이미지를 생성하기 때문에 noise 통계를 명시적으로 정규화할 것을 제안.

 

 

제안된 모델인 PITI(pretraining-based image-to-image translation)는 다양한 다운스트림 작업에서 전례없는 품질을 달성한다.

 

 

 

Approach

 

Generative pretraining

텍스트 조건화되고 거대하고 다양한 텍스트-이미지 쌍에 대해 훈련되는 GLIDE를 사전 훈련 모델로 채택.

또한 GLIDE는 여러 해상도의 모델이 포함된 계층적 구조의 cascade 생성 방법을 사용한다.

기본 모델과 업샘플러 모델을 각각 미세 조정하여 다양한 다운스트림 이미지 합성 작업에 적용할 수 있다.

 

Downstream adaptation

Base model finetuning

저해상도 모델의 생성은 훈련된 각각의 인코더 ε̃, 디코더 D̃, 사전 훈련에 사용된 조건 y에 대해

 

사전 훈련된 임베딩 공간에 조건부 입력을 매핑하도록 다양한 작업별 헤드 ε_i를 훈련한다.

(인코더에 대한 자세한 설명은 따로 없음.)

 

인코더만 먼저 훈련한 다음 전체를 미세 조정하는 2단계 미세 조정 방식 제안. 이러한 방식은 사전 훈련된 지식을 최대한 활용하고 품질 향상을 위해 매우 중요함.

 

Adversarial diffusion upsampler

고해상도 모델을 추가로 미세 조정.

훈련 이미지와 지나치게 부드러운 저해상도 샘플 간의 간격을 줄이기 위해 BSR 열화와 L₀ 필터 도입.

 

그럼에도 불구하고 지나치게 부드러운 샘플이 생성돼 확산 모델의 Gaussian noise assumption이 문제라고 추측, 이를 해결하기 위해 지각 손실과 적대적 손실 사용.

각각의 손실은 이미지 예측 x̂₀^t에 대해 다음과 같이 공식화

(D_θ는 적대적 손실을 최대화 하려는 판별자, ψ는 VGG 모델)

 

Normalized classifier-free guidance

모델이 조건부 입력을 무시하고 조건과 상관없는 결과를 생성하는 것을 막기 위해 샘플링 동안 p(x_t|y)와 함께 p(y|x_t)를 고려하는 classifier-free guidance를 도입한다.

 

샘플링 동안 조건에 null 토큰을 입력하고 무조건인 채로 noise를 추정한 뒤 조건부 모델과의 선형 결합으로 샘플을 생성할 수 있다. (가이드 강도 계수 w)

 

가이드는 조건에 대해 무조건 모델이 추정한 분포에서 평균과 분산을 이동시킴으로써 수행되는데, 이러한 통계 이동은 모든 T 확산 denoising 단계를 통해 누적되어 세부 정보가 씻겨나간 지나치게 포화된 이미지로 이어진다.

 

이를 해결하기 위해 원래 추정 ε_θ(x_t|y)에 guided noise 샘플 ε̂_θ(x_t|y)의 통계를 명시적으로 일치시키는 normalized classifier-free guidance를 제안한다.

Normalized classifier-free guidance는 특히 w가 클 때 샘플링 품질을 더욱 향상할 수 있다고 한다.

 

 

 

Experiments

 

디코더도 미세 조정 해야함.

 

NCF = Normalized Classifier-Free guidance