[논문리뷰] DreamBooth: Fine Tuning Text-to-Image Diffusion Models for Subject-Driven Generation

DreamBooth: Fine Tuning Text-to-Image Diffusion Models for Subject-Driven Generation
Nataniel Ruiz, Yuanzhen Li, Varun Jampani, Yael Pritch, Michael Rubinstein, Kfir Aberman
CVPR 2023 [paper] [github]

---

Contents

1. Introduction

2. Method

3. Experiments

4. Limitations

5. Summary

---

1. Introduction

본 논문은 Textual Inversion과 마찬가지로 pretrained text-to-image 모델을 personalization하는 기술인 DreamBooth를 제시한다. Personalization은 pretrained된 모델에 사용자가 원하는 새로운 컨셉을 주입시켜 원하는 이미지를 생성할 수 있도록 하는 task이다.

 

DreamBooth는 3~5장의 이미지를 이용하여 모델을 fine-tuning하고, 주어진 컨셉을 나타내는 unique identifier [V]를 사용하여 이미지를 합성한다. 아래 사진은 DreamBooth를 이용해 주어진 컨셉에 대한 이미지를 합성한 예시이다.

기존 연구는 주어진 이미지의 디테일을 표현하지 못하거나 prompt에 맞게 컨셉을 변형하는 능력에 한계가 존재한다. 아래 사진에서 DreamBooth의 subject와 prompt에 대한 높은 fidelity를 확인할 수 있다.

---

2. Method

DreamBooth는 text-to-image diffusion model을 base로 한다. Conditioning text prompt $\mathbf{P}$는 text encoder $\Gamma$를 이용하여 conditioning vector $\mathbf{c}$로 인코딩된다. Diffusion model의 objective function은 아래와 같으며, 자세한 설명은 생략한다. $$\mathbb{E}_{\mathbf{x}, \mathbf{c}, \boldsymbol{\epsilon}, t}\left[w_t\left\|\hat{\mathbf{x}}_\theta\left(\alpha_t \mathbf{x}+\sigma_t \boldsymbol{\epsilon}, \mathbf{c}\right)-\mathbf{x}\right\|_2^2\right]
$$

 

Personalization의 목표는 주어진 컨셉을 모델의 도메인에 주입하는 것이다. 연구진들은 subject를 표현하기 위한 prompt로 "[identifier] [class noun]"의 형태를 사용하였다. 이때 [identifier]는 subject를 의미하는 unique identifier이고, [class noun]은 subject의 class를 대략적으로 설명하는 단어이다. 예를 들어 cat, dog, watch 등이 [class noun]으로 사용된다. 이는 모델의 사전 지식을 활용하여 훈련 시간을 줄이고 성능을 향상시키는 역할을 한다. Unique identifier로는 특별한 의미를 가지거나 자주 사용되는 단어를 사용해서는 안된다. 이 경우에는 모델이 단어를 원래의 의미로부터 분리하고 컨셉과 연결되도록 학습해야 한다. 따라서 특별한 의미를 가지지 않는 단어를 활용하는 것이 효율적이다. 무작위 알파벳으로 구성된 "xxy5syt00"와 같은 단어는 tokenizer에 의해 각 글자별로 토큰화되기 때문에 마찬가지로 좋지 않은 선택이다. 연구진들은 모델의 vocabulary에서 희귀 토큰을 찾아 text space로 변환하는 방식으로 unique token을 선택한다. 희귀 토큰 $f(\hat{\mathbf{V}})$를 찾고 대응하는 단어 $\mathbf{V}$를 unique identifier로 사용한다.

 

모델의 fine-tuning 과정에서는 크게 두 문제가 발생할 수 있다. 첫 번째 문제는 language drift이다. Language drift는 모델을 fine-tuning하는 과정에서 언어의 문법적, 의미론적 능력을 상실하는 것이다. 예를 들면 모델이 subject와 동일한 class에 해당하는 다른 물체를 생성하지 못하는 문제가 해당한다. 두 번째 문제는 다양성 감소이다. Fine-tuning 과정에서 모델이 overfitting 되어 생성된 이미지가 학습한 이미지와 같은 자세와 각도로만 고정되는 현상이 발생할 수 있다.

 

