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

DreamBooth: Fine Tuning Text-to-Image Diffusion Models for Subject-Driven Generation

paper: https://arxiv.org/abs/2208.12242 github: https://github.com/XavierXiao/Dreambooth-Stable-Diffusion

github은 stable diffusion에 접목된 버전입니다.

DreamBooth: Fine Tuning Text-to-Image Diffusion Models for Subject-Driven Generation 0. Abstact 1. Introduction 2. Related Work 3. Preliminaries 4. Method 4.1. Representing the Subject with a Rare-token Identifier 4.2. Class-specific Prior Preserve Loss 4.3. Personalized Instance-Specific Super-Resolution 5. Experiments 5.1. Applications 5.2. Ablation Studies 5.3. Comparisons 5.4. Limitations Reference

0. Abstact

Large text-to-image 모델은 high-quality image genration을 가능하게 했지만, 주어진 reference set의 subjects(주제, 피사체)를 모방하는 능력이 부족했다. 본 논문에서는 (표현하고자 하는) subject에 대한 few 이미지만으로 text-to-image 모델을 fine-tune하는 접근방식을 제시한다. 특정 subject에 대한 고유한 identifier를 binding하고 모델에 embeded되어 있는 semantic prior를 (class-specific) prior preservaion loss와 함께 레버리지하여, subject를 레퍼런스 이미지에는 없었던 다양한 장면, 포즈, 구도, 조명 등으로 합성할 수 있도록 fine-tuning한다.

1. Introduction

Imaginery scene 합성은 그간 challenging한 task였으나, 최근 Large text-to-image 모델이 주목할만한 성과를 보여줬다. 가장 중요한 adventage는 semantic pair(”dog”라는 단어를 여러 개에 대한 instance와 바인딩 할 수 있는)인데, 이 Large text-to-image 모델들은 주어진 subject유지가 안된다.

근본적인 이유는 output domain이 제한되어 있기 때문이다. 심지어 아주 강하게 text decription을 주더라도 다른 모습으로 나온다. 이는 Text-Object pair가 Fig2.의 다른 모델에서 표현되는 것과 같이 근본적으로 다양하기 때문일 것이다.

본 논문에서는 text-to-image diffusion model의 “personalization(개인화)”에 대한 새로운 접근법을 제안한다. 목표는 모델의 language-vision dictionary를 확장하는 것이며, 이를 통해 유저들이 특정 subject를 자신이 원하는 새로운 단어와 결합하여 생성할 수 있게하기 위함이다. 보다 엄밀하게는, 주어진 소량(3~5장)의 subject 사진에 대해, 논문의 objective가 subject를 모델의 output domain에 심을 것이고, 유저는 subject에 바인딩된 unique identifier를 활용해 이미지를 합성할 수 있다.

본 논문에서는 다음의 2가지 새로운 기법을 제시한다.

주어진 subject를 rare token identifier로 표현하는 기법

diffusion based text-to-image 프레임워크 모델을 fine-tune하는 기법

2번 기법의 경우 diffsuion 모델의 생성 프레임워크는 2단계(text로 부터 low-resolution 이미지를 생성 → super-resolution model을 통한 upscale)로 구성되며, Fine-tuning 또한 2단계로 나뉘어진다.

unique-identifier를 포함하고 있는 text 프롬프트(eg. “A [V] dog”)에 대해 low-resolution text-to-image 모델을 fine-tune

overfitting과 languate drift를 방지하기 위해 class-specific prior preservation loss 적용

super-resolution 모델을 fine-tune

subject의 중요하지 않은 작은 디테일까지도 높은 fidelity로 보존하기 위함

요약하자면, 본 논문이 기여한 바는 다음과 같다.

새로운 문제를 정의: subject-driven generation(subject에 대한 몇장의 사진이 주어지면, 다양한 context로 합성하되, 동시에 subject에 대한 key feature를 high fidelity로 유지하는 것)

새로운 기법을 제안: few-shot setting에서 text-to-image diffusion model을 기존의 semantic knowledge를 유지하면서도 fine-tune하는 새로운 기법 제안

2. Related Work

Related Work에서는 Fine-Tuning 하지 않는 방식으로 subject를 보존하며 새롭게 생성하는 문제(subject-driven novel image generation)를 풀 수는 없었는지 다른 방식들에 대해 정리되어 있다. 기존 방법론들이 어떤 결과를 가졌었는지를 간략하게 정리해보았다.

