[논문 리뷰] One-Step Image Translation with Text-to-Image Models

arxiv paper code
Gaurav Parmar, Taesung Park, Srinivasa Narasimhan, Jun-Yan Zhu
Garnegie Mellon University
18 Mar 2024

Abstract

Fig. 1 : 태스크 예시

본 논문에서는 기존의 Diffusion Model에 대한 다음 두 가지 한계점에 대해 언급한다.

1. noising process를 반복해야해서 발생하는 느린 생성 속도
2. Fine-tuning을 위해서는 반드시 Paired dataset을 구성해야하는 이슈

위 두 이슈에 대해 다음 해결법을 제시하였다.

• Adversarial learning objective를 이용하여, 새로운 태스크 및 도메인에 single-step에 general하게 adapt 가능한 방법

우선, vanilla ldm을 하나의 end-to-end 생성 모델로 통합하였다. 이때, 과적합을 방지하며 입력 이미지의 구조를 보존하기 위한 작은 trainable weights를 결합하였다.

이렇게 제안한 모델 CycleGAN-Turbo는 scene translation task(day-to-night, 눈,안개,비 등의 날씨 추가/제거)에서 기존의 GAN 방식과 Diffusion 방식 모두의 성능을 능가한 성능을 보였다. 또한, ControlNet의 Sketch2Photo, Edge2Image 와 같은 paried data를 적용한 제안 모델 pix2pix-Turbo은 단 한 스텝으로 작동한다.

저자는 single-step diffusion 모델은 다양한 GAN 학습을 위한 강력한 Backbone이 된다고 언급하였다.

(참고로 SD-Turbo라는 논문이 2023년에 이미 등장하였으며 본 논문은 해당 논문과 구조적으로 비슷함 : paper code weight)

Introduction

저자는 먼저, 실험을 통해 ControlNet와 같은 Adapter는 one-step으로 이미지를 생성하기에 효과적이지 않았다고 한다. 그리고 one-step model에서는 기존의 Diffusion 모델과 달리 noise map이 생성 이미지의 구조에 영향을 끼치는 것을 확인할 수 있었다고 한다. 따라서 추가 adapter에 noise map과 input condition을 같이 입력하게 되면 서로 충돌이 일어난다고 주장한다.

저자가 강조하는 부분 중 하나는 image2image를 할 시, Multi-stage pipeline (Endoer-UNet-Decoder) 구조로 인해 입력 이미지의 디테일을 잃게 된다고 하였다. 이러한 요소는 '낮을 밤으로 바꿔라'와 같은 현실 이미지에서 눈에 띄는 문제를 발생시킨다. (이에 관한 업무를 경험한 적이 있어 관심이 많이 가는 부분이다.)

Fig. 2 : 저자가 제안하는 아키텍쳐

저자는 이러한 문제를 해결하기 위해 기존의 SD-Turbo weight를 유지하면서 입력 이미지의 구조를 보존하는 새로운 생성 아키텍쳐를 제안하였다.

1. noise map과 input control의 충돌을 피하기 위해 UNet의 noise encoder branch에 conditioning 정보를 다이렉트로 입력
2. Encoder, UNet, Decoder를 single end-to-end trainable architecture로 통합.
   -> fine-tuning 시간을 줄이고 overfitting을 해소하기 위해 LoRA를 사용하여 original weight를 유지한 채 새로운 control과 domain에 적용 가능하게 한다.
3. 입력 이미지의 high-frequency detail을 보존하기 위해 encoder와 decoder 사이에 zero-conv를 사용한 skip-connection을 추가

이 새로운 아키텍쳐는 다목적이며 CycleGAN과 pix2pix와 같은 conditional GAN의 objective를 위한 plug-and-play 역할을 한다.

• CycleGAN-Turbo : unpaired image translation model
• pix2pix-Turbo : paired image translation model such as ControlNet

Related Work

Image-to-Image translation

최근 Diffusion 모델에 사용 가능한 대표적인 모델(GLIGEN, T2I-Adapter, ControlNet)들은 Adapter 방식으로 효과적인 Fine-tuning 기술이지만, 여전히 large paired dataset이 필요하다는 조건이 있다. 그에 반해 본 논문에서 제시하는 모델은 image pair가 필요없으며 생성 속도도 굉장히 빠르다.(512*512 기준 0.3초; GPU 사양 확인해야함)

