[논문] Adding Conditional Control to Text-to-Image Diffusion Models

https://arxiv.org/abs/2302.05543

Adding Conditional Control to Text-to-Image Diffusion Models

해당 논문을 보고 작성했습니다.

 

Abstract

저자들은 large, pretrained text-to-image diffusion model에 spatial conditioning control을 추가하는 ControlNet이라는 neural network architecture를 제안합니다. ControlNet은 large diffusion model은 고정하고, 수십억 개의 image로 pretrain 된 deep and robust encoding layer를 strong backbone으로 재사용하여 다양한 condition을 control 합니다. neural architecture는 "zero convolution" (zero-initialized convolution layer)로 연결되며, 이는 parameter가 0에서 점진적으로 커지도록 하여 harmful noise가 finetuning에 영향을 주지 않도록 보장합니다. 저자들은 다양한 conditioning control (e.g., edges, depth, segmentation, human pose, etc.)로 실험을 진행했습니다. Stable Diffusion에 단일 또는 다중 condition, text prompt 유무에 따라 실험을 진행했습니다. 저자들은 ControlNet의 학습이 small (<50k) and large (>1m) dataset 모두에서 robust 하다는 것을 보여줍니다.

 

Introduction

많은 사람들이 독특한 image를 capture하고자 하는 시각적 영감을 경험한 적이 있을 것입니다. text-to-image diffusion model의 출현과 함께, text prompt를 입력하여 시각적으로 놀라운 image를 만들 수 있게 되었습니다. 하지만 text-to-image model은 image의 공간적 구성을 제어하는 데 한계가 있습니다. 복잡한 layout, pose, shape, form을 text prompt만 사용하여 표현하는 것은 어렵습니다. 우리가 머릿속으로 상상한 image와 정확하게 일치하는 image를 생성하기 위해서는, prompt를 수정하고 생성된 image를 확인하고 다시 prompt를 수정하는 여러 번의 시행착오가 필요할 수 있습니다.

large text-to-image diffusion model을 위한 conditional control을 end-to-end 방식으로 학습하는 것은 어렵습니다. 특정 condition에 대한 training dta의 양이 일반적인 text-to-image training을 위해 사용 가능한 data보다 훨씬 적습니다. 예를 들어 다양한 특정 문제(e.g., object shape/normal, human pose extraction, etc.)에 대해 가장 큰 dataset은 보통 100K 크기이며, 이는 Stable Diffusion model을 학습할 때 사용된 LAION-5B dataset 보다 50,000배 작은 크기입니다. 제한된 data로 large pretrained model을 fine-tuning 하거나 이어서 학습하는 것은 overfitting 및 치명적인 forgetting을 초래할 수 있습니다. 연구자들은 학습 가능한 parameter의 수나 rank를 제한함으로써 이러한 forgetting을 완화할 수 있음을 보여주었습니다. 저자들의 경우엔, wild conditional image를 복잡한 shape와 다양한 high-level semantic을 가진 채로 처리하려면, 더 deep 하거나 customized neural architecture가 필수적일 수 있습니다.

이 논문에서는 large pretrained text-to-image diffusion model을 위한 conditional control을 학습하는 end-to-end neural network architecture인 ControlNet을 제안합니다. ControlNet은 large model의 parameter를 고정하여 model의 quality와 capability를 보존하면서 encoding layer의 학습 가능한 복사본을 만듭니다. 이 architecture는 large pretrained model을 다양한 conditional control을 학습하는 strong backbone으로 활용합니다. 학습가능한 copy와 원래 고정된 model은 zero convolution layer로 connect됩니다. 여기서 zero convolution layer는 0으로 weight가 초기화되며, 학습 과정에서 weight가 점진적으로 커집니다. 이 구조는 학습 초기에 harmful noise가 large diffusion model의 deep feature에 더해지지 않도록 보장하며, 학습 가능한 copy의 large-scale pretrained backbone이 이러한 noise에 의해 damage 받지 않도록 보호합니다.