(1) Compositing Objects into Scenes (Object를 Scene으로 합성)

주어진 subject를 새로운 배경에 클론하는 방식인데, 이러한 방식으로는 배경의 조명이나 사물들, 그림자 등을 자연스럽게 적용할 수 없다.

(2) Text-Driven Editing (텍스트 기반 편집)

CLIP 베이스로 editing하는 GAN들이 대표적인데 highly-curated dataset(예를들어 사람 얼굴 등)에 대해서는 editing이 성공적이지만, 여러 subject 타입을 가진 diverse dataset에 대해서는 잘 못한다. 또한 전역적인 특성이나 local editing은 잘하긴 하는데, 특정 subject를 새로운 맥락으로 생성해내는 것은 불가능했다.

(3) Text-to-Image Syntesis (텍스트 기반 이미지 생성)

Imagen, DALL-E2, Parti, CogView2 등 텍스트 기반 이미지 생성 모델들은 최고 성능을 뽐냈지만 text guidance만을 사용하여 섬세한 컨트롤을 제공하지 못한다. 이를 극복하기 위해 몇몇 연구들에서 mask로 area를 제한하는 방식을 사용했으나, subject를 새로운 context로 합성하지는 못했다.

(4) Inversion of diffusion models (DM의 Inversion을 조작하는 방식)

이건 diffusion model에 input으로 들어가는 latent를 제어하고자 하는 방식(latent를 요리조리 가지고 노는 방식, w, w+, S space 등에서의 조작)으로 내가 GAN으로 Face Editing 다룰 때 연구했던 분야다. Diffusion model에서는 latent인 nosie map이나 원하는 conditioning vector를 찾아서 대응하는 image를 생성해야 하는데, 향후 본 문제(subject-driven novel image generation)를 풀 수 있는, 가능한 방식 중에 하나이다. 하지만, 근본적으로 Textual Modality(텍스트로 제어하는 경우) 표현력의 한계나 기존 모델 output domain의 한계로 인해 현재까지는 제한된다. (관련 방법론들은 아래 논문 참고)

📑

[논문리뷰] Prompt-to-Prompt Image Editing with Cross-Attention Control

(5) Personalization (개인화; 특정 도메인에 대한 fine-tuning 방식?)

개인화는 ML 씬에서 자주 다루어지는 방식인데, 추천시스템, 언어모델에서는 많이 다루고, 이 문제에 대해 적용된 사례는 GAN을 이용한 방식이 있다. MyStyle의 경우 face 합성을 위해 특정 identity로 fine-tune했다. MyStyle이 얼굴에 관한 개인화 문제를 풀기 위해서만 100장이 필요한데, DreamBooth는 다양한 피사체(animal, etc..)에 대해 3~5장이면 충분하다.

Dreambooth는 Personalization에 해당!

3. Preliminaries

(1) Cascaded Text-to-Image Diffusion Models

다음의 간단한 식으로 conditional diffusion model 가 noise가 추가된 이미지 를 denoising하도록 squared error loss로 학습된다.

는 ground-truth, 는 conditioning vector(text prompt 등으로 부터 얻어진), 은 noise다. , , 는 noise schedule과 샘플 품질을 조절하는 term이며, diffusion process time t는 을 따른다. Inference 시에는 을 반복적으로 denoising하여 샘플링하며, deterministic DDIM이나 stochastic ancestral sampler를 사용한다. 는 x를 예측하는 function이다.

최근의 SOTA text-to-image diffusion model들은 cascaded diffusion model을 사용하여 high-quality image를 생성한다. 구체적으로, 출력 해상도가 64×64인 기본 text-to-image 모델을 사용하고, 2개의 text-conditioning SR 모델이(6464 → 256256, 256256 → 10241024)사용된다.

(2) Vocabulary Encoding

text-conditioning의 detail는 Text-to-image diffusion model에서 visual fidelity와 semantic fidelity에 매우 중요하다. 기존 논문들은 학습된 prior를 사용하여 이미지 임베딩으로 변환되는 CLIP 텍스트 임베딩을 사용하거나, 사전 학습된 T5-XXL을 사용한다.

