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

오늘 소개할 논문은 "Adding Conditional Control to Text-to-Image Diffusion Models" 논문으로 2023년 ICCV에 등록된 논문이다. Diffusion model의 방법론 자체를 발전시킨 논문은 아니고, 응용 논문이라고 할 수 있다. 응용 논문인 만큼 수학적 지식 없이도 충분히 읽을 수 있으며, 방법은 매우 간단하지만 ControlNet이라는 이름으로 다른 논문에서도 종종 언급될 정도로 꽤 영향력 있는 논문이라 읽어두면 좋을 듯하다.

 

Introduction

최근 생성형 모델이 빠르게 발전하면서 특히 Text-to-Image model들이 엄청난 속도로 연구가 되고 있다. 그러나 단순히 텍스트만으로는 복잡한 레이아웃이나 포즈, 형태 등을 정확하게 제어하기 어려우며, 원하는 그림을 정확히 설명하기 위해서는 꽤나 긴 텍스트가 필요하다. 따라서 기존의 Text-to-Iage model만으로는 사용자의 의도에 맞는 이미지를 생성하기 위해 많은 시도와 수정을 필요로 하게 된다.

 

따라서 이 논문에서는 Text-to-Image Diffusion model에 다양한 conditional control을 추가하는 ControlNet이라는 architecture를 제안한다.

 

Method

ControlNet의 구조는 매우 간단하다. 어떠한 모델에도 적용이 가능하며, 그 구조는 아래와 같다. 

ControlNet의 구조

기존의 네트워크를 그대로 하나 더 복사해서 하나는 얼려두고, 하나는 condition을 학습시키는 방법이다. 이때, ControlNet은 condition의 특징을 학습하는 것이 목표이기 때문에 기존 네트워크의 encoder 부분만 복사하여 사용하며, decoder 부분은 zero convolution으로 초기화하여 사용한다.

 

Diffusion model의 경우, 기존의 네트워크가 Unet 형식이기 때문에, 아래의 그림을 참고하여 보면 조금 더 이해가 쉬울 것이다.

Stable Diffusion에서의 ControlNet

왼쪽은 기존의 Stable Diffusion이다. 매 t-step마다 t-step에서의 노이즈 이미지 $z_t$와 prompt embedding, time embedding이 input으로 주어지고, 노이즈를 제거한 다음 step의 이미지를 output으로 내놓는 Unet 구조이다.

여기에 ControlNet을 적용한다면 우선 기존 네트워크의 encoder 부분을 그대로 복사하여 가져오고, condition을 embedding 하는 부분과 decoder 부분은 0으로 초기화한 zero convolution으로 쌓아주고 ControlNet 부분만 학습을 시켜준다.

 

위와 같이 이 논문에서는 모델의 디테일로 trainable copy와 zero convolution 방법을 사용하였는데, 각각의 장점은 다음과 같다.

• trainable copy
  • 사전 학습된 모델을 복사하여 사용하였기 때문에 기존 모델의 성능을 보존
  • finetuning 시, 기존의 데이터에 대한 정보를 잃는 Catastrophic forgetting이 발생하는데, 기존의 네트워크를 그대로 두고, 추가로 복사한 네트워크를 학습시키기 때문에 Catastrophic forgetting 방지
  • Encoder 부분만 그대로 가져와 효율적으로 학습 가능
• zero convolution
  • 학습 안정성
    random 초기화 시, 학습 초기에 condition embedding으로부터 random noise가 추가될 수 있음.

 

아래 Experiments를 읽다 보면 여러 condition을 줬을 때도 좋은 성능을 내고 있다고 주장하는데, 이렇듯 ControlNet의 장점 중 또 하나는 따로 학습시킨 여러 condition들을 같이 줄 수 있다는 것이다. ControlNet 구조 그림에서 기존 네트워크의 output과 ControlNet의 output을 단순히 합쳤는데, 여러 condition을 주고 싶은 경우, 각 condition을 학습시킨 ControlNet들을 단순히 합쳐주면 된다.

 