저자들의 실험은 ControlNet이 다양한 conditioning input (canny edge, hough lines, user scribble, human key point, segmentation maps, etc.)을 통해 Stable Diffusion을 control 할 수 있음을 보여줍니다. 

 

저자들은 text prompt의 유무에 상관 없이 single conditioning image로 실험을 진행했으며, 어떻게 저자들의 approach가 multiple condition 구성을 지원할 수 있는지 보여주었습니다. 추가적으로 저자들은 ControlNet의 학습이 robust 하고 다양한 size의 dataset에 맞춰 scalable 할 수 있다는 것을 보였으며, 좋은 성능을 달성할 수 있음을 보였습니다. 마지막으로 model의 각 구성 요소의 contribution을 조사하는 ablation study을 수행했으며, 여러 strong conditional image generation baseline과 비교를 진행했습니다.

요약하자면 (1) 저자들은 ControlNet을 제안합니다. 이는 공간적으로 localize된 input condition을 pretrained text-to-image diffusion model에 더할 수 있는 neural network입니다. (2) 저자들은 Stable Diffusion을 여러 condition으로 control 할 수 있는 pretrained ControlNet을 보입니다. (3) 저자들은 여러 대체 architecture와 비교한 ablation study을 진행하여 효율성을 입증했습니다.

 

Related Work

 

Finetuning Neural Networks

neural network를 finetuning하는 방법 중 하나는 추가적인 training data로 이어서 학습하는 것입니다. 하지만 이러한 방식은 overfitting, mode collapse, catastrophic forgetting을 유발할 수 있습니다. 이러한 문제를 피하기 위한 fine-tuning strategy 관련 연구들을 진행하고 있습니다.

 

- HyperNetwork

자연어 처리에서 처음 등장한 방식으로 small recurrent neural network을 학습하여 더 큰 network의 weight에 영향을 주는 것을 목표로 합니다. 이 방법은 generative adversarial networks (GANs)을 사용한 image generation에도 적용되었습니다. 그리고 HyperNet를 Stable Diffusion에 사용하여 output image의 artistic style을 변경하기도 했습니다.

 

- Adapter

이 method는 pretrained transformer model을 다른 task로 customize 하기 위해 새로운 module layer를 추가하는 방식이며, NLP에서 널리 사용되었습니다. computer vision에서는, incremental learning과 domain adaptation을 위해 adaptor를 사용했습니다. 이 기술은 CLIP과 함께 pretrained backbone model을 다른 task로 transfer 하기 위해 자주 사용됩니다. 더 최근에는 adapter가 vision transformer에서도 성공적은 결과를 내고 있습니다.

 

- Additive Learning

original model의 weight를 freeze하고 학습된 weight mask, pruning, hard attention을 사용하여 적은 수의 새로운 parameter를 추가하여 forgetting을 방지하는 방식입니다. Side-Tuning은 frozen model의 output과 더해진 network의 output을 선형으로 혼합하여 추가적인 기능을 학습하는 side branch model을 사용합니다.

 

- Low-Rank Adaptation (LoRA)

LoRA는 low-rank matrices를 사용하여 parameter의 offset을 학습함으로써 catastrophic forgetting을 방지합니다. 이는 over parameterized model이 낮은 intrinsic dimension space에 존재한다는 사실을 기반으로 합니다.

 

- Zero-Initialized Layers

network block을 connect하기 위해 ControlNet에서 사용합니다. neural network에 대한 연구들은 초기화 및 조작에 대한 논의가 많이 이루어져 왔습니다. 예를 들어 weight의 gaussian initialization은 0으로 초기화하는 것보다 덜 위험할 수 있다는 연구가 있었습니다. 더 최근에는 diffusion model의 convolution layer의 weight 초기화를 어떻게 scale 하는 것이 학습을 향상시킬 수 있는지에 대한 논의가 있었으며, 해당 연구에서의 "zero-module"은 weight를 0으로 scale 하는 극단적인 경우를 의미합니다. 이 외에도 convolution weight의 초기화를 조작하는 다양한 연구들이 있었습니다.

 

Method

ControlNet은 spatially localized, task-specific image condition을 large pretrained text-to-image diffusion model에 추가하는 neural network architecture입니다.

 