본 논문에서는 pretrained LM(T5-XXL; Language Model)을 사용한다. 이 LM은 tokenized text prompt의 embedding을 생성하며, vocabulary encoding은 프롬프트 임베딩을 위한 전처리 단계에서 굉장히 중요하다. Text Prompt 를 conditioning embedding 로 변환하기 위해서는 먼저 텍스트가 학습된 vocabulary를 사용하는 tokenizer 로 토큰화되어야 한다. 본 논문에서는 SentencePiece tokenizer를 사용하였다.

tokenizer 로 를 tokenize하고나면 고정된 길이의 vector인 를 얻을 수 있다. LM 는 이 vector로 를 생성한다. diffusion model에는 가 conditioning된다.

4. Method

목적은 주어진 subject의 새로운 이미지(피사체가 있는 배경 변경 or 피사체의 색깔, species, 모양 변경 or 피사체의 포즈, 표현, 매질 or 다른 semantic 변경)를 만들어 내는 것이다. Method의 High-Level Overview는 아래와 같다.

이를 위한 첫번째 task는 subject를 output domain에 심는 것, subject를 unique identifier와 바인딩하는 것이다. 이때의 key problem은 subject를 담은 소수의 이미지셋으로 fine-tuning을 시도할때 주어진 사진에 대한 overfitting에 취약하다는 것이다. 또한 language drift 문제도 있는데 모델이 같은 클래스의 다른 피사체들을 어떻게 생성하는지 까먹는 문제로, 해당 클래스의 다양하고 자연스러운 variation에 대한 embedded knowled를 다 잃어버리는 문제다.

이 문제들을 해결하기 위해 autogenous class-specific prior preservation loss를 도입했는데, subject와 같은 클래스의 다양한 instance를 계속해서 생성할 수 있도록 함으로서 overfitting을 경감시키고 language drift를 방지한다.

detail의 보존을 강화하려면 모델의 super-resolution 부분도 fine-tuning해야 한다. 그러나 naive한 방식으로 대상 인스턴스를 생성하도록 fine-tuning된 경우 instance의 중요한 details를 복제할 수 없다. subject의 detail을 더 잘 보존하기 위해 이러한 SR 모듈을 학습시키고 테스트할 수 있는 통찰력을 제공하여 recontextualization에서 전례 없는 성능을 달성하였다. Training process는 Fig 4.에 나와 있다. 본 논문에서는 사전 훈련된 Imagen 모델을 기본 모델로 사용한다.

4.1. Representing the Subject with a Rare-token Identifier

(1) Designing Prompts for Few-Shot Personalization

본 논문의 Goal은 새로운 (key, value) pair를 모델의 dictionary에 추가하는 것으로 한 가지 방법은 few-shot fine-tuning을 하는 것이다. 문제는 ‘이 과정을 어떻게 supervise할 것인가?’이다. Text prompt는 사람에 의해 쓰여졌고, data는 온라인으로부터 sourced되었다. 주요한 한계는 fine-tuning에 사용할 이미지셋에 대해 일일히 디테일한 description 작성이 필요하다는 점, 그리고 사람이 작성한 caption은 높은 variance와 주관성을 가진다는 점이다.

때문에 본 논문에서는 ‘a [identifier] [class noun]’으로 라벨링하는 방식을 선택했다. 여기서 identifier는 subject(피사체)의 고유한 식별자이고, class noun은 classifier로 부터 얻어진 cat, dog, watch 등과 같은 굵직한 class 서술자이다. class noun 없이 identifier로만 학습시, 학습시간이 늘어나고 퍼포먼스도 떨어짐을 확인했다고 한다. 이 기법을 통해 목적이었던 ‘subject(피사체)를 새로운 포즈 or new context로 생성해내기’가 가능해진다.

identifier를 구성하는 가장 단순한 방법은 ‘unique’, ‘special’과 같은 존재하는 단어를 사용하는 것이다. 그러나 이는 기존의 의미를 disentangle하고 새로운 의미를 entangle 해야해서, 학습시간이 늘어나고 퍼포먼스가 떨어진다. 극단적인 경우 ‘blue’를 identifier로 두고, subject(피사체)가 회색을 띄고 있을 때, blue와 gray가 섞인 이미지가 나오게 된다.