Text-to-Image models

기존의 방식들(finetuning network/text embedding/ inversion)은 복잡한 이미지 생성에 어려움을 갖지만 저자의 방식은 새로운 도메인 및 태스크에 대해 paired/unpaired data로 이를 보강햇다고 한다.

One-step generative models

최근, diffusion model의 sampling step을 줄이기 위해 fast ODE solvers와 distillation 등 제안된 바 있다. 하지만 이러한 방법은 blurry result를 가질 수 있으며 이를 해소하기 위해 다양한 방법들이 제안되어오고 있다. 저자는 sampling step을 줄이는 위 방법들과 달리 one-step에 초점을 두고 접근하였다. 결과적으로, one-step distilled ControlNet를 능가했다고 한다.

Fig. 3

Method 1 : Adding Conditioning input

문제정의 : Conflicts between noise and conditional input

저자는, 기존의 Control Adapter 방식은 UNET Encoder와 Condition Encoder 두 개를 사용하여 context상에서 noise map과 임베딩된 control map 간의 마찰이 빚어진다고 주장한다. 또한, multi-step diffusion 모델과 달리, one-step 모델에서는 noise map이 생성 이미지의 layour이나 pose에 직접적인 영향을 끼치게 되어 control 이미지와 모순되는 경우가 많다고 한다. 이 때문에 UNet의 Decoder에서는 모순되는 두 residual feature를 받기에 학습 프로세스가 더 어렵다고 주장한다.

해결방안 : Direct conditioning input

위 그림(Fig. 3)에서 noise map $z$는 이미지 생성 결과의 구조에 큰 영향을 끼치는 것으로 보여진다. 저자는 이것에 인사이트를 얻어, conditioning image를 다이렉트로 Network의 입력으로 사용하는 방법을 제안하였다. 이는 추가적인 adapter를 사용하는 것보다 더 효과적임을 Figure 7과 Table 4에서 증명했다. 또한, backbone 모델이 다양한 conditioning에 적응할 수 있도록 UNet의 다양한 Layer에 LoRA를 적용하였다.

Method 2 : Preserving Input Details

Fig. 4

Why details are lost

Latent Diffusion Models(LDMs)의 Encoder(VAE)는 입력 이미지의 채널을 3 -> 4로 증가시키고 해상도를 8배 감소하여 결과적으로 이미지를 압축시킨다. 이는 학습 및 생성 속도를 증가시키기 위한 방법이다. 하지만 image translation task에서는 입력 이미지의 디테일을 보존하는것이 요구된다. 위 이미지(Fig. 4)는 낮의 드라이브 이미지(왼쪽 이미지)를 밤으로 바꾸는 태스크에 대한 예시이다. skip connection을 사용하지 않은 중간 이미지의 경우, 글자와 간판, 자동차 등에서 디테일이 바뀌는 모습을 볼 수 있다. 그러나, skip connection을 적용한 architecture(오른쪽 이미지)에서는 복잡한 디테일들이 보존되는 것을 볼 수 있다.

Connecting first stage encoder and decoder

저자는 입력 이미지의 디테일 정보를 보존하기 위해 Encoder와 Deocder 사이에 skip connection을 적용하였다. 구체적으로, Encoder의 Downsampling block에서 4개의 activation을 $1 \times 1$ zero-convolution layer에 통과시킨 후 Decoder의 upsampling block에 연결시켰다.

Method 3 : Unpaired Training

unpaired dataset 예시 (실제 학습 데이터)

• 프롬프트 예시 : picture of a hores

우선, 본 논문의 모든 실험은 one-step, SD2.1-Turbo를 기본으로 진행되었다. unpaired translation에는 수정된 CycleGAN 수식이 사용되었다. 구체적으로, $unpiared dataset X = { x \in X }, \quad Y = { y \in Y }$이 주어졌을 때, source domain $X \subset \mathbb{R}^{H \times W \times 3}$을 target domain $Y \subset \mathbb{R}^{H \times W \times 3}$로 변환하는 것을 목표로 한다.