ControlNet

 

ControlNet은 추가적인 condition을 neural network의 block에 inject합니다. 여기서 network block이라는 용어는 resnet block, conv-bn-relu block, multi-head attention block, transformer block과 같이 일반적으로 함께 묶여 neural network의 single unit을 형성하는 neural layer set을 의미합니다. F(· ; Θ)를 parameter Θ인 trained neural block이라 하겠습니다. 그리고 이는 input feature map x를 또 다른 feature map y로 변환한다고 하겠습니다.

 

저자들은 x와 y를 2D feature map으로 사용했으며, x ∈ R^{h x w x c}이고, x는 height, w는 width이고 c는 channel 수를 나타냅니다. 

ControlNet을 이러한 pre-trained neural block에 추가하기 위해, 저자들은 original block의 parameter Θ를 lock (freeze)하는 동시에 학습 가능한 parameter Θ_c를 사용하는 block 복사본을 준비합니다. 학습 가능한 복사본은 추가적인 condition vector c를 input으로 받습니다. 이 구조는 Stable Diffusion과 같은 대규모 model에 적용하기 위해, locked parameter는 수십억 개의 image로 학습된 준비된 model을 유지하면서 trainable copy는 다양한 input condition을 다루기 위해 large-scale pretrained model을 재사용하여 deep, robbust and strong backbone을 구성합니다. 

trainable 복사본은 locekd model에 zero convolution layer인 Z(· ; ·)로 연결됩니다. 구체적으로 Z(· ; ·)는 weight와 bias 모두 0으로 초기화된 1x1 convolution layer입니다. ControlNet을 설계하기 위해, 저자들은 각각 parameter Θ_{z1}, Θ_{z2}를 가지는 두 가지 instance를 사용합니다. 최종 ControlNet은 다음과 같이 계산됩니다.

 

y_c는 ControlNet block의 output입니다. 첫 training step에서 zero convolution layer의 weight와 bias가 0으로 초기화되어 있기 때문에, Z(· ; ·)은 0이 되며, y_c = y가 됩니다. 이 과정에서 harmful noise는 학습 시작할 때 trainable copy에 있는 neural network layer의 hidden state에 영향을 끼칠 수 없습니다. 그리고 Z(c ; Θ_{z1}) = 0이고 trainable copy는 input image x를 받기 때문에, trainable copy는 large pretrained model의 capability를 완전히 유지하여 이후에 있을 학습에 대한 strong backbone으로 작용할 수 있습니다. zero convolution은 이 backbone의 초기 학습 단계에서 random noise를 graident로 작용하는 것을 방지합니다.

 

ControlNet for Text-to-Image Diffusion

저자들은 ControlNet이 어떻게 large pretrained diffusion model에 conditional control를 추가할 수 있는 지에 대해 보여주기 위해 Stable Diffusion을 사용합니다. Stable Diffusion은 본질적으로 encoder, middle block, skip-connected decoder로 구성된 U-Net입니다. encoder와 decoder 모두 12개 block을 포함하며, full model은 middle block을 포함하여 25개 block으로 구성됩니다. 25개 block 중 8개 block은 down-sampling 또는 up-sampling convolution layer이며, 나머지 17개 block은 각각 4개 resnet layer와 2개 Vision Transformers (ViTs)를 포함하는 main block입니다. 각 ViT는 여러 cross-attention과 self-attention mechanism을 포함합니다.

 

예를 들어 위 그림의 (a)와 같이, "SD Encoder Block A"는 4개 resnet layer와 2개 ViTs를 포함하고 있으며, "x3"는 이 block을 3번 반복한다는 것을 나타냅니다. text prompt는 CLIP text encoder를 이용하여 encode 되며, diffusion timestep은 positional encoding을 사용하는 time encoder로 encode 됩니다. 

