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

CVPR 2023. [Paper]
Nataniel Ruiz, Yuanzhen Li, Varun Jampani, Yael Pritch, Michael Rubinstein, Kfir Aberman
Google Research | Boston University
25 Aug 2022

 

Abstract

Fine-tunning 이란?
기존에 학습된 모델(Pretrained model)을 기반으로 추가적으로 새로운 데이터셋 학습을 통해 모델의 파라미터를 미세조정하는 것이다. Latent Diffusion 모델에서 학습가능한 파라미터의 영역은 크게 텍스트 인코더와 U-net 영역이 있으며, 드림부스는 두가지 파라미터 모두 학습한다.

 

본 논문에서는 text-to-image diffusion model의 "개인화"를 위한 새로운 접근 방법을 제안한다. 피사체의 이미지 몇 개만 입력으로 주어지면, 사전 훈련된 텍스트-이미지 모델을 미세 조정할 수 있다.

 

Introduction

최근 개발된 대형 text-to-image 모델은 고품질의 다양한 이미지 합성을 가능하게 함으로써 전례 없는 기능을 보여준다. 이러한 모델의 주요 이점은 이미지-캡션 쌍의 대규모 컬렉션에서 학습한 강력한 semantic prior이다.

이러한 모델은 주어진 레퍼런스에서 피사체의 모양을 모방하고 다양한 맥락에서 동일한 피사체의 새로운 표현을 합성하는 능력이 부족하며, 이는 출력 도메인의 표현력이 제한되어 있기 때문이다.

또한 텍스트 임베딩이 language-vision space에 있는 모델도 주어진 대상의 모습을 정확하게 재구성할 수 없고, 이미지 콘텐츠의 변형만 생성할 수 있다. 

본 논문에서는 text-to-image diffusion model의 사용자의 이미지 생성 요구 사항에 맞게 모델이 조정되는, 즉 “개인화”를 위한 새로운 접근 방식을 제시한다. 

본 논문이 기여한 바는 다음과 같다.

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

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

 

Method

본 논문의 목표는 피사체의 이미지 몇 개(일반적으로 3-5)가 주어지면 모델의 출력 도메인에 피사체를 이식하여 고유 식별자로 합성할 수 있도록 하는 것이다. 높은 디테일 fidelity로 텍스트 프롬프트에 의해 guide된 피사체의 새로운 이미지를 생성하는 것이 본 논문의 목표이다. 입력 이미지 캡처에 제한을 두지 않으며 피사체 이미지는 다양한 컨텍스트를 가질 수 있다. 방법에 대한 개요는 다음과 같다.

위 그림을 보면, Pretrained Text-to-Image model, 몇장의 이미지, 대상에 대한 class 이름이 주어지고 이를 Dreambooth를 통해 Fine-tuning 하면, unique identifier [V]에 대한 학습이 된 Personalized Text-to-Image model 이 Output으로 나온다. 

첫 번째 작업은 피사체 인스턴스를 모델의 출력 도메인에 이식하고 피사체를 고유 식별자로 바인딩하는 것이다. 중요한 문제는 피사체를 보여주는 작은 이미지들을 finetuning하면 주어진 이미지에 overfitting되는 경향이 있다는 것이다. 또한 language drift는 언어 모델에서 흔히 발생하는 문제이며 text-to-image diffusion model에서도 나타난다. 이를 위해 저자들은 diffusion model이 주제와 동일한 클래스의 다양한 인스턴스를 계속 생성하도록 하여 overfitting을 완화하고 language drift를 방지하는 autogenous class-specific prior preservation loss를 제시한다.

디테일의 보존을 강화하려면 모델의 super-resolution 부분도 finetuning해야 한다. 그러나 naive한 방식으로 대상 인스턴스를 생성하도록 finetuning된 경우 인스턴스의 중요한 디테일들을 복제할 수 없다. 저자들은 주제의 디테일을 더 잘 보존하기 위해 이러한 SR 모듈을 학습시키고 테스트할 수 있는 통찰력을 제공하여 재맥락화(recontextualization)에서 전례 없는 성능을 달성하였다. 제자들이 제안한 학습 절차의 스케치는 아래 그림에 나와 있다. 본 논문에서는 사전 훈련된 Imagen 모델을 기본 모델로 사용한다.