이러한 문제를 해결하기 위해 연구진들은 class-specific prior preservation loss를 제안한다. Fine-tuning 전에 기존 모델을 이용해 이미지를 생성하고, 생성한 이미지를 fine-tuning 과정에서 학습하여 모델이 사전 지식을 유지할 수 있도록 한다. 이를 위해 frozen된 pre-trained 모델로부터 $\mathbf{x}_{\mathrm{pr}}=\hat{\mathbf{x}}\left(\mathbf{z}_{t_1}, \mathbf{c}_{\mathrm{pr}}\right)$을 생성한다. 이때 conditiong vector는 다음과 같다: $\mathbf{c}_{\mathrm{pr}}:=\Gamma(f($ "$\mathrm{a}$ [class noun]" $))$

따라서 전체 loss는 다음과 같다.$$
\begin{aligned}
\mathbb{E}_{\mathbf{x}, \mathbf{c}, \boldsymbol{\epsilon}, \boldsymbol{\epsilon}^{\prime}, t}\left[w_t\left\|\hat{\mathbf{x}}_\theta\left(\alpha_t \mathbf{x}+\sigma_t \boldsymbol{\epsilon}, \mathbf{c}\right)-\mathbf{x}\right\|_2^2+\right. \left.\lambda w_{t^{\prime}}\left\|\hat{\mathbf{x}}_\theta\left(\alpha_{t^{\prime}} \mathbf{x}_{\mathrm{pr}}+\sigma_{t^{\prime}} \epsilon^{\prime}, \mathbf{c}_{\mathrm{pr}}\right)-\mathbf{x}_{\mathrm{pr}}\right\|_2^2\right]
\end{aligned}
$$ 식에서 두 번째 항이 prior-preservation 항을 나타낸다.

아래 그림에서 모델의 학습 과정을 확인할 수 있다. 위에서는 few-shot dataset으로 subject-identifier를 학습하고 아래에서는 기존 모델이 생성한 이미지를 학습하여 사전 지식을 유지한다.

학습은 A100 GPU, Stable Diffusion 기준으로 5분 정도가 소요된다. 

---

3. Experiments

3.1. Dataset and Evaluation

연구진들은 DreamBooth의 성능 평가를 위해 데이터셋을 구축했다. 데이터셋은 30개의 subject에 대한 이미지들과 25개의 prompt로 구성되어 있고, 각 (subject, prompt)에 대해 4개의 이미지를 생성하여 전체 3000개의 이미지로 성능을 평가한다. 아래 그림은 데이터셋의 일부이다.

 

모델의 성능을 평가하는 요소 중 하나는 subject fidelity이다. 모델이 생성한 이미지가 subject의 디테일을 얼마나 잘 보존하고 있는지 평가하기 위해서, CLIP-I와 DINO를 사용한다. CLIP-I는 실제 이미지와 생성된 이미지의 CLIP embedding 사이의 average pairwise cosine similrarity이다. CLIP-I는 기존 연구에서도 사용되던 지표이지만, 비슷한 text description을 가질 수 있는 서로 다른 subject에 대한 평가가 어렵다. 따라서 연구진들은 실제 이미지와 생성된 이미지의 ViT-S/16 DINO embedding 사이의 average pairwise cosine similrarity인 DINO 평가 지표를 추가로 사용한다. DINO의 self-supervised training objective는 subject의 unique feature을 구별하도록 학습한다.

모델을 평가하는 또 다른 요소는 prompt fidedlity이다. 이는 prompt와 생성된 이미지의 CLIP embedding 사이의 average cosine similarity인 CLIP-T로 평가한다.

3.2. Comparisons

아래 표는 DreamBooth와 Textual Inversion을 위에서 설명한 세 개의 지표로 평가한 결과이다. Baseline 모델은 Imagen과 Stable Diffusion이다. subject fidelity와 prompt fidelity 모두 DreamBooth가 Textual Inversion보다 높은 성능을 보이는 것을 확인할 수 있다.