ControlNet structure는 U-Net의 각 encoder level에 적용됩니다. 특히 저자들은 ControlNet을 사용하여 Stable Diffusion의 12개 encoding block과 1개 middle block을 copy하여 trainable copy를 생성합니다. 12개 encoding block은 4개 resolution (64x64, 32x32, 16x16, 8x8)이고, 각각 3번 반복됩니다. output은 12개 skip-connection에 더해지고 U-Net의 1개 middel block에 더해집니다. Stable Diffusion은 일반적으로 U-Net 구조이기 때문에, 이러한 ControlNet 구조는 다른 model에도 적용될 가능성이 큽니다.

ControlNet을 connect 하는 방법은 연산 효율적입니다. 왜냐하면 locked copy parameter는 frozen이고 finetuning과정에서 locked encoder는 기울기를 계산할 필요가 없기 때문입니다. 이러한 방식은 학습 속도를 향상시키고 GPU memory를 save 할 수 있습니다.

Image diffusion model은 점진적 denoise image를 학습하고 training domain으로부터 sample을 생성합니다. denoising process는 pixel space 또는 training data로부터 encode된 latent space에서 수행될 수 있습니다. Stable Diffusion은 latent image를 training domain으로 사용하고, 이 space에서 학습을 진행하면 안정적이라는 것이 발견되었습니다. 구체적으로 Stable Diffusion은 VQ-GAN처럼 512x512 pixel-space image를 smaller 64x64 latent image로 변환하는 pre-processing method를 사용합니다. ControlNet을 Stable Diffusion에 추가하기 위해, 저자들은 각 input conditioning image를 512x512에서 64x64 feature space vector로 변환했습니다. 저자들은 4x4 kernel과 2x2 stride를 가진 4개 convolution layer로 구성된 tiny network ε()를 사용하여 image-space condition c_i를 feature space conditioning vector c_f로 encode 합니다.

 

이 conditioning vector c_f는 ControlNet에 pass됩니다.

 

Training

input image z_0가 주어지면, image diffusion algorithm은 image에 noise를 점진적으로 더하여 noisy image z_t를 생성하며, 여기서 t는 noise가 더해진 횟수를 나타냅니다. time step t, text prompt c_t 뿐만 아니라 task-specific condition c_f를 포함한 condition set이 주어지면, image diffusion algorithm은 noisy image z_t에 더해진 noise를 예측하는 network ε_θ를 학습합니다. 

 

L은 entire diffusion model의 learning objective를 나타냅니다. 이 learning objective는 ControlNet을 사용하는 finetuning diffusion model에 바로 사용될 수 있습니다.

학습 과정에서, 저자들은 random 하게 50% text prompt c_t를 empty string으로 대체합니다. 이 방식은 ControlNet이 prompt 대신 input conditioning images (e.g., edges, poses, depth, ... )에서 직접 segment를 인식하는 능력을 높입니다.

학습 과정에서 zero convolution이 network에 noise를 더하지 않기 때문에, model은 항상 high-quality image를 예측할 수 있습니다. 저자들은 model이 점진적으로 condition을 control 하는 것을 학습하는 것이 아니라, 갑자기 input conditioning image에 맞춰 동작하는 것을 성공한다는 걸 관측했습니다. 이는 주로 10K optimization step 이전에 발생합니다.

 

위 그림과 같이 동작하며, 저자들은 이를 "sudden convergence phenomenon"이라 부릅니다. 

 

Inference

저자들은 ControlNet의 추가적인 condition이 denoising diffusion process에 미치는 영향을 여러 가지 방식으로 제어할 수 있습니다. 

 

- Classifier-free guidance resolution weighting

Stable Diffusion은 high-quality image를 생성하기 위해 Classifier-Free Guidance (CFG)라는 기술에 의존합니다. CFG는 

위 식으로 나타낼 수 있으며, ε_{prd}는 model의 final output, ε_{uc}는 unconditional output, ε_{c}는 conditional output β_{cfg}는 user-specified weight입니다. ControlNet을 통해 conditioning image를 추가할 때,  ε_{uc}와 ε_{c} 모두에 추가하거나, ε_{c}에만 추가할 수 있습니다. prompt가 없는 경우와 같이 어려운 상황에서는 ε_{uc}와 ε_{c} 모두에 추가하면 CFG guidance가 완전히 제거됩니다. 아래 그림의 b와 같습니다.

 

