공식 프로젝트 페이지
구현체가 깃헙에 나왔다!
Dreambooth-Stable-Diffusion
XavierXiao • Updated Sep 18, 2022
공식 구현체는 Imagen을 쓰지만, stable diffusion에 맞게 수정된 코드가 올라왔다.
→ 실제로 학습해보니 학습이 잘 된다!
이슈
- 기본 lr = 1e-6 (DreamBooth의 1e-5는 나쁜 결과를 보임)
- 2개 A6000 GPU로 학습, 약 15분 소요
- 학습 위해서는 약 ~33GB의 VRAM이 필요 → 3090에서는 학습할 수 없음 ㅜㅜ
어떤 문제를 풀기 위한 논문이지?
Stable Diffusion 등장 이후 Diffusion 모델에 대한 관심과 성장이 엄청나다.
한편, 해당 모델을 Text Input → Image Output으로 생성하는데, 그 과정에서 <특정한 Object>를 지속적으로 생성하게 하는 것은 사실상 불가능하다.
실제로, DALL-E2, Imagen에 단순히 Text Prompt만을 넣어주는 것 & 설령 Image 기반 가이드된 생성이라고 해도, 실 이미지와는 동떨어진 Object를 생성한다.
- 입력 이미지에 들어있는 노란색 시계와는 다른 느낌의 시계를 생성하거나
- 배경이 완전히 동일한 수준으로만 생성한다. (백그라운드 배경을 바꾸지 못함)
페이퍼의 목표
이 페이퍼에서 목표하는 것은 크게 두 가지이다.
- 입력한 Object의 특징을 유지하고
- 다양한 상황에서의 생성을 진행 (배경, 물체의 방향 등)
무엇을 하나?
이 두 가지를 위해서 Imagen와 같은 이미지 생성 모델을 fine-tune 하는 단계를 거친다.
위 그림과 같이 적은(3개~5개 정도) 수의 이미지를 넣어고, 해당 이미지에서 타겟으로 하는 “dog”와 같은 단어(논문에서는 ‘class’ 라고 지칭한다.)를 Prompt에 함께 넣고, Unique ID(논문에서
[v] 로 표현하는데, 실제로는 3-4개 토큰으로 쪼개지는 이상한 단어(ex: xxy5syt00)를 쓴다.)를 함께 넣어 Fine-Tune을 진행한다.즉, 핵심적으로는 “신규 토큰"을 넣어주되, 해당 토큰은 ‘충분히 사용되지 않아서, 타 LM과 임베딩에 영향을 주지 않는' 수준의 단어를 써야 한다.
- “Rare-token Identifier”
- 스페이스 없이, 쭉 이어쓴 랜덤유니코드3개문자 가 합쳐진 문자열
- 혹은 T5-XXL 토크나이저 기준 대략 5000번대~1000번대 토큰들을 사용
- 실제로 모델의 토크나이저에서 해당 단어를 분리할 수도 있지만(ex:
xxy / 5sy / t00), 그래도 잘 동작한다.
모델 학습 과정
좀 더 세분화된 학습 과정은 크게 세 단계로 나눠진다.
- 적은 수의 이미지를
[v]텍스트와 함께 학습하는 것. 이때는 Text → 64x64 크기의 Text to Image 모델을 학습한다. 이때는 “A [v] dog” 라는 텍스트를 입력으로 사용한다.
- 기존 “dog” 클래스에 대해 모델이 Cascade Forgetting(Overfitting)/Language Drift 가 발생하는 것을 방지하기 위해, 기존 Text to Image모델이 생성한 이미지와 ↔ Fine-tune된 모델이 생성하는 이미지 사이의 간극을 줄이기 위한 “Class-Specific Prior Preservation Loss”를 도입한다.
- 이때는 만약 클래스가 “dog”이라면 “A dog” 와 같은 Text를 두 모델의 공통 입력으로 쓴다.
- 새로 학습하는 이미지를 → Super Resolution에도 학습한다. 이때는 신규 이미지를 Downsampling해서 생긴 저해상도(64x64)를 고해상도(1024x1024)로 학습하고, “Reconstruction Loss”를 사용한다.
Prior-Preservation Loss
기존의 Text to Image 모델이 생성하는 다양한 이미지 생성을 유지하기 위해 쓰는 loss.
fine-tune 하는 모델이 기존과 비슷한 다양성을 확보하기 위해서 기존 모델이 생성한 Image == Label 처럼 사용해서 Supervised Learning을 진행한다.
Initial Noise ~ N(0,1) 와 Text condition vector “a dog”와 같이 텍스트 입력을 넣어준다.
굉장히 단순히 두개의 차이만을 계산하는 Loss다. (MSE급)
- 200epochs 정도를 lr=1e-5, =1 로 두는게 좋은 결과를 가져왔다고 한다.
- 이때, 기존 모델에서 이미지를 샘플링하는 데이터 갯수는 200 x N개 (N=신규 데이터 이미지 수)를 썼다.
- TPUv4 기준 15분 내외의 시간이 소요된다.
이 Loss를 통해서 학습된 효과는 아래 그림에서와 같이 “class + [v]” 경우에 가장 잘 동작하는 모델을 만들어주고, 해당 토큰이 없거나 혹은 다른 상황에 대해서는 모델이 해당 부분을 무시하도록 학습한다.
prior-preservation loss 없이 단순 fine-tune한 모델은 “A dog”라는 입력에 대해서 모두 신규 학습한 이미지만을 생성해, 기존의
dog 클래스에 대한 정보를 잊어버린다는 것을 알 수 있다.Personalized Instance-Specific Super-Resolution
Text → 64x64로 생성하는 모델도 중요하지만, 이후에 1024*1024로 UpScale하는 모델은 신규 오브젝트/이미지에 대한 정보가 없기 때문에, 사전 학습된 모델을 그대로 사용하면 뿌연 이미지 혹은 Halluciation을 생성하기도 한다.
아래 이미지에서 볼 수 있는 것과 같이, Normal Noise를 쓴 경우는 최종 이미지에서 Blurry한, 뿌연 이미지가 생성된다. 그리고 Finetune없이 진행할 경우에는 주요한 Feature가 완전히 다르게 생성되거나, 기존에 없던 부분이 생성되는 등의 이슈가 있다.