아래는 DreamBooth와 Textual Inversion의 subject fidelity와 prompt fidelity를 설문을 통해 비교한 결과이다. DreamBooth가 훨씬 높은 선호도를 보인다.

아래 이미지는 꽃병 사진을 학습한 결과를 비교한 것이다. 사용된 Prompt는 왼쪽부터 순서대로 “a [V] vase in the snow”, “a [V] vase on the beach”, “a [V] vase in the jungle”, “a [V] vase with the Eiffel Tower in the background”이다. Subject fidelity와 prompt fidelity 모두 DreamBooth가 더 좋은 결과를 생성하는 것을 볼 수 있다.

3.3. Ablation Studies

Prior preservation loss(PPL)의 효과를 확인하기 위해 PPL을 사용한 모델과 사용하지 않은 모델을 학습했다. PPL의 목적은 language drift를 방지하고, 모델의 prior를 보존하는 것이다. Prior 보존을 평가하기 위한 지표로 subject가 해당하는 prior class의 random subject에 대해 생성한 이미지와 subject 훈련 이미지의 average pairwise DINO embeddings인 PRES를 사용한다. PRES가 높을수록 class의 random subject가 우리의 subject과 유사한, prior가 잘 보존되지 않고 langauge drift 문제가 발생했다는 의미이다.

또한 모델의 diversity를 평가하기 위해 동일한 subject와 prompt로 생성된 이미지들의 average LPIPS consine similrarity인 DIV를 사용한다. 결과는 아래 표에서 확인할 수 있다.

PPL을 이용해 학습한 모델이 약간의 subject fidelity 감소와 함께 높은 diversity를 가진다는 것을 확인할 수 있다.

 

Prompt에서 identifier [V] 뒤에 위치한 [class noun]의 영향을 확인하기 위해 class가 없거나 잘못된 class이 적힌 prompt로 모델을 학습했다. 잘못된 class로 학습된 모델은 주어진 subject의 형태를 변형한다. Class가 없이 학습된 모델은 subject를 학습하는데 어려움을 겪었다. 결과는 아래 표에서 확인할 수 있다.

3.4. Applications

아래 그림과 같이 상황을 서술하는 prompt("a [V] [class noun] [context description]")를 통해 subject에 대한 새로운 이미지를 생성할 수 있다.

Prompt를 적절히 조정하여 아래 그림과 같은 작업도 가능하다. 물체를 이전에 바라본 적 없는 방향에서 바라보거나, style transfer처럼 사용하거나, 물체의 특징을 변형할 수 있다.

---

4. Limitations

위 그림에서 DreamBooth가 실패한 결과물들을 확인할 수 있다. 첫 번째 문제는 prompt를 잘 반영하지 않는 문제이다. 기존 모델의 사전 지식이 부족하거나, 흔하지 않은 subject-prompt 쌍의 경우가 해당된다. 두 번째 문제는 prompt에 의해 subject의 외형이 변하는 문제이다. 세 번째 문제는 훈련에 사용한 이미지와 비슷한 prompt를 사용하는 경우에 overfitting된 이미지가 생성되는 문제이다. 또한 특정 subject에 대해서는 학습이 잘 되지 않거나 희귀한 subject에 대해 다양성이 감소하는 문제 등이 존재한다.

---

5. Summary

연구진들은 적은 수의 이미지와 prompt만으로 subject에 대한 다양한 이미지를 생성할 수 있는 DreamBooth 기술을 제시했다. DreamBooth의 핵심은 subject를 unique identifier [V]와 결합함으로써 모델의 도메인에 임베딩하는 것이다. 또한 prior preservation loss를 도입하여 모델이 subject를 fine-tuning하는 동안 language drift가 발생하지 않도록 하였다. Textual Inversion보다 훨씬 빠른 학습 속도를 가지면서도 subject와 prompt를 보존하는 능력이 더 뛰어남을 확인했다. 다만, subject에 따라 fine-tuning이 실패하거나 prompt 무시, 과적합 등의 문제가 발생할 수 있다는 한계가 존재한다.