Adding Conditional Control to Text-to-Image Diffusion Models (ControlNet)

https://arxiv.org/abs/2302.05543

Adding Conditional Control to Text-to-Image Diffusion Models

 

# Introduction

Text-to-Image diffusion 모델이 발전하면서 많은 이미지를 text prompt로 생성할 수 있게 되었다. 그러나 text-to-image 모델은 image의 공간 구성에 대한 control이 잘 되지 않는 문제점이 있다. (예를 들어 complex layout이나, poses, form에 대한 모양 같은 것들) 따라서 이는 text에 의존하는 경향이 크다. 하지만 이는 우리가 상상하는 것을 text로 표현해야된다라는 문제점도 발생한다. 그래서 계속 prompting 을 반복해야된다는 문제점이 있다.

 

그러면 추가적인 이미지를 주어 내가 원하는 방향대로 만들 수 있으면 좋지 않나? 그래서 기존에도 image to image translation 분야가 발전해왔다. 하지만 여전히 full training (end-to-end training) 을 해야된다라는 단점이 존재한다. 

 

diffusion 모델을 end-to-end로 condition을 학습한다는 것은 매우 어려운 일이다. 왜냐하면 condition에 대한 많은 정보가 없을 뿐만 아니라, 설사 학습한다해도 forgetting 문제가 발생할 수 있기 때문이다. 따라서 이를 해결하기 위해, controlnet을 제안한다.

 

controlnet은 모델을 전체적으로 학습하는 것이 아닌 conditional control 부분만 학습하는 새로운 architecture를 제안한다. (물론 gpu 있다면 전체 학습 가능). ControlNet은 large model의 encoding 부분을 frozen 시켜 image의 quality를 보장하고, 이를 trainable copy (모델 architecture 그대로 복사)를 만들어 학습한다. 이를 통해 다양한 condition을 학습시킬 수 있다.

또한, zero convolution을 이용하여 이를 연결하여 training 과정동안 학습하도록한다. 이는 harmful한 noise를 추가하지 않고, large-scale의 pretrained backbone을 보존하는 역할을 한다.

 

singe gpu (3090)으로 학습이 가능하여 실험하였다. (사실 거짓말은 아니지만 A100도 씀)

 

그래서 이 논문의 contribution은 다음의 3가지이다.

 

1. 효율적인 fine tuning을 통해 text-to-image diffusion 모델에 공간으로 한정된 정보 (이미지 등등)가 있는 조건을 추가할 수 있는 controlnet architecture를 제안한다.
2. canny edege, hough line, user scribble, human keypoints, segmentation map, shape normal, depth, cartoon line drawing 을 condition으로 하는 stable diffusion으로 pretrained 한 controlnet을 공개한다.
3. ablative 실험을 통해 다른 architecture들과 비교하여 검증하였고, 다양한 task에 걸쳐 user-study 를 수행하였다.

 

# Method

ControlNet

기존의 neural network 는 x라는 input을 neural network block에 넣어 output y를 뽑아낸다.

반면에 controlnet은 condition c를 zero convolution에 태워 input x 와 concat을 한다. 그리고 trainable copy 즉, neural network block을 trainable 하게 바꾼 것에 태우고 이에 대한 output을 다시 zero convolution (1x1 zero conv)에 태워 이를 frozen 한 block에 input을 넣은 x의 결과와 concat 해서 \(y_c \) output을 내뱉는 것이다.

그럼 여기서 왜 weight initalization 된 zero convolution을 하는 것일까? noise로 initalization 했을 땐 초반 몇 step 만에 trainable copy 가 망가질 가능성이 높다. 따라서 아예 이러한 noise를 제거하는 효과를 가질 수 있게 zero convolution을 사용하는 것이다. 

 

그럼 여기서 한가지 의문이 드는 것은 zero-convolution 을 하면 학습이 안되지 않을까? 라는 생각이다. 이는 간단한 미분식으로 그게 아니라는 것을 증명한다.

 

neural network : \( y = wx + b \) 라고 표현되니, 이걸 한 번 각각 w로 미분을 해보자

$$ \delta y / \delta w  = x  $$

$$ \delta y / \delta x  = w $$

$$ \delta y / \delta b  = 1 $$

 

그럼 여기서 w=0 이고 \( x \neq 0 \) 일 때, 모두 0 이 되지 않는다. 이와 같은 직관으로 0이 되지 않는다는 것을 다시금 알려주고, 이를 좀 더 controlnet에 적용을 하면

$$ y_c = \mathcal{F} (x ; \Theta ) + \mathcal{Z} (\mathcal{F}(x + \mathcal{Z} ( c, \Theta_{z1}); \Theta_c); \Theta_{z2}) $$

where, \( \Theta_{z1} \) 과  \( \Theta_{z2} \)는 zero convolution의 parameter이고, \( \mathcal{Z}(\cdot , \cdot) \) 은 zero convolution layer를 의미한다. 그리고 \( \mathcal{F} \)는 trained neural block 이다. 

