Adding Conditional Control to Text-to-Image Diffusion Models 논문 요약

1. 스탠포드 대학교 논문 아니랄까봐 굉장히 논리적으로 잘 쓴 논문의 전형이면서 동시에 2. 논문 결론이 "했더니 되던데?" 라는 느낌을 지울 수 없는...

 

적은 데이터로 pretrained model을 finetuning 및 continued training 하면 overfitting이나 catastrophic forgetting의 위험이 있다.

Stable Diffusion 같은 것을 train할 때, LAION-5B 데이터셋이 사용되었는데, 이는 이름 그대로 50억 개 분량의 데이터셋이다. 그런데 conditioned diffusion을 train하려면 데이터셋이 10만개 수준이기 때문에, 이를 통해 Stable Diffusion 모델을 fine-tuning하거나 continued training 하려고 하면 신중을 기해야 한다.

 

왜냐하면 그렇게 적은 데이터로 그렇게 큰 모델을 학습시키면 overfitting, catastrophic forgetting의 위험이 있기 때문이다. catastrophic forgetting이란 추가 학습의 과정에서 alignment에 문제가 생겨 사전학습 때 학습한 지식을 다 잊어버리는 현상을 말한다. 이를 막기 위해 최근에는 LoRA 처럼, 추가 학습 시에 행렬의 rank나 혹은 파라미터 수를 제한함으로써 이를 해결하고는 한다.

 

이 논문에서는 조금 다른 접근법으로 그 문제를 해결한다. 이미 충분히 잘 훈련된 Stable Diffusion 등 기존 모델의 파라미터를 freeze하되, 대신 encoding layer를 그대로 복사본(trainable)을 떠온다. 그 복사본과 원래 모델은 zero convolution으로 연결되어 있고, 이들은 커널과 bias가 전부 0으로 초기화되어 학습이 점진적으로 진행된다.

 

이를 ControlNet 구조라고 한다.

 

ControlNet

before

after

 

이 논문에서는 U-Net이 활용된 Stable Diffusion 모델을 활용했지만, Neural Network를 활용하는 모델이라면 아무거나 이 ControlNet을 활용할 수 있다고 이 논문은 주장한다. 

 

Stable Diffusion에서 각 step마다 노이즈가 낀 이미지 z_t를 받아서 ε(z_t, t, c_t, c_f)를 예측하는 것이다. ε()는 이미지에서 없애야 하는 노이즈를 의미하며 (ε-prediction), c_t는 텍스트 프롬프트, c_f는 기타 condition을 의미한다. 이 논문은 ControlNet을 설명할 때 Stable Diffusion만을 예시로 들었으나 꼭 U-Net이 아니어도 Neural Network를 활용하는 모든 모델에서 이를 활용할 수 있다고 한다. 

 

텍스트 프롬프트가 아닌 기타 condition은 segmentation, sketch, depth map 등 아무거나 될 수 있다.

 

꼭 zero conv여야 함?

 

Gausian으로 초기화된 conv보다 zero conv가 훨씬 효과가 좋았음이 실험을 통해서도 나타났다. 큰 이유는 다음과 같다:

 

1. 유해한 노이즈 방지

학습 초기 단계에서 원래 사전학습된 모델의 특징 맵에 유해한(랜덤한) 노이즈를 추가하지 않게 되므로, 기존 모델의 품질과 기능을 보존할 수 있다.

 

2. 점진적 학습

가중치를 0으로 시작하여 모델이 갑작스럽게 사전학습된 파라미터로 생성되는 것에 큰 변화를 가하지 않는다.

 

3. 사전학습된 모델 보호

이미 사전학습된 모델은 뛰어나다. 그러므로 충분히 추가 컨디션에 대한 학습이 이루어지지 않은 상황에서 사전학습된 모델에 추가되는 것을 방지한다.

 

Multi-conditions는?

했더니 되더라

 

 

수식에서도, 아키텍처에서도 c_f는 하나이다. 그런데 실험에서는

 

이렇게 여러 컨디션이 조건으로 걸린 실험 결과를 보고한다. 그러고서는 하는 말이 "하면 된다."이다. 아키텍처 상으로는 어떻게 컨디션이 두 개가 걸렸는지 확실하지 않은데, 그냥 진짜 ControlNet 하나 더 추가하면 된다는 건가...