따라서, Language Model과 Diffusion Model 모두에 weak prior를 가진 identifier가 필요하다. 가장 먼저 생각나는 방식은 랜덤 문자열(eg. “xxy5syt00”)을 생성해 사용하는 것인데, 이 경우 문자 그대로가 그림으로 묘사되거나, 조금이나마 연관된 이미지가 나오는 등 위험하고, 기존 단어를 사용하는 것과 마찬가지로 성능이 나쁘다

(2) Rare-token Identifiers

간단히 말해서 본 논문의 접근 방식은 vocabulary에서 상대적으로 rare한 token을 찾고, 이 token을 text space로 invert시키는 것이다. 먼저, rare-token lookup을 수행해 rare-token identifier 를 얻는다. de-tokenizer를 사용해 로부터 디코딩된 text 를 얻는다. 토큰은 공백없이 3자 혹은 더 적은 수의 Unicode 문자로 구성되는 것이 성능이 가장 좋다. (eg. , sks는 9258길이의 vocabulary의 9061번째 index로 rare-token identifier이다)

4.2. Class-specific Prior Preserve Loss

(1) Few-shot Personalization of a Diffusion Model

target 피사체를 묘사한 이미지 set 과 text prompt “a [identifier] [class noun]”으로 부터 얻어진 동일한 conditioning vector 에 대해 기존 loss ()로 fine-tune 했더니 두가지 문제가 발생했다. Overfitting과 Language-drift다.

(2) Issue-1: Overffiting

주어진 image set이 꽤 적었기 때문에, large image generation 모델을 fine-tuning하는 것은 주어진 image set의 context나 모습에 대해 overfit 될 수 있었다. Figure 12.의 위쪽 경우가 input image에 대해 overfit된 케이스 이다. 이 문제에 대해 접근할 수 있는 regularization이나 모델의 부분만 선택적으로 fine-tuning 하는 등의 기법이 있는데, fine-tuning의 경우 모델의 어느 파트가 frozen 되어야 subject-fidelity와 semantic modification의 유연성을 얻을 수 있는지 불확실했다. 경험에 따르면 모델의 모든 layer에 대해 fine-tuning하는것이 성능이 좋으나 language-drift문제를 일으켰다.

(3) Issue-2: Language Drift

Language Drift문제는 Language Model 문헌들에 나타나있는데, 특정 task를 위해 language model을 fine-tuning할 때, 해당 task에 대한 성능을 높이기 위해 모델이 언어에 대한 syntactic and semantic knowledge를 점진적으로 까먹는 현상이다. 본 논문에서는 text prompt가 [identifier]와 [class noun]을 갖는데, diffusion 모델이 피사체에 대한 약간의 image set을 fine-tuning한 뒤에 점차 class-specific prior를 까먹으면서 같은 class에 속하는 다른 instance들을 생성해내지 못함을 발견했다. Figure 13의 중간에 이러한 현상을 볼 수 있다.

(4) Prior-Preservation Loss

본 논문은 overfitting과 language-drift 문제를 해결하기 위해 class-specific prior-perserving loss를 제안한다. 이 method는 모델을 기존에 생성되던 sample들로 supervise하는데, few-shot fine-tuning이 시작된 뒤에도 기존의 prior를 유지하기 위함이다. 특히나, data를 frozen된 diffusion model의 기존(ancestral) sampler로 생성한다. 이때 랜덤 초기 noise 이고, conditioning vector 이다.

는 prior-preservation term을 컨트롤한다. Figrue 4의 이 loss로 학습하는 과정을 보인다. 간단한데도 불구하고 이 prior-preservaiton loss가 overfitting과 language-drift 문제를 해결하는데 효과적임을 발견했다. learning rate 로 epoch 정도가 가장 좋은 결과를 내고, TPUv4 1개로 15분 정도만 학습하면 된다.

4.3. Personalized Instance-Specific Super-Resolution

text-to-image 모델이 가장 중요한 visual semantics를 컨트롤 한다고 할 때, SR(Super-Resolution)모델들은 이미지에서 피사체의 detail을 보존하고 phorealistic하게 만든다. 만약 SR모델이 함께 fine-tuning되지 않는다면 생성된 새로운 피사체 이미지에서 artifacts가 발견되거나, 부정확한 feature묘사, 디테일 상실 등이 있을 수 있다. Figrue 14.의 아랫줄에서 이런 현상들을 확인할 수 있다. 64x64 → 256x256 SR 모델을 fine-tuning하는 것은 필수적이고, 256x256 → 1024x1024 SR 모델의 경우는 높은 수준의 fine-grained deatil에 대해 좋다.

