[논문 리뷰] Dreambooth의 원리와 기본적 이해

Stable Diffusion에 대한 인기가 높아지면서 모델의 파인 튜닝에 대한 관심도 같이 높아지고 있다. 그런데 여기서 간편한 모델 파인 튜닝 방법 중 하나인 Dreambooth(드림부스)라는 것이 주목 받고 있다. 이 글에서는Dreambooth 논문을 요약하고, 어떻게 쓰는 것이 효율적인지에 대하여 연구한 내용을 다룰 것이다.

 

이 글을 읽기 전에 Stable diffusion에 대한 글을 먼저 읽는것을 추천한다.

 

Fine-tuning 이란?

먼저, Dreambooth가 하는 일인 Fine-tuning이 뭔지부터 설명해보겠다. Fine-tuning은 기존에 학습된 모델(pretrained model)을 기반으로 추가적으로 새로운 데이터셋을 학습하는 것을 통해 모델의 파라미터를 미세조정하는 것이다. Fine-tuning을 하면 내가 원하는 결과물을 처음부터 모델을 학습할 필요없이 기존 모델에 이어서 빠르게 학습시킬 수 있기 때문에 더욱 효율적이다. Latent Diffusion 모델에서는 학습가능한 파라미터의 영역은 크게 텍스트 인코더와 U-net 영역이 있으며, 드림부스는 두가지 파라미터 모두를 학습한다.

 

Dreambooth란?

개요

Dreambooth의 효과

Dreambooth는 구글 연구진이 발표한 논문에서 사용된 학습 방법의 이름으로 Imagen이라는 text-to-image 생성 모델을 어떤 subject에 대한 몇 장의 사진으로 Fine-tuning 하여  개인화된 text-to-image 생성 모델을 만들고 해당 subject를 새로운 context에서 높은 fidelity로 이미지를 생성할 수 있게 해주는 학습 방법이다.

 

Dreambooth는 기존의 Fine-tuning에서 발생하는 다음 두 가지 문제를 해결하는 것을 목표로 제안된 방법이다.

1. 주제 기반의 이미지 생성. 몇 장의 간단한 사진만으로  주제의 시각적 특징에 대한 높은 충실도를 유지하면서 새로운 맥락의 사진을 합성하는 것.
2. 몇 장의 이미지로 text-to-image diffusion모델을 fine-tuning 하면서도, 기존 모델의 Semantic knowlege를 보존하는 것.

위 내용을 통해 알 수 있듯이, Dreambooth는 기존의 Fine-tuning과는 달리 적은 수의 이미지만으로 모델의 오염(overfitting, language drift) 없이 학습이 가능하기 때문에 개인이 더 손쉽게 fine-tuning이 가능하다는 장점이 있다.

 

 

Dreambooth의 개요

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

 

예를 들어 설명해보겠다. 이미 여러 동물에 대해 학습된 어떤 모델이 있을때, 나는 내가 키우는 개인 [V]를 학습시키고 싶다. 이때, 학습 images는 내가 키우는 개 [V]에 대한 사진이고, [V]는 개 이므로 class name은 dog이다.

 

이 예제 처럼 class name은 [V]의 상위 개념을 써주면 된다. dog 말고도 더 상위의 개념인 animal로도 정할수 있지만, 너무 포괄적인 범위인 animal로 정한다면 모든 animal의 특성이 [V]와 섞이게 되어 아마 좋지 못한 결과물이 나올 가능성이 높아질 것이다. 얼마나 상위 개념을 class name으로 정할 것이냐는 어떤 모델과 어떤 대상을 학습하느냐에 따라 달라지므로 개인의 선택이다. 또, class name을 어떻게 정하느냐에 따라서 모델 오염의 정도가 달라질 수 있으므로 잘 감안하여 선택하는 것이 좋을 것이다. 

 

원리

그런데 드림부스는 어떤 원리로 기존 모델의 내용을 유지하면서 새로운 것을 학습할 수 있을까?

이것에 대한 비밀은 class name에 있다.

 

Dreambooth 원리

Dreambooth는 이런 문제를 class image(Regularization image, 정규화 이미지)를 통하여 해결했다.

 

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

 

구체적인 방법으로는,  학습할 때 사용하는 손실함수로 class-specific prior preservation loss 라는 자체적인 손실 함수를 이용하였다.

 

Prior-Preservation Loss

 

기존 Diffusion model의 손실함수와의 차이점은 뒤에 ${\lambda} {\omega _{t^\prime}} \left\| {\hat{x_{\theta}} \left( {\alpha _{t^\prime} {x_{pr}} + \sigma_{t^\prime}\epsilon ^\prime}, {c_{pr}} \right)} -{x_{pr}} \right\| $ 가 추가된 것인데, $\lambda$ 와 class image에 대한 손실함수를 곱하였기 때문에 $\lambda$값을 조절하여 가중치에 어느정도로 class image의 손실함수를 추가할지 조절할 수 있는 것이다.

논문에서는  $\lambda$가 1일때 가장 좋은 결과를 가져왔다고 말하지만 이는 학습할 이미지 개수나 자신이 원하는 결과물이 무엇인지에 따라서 다르므로 잘 조절하여야 한다. 왜냐하면 정규화 이미지가 과적합 방지에는 도움이 되지만 조금 더 학습할 이미지와 유사하게 학습하기 원할 경우에는 정규화 이미지가 오히려 방해할 수 있기 때문이다.

결론

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

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

 

초기에는 공개된 코드가 몇개 없었지만 현재는 계속해서 Stable diffusion 버전 dreambooth가 (Joepenna, Shivam shirao등) 여러 변형버전으로 공개되어 다양한 방식의 training 코드가 있는 상황이다.

 

드림부스 학습, 파인튜닝에 대한 추가적인 궁금증이 있을경우 오픈 카톡방 https://open.kakao.com/o/g3yx5sOe 에서 질문할 수 있다. 대리학습도 시간나면 받는다

 

 

 

참고자료

• https://dreambooth.github.io/