일단 글 더 싸재끼기 전에, 드림부스 써본 적도 없고 코드도 안봐서 다른 정보글만큼 확실하진 않음
이 글은 드림부스 논문이랑 관련 정보 몇개 찾아보고 그걸 토대로 쓰는거임
일단, 드림부스는 Stable Diffusion 모델을 Fine-tuning하는 기법이다. Fine-tuning이란 이미 학습이 완료된 모델과 그 가중치를 그대로 가져와서, 새로운 학습 데이터를 훈련시켜 그 데이터에 대해 결과가 잘 나오도록 바꿔주는 과정이다.
근데 방대한 자료로 학습해왔던 모델에다가 사진 몇 장 던져주고 '이제 너는 이것들의 정답만 맞추면 된다'라고 하면 어마어마한 속도로 과적합해버린다... 그러면서 방대한 자료로부터 학습해왔던걸 점차 까먹기 시작한다.
'A dog'라고 쳤을때의 결과. 맨위: 기존 모델, 중간: 과적합 모델, 아래: 드림부스
드림부스는 어떻게 과적합을 방지할까? 정규화 이미지를 통해서다. 훈련시키고 싶은 대상의 조금 더 큰 범주인 class의 정규화 이미지들을 통해서, 모델은 훈련 이미지의 특성과 함께 class의 다양성을 같이 학습하려고 한다.
사진의 윗부분인 훈련 이미지와 생성 이미지 간의 MSE loss를 구하는 방식은 Textual Inversion과 하이퍼네트워크랑 똑같다! 단지 차이가 있다면 드림부스같은 fine-tuning에서는 U-NET의 gradient를 전부 계산하도록 해서 모든 layer의 weight가 바뀌도록 조절하는 것이다. 이러니까 8기가 VRAM으로는 택도 없다는거다. 아 3070 신품 괜히샀음...
드림부스의 가장 큰 특징은 사진의 아랫부분에 있는데, 훈련 전에 먼저 클래스에 대한 정규화 사진을 뽑은 후에, 그걸 훈련할때 프롬프트에 클래스 이름을 넣어 생성한 이미지 간 비교하여 MSE loss를 추가로 구한다는 것이다. 그리고 그걸 그냥 더해서 새로운 loss function으로 써먹는다.
예를 들어 훈련시키고 싶은 캐릭터의 이름이 'mia'이고 클래스 이름이 'girl'이라면,
모델에서 'A mia girl'이라는 프롬프트로 생성한 이미지와 훈련 이미지를 비교함과 동시에
모델에서 'A girl'이라는 프롬프트로 생성한 이미지와 정규화 이미지를 비교하는 것이다.
만약에 모델이 'mia'에 대해서 너무 많이 배웠다면, 아래의 정규화 이미지와 비교할 때의 loss값이 커질 것이기 때문에, 모델이 '아! mia만 너무 많이 배웠구나! 다른 정규화 이미지도 좀 봐야겠다' 같은 방식으로 과적합을 방지할 수 있게 된다.
논문의 설명대로라면 정규화 이미지는 훈련시키고자 하는 모델에서 생성되어야 한다.
왜 이런지 설명하자면... 모델이 훈련을 진행하면 진행할수록 프롬프트에 'mia' 없이 그냥 'A girl' 만 들어가 있어도 생성하는 이미지가 훈련 이미지와 비슷해지게 된다. 그렇기에 훈련 전에 만들었던 정규화 이미지들과 비교하면서 '너 이런것도 만들줄 알았으면서!! 기억해내!!' 하면서 과적합할뻔한 모델을 바로잡게 되는 것이다.
논문에서는 epoch마다 정규화 이미지 200장 정도를 생성해서 사용하는데, 그만큼 많고 다양해야 드림부스 모델도 어느정도 다양해지는 듯하다.
근데 씹덕 그림만 뽑아낼거라면 정규화 이미지도 어쨌든 다양하게 씹덕스러운거로 뽑아야지, masterpiece best quality로 채워넣으면 맨날 보던것만 보게 될거다.
---
그리고 논문보다가 알게 된건데, 드림부스에서는 기본적으로는 텍스트 모델은 훈련하지 않아서, 훈련할 대상을 불러내기 위해 사용하는 프롬프트 단어는 문자열 1~3개 정도의 랜덤한 문구 (논문에서는 'sks')를 사용했다고 함. 근데 왜 문자열 4개 이상은 안되는지는 안써놨는데, 왠지 알것 같음. CLIP은 문자열 4개부터 강제로 토큰 2개로 쪼개기 때문이다. 즉 지금까지 잘 돌아가는 드림부스 모델들은 1토큰으로 원하는걸 다 뽑아냈다는 것이다.
토큰 얘기 나온김에 추가하자면 이 훈련할 대상을 불러오는 문자열에 이미 다른 의미가 있다면 (예를 들면 'red'를 쓴다던가...) 훈련이 완전히 망가지기 때문에, 랜덤한 문자열을 쓰는 것이다. 숫자 넣어도 안됨.
드림부스 라이스 샤워랑 그림체 한번에 학습할려다 망했다 - AI그림 학습 채널 (arca.live)
이친구는 rice_shower로 학습할려고 했다가 rice와 shower라는 토큰 안에 말딸 데이터가 들어가버렸다...
---
Training Stable Diffusion with Dreambooth – Weights & Biases (wandb.ai)
이 링크에 드림부스로 실험한 내용들이 있는데, 재밌는 내용들이 몇가지 있다.
먼저 5e-6으로 돌리면 과적합나서 2e-6으로 해야 겨우 쓸만해진다는 부분이나 (...)
Text Encoder(텍스트 모델)에 대해서도 학습을 진행했더니 퀄리티가 훨씬 높아졌다는 것이다. 근데 당연하지 그건... 지금까지 sks라는 의미없는 임베딩에 대해서 어떻게든 훈련 이미지랑 대응시키고 있었는데, 텍스트 모델도 학습시키게 되면 임베딩 안에까지 자연스럽게 정보가 들어가게 되기 때문이다. 그리고 텍스트 모델까지 100% 학습시킬려고 하면 12기가 16기가로는 택도 없고 무조건 24기가 써야함.
이건 텍스트 모델 고정
이건 텍스트 모델 훈련된거. 퀄리티 차이가 정말 장난 아니다.
제일 재밌는건, 여기서 sks에 대해서 Textual Inversion을 통해 임베딩을 따로 학습한 뒤에, 그걸 드림부스에 넣어서 추가로 학습했더니 결과가 딱 중간만큼 나왔다는 것이다.
그랬더니 조커가 증발하긴 했는데, 어쨌든 얼굴은 텍스트 모델 고정한 드림부스보다 더 잘 살렸다. 텍스트 모델 전체를 훈련시키지 않고도 이런 성능이 나온다는 것이다.
코랩에 이걸 구현할려면 따로 추가코드를 만들어야 할거고, WEB UI 확장에 이식하는건 그렇게 어렵진 않을듯 (아니면 이미 되어있나?)
그러니 임베딩+하이퍼, 임베딩+드림부스 이런게 생각보다는 쓸모있을듯 (하이퍼+드림부스는 하면 안됨)
