Adding Conditional Control to Text-to-Image Diffusion Models (ControlNet)

End-to-End 방식으로 확산 모델에 조건 입력

 

Github

arXiv

 

 

 

Abstract

확산 모델에 조건부 입력을 가능하게 하는 end-to-end 방식의 ControlNet 제안

 

 

 

Introduction

• 작업별 도메인에서 사용 가능한 데이터셋은 크지 않다. 따라서 일반화 능력을 보존하기 위한 특별한 훈련 방법이 필요하다.
• 시간과 메모리의 제약으로 인해 fine tuning 전략이 필요하다.
• 다양한 형태의 이미지 처리 문제에 대해 denoising process 제한, attention의 편집 등 절차적 방식으로 규제했지만 개체 수준의 본질적인 이해를 위해서는 end-to-end 학습이 필수불가결하다.

ControlNet은 확산 모델의 가중치를 trainable copy, locked copy로 복제하며 locked copy는 네트워크의 feature를 보존하고 trainable copy는 조건부 통제를 학습한다.

 

ControlNet의 학습은 fine tuning 만큼 빠르며, 개인용 GPU에서도 학습할 수 있다.

 

 

 

Method

ControlNet

'네트워크 블록'은 신경망에서 자주 사용되는 단위.

 

예를 들어 2D feature는 다음과 같이 feature map x를 받아 블록 F에서 피라미터 Θ를 통해 다른 feature map y로 변환된다.

 

ControlNet에서는 사전 훈련된 대형 모델의 생성 품질을 보존하고 fine tuning의 과적합을 방지하기 위해 Θ를 둘로 나누고 trainable copy만 훈련한다.

 

이때, 가중치와 편향을 모두 0으로 초기화한 1x1 convolution인 'zero convolution'을 사용하며 Z로 표기한다.

 

ControlNet의 연산 구조

 

Zero convolution으로 인해 학습 초기에는 사전 훈련 모델과 완전히 같음.

---

가중치 W, 편향 B, input I에 대해 zero convolution의 gradient :

 

W, B = 0 인 초기 상태에서의 gradient :

 

Input이 0이 아닌 한 zero convolution의 W, B의 gradient는 0이 되지 않고 경사 하강법이 가능하고, 손실 함수 L에 대해 outside gradient를 이용한 경사 하강으로 W*을 구하면

 

0이 아닌 W가 구해지고 신경망의 학습을 시작한다.

---

ControlNet in Image Diffusion Model

Stable Diffusion(LDM)의 잠재공간을 위한 인코딩 네트워크를 사용하고

 

다음 그림과 같이 U-Net을 제어

---

Training

Noise image z_t, 시간 단계 t, 조건 텍스트 c_t, 사용자 입력 조건 c_f에 확산 과정 최적화

 

입력 조건에서 의미론적 내용을 인식하는 ControlNet의 기능을 최대화하기 위해 훈련 중 조건 텍스트의 50%를 빈 문자열로 대체.

---

Improved Training

Small-Scale Training

계산 리소스가 제한적일 때는 SD와의 연결 중 일부를 끊는 것으로 훈련 속도를 1.6배 정도 향상할 수 있다.

 

Large-Scale Training

대규모 훈련이 가능한 환경의 경우 ControlNet을 충분히 훈련한 뒤 SD의 가중치까지 공동훈련하여 더 최적화할 수 있다.

---

 

 

Experiment