Introduction
이 논문은 흔히 ControlNet이라고 불리고 제목 그대로 추가적인 input condition을 지원하여 large diffusion model를 제어하기 위한 모델이다. ControlNet은 end-to-end 방식으로 학습하며, 학습 데이터 세트가 적은(<50k) 경우에도 좋은 성능을 보이고 학습 속도 또한 diffusion model의 fine tuning만큼 빠른 장점을 가진다.
저자는 기존 large text-to-image 모델의 prompt에 대한 의존성과 특정 task에서 활용함에 있어서 발생할 수 있는 현실적인 문제를 언급했다. 이와 관련해서 세 가지 측면에서 검토하고 제안했다.
1. task-specific domain의 경우 일반적인 text-to-image 데이터 스케일만큼 크지 않으므로 large model을 특정 문제에 대해 학습시킬 때는 과적합을 방지하고 일반화 능력을 보존할 수 있는 방법이 필요하다.
2. large model를 항상 large computation cluster에서 사용할 수 없기 때문에 제한된 시간 및 메모리 공간내에서 사용가능한 학습 전략이 필요하다. 이를 위해 fine-tuning이나 trasfer learning 뿐만 아니라 pretrained weight의 활용이 필요하다.
3. 다양한 image processing은 denoising process 제한, multi-head attention activation 편집등과 같은 hand crafted 룰을 따른다. 이러한 문제는 필수적으로 raw data에 대한 object-level이나 scene-level에서 이해가 필요하므로 실현 가능성이 떨어지게 되어 end-to-end 학습이 필요하다.
이러한 질문들에 대답하기 위해, 저자는 다양한 image processing applications들을 조사했고, 세가지 발견을 할 수 있었다.
3가지 발견
- Task-specific 도메인에서 사용가능한 data 규모가 항상 general text2img data의 규모만큼 크지 않다.
- Image processing task을 data-driven solution으로 수행하고자 할때, 항상 large computation이 사용가능하지는 않다.
- 다양한 image processing 문제들은 “문제 정의”, “유저 컨트롤”, “이미지 annotation”등에 있어 다양한 form을 가진다.
ControlNet은 end-to-end neural network 구조로 Stable Diffusion 같은 large image diffusion model들을 특정 input condition으로 다룰 수 있게 한다. ControlNet은 large diffusion model의 weights를 “trainable copy”와 “locked copy”로 복제하는데, “trainable copy”가 conditional control을 배우기 위해 task-specific 데이터셋에 대해 학습된 것에 반해 “locked copy”는 기존 network의 capability를 보존한다. 이 두가지 neural net block은 특별한 종류의 convolution layer인 “zero convolution”와 연결된다. production-ready weight가 보존되기 때문에, 학습은 다른 규모의 dataset에 대해서도 robust하다. 또한, zero convolution이 새로운 노이즈를 더하지 않기 때문에, 학습 역시 diffusion 모델의 fine-tuning만큼 빠르다.
1x1 convolution의 효과는 학습 파라미터 감소, 비선형성 추가, 객체 정보와 위치 정보의 독립적 연산등이 있다.
Method
ControlNet
ControlNet은 신경망 블록의 입력 조건을 조작하여 전체 신경망의 전반적인 동작을 추가로 제어한다. 여기서 “네트워크 블록”은 신경망을 구축하기 위해 자주 사용되는 단위로 한데 모인 레이어의 집합을 의미하며, 예를 들어 “resnet” 블록, “conv-bn-relu” 블록, multi-head attention 블록, transformer 블록 등이 있다.
예를 들어 2D feature를 사용하면 {ℎ,𝑤,𝑐}가 높이, 너비, 채널 수 feature map 𝑥∈𝑅^ℎ×𝑤×𝑐가 주어지면 신경망 블록 𝐹(⋅;Θ)는 다음과 같이 𝑥를 다른 feature map 𝑦로 변환한다. (위 그림의 (a))
𝑦=𝐹(𝑥;Θ)
모든 파라미터를 Θ로 잠근 다음 Θ𝑐에 복제한다. 복사된 Θ𝑐는 외부 조건 벡터 𝑐로 학습된다. 본 논문에서는 원래 파라미터와 새 파라미터를 locked copy와 trainable copy라고 한다. 원래 가중치를 직접 학습하는 대신 이러한 사본을 만드는 목적은 데이터 셋이 작을 때 overfitting을 방지하고 수십억 개의 이미지에서 학습한 대형 모델의 품질을 보존하는 것이다.
신경망 블록은 “zero convolution”이라고 하는 고유한 유형의 convolution layer, 즉 가중치와 바이어스가 모두 0으로 초기화된 1×1 convolution layer로 연결된다.
일반적으로 Convolution filter의 값이 0이면 결과값(Destination pixel)은 0이 된다. 그러나 ControlNet에서는 이를 역전파를 이용해 학습시키는데, 이를 통해 이미지의 실루엣처럼 남아서 가중치 부분이 조건부 생성을 지원한다. 이는 ControlNet에서 이미 학습된 모델의 의미론을 보존하면서 새로운 조건을 학습시키는 핵심 기술 중 하나이다. 이러한 ControlNet 구조는 다양한 이미지 생성 작업에서 사용될 수 있으며, 높은 수준의 이미지 생성 결과를 보장한다.
Zero convolution 연산을 𝑍(⋅;⋅)로 표시하고 파라미터의 두 인스턴스 {Θ_𝑧1, Θ_𝑧2}를 사용하여 다음과 같이 ControlNet 구조를 구성한다. (위 그림의 (b))
𝑦𝑐=𝐹(𝑥;Θ)+𝑍(𝐹(𝑥+𝑍(𝑐;Θ𝑧1);Θ𝑐);Θ𝑧2)
여기서 𝑦𝑐는 이 신경망 블록의 출력이 된다.
Zero convolution layer의 가중치와 바이어스는 모두 0으로 초기화되기 때문에, 첫 번째 학습 step에서 다음과 같다.
𝑍(𝑐;Θ𝑧1)=0
𝐹(𝑥+𝑍(𝑐;Θ𝑧1);Θ𝑐)=𝐹(𝑥;Θ𝑐)=𝐹(𝑥;Θ)
𝑍(𝐹(𝑥+𝑍(𝑐;Θ𝑧1);Θ𝑐);Θ𝑧2)=𝑍(𝐹(𝑥;Θ𝑐);Θ𝑧2)=0
그리고 이는 다음과 같이 변환될 수 있다.
𝑦𝑐=𝑦
첫 번째 학습 step에서 신경망 블록의 trainable copy와 locked copy의 모든 입력 및 출력은 마치 ControlNet이 존재하지 않는 것과 일치한다. 즉, ControlNet이 일부 신경망 블록에 적용될 때 최적화 전에는 feature에 영향을 미치지 않는다. 모든 신경망 블록의 능력, 기능 및 결과 품질은 완벽하게 보존되며 추가 최적화는 fine-tuning만큼 빠르다 (해당 레이어를 처음부터 학습하는 것과 비교할 때).
Zero convolution layer의 기울기 계산을 간단히 추론해보자. 입력 map 𝐼∈𝑅ℎ×𝑤×𝑐가 주어지면 임의의 공간적 위치 𝑝와 채널별 인덱스 𝑖에서 가중치 𝑊와 바이어스 𝐵를 갖는 1×1 convolution layer를 고려하면 forward pass는 다음과 같이 쓸 수 있다.
𝑍(𝐼;{𝑊,𝐵})𝑝,𝑖 = 𝐵_𝑖 + ∑(𝑗 to 𝑐)𝐼_𝑝,𝑐𝑊_𝑖,𝑗
Zero convolution은 최적화 전에 𝑊=0과 𝐵=0이므로 𝐼𝑝,𝑖가 0이 아닌 모든 위치에 대해 기울기는
∂𝑍(𝐼;{𝑊,𝐵})𝑝,𝑖 / ∂𝐵_𝑖=1
∂𝑍(𝐼;{𝑊,𝐵})𝑝,𝑖 / ∂𝐼_𝑝,𝑖=∑(𝑗 to 𝑐)𝑊𝑖,𝑗=0
∂𝑍(𝐼;{𝑊,𝐵})𝑝,𝑖 / ∂𝑊𝑖,𝑗=𝐼_𝑝,𝑖≠0
가 된다.
Zero convolution으로 인해 feature 항 𝐼의 기울기가 0이 될 수 있지만 가중치와 바이어스의 기울기의 영향을 받지 않는다는 것을 알 수 있다. Feature 𝐼가 0이 아닌 한 가중치 𝑊는 첫 번째 gradient descent iteration에서 0이 아닌 행렬로 최적화된다. 특히 feature 항은 데이터셋에서 샘플링된 입력 데이터 또는 조건 벡터이며 자연스럽게 0이 아닌 𝐼를 보장한다.
ControlNet for Text-to-Image Diffusion
Stable Diffusion을 예로 들어 ControlNet을 사용하여 task별 조건으로 대규모 diffusion model을 제어하는 방법을 소개한다.
Stable Diffusion은 수십억 개의 이미지에 대해 학습된 대규모 text-to-image diffusion model이다. 이 모델은 본질적으로 인코더, 중간 블록, skip-connection으로 연결된 디코더가 있는 U-net이다. 인코더와 디코더 모두 12개의 블록이 있고 전체 모델에는 25개의 블록 (가운데 블록 포함)이 있다. 이 블록에서 8개의 블록은 다운샘플링 또는 업샘플링 convolution layer이고, 17개의 블록은 각각 4개의 resnet 레이어와 2개의 ViT (Vision Transformers)를 포함하는 기본 블록이다. 각 ViT에는 몇 가지 cross-attention과 self-attention 메커니즘이 포함되어 있다. 텍스트는 OpenAI CLIP으로 인코딩되고 diffusion timestep은 위치 인코딩으로 인코딩된다.
Stable Diffusion은 VQ-GAN과 유사한 전처리 방법을 사용하여 안정화된 학습을 위해 512×512 이미지의 전체 데이터샛을 더 작은 64×64 크기의 “latent 이미지”로 변환한다. 이를 위해서는 convolution 크기와 일치하도록 이미지 기반 조건들을 64×64 feature space로 변환하는 ControlNet이 필요하다. 본 논문은 4×4 kernel과 2×2 stride를 가지는 convolution layer로 구성된 작은 네트워크 𝐸(⋅)를 사용한다. 다음과 같이 이미지 조건 𝑐𝑖를 feature map으로 인코딩한다.
𝑐𝑓=𝐸(𝑐𝑖)
여기서 𝑐𝑓는 변환된 feature map이다. 이 네트워크는 512×512 이미지 조건을 64×64 feature map으로 변환한다.
위 그림과 같이 ControlNet을 사용하여 U-net의 각 레벨을 제어한다. ControlNet을 연결하는 방법은 계산적으로 효율적이다. 원래 가중치가 잠겨 있기 때문에 학습을 위해 원래 인코더에 대한 기울기 계산이 필요하지 않다. 이렇게 하면 원본 모델에서 기울기 계산의 절반을 피할 수 있으므로 학습 속도가 빨라지고 GPU 메모리가 절약된다. ControlNet으로 Stable Diffusion 모델을 학습하려면 GPU 메모리가 약 23% 더 필요하고 각 학습 iteration에서 34% 더 많은 시간이 필요하다.
구체적으로, ControlNet을 사용하여 Stable Diffusion의 12개 인코딩 블록과 1개 중간 블록의 trainable copy를 만든다. 12개의 블록은 각각 3개의 블록을 갖는 4개의 해상도 (64×64, 32×32, 16×16, 8×8)로 되어 있다. 출력은 U-net의 skip connection 12개와 중간 블록 1개에 추가된다. Stable Diffusion은 일반적인 U-net 구조이므로 이 ControlNet 아키텍처는 다른 diffusion model에서 사용할 수 있다.
Training
이미지 diffusion model은 이미지를 점진적으로 noise를 제거하여 샘플을 생성하는 방법을 학습한다. Noise 제거는 pixel space 또는 학습 데이터에서 인코딩된 latent space에서 발생할 수 있다. Stable Diffusion은 latent를 학습 도메인으로 사용한다.
이미지 𝑧0가 주어지면 diffusion 알고리즘은 점진적으로 이미지에 noise를 추가하고 noise가 추가된 횟수를 𝑡로 하여 noisy한 이미지 𝑧𝑡를 생성한다. 𝑡가 충분히 크면 이미지는 순수한 noise에 가깝다. Timestep 𝑡, 텍스트 프롬프트 𝑐𝑡, task별 조건 𝑐𝑓를 포함한 일련의 조건이 주어지면, 이미지 diffusion 알고리즘은 다음의 목적 함수로 네트워크 𝜖𝜃를 학습하여 𝑧𝑡에 추가된 noise를 예측한다.
이 목적 함수는 fine-tuning diffusion model에서 직접 사용할 수 있다.
학습 중에 50%의 텍스트 프롬프트 𝑐𝑡를 빈 문자열로 랜덤하게 바꾼다. 이는 입력 조건 map에서 의미론적 내용을 인식하는 ControlNet의 능력을 용이하게 한다. 이것은 주로 Stable Diffusion 모델에 대한 프롬프트가 표시되지 않을 때 인코더가 프롬프트 대신 입력 제어 map에서 더 많은 semantic을 학습하는 경향이 있기 때문이다.
Improved Training
GPU 사용이 매우 제한적이거나 매우 강력한 경우 ControlNet의 학습을 개선하기 위한 몇 가지 전략을 사용할 수 있다.
Small-Scale Training
GPU 사용이 제한될 때 ControlNet과 Stable Diffusion 간의 연결을 부분적으로 끊으면 수렴이 가속화될 수 있다. 기본적으로 ControlNet을 “SD Middle Block”과 “SD Decoder Block 1,2,3,4”에 연결한다. 중간 블록을 연결하면 학습 속도를 약 1.6배 향상시킬 수 있다 (RTX 3070TI 노트북 GPU 기준). 모델이 결과와 조건 사이에 합리적인 연관성을 보여주면 정확한 제어를 용이하게 하기 위해 학습에서 끊긴 연결을 다시 연결할 수 있다.
Large-Scale Training
여기서 대규모 학습이란 강력한 계산 클러스터 (최소 8개의 Nvidia A100 80G 또는 동급)와 대규모 데이터셋 (최소 100만 개의 학습 이미지 쌍)을 모두 사용할 수 있는 상황을 말한다. 이는 일반적으로 데이터를 쉽게 사용할 수 있는 task에 적용된다. 이 경우 overfitting의 위험이 상대적으로 낮기 때문에 먼저 충분히 많은 iteration (보통 50,000 step 이상) 동안 ControlNet을 학습한 다음 Stable Diffusion의 모든 가중치를 잠금 해제하고 전체 모델을 공동으로 학습시킬 수 있다. 이것은 보다 문제에 특화된 모델로 이어질 것이다.
Experiment
Ablative Study
우선 4가지 종류의 프롬프트에 zero conv를 일반 conv로 변경한 실험, 각 학습가능한 copy를 conv로 변경한 실험을 하여 한 Input에 대하여 12가지 테스트를 하였다.
zero convolution의 성능이 대단하다는 것을 알 수 있었고 prompt engineering의 중요성도 알 수 있었다.
Quantitative Evaluation
1에서 5점사이의 Average Human Ranking을 사용했다.
Comparision to Previous Methods
Extra