Experiments

이 논문에는 다른 논문들에 비해 꽤 많은 실험들이 소개되어 있다. 한번 같이 하나하나 살펴보자!

 

Canny edge와 pose 정보를 condition으로 준 ControlNet 결과

논문의 맨 앞에 나오는 실험 결과이다. 간단한 text만을 줬을 때에는 다양한 이미지들이 나올 수 있지만 canny edge와 pose 정보를 같이 주면 원하는 형태에 맞게 이미지가 생성되는 것을 볼 수 있다.

 

다양한 condition에 대한 ControlNet 결과

condition으로는 총 8가지의 condition으로 학습을 시켜고 실험을 돌렸으며, 위의 그림과 같이 8가지에 대해 성공적으로 학습이 이루어졌음을 확인할 수 있다.

 

2개의 condition(pose & depth)을 동시에 준 결과

앞서 언급했던 대로 따로 학습시킨 여러 개의 condition을 동시에 줄 수도 있는데, 위 사진이 이에 대한 실험 결과이다. 개별 condition에 대한 이미지에 비해 약간 부자연스러워 보이기는 하지만 두 condition을 모두 잘 나타내고 있는 모습을 볼 수 있다.

 

sudden convergence phenomenon

이 논문에서는 sudden convergence phenomenon이라는 현상이 나타났다고 주장하고 있는데, 이는 학습 시 초반에는 condition에 대한 정보를 잘 모르다가 10K가 되지 않는 어느 시점에서 갑자기 성능이 좋아진다는 것이다.

 

trainable copy와 zero convolution에 대한 실험 결과

Method 부분에서 여러 장점들을 설명하면서 trainable copy와 zero convolution을 사용하였다고 언급하였는데, 이를 설명하는 실험이라고 볼 수 있다. 우선 가장 위가 이 논문에서 주장하고 있는 방법이다. 사진들을 보면 당연히 잘 되고 있는 것을 볼 수 있다. 여기서 zero convolution을 사용하지 않고 decoder 부분과 condition encoder 부분을 random 하게 초기화한 경우와 기존 네트워크를 복사하여 사용하지 않고 그냥 convolution network를 쌓은 경우에 대해서는 prompt를 완벽하게 주어야만 성능이 어느 정도 나오는 모습을 볼 수 있다.

 

학습 데이터 셋 크기에 따른 성능 결과

당연하게도 학습 데이터가 많을수록 성능이 좋았다는 결과를 보여주는 사진이다.

 

애매한 condition에 대한 실행 결과

이 실험이 살짝 재밌는데 왼쪽 Input을 보면 물병이 먼저 보이냐 사람 옆모습이 먼저 보이냐에 대한 애매한 표현을 나타낸 이미지이다. 이렇게 애매한 조건을 주었을 경우 생성할 때마다 condition은 만족시키지만 다 다른 결과가 나타나는 것을 볼 수 있다.

 

이 외에도 몇 가지 실험 결과가 더 있는데, 직관적으로 잘 쓰여 있어 직접 읽어보면 좋을 듯하다. (절대 귀찮아서가 아니다 ㅎ)

 

Conclusion

• Text-to-Image diffusion model에 공간적인 조건(spatial conditional control)을 추가하여 이미지 생성에 대한 제어성 향상
• 다양한 conditional input을 통한 이미지 생성
• 소규모 데이터셋에서도 안정적인 학습 성능

 

마치며...

이 논문을 읽으면서 물론 저자들은 열심히 이것저것 실험한 끝에 얻어낸 결과물이겠지만 꽤나 간단한 방법으로도 diffusion model을 활용할 수 있다는 생각이 들었다. 또한 나처럼 수학을 못하는 사람들은 diffusion의 수식을 달달 공부하는 것보다 빨리 여러 응용 실험을 해야 좋은 연구를 할 수 있지 않을까라는 생각이 들었다.