[paper] ControlNet, Adding Conditional Control to Text-to-Image Diffusion Models

Adding Conditional Control to Text-to-Image Diffusion Models, 2023, arxiv

ControlNet은 Stable Diffusion 모델에 조건을 추가하여 더 정밀한 이미지 생성을 가능하게 하는 신경망 아키텍처입니다. Stable Diffusion은 원래 텍스트 입력을 바탕으로 이미지를 생성하는 모델로, 입력된 텍스트에 따라 이미지를 생성할 수 있는 강력한 능력을 가지고 있습니다. 텍스트만으로는 원하는 형태나 레이아웃을 반영하기 어려울 뿐 아니라, UNet의 각 레이어에 원하는 컨디션을 concat해주더라도 충분한 퀄리티가 나오기까지의 cost가 굉장히 많이 필요합니다.

Stable diffusion (latent diffusion model)

ControlNet

ControlNet은 이러한 문제를 해결하기 위해 Stable Diffusion에 조건적 제어(conditional control)를 추가하는 방안을 고안했는데요,
ControlNet은 텍스트와 함께 추가적인 condition 이미지를 입력받아, 이 컨디션에 맞는 이미지를 생성합니다. 예를 들어, 사용자가 제공하는 Canny edge, pose, segmentation map, depth map 등을 추가적으로 입력받아, Stable Diffusion의 이미지 생성 과정을 세밀하게 조정할 수 있게 됩니다. 이렇게 하면, 이미지 생성 시 사용자가 원하는 정확한 형태나 구성 요소를 반영하는 것이 가능합니다.

기존 Stable Diffusion 모델의 UNet 부분을 고정된 부분(frozen layers)으로 두고, condition을 학습하는 복제본(trainable copy)을 두어 두 브랜치로 이루어져 있습니다.

피규어를 살펴보면, Stable Diffusion 모델의 원래 파라미터는 frozen되어 그대로 사용됩니다. 즉, Stable Diffusion의 대형 파라미터는 그대로 유지됩니다.
frozen 모델에 복제본(trainable copy)을 만들어 이 부분을 학습합니다. 복제본은 입력된 조건적 이미지(예: Canny edge, pose, segmentation map 등)로부터 특징을 학습하며, 이를 통해 이미지를 생성할 때 더 정교한 제어가 가능합니다.

Denosing process in ControlNet

아키텍처를 자세히 들여다보면,

denoising process에서 노이즈 $\epsilon$를 예측하는 UNet은 두 브랜치로 이루어져 있고, condition이 들어가는 ControlNet 브랜치는 zero convolution layer를 사용하여 연결됩니다.

Zero-initialized Layers

zero convolution 레이어들은 bias와 weight들이 모두 0으로 설정되어 있어, 학습이 전혀 되지 않은 상태에서는 기존 stable diffusion의 output과 같습니다. 이 방식은 학습 초기 단계에서 불필요하거나 유해한 노이즈가 모델에 영향을 미치지 않도록 하고, 이는 모델의 기존 성능(사전 학습된 Frozen Model)을 보호하면서 새로운 condition을 점진적으로 학습할 수 있도록 돕습니다.

Sudden Convergence Phenomenon

기존 Stable Diffusion 모델을 그대로 두고, 그 위에 ControlNet을 덧붙이는 방식이므로, finetunning시 학습 속도가 빨라지고 더 적은 데이터로도 효과적인 학습이 가능합니다.

또 그러한 방식을 사용하기 때문에, 논문에서는 ControlNet이 학습 초기에도 high quality의 이미지를 생성하는데, 그 과정에서 "sudden convergence phenomenon"을 경험한다고 언급합니다. 이는 학습 초기에 점진적으로 조건을 따라가지는 않지만, 특정 시점에 갑자기 퀀텀 점프를 하여 빠르게 조건을 따르게 되는 현상입니다. 이는 기존 모델을 보호하면서도 새로운 조건을 빠르게 학습할 수 있도록 도와주는 중요한 특징입니다.