source domain과 target domain을 상호 변환할 수 있도록 하였으며 이를 아래와 같이 표기하였다.

• $G(x, c_Y): X \rightarrow Y$ : 이미지 X와 캡션 Y를 입력하여 Y이미지가 출력되는 것
• $G(y, c_X): Y \rightarrow X$ : 이미지 Y와 캡션 X를 입력하여 X이미지가 출력되는 것
  이 때, 위 두 $G$는 같은 Network를 공유한다.

Cycle consistency with perceptual loss

• $L_{cycle}$ : cycle consistency loss. 각 source image $X$가 $Y$로 변환된 후, 다시 $X$가 되도록 하는 loss. (X -> fake Y -> X)
• $L_{rec}$ : L1 difference와 LPIPS의 조합 (Appendix D 참고) -> 근데 Appendix에서는 L2 사용중. 확인필요.

$$L_{\text{cycle}} = \mathbb{E}_x [L_{\text{rec}}(G(G(x, c_Y), c_X),x)] + \mathbb{E}_y [L_{\text{rec}}(G(G(y, c_X), c_Y),y)] \tag{1}$$

Adversarial loss

$D_X$와 $D_Y$를 통해 해당 domain의 이미지가 맞는지를 분류하는 판별기를 사용한다. 이때, Vision-Aided GAN에서 제안된대로 CLIP model을 backbone으로 사용한다.

$$\begin{align}
L_{\text{GAN}} = \mathbb{E}_y [log D_Y(y)] + \mathbb{E}_x [log(1-D_Y(G(x, c_Y)))]
+ \mathbb{E}_x [log D_X(x)] + \mathbb{E}_y [log(1-D_X(G(y, c_X)))] \tag{2}
\end{align}$$

Full objective

Full objective에는 위 두 Loss 외에 identity regularization loss $L_{idt}$가 포함된다.

• $L_{idt} = \mathbb{E}_y[L_{rec}(G(y,c_Y),y)] + \mathbb{E}_x[L_{rec}(G(x, c_X),x)] $

다음은 최종 objective이다.

$$\underset{G}{argmin} \quad L_{\text{cycle}} + \lambda_{\text{idt}} L_{\text{idt}} + \lambda_{\text{GAN}} L_{\text{GAN}} \tag{3} $$

• $\lambda_{idt}$, $\lambda_{gan}$ : just weight

Extensions

Paired training

canny edge나 lineart에서 이미지를 생성하는 태스크에 대해 위에서 제안한 network $\G$를 사용하였다.
$$G(x,c) : X \rightarrow Y$$

• $X$ : source domain(예 : sketch)
• $Y$ : target domain(예 : output image)
• $c$ : caption

objective function에 사용된 loss는 다음과 같다.(자세한 사항은 Appendix D 참고)

• perceptual loss와 pixel-space reconstruction loss를 조합한 reconstruction loss
• 수식 (2)에서 target domain에 대한 유사도 판별을 진행하는 GAN loss
• CLIP text-image alignment loss인 $L_{CLIP}$

Generating diverse outputs

저자는 interpolation coefficient $\gamma$가 주어졌을 때, 다음 세 가지를 변경한다.

1. Gaussian noise와 encoder output을 결합(combine)한다.
   Generator $G(x,z,r)$는 총 세개의 입력을 취한다. image $x$, noise map $z$, coefficient $\gamma$. 우선, 다음 수식을 통해 noise $z$와 encoder output을 결합한다. $\gamma G_{enc}(x)+(1-\gamma )z$. 결합된 signal은 UNet에 입력된다.
2. Skip-connection 부분의 LoRA에 scale을 적용한다.
   $$\theta = \theta_0 + \gamma \cdot \Delta \theta$$

• $\theta_0$ : pre-trained weights
• $\Delta \theta$ : 새로 추가된 weights

3. reconstruction loss에 scale을 적용한다.
   $$L_{diverse} = \mathbb{E}_{x,y,z,\gamma}[\gamma L_{rec}(G(x,z,\gamma ),y)] \tag{4}$$

