[2023 arXiv] ControlNet : Adding Conditional Control to Text-to-Image Diffusion Models

[gusdlf93]

사전에 학습된 Large Model을 어떻게 해야 적은 데이터로 원하는 도메인에 적합하게 학습시킬 수 있을까요?

이에 대한 기본적인 해결방법은 Large Model을 최소한으로 수정해서, 도메인에 적합하게 네트워크에서 수정되야 하는 부분들을 학습시키는 겁니다.
그럼 어느 부분이 수정되어야 하나요? 라고 물어보면 선뜻 대답하기는 쉽지 않습니다.

이 논문에서는 큰 모델을 Fine tuning하는 과정에서 발생하는 문제점들이 3가지가 있다고 주장하고, 이에 대한 해결법을 제시합니다.

1. 큰 모델을 다른 도메인에 Fine tuning하기 위해서는 그에 상응하는 많은 데이터가 필요합니다.
   데이터가 부족하다면 과적합되기 쉽습니다. 하지만, 대부분의 연구 환경에서는 5만개 이하의 데이터를 사용할 수 있는데, 이는 Large Model을 재학습시키기에 충분한 데이터량은 아닙니다.)
2. Large Model은 Finetuning 할 때도 오래걸리는데, 대부분의 연구환경에서 이를 감당하기란 쉽지 않습니다.
   즉, Fine tuning을 할 수 있는 연구 환경을 갖춘 곳도 그리 많지 않습니다.
3. 특정 목적에 맞게 Regularization을 주는 것도 쉽지 않습니다.
   모델을 수정하거나, Regularization을 주거나, 하여튼, 원하는 값을 얻기 위해 튜닝하는 행동들이 원하는 결과를 도출하게 만들기도 쉽지 않죠.

하지만, 최근 나온 많은 방법론들은 이걸 personal computer에서 가능하도록 만들었습니다.
특히 ControlNet은 Stable Diffusion을 Single GPU - 3090Ti 환경에서 fine tuning을 할 수 있도록 만들었으며, Condition을 사용자가 줌으로써, 결과물을 원하는 방향으로 유도할 수 있음을 보입니다.

Control Net은 Hyper Network로부터 영감을 받았습니다.
Hyper Network란 초기 RNN의 레이어들이 가중치를 공유하는 문제점을 해결하기 위해서 제안되었는데, 레이어별 가중치를 다르게 주기 위해, 필요한 순간에 레이어를 생성하는 네트워크입니다.

그림에서 볼 수 있듯이, 특정 레이어에 가중치(weight, bias)를 부여해야 할 상황이 오면,
레이어에 임의의 벡터를 부여하고, 벡터를 레이어의 가중치로 만들어 주는 HyperNetwork(두번의 Fully Connected Layer)는 입력받은 벡터를 해당 레이어에 맞는 가중치로 만들어줍니다.

이게 유명해지게 된 계기는 2021년에 Novel AI에서 Transformer의 cross attention 단계에서의 k,v 벡터를 HyperNetwork로 생성했더니, 성능이 너무 잘 나오게 되면서 유명해지게 된 것입니다.

어떻게 사용되냐면 Stable Diffusion에서 이미지 정보에 해당하는게 Q이고, Condition 정보(Text, Image, Semantic Map)에 해당하는게 K, V라고 볼 수 있습니다.
여기서 Condition( 모나리자의 화풍이나, 특정 작가의 그림체)을 이미지에 대입하게 되면, 이미지가, Condition의 영향을 받도록 강제되서 해당 Condition에 오버피팅되는 것처럼 학습됩니다.
원래 가지고 있는 정보들을 상당부분 잃어버리고, 학습 데이터와 유사해지도록 만들어줍니다.
이미지와 Condition은 사전에 학습된 Stable Diffusion의 임베딩 레이어를 통해서, 임베딩 벡터로 변환이 되고,
Condition에 해당하는 임베딩 벡터가, Cross attention하기 전에 HyperNetwork를 통과해서 우리가 원하는 결과값을 나타낼 수 있도록 학습이 되게 되는 것입니다.
여기서 HyperNetwork는 단순하게 (FC layer + activation function + Layer Norm)로 구현이 되고, 데이터 수천장 정도를 학습하게 되면, 임베딩 벡터가 우리가 원하는 Condition에 어울리는 이미지를 만들도록 학습이 되게 됩니다.

그럼, 이제 ControlNet에 대해서 알아보겠습니다.
ControlNet은 Stable Diffusion 모델을 학습 가능한 버전과, 가중치를 잠근 버전 두개로 복사해줍니다.
그리고 이들에 zero convolution이라는 레이어를 붙이고 학습을 진행하게 됩니다.
학습 과정에서는 Condition으로 (Canny edges, Hough lines, user scribbles, human key points, segmentation maps, shape normals, depths)등의 조건을 부여할 수 있습니다.

Condition이란 무엇이고, 어떻게 사용되는지에 대해서 잠깐 설명하고 넘어가자면,
Source Image에 Canny edge 필터를 씌우고, 이를 Stable Diffusion에 Condition으로 넣어줍니다.
이미지를 생성할 때, prompt로 “a high-quality, detailed, and professional image”를 넣었을 때, 사슴에 해당하는 prompt를 주지 않았음에도 사슴에 해당하는 색상과 테스쳐를 아주 자연스럽게 넣어주는 것을 볼 수 있습니다.

이번에는 zero convolution에 대해서 설명드리겠습니다.
Zero convolution은 weight랑 bias를 0으로 초기화한 point wise convolution입니다.
C는 우리가 학습하길 원하는 condition입니다.
Weight랑 bias가 0이면, zero convolution이 학습이 될 수가 없겠지요

zero convolution에서 weight랑 bias가 0이면 원래는 학습이 안되는게 정상이지만, chain rule에 의해서 가중치에 상관없이 Learning rate, Loss, Weight만 존재하면 gradient를 계산할 수 있고, zero convolution layer는 학습을 진행할 수 있게 됩니다.

간단하게 생각하자면, 원래는 원본 입력과 동일한 입력값을 내뱉게 되지만, 우리가 주는 Condition에 적합한 정답을 내뱉을 수 있도록 모델의 가중치를 아주 조금씩 깍아나가기 위해서 zero convolution을 넣어준 겁니다.
큰 모델을 적은 데이터에 튜닝하기 위해서는 가중치의 수정이 매우 조금씩 이루어져야 하는데, 이를 위해서 weight랑 bias를 0으로 초기화했다고 보시면 됩니다.

결국, ControlNet은 Stable Diffusion을 적은 데이터에 적합하도록 미세조정하는 네트워크라고 볼 수 있습니다.
자세한 구조는 밑에 그림과 같습니다.

딱히 비교할만한 성능은 없고, 논문의 Result 11페이지부분을 보시면 Condition으로 headpose를 줬을 때, 훌륭하게 사람을 생성해내는 것을 볼 수 있습니다.
상당히 재밌는게, 사람이 주는 Condition이라는 것을 Stable Diffusion이 이토록 잘 이해한다는게 놀랍네요.
이런 Head Pose를 이어붙이면, 하나의 동영상을 만들 수도 있습니다.