드림부스는 class name을 통하여 기존 모델의 내용을 유지하면서 새로운 내용을 학습하도록 하였다.

학습하고 싶은 이미지인 A [v] dog에 대한 input image를 학습함과 동시에, 기존 모델을 이용하여 출력한 A dog 라는 class name에 대한 이미지인 class image를 같이 학습하여 기존 모델이 가진 class name에 대한 지식을 잊지 않도록 하는 것이다.

Designing Prompts for Few-Shot Personalization

저자들의 목표는 새로운 (key, value) 쌍을 diffusion model의 사전에 삽입하여 주제에 대한 key가 주어지면 텍스트 프롬프트로 guide된 유의미한 의미 수정을 통해 이 특정 주제의 완전히 새로운 이미지를 생성할 수 있도록 하는 것이다.

주제에 대한 식별자를 구성하는 naive한 방법은 기존 단어를 사용하는 것이다. 예를 들어 “unique” 또는 “special”과 같은 단어를 사용하는 것이다. 한 가지 문제는 기존 영어 단어가 text-to-image diffusion model의 학습 데이터셋에 존재하기 때문에 prior에 더 강한 경향이 있다는 것이다.

따라서 언어 모델과 diffusion model 모두에서 weak prior를 갖는 식별자가 필요하다. 이를 수행하는 위험한 방법은 영어에서 임의의 문자를 선택하고 연결하여 희귀한 식별자 (ex. “xxy5syt00”)를 생성하는 것이다. 실제로 tokenizer는 각 문자를 개별적으로 토큰화할 수 있으며 diffusion model의 사전은 이러한 문자에 대해 강력하다.

Class-specific Prior Preservation Loss

목표 대상을 묘사하는 작은 이미지 세트와 텍스트 프롬프트 “a [identifier] [class noun]”에서 얻은 동일한 조건 벡터 𝑐𝑠를 사용하여 오리지널 diffusion model의 denoising loss로 text-to-image model을 finetuning한다. 이러한 naive한 finetuning 전략에는 overfitting과 language drift라는 두 가지 주요 문제가 발생한다.

저자들은 autogenous class-specific prior-preserving loss를 통해 overfitting 문제와 language drift 문제를 동시에 해결하고자 한다. 본질적으로 본 논문의 방법은 모델을 자체 생성한 샘플들로 supervise하며 few-shot finetuning이 시작한 후 prior를 유지한다.

Loss는 다음과 같다.

𝜆는 prior-preservation 항의 상대적 가중치를 조절하는 값이다. 저자들은 loss가 단순함에도 불구하고 overfitting과 language-drift 문제를 극복하는 데 효과적이라는 것을 발견했다. 200 이하의 epoch, learning rate 10−510−5, 𝜆=1으로 두는 것이 좋은 결과를 달성하는 데 충분하다고 한다. 학습 과정에서 200×𝑁개 이하의 “a [class noun]” 샘플들이 생성되며, 𝑁은 주제 데이터셋의 크기로 일반적으로 3~5이다. 학습 과정은 1개의 TPUv4에서 15분이 걸린다.

 

Experiments

(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과 동일하지 않다. 대신 예술적 스타일에 따라 피사체 인스턴스 디테일과 정체성 보존을 통해 장면에서 의미 있는 변화를 얻을 수 있다.

(3) Expression Manipulation

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

 

Limitations

본 논문의 방법은 몇가지 제한점이 존재한다. 먼저, 다음과 같이 3가지 주요 failure mode가 존재한다.

1. 프롬프트 맥락과 다른 이미지 생성
2. Context-appearance entanglement
3. 보여준 원본 대상과 비슷한 프롬프트에 대하여 overfitting 발생

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

Conclusions

드림부스는 이런 방식으로 적은 이미지 수로 학습하면서도 Fidelity를 유지하고, 기존 모델의 정보를 잊지 않게 하는 것에 성공했다.

이러한 방법을 응용하여 현재는 Imagen 말고도 최신 diffusion 모델인 Stable diffusion에 적용할 수 있도록 만든 코드가 오픈소스로 github에 공개되어있고, 이를 통해 Stable Diffusion을 Dreambooth로 Fine-tuning 할 수 있게 되었다.