• $\gamma = 0$인 경우, pre-trained 모델의 기본 strochastic 행동에 해당되며 reconstruction loss는 반영되지 않음.
• $\gamma = 1$인 경우, deterministic translation에 해당되며 reconstruction loss가 반영됨.
  저자는 다양한 interpolation coefficient $\gamma$로 fine-tuning을 하였다고 한다.

Experiments

저자는 실험을 통해 다음 세 결과를 보였다.

1. image translation task에 대해 GAN-based, Diffusion-based와 비교하여 양적으로나 질적으로나 본 논문에서 제안한 것이 결과가 우수함
2. unpiared method인 CycleGAN-Turbo에 대해, 여러 요소들에 대한 조합을 실험하여 어떤 요소의 결합이 효과적인지 분석하였음
3. paired 환경에서 어떻게 다양한 결과를 생성할 수 있는지를 보임

CycleGAN-Turbo Training Details

• 학습 파라미터 수 : 330MB(LoRA, zero-conv layer, first conv layer of U-NET)
• 학습 및 추론 해상도 : 256*256(얼룩말, 날씨), 512*512(교통)

Table 1 :  Evaluation on standard CycleGAN datasets (256 × 256)

세 데이터 셋에 대한 FID score와 DINO Struct에 대한 비교 결과이다. CycleGAN-Turbo가 Horse -> Zebra 데이터에서 FID score가 조금 더 높게 나왔지만 그 외 모든 데이터에 대해 우수한 성과를 보였다.(Horse 데이터에서도 DINO Struct 점수는 가장 우수하게 나왔다.)(GPU : A6000)

개인적인 견해로, Image Quality를 측정하는 평가지표는 불안정해서 다양한 평가지표가 사용되고 있다. 그런측면에서 4가지 데이터 * 2개의 평가지표 = 총 8개의 평가 결과에서 7개의 결과에서 가장 좋게 나온 것으로 볼 때, 양적 평가 결과로는 이미 충족된다고 판단된다.

Evaluation Protocol

저자는 효과적인 image translation 방법은 두 가지가 충족되어져야 한다고 주장한다.

1. target domain의 데이터 분포와 매칭해야 함
2. 변환된 결과에서 입력 이미지의 구조가 보존되어야 함

저자는 위 두 가지 측면에서의 평가를 위해 clean-FID score를 통해 분포를 비교하였으며, 입출력 이미지 feature space에서의 구조적 유사성을 평가하기 위해 DINO-Struct-Dist를 사용하였다.

만약, FID score가 낮고 DINO-Struct-Dist가 높은 경우는 입력 이미지의 구조가 보존되지 않는 것을 의미한다. 반대로, FID score는 높고 DINO-Struct-Dist가 낮은 경우는 입력 이미지로 부터 거의 변하지 않았다는 것을 의미한다. 따라서, image translation 태스크에서는 위 두 평가지표를 같이 써야 유의미한 평가가 가능하다.

Comparison to Unpaired GAN-based Methods

저자는 우선 GAN-based 모델인 CycleGan, CUT에 대한 비교 실험을 진행하였다. default hyperparameter를 사용하여 100,000 steps을 학습하고 그 중 best ckpt를 채택하였다. 객체 중심 데이터에서 두 모델은 낮은 FID score와 DINO-Structure score를 달성하는 모습을 보였다.(CycleGAN-Turbo는 위 두 모델보다 약간 더 낮은 점수를 기록하였음.) 하지만, 좀 더 복잡한 Scene(night -> day)에서는 CycleGAN-Turbo보다 꽤나 높은 FID score 값을 보였다. 그림을 보아도 조금씩 이상한 artifact를 생성하는 결과를 보였다.

Comparison to Diffusion-based Editing Methods

Diffusion-based 모델들은 객체 중심 데이터에서 그럴듯한 얼룩말을 그려냈지만, 다소 큰 DINO-Structure score를 기록하며 구조적으로 다른 이미지를 생성했다. driving dataset에서도 좋지 못한 성능을 보였는데, 저자는 그 이유에 대해 세 가지를 언급하였다.

1. 다양한 객체가 포함된 복잡한 이미지를 생성하는데 어려움을 겪음
2. instruct-pix2pix를 제외한 모델들은 이미지를 noise map으로 바꾸어 잠재적인 artifact가 발생
3. pre-trained 모델은 driving dataset의 하나의 캡쳐 이미지와 유사한 거리 이미지를 생성할 수 업음