그러면 첫번째 training step을 살펴보자.

\( \mathcal{Z} (c; \Theta_{z1} ) = 0 \) 이고 (zero convolution이니)

\( \mathcal{F} ( x + \mathcal{Z} ( c ; \Theta_{z1} ); \Theta_c ) = \mathcal{F} (x ; \Theta_c ) = \mathcal{F} (x ; \Theta) \)

\( \mathcal{Z} (c , \Theta_{z1}) ; \Theta_c) ; \Theta_{z2} ) = \mathcal{Z} ( \mathcal{F} ( x ; \Theta_c ) ; \Theta_{z2} ) = 0 \)

 

이다. 여기서 0이 되지 않음을 보였고, \( y_c = y \)가 된다. 

 

좀 더 일반화해서 설명을 하면

forward process

zero convolution의 forward process를 이렇게 표현하면 (where, W = weight, B = Bias) 여기서 gradient 는

gradient

가 된다.

여기서

update

이렇게 되기 때문에 0으로 update가 되는 게 아닌 어떠한 다른 값으로 update 가 된다는 것 알 수 있다.

 

그래서 이제 다시 controlnet에 대해서 살펴보도록 하자.

ControlNet 부분은 Stable diffusion encoder block을 그대로 복사하여 encoder 부분만 학습을 시켜서 condition이 포함된 encoding을 하는 것이다.

 

그럼 여기서 왜 Stable diffusion encoder block 을 복사해서 사용할까? 이를 위해 논문에는 나와있지 않지만, github disscussion에 ablation study를 한 결과가 나타나있다. https://github.com/lllyasviel/ControlNet/discussions/188

Ablation Study: Why ControlNets use deep encoder? What if it was lighter? Or even an MLP? · lllyasviel/ControlNet · Discussion

실험 셋팅은 다음과 같다. 

controlnet-self , controlnet-lite , controlnet-mlp 를 셋팅하여 self는 본 논문에서 제시된 stable diffusion을 그대로 복사한 것이고, lite 는 stable diffusion encoder 대신 cnn 을 사용한 것이고, mlp 는 mlp를 사용한 것이다.

그러고 이미지 생성에 대한 효과를 보기 위해, prompt를 매우 상세하게 주어 scrrible image를 condition으로 주고 생성하라고 해본다.

prompt는 "Professional high-quality wide-angle digital art of a house designed by frank lloyd wright. A delightful winter scene. photorealistic, epic fantasy, dramatic lighting, cinematic, extremely high detail, cinematic lighting, trending on artstation, cgsociety, realistic rendering of Unreal Engine 5, 8k, 4k, HQ, wallpaper" 

scrrible image

그리고 각각의 결과를 보자

Controlnet self results

controlnet-lite results

controlnet-mlp

다들 모두 잘 생성되는 것을 볼 수 있다. 그러면 controlnet을 안쓰고 stable diffusion에 적용해보면 어떨까?

stable diffusion resutls

매우 잘 생성이 되지만 scrrible image의 정보는 거의 반영을 못한 것을 볼 수 있다.

그러면 굳이 trainable copy를 안해도 되지 않을까? 라는 의문을 해결하기 위해 다른 방식으로 실험을 하였다. 

아예 prompt를 주지 않음으로써 text에 대한 효과를 없애는 것이다. 그러면 결과가 어떻게 될까?

controlnet-self

controlnet lite

controlnet-mlp

prompt를 없애니 condition image에 대한 영향을 잘 못받는 것을 확인이 되었다. 따라서 이래서 trainable copy를 한 것이라고 하는 것 같다.

 

# Experiments

다양한 condition을 기반으로 실험을 하였고, 중요한 것은 training detail과 데이터셋이라 이것만 설명한다.

## training detail

 

1. random으로 50% 의 text prompt를 empty string으로 교체하여 input condition map으로 부터 semantic content를 인지하는 능력을 향상 시킨다.

2. small scale gpu (1-gpu) 일 때: stable diffusion이랑 연결되는 부분을 유동적으로 변경시킨다. 초반에는 middle block 부분만 연결시켜서 training speed 를 높히고, 어느 정도 결과가 align 되면 decoder도 연결시켜서 quality를 향상시킨다.

3. large scalp gpu (cluster 단위) 일 때 : 처음에 controlnet을 50k step만 학습시키고 추후에 stable diffusion까지 joint training을 시키는게 훨씬 퀄이 좋았다.

 

## dataset

기존 데이터셋마다 captioner (BLIP)을 붙여 condition-image-caption pair를 제작한다. 또한, condition 을 다양하게 하기 위해 pretrained model과 다양한 library를 활용하여 직접 제작하기도 하였다. training time은 A100 으로 100~600 시간 정도 소요했다.