[딥러닝 연구] Classifier-Free-Guidance(CFG)와 Stable diffusion 파인튜닝

동기

지금까지 text-to-image 생성 모델인 여러 Latent Diffusion Model을 파인튜닝 해왔는데, 이미지 생성 모델에 대한 논문을 다시 정독하다가 CFG에 대한 부분을 전혀 고려하지 않고 있었다는 것을 깨달았다. 그래서 이걸 고려해서 새로운 기법을 적용하여 파인튜닝 해 본 결과 꽤 재미있는 결과가 나와서 글을 쓰게 되었다. 

 

Text Conditioning의 중요성

text-to-image 생성 모델에서는 text conditioning이 매우 중요한 역할을 한다. 자신이 원하는 이미지를 출력하기 위해서는 입력한 프롬프트를 이해하고 정확하게 Sampling을 해야 하기 때문이다.

Stable Diffusion 에서는 샘플링 과정에서 text condition이 적용된 샘플의 퀄리티를 높이기 위해서 Diversity를 줄이고 자신이 원하는 이미지를 얻을 수 있도록 하는 classifier-free-guidance를 적용하였는데 이는 실제로 큰 효과를 보여줬다.

High-Resolution Image Synthesis with Latent Diffusion Models 논문에 나온 모델 평가 비교

논문의 표의 맨 아래에 있는 LDM-4-G (CFG를 적용한 모델)에서 FID와 Precision값이 훨씬 더 좋게 나온것을 알 수 있다. 그리고 Diversity는 감소하여 Recall 수치는 내려갔다.

Classifier Guidance란 무엇인가?

지금 알아보고자 하는 개념인 Classifier-free-guidance에 대해 알아보려면 우선 classifier guidance가 무엇인지 알아야 한다.

Diffusion Model이 대두되기 전, 기존에 사용되던 GAN 생성 모델에서는 Truncation 이나, Low temperature Sampling으로 Diversity를 줄여 샘플의 퀄리티를 개선해왔다. 그러나 이런 기법들은 Diffusion model에서는 효과적이지 못했고 이로 인하여 Diffusion 모델이 GAN보다 성능이 떨어지는 것으로 생각되었으나, Diffusion Models Beat GANs on Image Synthesis이라는 논문에서 classifier guidance라는 방법을 제시하면서 판이 뒤집혔다.

 

Classifier guidance는 conditional diffusion 모델의 학습 후 과정에서 샘플의 divesity와 fidelity를 trade-off 하는 방법으로 샘플링 과정의 score(negative log likelihood 함수의 1차 미분) 추정 함수 뒤에 gradient of the log likelihood of an auxiliary classifier model를 더하였다. 

기존 score 함수

위 함수에서 $z_{\lambda}$는 노이즈가 추가된 데이터를 의미한다.

기존 함수에 특정 class에 대한 classifier의 negative log likelihood 를 더함(w는 조절가능한 가중치)

이 방법으로 classifier를 이용하여 guide를 제시하여 더 높은 퀄리티의 이미지를 얻을 수 있다.

하지만 classfier guidance를 쓰기 위해서는class에 대한 label이 필요하며 classifier를 따로 학습해주어야 한다는 큰 단점이 있다.

 

 Classifier-Free-Guidance란?

CLASSIFIER-FREE DIFFUSION GUIDANCE 논문에서는 Classifier Guidance가 noise level에 따른 classifier를 따로 학습시켜야 하고 classifier에 기반하여 모델을 평가하는 수치인 IS와 FID를 편법으로 올리는 방법(adversarial attack) 이라고 주장한다.

연구진들은 그래서 모델 자체만으로 작동하고 매우 간단하면서도, 효과가 확실한 방법인 Classifier-Free-Guidance를 제시했다.그 방법은 학습 코드 한 줄만 바꾸기만 해도 바로 적용가능하기 때문에 매우 간단하다.

classfier free guidance($w$는 조절 가능한 가중치)

위 식에서 $\epsilon_{\theta}(z_{t},c)$는 $c$라는 텍스트 조건에 의해 conditioned(조건화) 된 것이고, $\epsilon_{\theta}(z_{t})$는 unconditioned 된 것이다.

이 식을 $w$에 대하여 정리하면, $$(1+w)\epsilon_{\theta}(z_{t},c)-w\epsilon_{\theta}(z_{t}) = w(\epsilon_{\theta}(z_{t},c)-\epsilon_{\theta}(z_{t}))+\epsilon_{\theta}(z_{t},c)$$

위 식을 보면 이 방법에서 guidance가 어떻게 작동 하는지 직관적으로 이해할 수 있다. sample의 conditional likelihood에서 unconditional likelihood를 뺀 값을 늘린다는 간단한 구조 덕분이다.

 

파인튜닝과의 연관성

Classifier-free-guidance를 적용하기 위해서는 위에서 설명했듯이 conditional likelihood에서 unconditional likelihood를 빼야하므로 unconditioned 되는 경우도 학습해야 한다. 지금까지는 CFG를 고려하지 않고 unconditioned되는 경우는 학습하지 않았다. 그래도 학습 결과에 치명적인 악영향을 체감하지는 못했지만, classfier-free-guidance를 적용하기 위해서는 텍스트 조건이 없는 경우도 학습을 해줘야 한다.

 

그래서 학습하는 데이터셋에서 일부 이미지의 캡션을 랜덤하게 비워버리는 caption dropout 기능이 있는 파인튜닝 repository를 이용하는 것을 추천한다.

또한, caption dropout rate(데이터셋에서 캡션을 비우는 비율)는 0.1(10%)를 추천한다. 왜 0.1인지는 모르겠으나 보통 학습할때 0.1이 기본값인듯 하다. 참고로 너무 높이면 오히려 학습한 캐릭터와 유사도가 떨어지는 등 학습이 잘 되지 않았다.

For classifier-free guidance, we joint-train unconditionally via zeroing out the text embeddings with 10% probability for all three models. We train on a combination of internal datasets, with 460M image-text pairs, and the publicly available Laion dataset [61], with 400M image-text pairs. - Imagen 논문 에서 언급된 내용

Imagen 논문에 의하면 Imagen 모델도 10%의 비율로 caption dropout를 적용하였다고 한다.

 

실제로 Caption Dropout을 적용하여 학습해본 결과 CFG값에 따라 diversity가 증가하는게 체감되었다.

그리고 적용 전에는 CFG값이 15이상으로 올라가면 이미지가 완전히 망가졌는데, 적용 후에는 값을 20까지 올려도 별 문제가 없었다.

이 방법은 학습하는 class가 여러개인 경우에 더 큰 도움이 될것 같다.

 

이 기능을 지원하는 repo는 kohya-ss의 sd-scripts와 victorchall의 EveryDream2trainer가 있다.

참고자료

CLASSIFIER-FREE DIFFUSION GUIDANCE / https://arxiv.org/abs/2207.12598

Diffusion Models Beat GANs on Image Synthesis / https://arxiv.org/abs/2105.05233

 

다음 딥러닝 글은 파인튜닝하면서 너무 어둡거나 밝은 이미지를 학습하지 못하는 이유와 해결법에 대하여 쓸 예정이다.

 

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