ControlNet 논문 리뷰

Adding Conditional Control to Text-to-Image Diffusion Models

---

Summary

본 논문은 텍스트만으로는 구조·구도 제어가 약한 Text-to-Image 확산모델의 한계를 개선한다. 대규모 백본(예: Stable Diffusion)을 고정한 채, zero-convolution으로 연결된 학습 가능한 인코더/미들 블록 복사본(ControlNet)에 edge·depth·segmentation·pose 등 조건 이미지를 태워 스킵 연결과 미들 블록에 주입한다. 이 방식은 원 모델 성능을 해치지 않고(안전한 미세조정), 소·대규모 데이터 모두에 강건하며, 단일/다중 조건과 프롬프트 유무 등 다양한 시나리오에서 정밀 제어를 가능케 한다.

Method

ControlNet

여기에서 neural network block는 resnet block, conv-bn-relu block 등 다양한 네트워크 블록이 될 수 있다. $\Theta$를 학습된 neural network block, $x$, $y$를 각각 input feature map, output feature map라 할 때 (a)는 다음과 같은 식으로 나타낼 수 있다:

$$y =  \mathcal{F} (x ; \Theta) $$

여기에서 ControlNet은 기존의 학습된 패러미터 $\Theta$를 고정, 복사하여 학습 가능한 $\Theta_c$를 만든다. (a)에서 input을 $x$로 넣어준 것과는 다르게 $x$와 조건 이미지 $c$를 zero-convolution layer $\mathcal{Z}$을 거친 feature map과 합쳐주어 $\mathcal{F}$에 input으로 넣어준다. 이후 zero convolution을 한 번 더 거쳐서 feature map을 만들어주고, 이를 (a)에서 만들어진 $y =  \mathcal{F} (x ; \Theta) $에 더해준다. 즉 다음과 같다:

Zero-convolution은 초기에 weight, bias가 전부 0으로 초기화된 1×1 convolutional layer인데, 이것이 ControlNet의 핵심이라고 볼 수 있다. 보통 잘 알려진 initalization인 Xavier 또는 He initalization을 사용하지 않고 0으로 전부 초기화했는지 알아보자.

 

우선 논문은, 첫 번째 training step에서 $\mathcal{Z}$가 0으로 초기화되었기 때문에 그림 (a)와 (b)의 결과가 완전히 같다($y_c = y$)고 보여준다. 따라서 이 방식으로는 유해한 노이즈가 학습 가능한 복사본의 신경망 계층에 바로 영향을 미칠 수 없으며, 학습 과정에서 필요한 만큼만 residual을 더하는 방식으로 학습된다.

 

ControlNet feature $F_c$ (pretrain된 인코더의 feature와 초기에 동일)과 0으로 초기화된 weight, bias $W, b$를 보자. 이렇게 convolution을 적용한다고 하면

$$ Z = W \cdot F_c + b $$

가 되고, 초기에는 $Z=0$ 이기에 $F_c$의 신호가 전달되지 않는다. 다만

$$ \frac{\partial L}{\partial W} = F_c \cdot \frac{\partial L}{\partial Z} \neq 0$$

$$ \frac{\partial L}{\partial b} = \frac{\partial L}{\partial Z} \neq 0 $$

이므로 weight와 bias는 학습이 가능하다. 몇 스텝이 지나면 ContronNet 인코더 블록까지 gradient가 전달될 수 있어 학습이 가능하다.

이러한 과정에서 weight와 bias가 백본 모델의 특성에 adapt되고 나서 $F_c$의 학습이 시작되기 때문에 백본 모델을 해치지 않으면서 조건적 정보에 적응할 수 있게 되는 것이다.

반면 Xavier, He initalization은 처음부터 $F_c$가 학습 가능하도록 gate가 열리기 때문에 기존 모델의 패러미터로도 gradient update가 적용될 수 있다. 이 경우에는 기존 백본 모델을 보호하는 것을 어렵게 한다는 문제점이 있다.

 

이후 experiment를 참고하며 zero-initalization이 왜 유리한지 추가적으로 알아보자.

 

ControlNet in Image Diffusion Model

기본적으로 stable diffusion은 U-Net 구조를 사용하며, 인코더의 feature map이 디코더에 skip-connection으로 전달된다. Input condition $c_t$가 입력되는 경우 zero-convolution을 지나 controlnet에 맞게 학습 가능한 복사본으로 구성된 encoder block으로 들어가며, 각 skip connection layer에서 zero-convolution을 한 번 더 거쳐 만들어진 feature map이 기존의 encoder block의 feature map와 합쳐져서 decoder의 skip-connection으로 전달된다. Middle 레이어의 경우에도 zero-convolution을 거친 feature map이 전달된다.

 

