Fast High-Resolution Image Synthesis with Latent Adversarial Diffusion Distillation (SD3-Turbo, LADD)

 

Abstract

픽셀 기반인 Adversarial Diffusion Distillation과 달리 사전 훈련된 잠재 확산 모델의 generative feature를 활용하여 ADD의 한계를 극복한 Latent Adversarial Diffusion Distillation (LADD) 제안

 

[arXiv](2024/03/18 version v1)

 

 

 

Introduction

Adversarial Diffusion Distillation (ADD)에서는 판별기의 backbone으로 DINOv2 encoder를 사용하는데, 이는 단점이 있다.

• 판별기의 훈련 해상도가 518×518로 제한된다.
• 판별기의 피드백 수준을 제어할 수 없다.
• 판별기가 잠재 공간에서 훈련되지 않아 RGB 공간으로 디코딩해야 한다.

 

 

 

Background

Adversarial Diffusion Distillation

DiT 기반의 Diffusion Transformer인 Stable Diffusion 3

 

 

 

Method

• 잠재 공간에서 증류 수행
• 교사 모델을 판별기 backbone으로 사용
• 훈련을 위한 합성 데이터 생성

논문에 loss function에 대한 구체적인 언급이 따로 없다. ADD와 똑같은 loss를 사용하는 듯.

 

Unifying teacher and discriminator

판별기의 backbone으로 교사 모델을 사용하며 판별기 입력은 깨끗한 이미지가 아닌 noised latent이다.

교사 모델의 각 attention block 이후의 토큰 시퀀스를 추출하여 각각에 독립적인 discriminator head를 적용한다.

 

Latent Adversarial Diffusion Distillation (LADD)의 장점:

• RGB 공간으로 디코딩할 필요가 없으므로 시스템이 단순화된다.
• 확산 모델은 timestep에 따라 각각 다른 feature level에 집중하는데, backbone으로 timestep과 text를 조건으로 받는 교사 모델을 사용함으로써 이를 반영할 수 있다.
• 교사 모델은 다양한 종횡비의 이미지를 생성할 수 있기 때문에 해상도가 제한되지 않는다.
• CNN은 전체적인 모양보다 질감을 우선시하는 texture bias가 있다. 하지만 교사 모델(Transformer)은 인간과 유사한 형상 편향을 갖고 있어 인간 인식과의 정렬을 향상시킬 수 있다.

 

Discriminator head 아키텍처는 ADD와 같이 StyleGAN-T를 따르며 종횡비를 구별하지 못하는 1D conv 대신 2D conv를 사용.

 

Leveraging synthetic data

텍스트 정렬은 데이터셋에 따라 크게 다르며 최첨단 확산 모델은 더 높은 CLIP score를 달성할 수 있다.

예를 들어, COCO dataset은 평균 0.29의 CLIP score를 갖지만 SD3은 COCO prompt에서 0.35의 더 높은 CLIP score를 달성한다.

 

따라서 일정한 CFG 값에서 합성 데이터를 생성하여 사용한다. 이는 텍스트 정렬을 높이고 교사의 지식을 정제할 수 있다.

 

 

 

Experiments

이 섹션에서의  정성적 결과는 깊이를 줄인 SD3에서 수행되었으므로 실제 모델의 출력은 더 좋음.

Teacher noise distribution

깨끗한 이미지에 추가되는 noise level의 분포를 조절. 낮으면 global consistency ↓, 높으면 detail ↓.

m = 1, s = 1 설정 채택. 

 

Synthetic data

panda scientist mixing chemicals

 

Latent distillation approaches

Latent Consistency Model과 비교.

 

Scaling Behavior

데이터 품질과 교사 모델 크기는 포화되었으며 학생 모델의 크기는 매우 큰 영향을 끼친다.

 

Direct preference optimization

교사 모델의 모든 선형 레이어에 LoRA를 추가하고 Diffusion-DPO를 수행한다.

여기까지는 원래 SD3에도 있는 내용이다. 그러니까 교사 모델은 애초에 이미 DPO가 적용된 모델이다.

 

놀라운 점은, 학습된 학생 모델에 이 교사 모델의 DPO-LoRA를 그대로 추가하는 것만으로도 품질이 향상된다는 것이다.

바퀴, 손 등의 디테일이 향상되었다.

 

 

 

Comparison to State-of-the-Art

기본으로 4 step, 또는 2 step을 사용

Consistency Model의 sampler 채택 (그냥 t = 0 예측하고 노이즈 추가하는 게 다임.)

고해상도에서 훈련할 때, 훈련 초기에 낮은 noise level부터 시작하는 warm-up 단계가 매우 중요하다고 한다.

 

 

 

 

Image-to-Image

 

Inpainting