ε_c만 사용하는 것은 guidance를 매우 강하게 만듭니다. 결과는 위 c와 같습니다. 저자들의 solution은 conditioning image를 ε_c에 더하고 Stable diffusion과 ControlNet 사이 각 connection에 weight w_i를 곱하는 것입니다. 여기서 w_i는 각 bblock의 resolution에 따라 다르며, w_i = 64/h_i 형태입니다. 여기서 i는 block의 index를 나타냅니다. CFG guidacne 강도를 조절했을 때 결과는 위 d와 같고, 이를 CFG Resolution Weighting으로 부릅니다.

 

- Composing multiple ControlNets

 

여러 conditioning image (e.g., Canny edges, and pose)를 Stable Diffusion의 single instasnce에 더하기 위해, 저자들은 Stable Diffusion model에 대응하는 ControlNet의 output을 바로 더합니다. 추가적인 weight나 linear interpolation이 필수적이지 않습니다.

 

Experiments

저자들은 Canny Edge, Depth Map, Normal Map, M-LSD lines, HED soft edge, ADE20K segmentation, Openpose, user sketch를 포함한 다양한 test condition을 사용하는 Stable Diffusion with ControlNet으로 실험을 진행합니다.

 

Qualitative Results

 

위 그림은 여러 prompt setting에 따른 생성된 image를 보여줍니다. 

 

위 그림은 prompt 없는 다양한 condition으로 생성된 결과들을 보여줍니다. ControlNet은 다양한 input conditioning image에서 content semantic를 robust 하게 해석합니다. 

 

Ablative Study

저자들은 1) zero convolution을 Gaussian weight로 초기화된 standard convolution으로 대체하거나, 2) 각 block의 trainable copy를 1개 single convolution layer로 대체합니다(이를 ControlNet-lite라 부름).

저자들은 실제 사용자 행동을 test 하기 위해 4가지 prompt setting을 제시했습니다. 1) prompt 없음, 2) conditioning image의 object를 완벽하게 cover 할 수 없는 불충분한 prompt (이 논문에서의 default prompt는 "a high-quality, detailed, and professional image"입니다), 3) conditioning image의 의미를 변경하는 conflicting prompt, 4) 필요한 content semantic을 묘사하는 완벽한 prompt입니다.

 

위 그림은 ContrlNet이 성공적이게 4가지 setting을 모두 수행하는 것을 보여줍니다. 위 그림의 (c)는 lightweight ControlNet-lite를 보여주며, conditioning image를 번역하기에 충분히 strong 하지 않고 no prompt condition과 insufficient prompt condition에서 실패하는 모습을 보여줍니다. zero convolution이 대체되었을 때, ControlNet의 성능이 ControlNet-lite처럼 떨어지며, 이는 fine-tuning 과정에서 trainable copy의 pretrained backbone이 붕괴되었음을 나타냅니다. 위 그림의 (b)를 보면 알 수 있습니다.

 

Quantitative Evaluation

 

- User study

저자들은 20개 unseen hand-drawn sketch를 sample 하고 각 skecth에 5가지 method를 할당합니다: PITI's sketch model, Sketch-Guided Diffusion (SGD) with default edge-guidance scale (β=1.6), 높은 edge guide scale (β=3.2)을 적용한 SGD, ControlNet-lite, ContrlNet입니다. 저자들은 12명의 사용자에게 image quality와 sketch에 대한 fidelity 관점에서 평가를 진행했습니다. 결과는 아래와 같습니다.

 

Conclusion

ControlNet은 large pretrained text-to-image diffusion model을 위한 conditional control을 학습하는 neural network structure입니다. 이는 source model의 large-scale pretrained layer을 재사용하여 특정 condition을 학습하기 위해 dee pand strong encoder를 build 합니다. original model과 trainable copy는 "zero convolution" layer로 connect 되며, 이를 통해 학습 과정에서 harmful noise를 제거할 수 있습니다. 실험을 통해 ControlNet이 효과적으로 Stable Diffusion with single or multiple condition을 control 할 수 있다는 것을 보였습니다.