ControlNet은 Stable Diffusion의 인코더 및 미들 블록을 복사하지만, 원본 블록은 고정되어 있기 때문에 gradient 계산이 필요없으며, 학습 시 그래디언트가 전달되는 부분은 ControlNet 브랜치와 zero-conv에만 국한된다. 실제 논문에서는 A100 40GB 기준으로 ControlNet을 붙이더라도 학습 시 메모리 약 +23%, 연산 시간 +34% 증가로 보고하였으며, 전체를 다시 학습하는 것보다 GPU 메모리 사용량과 연산량을 크게 줄일 수 있음을 설명한다.

 

Training

Task-specific 조건인 $c_t$가 추가된 것을 확인할 수 있다. 또한 논문에서는 훈련 과정에서 텍스트 프롬프트의 50%가량을 빈 문자열로 대체하여 제공하는데, 이는 모델이 conditioning image을 곧바로 인지할 수 있는 능력을 부여한다고 설명한다.

Inference

$$\epsilon_{prd} = \epsilon_{uc} + \beta_{cfg}(\epsilon_c - \epsilon_{uc} ) $$

여기에서 prd는 final output, uc는 조건이 없을 때 모델이 예측한 노이즈, c는 조건이 존재할 때 모델이 예측한 노이즈, $\beta$는 조건을 얼마나 반영하느냐에 대한 하이퍼패러미터이다. 여기서 논문은

• 만약 $\epsilon_{uc}$ 와 $\epsilon_c$ 에 전부 ControlNet의 조건을 제공하는 경우 하이퍼패러미터 $\beta$로 모델 제어가 불가능해질 것이고,
•  $\epsilon_c$ 에만 조건을 제공하는 경우 가이던스 영역의 영향력을 지나치게 키울 가능성이 있다는 문제를 제기한다.

따라서 본 논문은 $\epsilon_c$ 에만 조건을 제공하되, 저해상도 블록에는 ControlNet의 영향을 크게, 고해상도 블록에서는 작게 주는 방향으로 이러한 문제를 해결하였다. 이렇게 하는 이유는 다음과 같다:

• 저해상도 단계에서는 이미지의 전반적인 구조 및 레이아웃이 더 주목되어, 조건을 강하게 반영해야 한다.
• 고해상도 단계에서는 디테일 및 텍스처를 채우는 단계이기 때문에 조건을 너무 강하게 주면 노이즈가 생겨 모델에 피해가 갈 수 있어 약하게 반영해야 한다.

Experiments

Zero-Convolution을 도입하는 것의 효과부터 확인해보자. 우선 본 논문에서는 sudden convergence phenomenon을 보여준다.

6100 step까지도 ControlNet의 조건과 거의 무관하게 학습되지만 갑작스럽게 조건을 잘 따르는 모델로 바뀌는 것을 확인할 수 있다.

Zero-initalization에서 시작하지 않는 경우 결과물이 엉뚱하거나 artifact가 많은 것은 확인할 수 있으며 부족하거나 모순되는 프롬프트에서는 condition 반영이 약하거나 오히려 모델에 안좋은 노이즈로 전달되며, 텍스트에 끌려가는 모습을 확인할 수 있다. 인코더를 복사하지 않고 conv 레이어만 adapt하는 것은 condition 해석이 거의 불가능한 것을 보여주며 엉뚱한 프롬프트를 제공하는 경우 텍스트 프롬프트에 따라가는 것을 확인할 수 있다.

반면 본 논문의 방법론은 기초적으로 전부 우수하며, 모순되는 프롬프트더라도 그 형태를 잘 따라가는 모습을 확인할 수 있다.

 

단순히 condtion만 잘 따라가는 것이 아니라, 사용자가 원하는 조건을 더욱 다채롭게 할 수 있는 것을 확인할 수 있다. 필자도 SD를 좀 많이 썼었는데, 기존의 SD에서 같은 포즈를 재현하는 것이 상당히 어려우며, 또한 어떤 포즈와 어떤 모습은 서로 트레이드오프의 관계여서 타협 지점들이 많았던 기억이 있다. 그러나 ControlNet은 이 부분을 잘 해결해준다.