위 이미지에서 볼 수 있듯이 퀄리티 적으로 좋지 아니하고(낮은 FID score), 입력 이미지와 다른 구조의 이미지를 생성하였다.(낮은 DINO-Struct score)

Human Preference Study

Table 3 : Human Preference Evaluation

저자는 Amazon Mechanical Turk (AMT)를 이용해 정성평가를 진행하였다. 평가자에게는 변환된 결과 이미지들과 다음의 프롬프트가 주어졌다. "Which images looks more like a real picture of a driving scene taken in the night?"

평가 모델은 각 GAN/Diffusion-based 모델들 중 가장 성능이 좋았던 CycleGAN과 Instruct-Pix2Pix 그리고, 본 논문에서 제안한 CycleGAN-Turbo이다. 한가지 태스크(Clear -> Foggy)를 제외하곤 Driving dataset에서 가장 좋은 평가를 받았다. Clear -> Foggy의 경우, Instruct-Pix2Pix 모델이 좀더 artistic한 안개 이미지를 생성했었기 때문이라고 주장하며 해당 모델이 본 논문에서 제안하는 모델보다 더 높은 DINO-Struct score를 보였음을 근거라고 한다. 또한, target fog dataset과 달리 상당히 다양한 이미지를 생성했다고 하며 FID score(170.8 > 137.0)를 근거로 내세웠다.

Ablation Study

Table 4 : Ablation with Horse to Zebra

• [Config A] text-to-image model에 대한 pre-trained weight를 사용하지 않은 경우 부자연스러운 결과를 생성(High FID)
• [Config B] ControlNet Encoder를 사용한 경우, Directly Input과 유사한 FID를 달성하지만 높은 DINO를 보임. 또한, 기존에 없던 새로운 객체를 생성하는 결과를 보임(예 : 빌딩, 부자연스러운 얼룩말 무늬).
• [Config C] 적은 weight를 사용하는 T2I-Adapter Encoder의 경우, 좋지 못한 FID, DINO를 보여줌. 결과 이미지에서도 여러 artifact가 생성되고 입력 이미지의 구조가 보존되지 않는 결과를 보임.
• [Config D] Encoder와 Decoder 사이에 skip-connection을 적용한 경우, FID가 약간 상승하지만 구조 보존에 있어서 상당히 효과를 보임.

추가적인 Ablation Study는 Appendix A, C에서 확인할 수 있다.(다양한 데이터 셋에 대한 실험 결과, 다양한 학습 데이터 수, encoder-decoder fine-tuning 방법)

Extensions : Paired translation (pix2pix-Turbo)

저자는 Latent Consistency Models (LCM), Stable Diffusion Turbo (SD-Turbo), Stable Diffusion (SD)에 ControlNet을 적용하여 비교 실험하였다. 실험 결과, 512*512 이미지 생성에 0.29s가 소요되었으며 미적으로도 준수한 결과를 보였다.

또한 저자는 동일한 프롬프트와 input sketch에 대해 noise map을 sampling하여 다양한 이미지를 생성할 수 있다고 강조하였다.

Discussion and Limitations

저자는 one-step pre-trained model이 다양한 downstream image synthesis task에서 강력한 backbone이 될 수 있다고 제안하였다. 또한, multi-step diffusion training 없이 다양한 GAN's objectives를 활용하여 새로운 도메인에 대한 학습이 가능하다는 것을 보였다. 그리고 이러한 학습에 적은 수의 trainable parameter면 충분하다는 것을 입증하였다.

저자는 본 논문에서 제안하는 몬델의 한계점으로 다음

1. 본 논문에서 제안한 Backbone인 SD-Turbo는 classifier-free guidance를 사용하지 않기 때문에 특정 Guidance를 지정할 수 없음(CFG-scale)
2. artifact를 편리하게 줄일 수 있는 Negative prompt를 사용할 수 없음
3. 학습에 사용되는 cycle-consistency loss와 high-capacity generator는 메모리 집약적임.
4. 끝으로 저자는, 다음 연구에서 one-step으로 고해상도 이미지를 생성하는 연구에 대한 필요성을 언급하였다.