Stanford 대학교에서 2023년에 발표한 이미지 생성 논문
해당 논문이 발표된 2023년에는 diffusion 모델을 활용하여 이미지 생성이 대중적으로 많이 사용되기 시작한 시기이다. 사용자의 프롬프트를 받아서 이를 이미지로 생성하는 text-to-image 모델은 매우 고품질의 이미지를 생성할 수 있다. 그러나 정확한 자세나 모양 등은 텍스트로만 표현하기에는 어렵고 이는 모델이 이미지를 컨드롤하는데 한계가 존재한다.
이러한 미세한 컨트롤을 모델이 할 수 있도록 학습하는 것은 어려웠는데 문제점이 존재한다. 물체 모양, 사람 포즈 같은 특정 문제만을 위해서 학습할 수 있는 데이터 수는 가장 기본적인 text-to-image모델의 데이터수보다 5만배 적다. 즉 한정적인 데이터만을 활용할 수 있다. 만약 이렇게 적은 데이터로 대규모 모델을 파인튜닝하면 과적합이나 catastrophic forgetting(파괴적 망각)에 빠질 수 있다.
그래서 문제를 해결하기 위해, 복잡한 형태와 다양한 고수준의 의미를 다룰 수 있는 신경망 아키텍쳐가 필요하다. 그래서 해당 논문에서는 stable diffusion와 같은 대규모 사전학습 모델에 대한 조건부 제어를 학습하는 아키텍처인 ControlNet을 소개한다. ControlNet은 입력에서 출력까지 한 번에 다루는 end-to-end 모델이며 대규모 모델의 파라미터를 잠그면서 성능을 보존하고 인코딩 레이어에 학습 가능한 복사본을 생성한다. 이러한 아키텍처는 조건부 제어를 위한 강력한 backbone이 된다.
실험 결과로 다양한 조건화 입력으로 stable diffusion을 제어할 수 있었고 GPU RTX3090ti 한 장으로도 대규모 클러스터로 학습한 산업용 모델과 경쟁력있는 결과를 달성하였다.
3.Methods
3.1에서는 ControlNet의 기본 구조를
3.2에서는 Stable diffusion에 ControlNet을 적용하는 방법을
3.3에서는 학습 방법에 대해서
3.4에서는 추론시 고려사항에 대해서 설명할 것이다.
3.1. ControlNet
ControlNet이 제어할 수 있는 이유는 신경망 블록에 추가적인 조건을 주입하기 때문이다. 아래 그림처럼 조건 c를 추가한다.
기존의 사전학습된 모델을 보존하기 위해서 기존 파라미터는 변하지 않게 잠근다. 대신 복사본을 만들어서 조건 c로 학습이 되도록하는 trainable copy가 위의 사진 오른쪽에 존재한다. 이를 통해 특정 테스크을 위해 사용하는 데이터셋이 작아서 발생하는 오버피팅을 방지할 수 있다.
수식으로 ControlNet을 표현하기 위해서 위의 사진의 왼쪽부터 먼저 나타내어 보면 입력 feature map x에서 출력 feature map y가 신경망 블록을 거친다는 것을 알 수 있다. 그래서 신경망을 F(;$\Theta$)라고 한다면
$y = F(x;\Theta)$
라고 할 수 있고 이를 zero convolution을 이용하여 잠겨져있는 원본 모델과 합칠 수 있다. 이를 식으로 표현하면
$ y_c = F(x;\Theta) + Z(F(x + Z(c;\Theta_{z1});\Theta_c );\Theta_{z2}) $
이다.
얼핏 복잡해 보이지만 zero convolution이 두 번, 그리고 복사된 trainable copy 신경망 한 번을 겹쳐서 원래 결과와 합치는 과정을 수식으로 표현한 것이다.
그리고 첫 번째 학습 스텝에서는 가중치와 편향이 모두 0으로 초기화되어 있기 때문에 zero convolution의 결과가 모두 0으로 나온다. 따라서 $y_c = y$이 가장 처음에는 성립한다. 여기에 또다른 의미가 존재하는데 노이즈가 학습 초반에 trainable copy에 영향을 못 미치게 막는 효과가 있다.
3.2. ControlNet for Text-to-Image Diffusion
Stable Diffusion을 사용하여 어떻게 ControlNet이 거대한 사전학습 diffusion 모델을 제어할 수 있는지 보여준다. 그래서 기존 stable diffusion과 여기에 ControlNet을 추가하는 모습은 위에 figure2를 확장하여 자세하게 보면 아래 아키텍처와 같다:
Stable Diffusion은 U-Net의 형태로 12개의 Encoder, 1개의 Middle 그리고 12개의 Decoder로 총 25개의 블록으로 이루어져 있다. feature map의 크기를 조절하는 down sampling과 upsampling을 하는 블록이 8개가 존재하고 나머지 17개의 블록은 가각 3개의 ResNet과 2개의 ViT(Vision Transformer)를 갖고 있다.
사용자가 입력한 프롬프트와 time step도 사용하는데 텍스트인 프롬프트는 CLIP을 이용해서 인코딩하고 time stpe는 positional encoding을 이용하여 인코딩한다.
위의 사진에서 오른쪽 ControlNet는 Stable Diffusion의 12개의 encoder와 1개의 middle 블록을 복사하였고 12개의 zero convolution과 1개의 middle 블록이 skip connection을 이루고 있다.
그렇다면 ControlNet은 얼마나 효율적일까? ControlNet를 사용하여 Stable Diffusion을 최적화하였을 때 GPU 메모리를 23% 더 쓰고 시간은 34%정도 더 걸렸다. 이것과 비교해야하는 것은 ControlNet없이 처음부터 다시 Stable Diffusion을 조건부 제어를 하려고 하였을 때이며 재학습에 걸리는 시간은 ControlNet을 사용하는 것보다 더 많이 걸린다는 것을 이해할 수 있다.
안정적인 학습을 위해서 이미지를 latent 형태로 압축한다. 따라서 512x512크기의 이미지를 64x64크기의 feature map으로 줄이기 위해서 4x4 kernel과 2x2 stride로 이루어진 4개의 합성곱인 작은 신경망 $\varepsilon(\cdot)$를 사용하여 $c_i$를 $c_f$로 변환한다.
$c_f = \varepsilon(c_i) $
이렇게 condition에 해당하는 $c_f$가 생성된다.
3.3. Training
Diffuison Model은 이미지에 조금씩 노이즈를 추가하고 반대로 노이즈를 조금씩 제거하면서 입력 이미지와 유사한 확률 분포를 가지도록 이미지를 생성하는 모델이다. ControlNet에서는 이미지$z_0$에서 노이즈를 t번 추가하여 $z_t$를 생성한다. 또한 텍스트 프롬프트 $c_t$, task별 조건 $c_f$를 포함하여서 노이즈를 예측하는 것이 목표이다. 노이즈를 잘 예측한다는 것은 반대 방향으로 노이즈를 제거할 때 깨끗한 이미지를 생성하는 것이 가능하기 때문에 Diffusion 모델 (Latent diffusion Model)에서 가장 핵심이다.
따라서 정규 분포 N(0,1)에서 샘플링한 노이즈 $\epsilon$를 모델은 최대한 같은 노이즈를 예측하는것이 $\epsilon_\theta$에 해당한다. 손실함수 수식으로 보면 다음과 같다.
$\mathcal{L} = \mathbb{E}_{\mathbf{z}_0, \mathbf{t}, \mathbf{c}_t, \mathbf{c}_f, \epsilon \sim \mathcal{N}(0,1)} \left[ \| \epsilon - \epsilon_\theta(\mathbf{z}_t, \mathbf{t}, \mathbf{c}_t, \mathbf{c}_f) \|_2^2 \right]$
따라서 최대한 같게 노이즈를 예측할수록 차이가 적게되며 loss가 작아지게 된다.
학습에서 50%확률로 프롬프트 $c_t$를 빈 문자열로 바꾸었는데 이를 통해서 ControlNet이 조건부 이미지(선, 자세 등)의 의미를 파악하는 능력을 키울 수 있게된다.
학습하는 과정에서 점진적으로 조건부 제어를 학습하는 것이 아니라 어느순간 갑자기 학습을 하는 것을 관찰하였다. 아래 그림처럼 10000번 이전에 이런 현샹아 나타났다.
3.4. Inference
위에서 Diffusion 모델은 노이즈를 제거함으로써 이미지를 생성하는 것을 알아보았다. ControlNet에서는 조건이 노이즈를 제거하는 과정에 영향을 미치는 방법을 해당 파트에서 설명한다.
1. Classifier-free guidance resolution weighting(CFG)
Stable Diffusion 계열에서 고품질 이미지를 사용하기 위해서 CFG라는 기술을 사용한다. 다음과 같은 수식으로 나타난다.
$\epsilon_{\text{prd}} = \epsilon_{\text{uc}} + \beta_{\text{cfg}}(\epsilon_{\text{c}} - \epsilon_{\text{uc}})$
- $ \epsilon_{\text{prd}} $ : 최종 출력
- $ \epsilon_{\text{uc}} $ : 무조건 출력
- $ \epsilon_{\text{c}} $ 조건 출력
- $ \beta_{\text{cfg}} $ 사용자 지정 가중치
ControlNet와 같은 추가 조건(conditioning image)를 넣을 때 $\epsilon_{\text{uc}}$나 $\epsilon_{\text{c}}$ 어디에 두는지에 따라 결과가 달라진다.
만약 $\epsilon_{\text{uc}}$와 $\epsilon_{\text{c}}$ 모두에 조건을 추가할 경우: CFG항이 소거가 되기 때문에 CFG 가이드 효과가 사라진다. 사진 5b가 해당한다.
만약 $\epsilon_{\text{c}}$에만 조건을 추가할 경우 : CFG항이 온전히 남아서 가이드 효과가 너무 강해져서 이미지가 5c처럼 왜곡이 될 수도 있다.
따라서 $ \beta_{\text{cfg}} $를 조절(억제)하면서 CFG가이드를 최적화하고자한다. 방법으로는 진행할수록 감소하는 방식으로 설정하여 점점 CFG 가이드가 줄어들게 설정하였다. 이를 통해 아래의 5d와 같은 이미지를 얻었다.
2. Composing multiple ControlNets
canny edge와 포즈처럼 여러 개의 조건부 이미지를 적용하여 단일 Stable Diffusion애 적용할 수 있다.
4. Experiments
4.1. Qualitative Results
Figure 1은 다양한 프롬프트로 생성한 결과물이다.
Figure 7에서는 프롬프트 없이 다양한 조건에서 생성한 결과물이고 ControlNet이 의미를 견고하게 잘 해석하는 것을 보여준다.
4.2. Ablative Study
ControlNet이 얼마나 효과적인지를 확인하기 위해서 zero convolution과 trainable copy를 각각 없애보면서 실험을 진행하였다.
(a)는 기존 ControlNet 그대로 사용하였고 (b)와 (c)는 각각 zero conv와 trainable copy없이 사용한 결과이다. (a)가 가장 적절하게 잘 나왔다는 것을 어느정도 확인할 수 있다.
추가로 (c)처럼 trainalbe copy없이 진행한 ControlNet을 ControlNet-lite로 지칭하였다.
4.3. Quantitative Evaluation
다른 5가지 방법 PITI 스케치 모델, Sketch-Guided Diffusion 두가지, 그리고 ControlNet과 ControlNet-lite를 평가한다.
12명의 유저들이 평가를 했고 그림의 퀄리티와 스케치를 충실히 구현하였는지 두 가지 요소로 1점에서 5점까지(낮으면 나쁨) 점수를 매겼다.
ControlNet의 평가가 가장 좋은 것을 확인할 수 있다.
이번에는 산업에서 실제로 사용하는 모델들과 비교를 해보자. Stable Diffusion V2 Depth-to-Image (SDV2-D2I)는 거대한 규모의 NVIDIA A100 클러스터를 이용하여 GPU를 수천시간 동안 학습하며 1200만장을 학습용 이미지로 사용하였다. 이에반해 ControlNet은 20만장을 학습으로 사용하였으며 단일 RTX 3090Ti를 사용하여 5일동안 학습을 하였다.
이 두가지 모델을 사용한 이미지를 다시 12명의 유저에게 보여주면서 어떤 모델을 사용하였는지 맞추어보라고 하였을 때 결과가 약 0.52으로 이는 두개를 구별하기 어려울 정도로 ControlNet이 성능이 좋다고 설명하고 있다.
다음으로는 조건화 충실도를 평가하기 위해서 데이터셋 ADE20K를 사용하여 평가하였다. FID(Frechet Inception Distance)를 이용하여 정답 이미지와 얼마나 차이를 나는지를 나타내었고 또한 CLIP 점수도 평가에 활용하였다.
Stable Diffusion을 제외한 나머지 모델과 비교할 때 ControlNet이 FID, CLIP 점수가 모두 우세하였고 Stable Diffusion와도 FID는 나쁘지만 CLIP 점수는 크게 차이가 나지않는다는 것을 확인할 수 있다.
4.4. Comparison to Previous Methods
해당 논문에서 소개한 Stable Diffusion + ControlNet 조합과 이전의 Sketch-Guided Diffusion 그리고 Taming Transformers와 비교하면 다음과 같다.
ControlNet이 견고하게 이미지를 잘 조절하여 깔끔한 결과를 만들어내는 것을 확인할 수 있다.
4.5. Discussion
다음 세가지의 현상을 언급한다.
1. 학습 데이터 수의 크기 : 데이터가 많을 수록 더 좋은 결과 생성
2. 해석하는 능력 : ControlNet이 의미를 잘 해석함
3. 다른 모델에 이식 가능 : Stable Diffusion에서 바꾸지 않았기 때문에 다른 모델에도 적용 가능함
5. Conclusion
ControlNet은 거대한 사전학습 text-to-image diffusion 모델을 제어하는 방법을 학습한다. 원본 모델과 trainable copy는 zero convolution 층과 연결되며 이는 초기학습 중 노이즈가 학습에 악영향을 끼치는 것을 방지한다.
또한 ControlNet은 프롬프트 유무의 상관없이 다양한 조건에서 Stable Diffusion을 효과적으로 제어할 수 있음을 확인할 수 있다.