5. Experiments

5.1. Applications

(1) Recontextualization

모델 가 주어지면 unique identifier와 class noun을 포함하는 문장을 학습된 모델에 입력하여 특정 주제 인스턴스에 대한 새로운 이미지를 생성할 수 있다. Recontextualization을 위해 일반적으로 “a [V] [class noun] [context description]”의 형식으로 문장을 구성한다.

(2) Art Renditions

“a painting of a [V] [class noun] in the style of [famous painter]”나 “a statue of a [V] [class noun] in the style of [famous sculptor]”와 같은 프롬프트로 대상 인스턴스의 예술적 연출을 생성할 수 있다. 특히, 이 task는 source scene의 semantic을 유지하고 다른 이미지의 스타일을 source scene으로 transfer하는 style transfer과 동일하지 않다. 대신 예술적 스타일에 따라 subject 인스턴스 디테일과 정체성 보존을 통해 장면에서 의미 있는 변화를 얻을 수 있다. 즉, 생성되는 일부 이미지에는 주제 ground truth 데이터에서 볼 수 없는 다른 포즈의 주제와 ground truth에도 없는 장면이 있다.

(3) Expression Manipulation

표정을 수정한 새로운 이미지 생성에 대한 예시

(4) Novel View Synthesis

새로운 시점으로 대상을 렌더링한 예시

(5) Accessorization

prior에 의해 가능한 대상에 액세서리를 추가하는 능력을 보여주는 예시

(6) Property Modification

대상의 속성을 수정한 예시이다. 본 논문의 방법을 사용하면 semantic하게 복잡한 속성을 수정할 수 있다.

5.2. Ablation Studies

(1) Class-Prior Ablation

잘못된 클래스에 대한 클래스 prior가 얽힌 상태로 남아 있고 모델이 이러한 방식으로 학습될 때 대상의 새로운 이미지를 생성할 수 없음을 관찰할 수 있다. Class noun 없이 학습시키는 경우 모델은 대상 인스턴스를 학습하는 데 어려워하며 클래스 prior와 대상 인스턴스를 얽지 못한다. 모델은 수렴하는 데 더 오래 걸리고 오차가 큰 샘플을 생성한다.

(2) Prior Preservation Loss Ablation

다음은 prior-preservation loss가 overfitting을 피하는 데 효과적인 것을 보여준다.

다음은 prior-preservation loss를 사용할 때 클래스 의미 prior가 보존되는 것을 보여준다.

(3) Super Resolution with Low-Noise Forward Diffusion

다음은 noise augmentation의 level을 낮게 두는 것의 효과를 보여준다. Noise augmentation의 level을 낮추었을 때 샘플 품질과 대상 fidelity가 개선된다.

5.3. Comparisons

다음은 An Image is Worth One Word 논문의 모델과 결과를 비교한 것이다. 본 논문의 접근 방식이 의미론적으로 정확한 이미지를 생성하며 대상의 feature를 더 잘 보존한다.

5.4. Limitations

몇가지 limitation이 존재한다. 먼저, 다음과 같이 3가지 주요 failure mode가 존재한다.

프롬프트 맥락과 다른 이미지 생성

Context-appearance entanglement

보여준 원본 대상과 비슷한 프롬프트에 대하여 overfitting 발생

또 다른 제한점은 몇몇 대상이 다른 대상에 비해 학습이 더 빠른 것이다. 흔한 대상에 대해서는 강한 prior가 있어 학습이 빠르고 희귀하거나 복잡한 대상에 대해서는 학습이 오래 걸린다. 마지막으로 주제의 fidelity에도 가변성이 있으며 일부 생성된 이미지에는 모델 prior의 강도와 의미론적 수정의 복잡성에 따라 주제에 대한 환각적인 feature가 포함될 수 있다.

Reference

https://arxiv.org/abs/2112.10752

< PREV

📑

[논문리뷰] LDM: High-Resolution Image Synthesis with Latent Diffusion Models

NEXT >

📑

[논문리뷰] Prompt-to-Prompt Image Editing with Cross-Attention Control