[논문리뷰] VICO : PLUG-AND-PLAY VISUAL CONDITION FOR PERSONALIZED TEXT-TO-IMAGE GENERATION

VICO : PLUG-AND-PLAY VISUAL CONDITION FOR PERSONALIZED TEXT-TO-IMAGE GENERATION
Shaozhe Hao, Kai Han, Shihao Zhao, Kwan-Yee K. Wong
[paper] [github]
7 Dec 2023

---

Contents

1. Introduction

2. Method

3. Experiments

4. Limitations

5. Summary

---

1. Introduction

Personalization은 text-to-image generation model이 unique object에 대한 이미지 합성이 가능하도록 하는 분야이다. 최근 DreamBooth, Custom Diffusion, Textual Inversion과 같은 연구에 의해 personlaization 분야가 발전되고 있다.

 

DreamBooth와 Custom Diffusion은 각각 전체 모델과 cross-attention layers를 fine-tuning한다. 두 모델 모두 모델을 fine-tuning하는 과정에서 사전 지식을 잃어버리는 language-drift 문제가 발생한다. 이러한 문제를 해결하기 위해서는 이미지를 사전생성하거나 대형 dataset에서 이미지를 검색해야하는 prior-preservation loss를 사용해야 한다. 기존의 모델은 물체의 정교한 디테일을 표현하지 못하며 CLIP의 표현력 한계에 의해 물체에 과적합될 가능성이 높다.

 

저자들은 모델의 사전 지식을 활용하기 위해 모델을 fine-tuning하지 않고 reference image를 통해 image conditiong이 가능한 기술인 ViCo를 제안한다. 이를 위해 기존 모델에 image cross-attention 모듈을 추가하며, fine-tuning 과정이 없으므로 language drift 문제가 발생하지 않는다.

 

또한 target object와 배경을 분리하기 위한 mask를 생성하는 메커니즘도 제안한다. 이는 추가적인 입력을 필요로 하지 않으며, 효율적으로 배경을 분리할 수 있다.

---

2. Method

Personalization에 사용되는 모델은 사용하는 token에 따라 두 종류로 분류할 수 있다. "[V] class" 형태의 class token을 사용하는 DreamBooth, Custom Diffusion과 single token $S_*$만을 사용하는 Textual Inverison으로 나뉜다. 저자들은 single token만으로 object를 임베딩하는 것이 personlaization의 중요한 요소라고 생각하여 single token $S_*$ 방식을 사용한다.

 

기존 논문들과 마찬가지로 ViCo의 베이스라인은 Latent Diffusion Model을 기반으로 하는 Stable Diffusion이다. 인코더 $\mathcal{E}$는 이미지 $x \in \mathcal{I}$를 latent $z = \mathcal{E}(x)$로 매핑하고, 디코더 $\mathcal{D}$는 latent로부터 다시 이미지 $\hat{x}=\mathcal{D}(\mathcal{E}(x))$를 복원한다. Text-to-Image synthesis를 위해 text prompt로 conditioning하는 것이 가능하며, 전체 loss는 다음과 같다.$$\mathcal{L}_{S D}=\mathbb{E}_{z \sim \mathcal{E}(x), y, \epsilon \sim \mathcal{N}(0,1), t}\left[\left\|\epsilon-\epsilon_\theta\left(z_t, t, c_\pi(y)\right)\right\|_2^2\right]$$

Image conditiong을 위한 기존 연구들은 feature concatenation이나 element-wise addition 방식으로 접근하였다. 이런 방식들은 이미지의 구조를 잘 포착하지만, 정교한 디테일을 보존하지 못한다. 저자들은 동일한 위치에 있는 patch만을 고려하는 대신 모든 patch의 pairwise한 상관관계를 고려한다. 이를 위해 기존 U-net의 text cross-attention layers와 동일한 구조를 가진 image cross-attention layers를 제안한다. 

 

추가적인 feature extractor를 학습하는 기존의 image conditiong 연구들과 다르게 ViCo는 pretrained된 autoencoder $\mathcal{E}$를 그대로 사용하여 reference image $x_r \in \mathcal{I}$를 latent $z_r = \mathcal{E}(x_r)$로 매핑한다. $\epsilon_\theta^l(\cdot)$를 U-net의 $l$번째 attetion layer라고 하면 다음과 같이 쓸 수 있다.$$c_I^l=\epsilon_\theta^l\left(z_r, t, c_T\right), l \in\{0,1, \cdots, L-1\}$$

$c_T=c_\pi(y)$는 text condition이다. 이렇게 계산된 visual conditiong은 다시 cross-attention을 통해 noisy image에 결합된다.$$n_t^l=\epsilon_{\theta, \psi}^l\left(z_t, t, c_T, c_I^l\right), l \in\{0,1, \cdots, L-1\}$$

정확한 계산 과정은 다음과 같다. Text cross-attention을 $\mathcal{A}_T(q, k v)$, image cross-attention을 $\mathcal{A}_I(q, k v)$라 하자. 기존의 cross-attention은 $n_t^{\prime}=\mathcal{A}_T\left(n_t, c_T\right)$로 표현할 수 있다. 이는 아래 식으로 대체된다.$$\hat{n}_t=\mathcal{A}_T\left(n_t, c_T\right), \ \hat{c}_I=\mathcal{A}_T\left(c_I, c_T\right), \ n_t^{\prime}=\mathcal{A}_I\left(\hat{n}_t, \hat{c}_I\right)$$

모델의 전체 구조는 아래와 같다. Image attention layer의 $\hat{M}$은 mask로, 아래에서 설명한다.

 

기존 연구에서는 공통적으로 training image의 배경에 대해 모델이 과적합되는 문제가 발생한다. ViCo는 mask를 통해 target object에서 배경을 분리하여 cross-attention을 계산함으로서 이러한 문제를 피한다. Text cross-attention의 식은 다음과 같다. $$\text { TextAttention }(Q, K, V)=\operatorname{softmax}\left(\frac{Q K^T}{\sqrt{d_k}}\right) V$$

Prompt-to-Prompt 논문에서는 attention map이 토큰에 대해 훌륭한 masking 성능을 보인다는 것을 보여주었다. 이에 영감을 받아 ViCo는 unique token $S_*$에 대한 attention map을 binarization하여 mask로 사용한다. Binarization에는 Otsu thresholding을 사용한다. Masked image cross-attention은 다음과 같이 쓸 수 있다.$$\text { ImageAttention }(Q, K, V)=\left(\hat{M} \odot \operatorname{softmax}\left(\frac{Q K^T}{\sqrt{d_k}}\right)\right) V$$

여기서 $\odot$은 element-wise product이다.

전체 masking 과정은 아래와 같다.

 

적은 수의 training set만으로 fine-tuning하는 과정에서 과적합은 흔하게 발생한다. 저자들은 transformer의 end-of-text token인 $\langle|\mathrm{EOT}|\rangle$가 unique object에 대해 consistency semantics를 유지하는 것을 실험적으로 확인하였다. 이를 이용하여 저자들은 text cross-attention map에서 $S_*$와 $\langle|\mathrm{EOT}|\rangle$의 similarity를 이용한 regularization을 제안한다. Cross-attention map은 다음과 같다.$$A:=\operatorname{softmax}\left(\frac{Q K^T}{\sqrt{d_k}}\right) \in \mathbb{R}^{B \times D_p \times D_t}$$

$S_*$가 i번째 토큰, $\langle|\mathrm{EOT}|\rangle$가 j번째 토큰이라고 가정하고 다음과 같은 regularization을 구성할 수 있다.$$\mathcal{L}_{\text {reg }}=\left\|A_{\star, i} / \max \left(A_{\star, i}\right)-A_{\star, j} / \max \left(A_{\star, j}\right)\right\|_2^2$$

Regularization을 통해 신뢰할 수 있는 mask를 얻을 수 있다.

 

ViCo의 전체 loss는 다음과 같다.

$$\mathcal{L}=\mathbb{E}_{z \sim \mathcal{E}(x), z_r \sim \mathcal{E}\left(x_r\right), y, \epsilon \sim \mathcal{N}(0,1), t}\left[\left\|\epsilon-\epsilon_{\theta, \psi}\left(z_t, t, z_r, c_\pi(y)\right)\right\|_2^2\right]+\lambda \mathcal{L}_{r e g}$$

실제로 training되는 부분은 $\psi$로 parameterized되는 image cross-attention layers과 text embedding $S_*$이다.

---

3. Experiment

실제 학습에서는 regularization weight $\lambda = 5 \times 10^{-4} $, learning rate of $S_*$ $5 \times 10^{-3}$, learning rate of image cross-attention layers $10^{-5}$를 사용한다. Training은 batch size 4로 400 steps 동안 진행된다.

 

기존 연구들(DreamBooth, Custom Diffusion, Textual Inversion)의 데이터셋을 균일하게 샘플링하여 20개의 unique concepts로 이루어진 데이터셋을 사용한다. ViCo의 목적은 unique object의 외형을 최대한 보존하는 것이므로 목적에 맞지 않는 prompt인 "A cube shaped $S_*$"는 사용하지 않는다.

 

기존 연구들과의 성능 비교를 위해 3개의 evaluation metric을 사용한다. DreamBooth와 마찬가지로 subject fidelity metric인 CLIP-I, DINO와 prompt fidelity인 CLIP-T를 사용한다.

 

위의 표를 보면 ViCo가 CLIP-I와 DINO에서 가장 높은 subject fidelity를 보여주는 것을 확인할 수 있다. 표를 보면 Textual Inversion과 ViCo가 DreamBooth와 Custom Diffusion에 비해 낮은 CLIP-T score를 가지는 것을 볼 수 있다. 저자들은 이를 DreamBooth와 Custom Diffusion가 "[V] class" 형태의 토큰을 사용하여 모델의 사전 지식을 사용하기 때문이라고 설명한다.

아래 표는 training과 inference 시간을 비교한 것이다. Inference 시간은 살짝 증가하였지만 training 시간이 훨씬 짧은 것을 알 수 있다. Prior-preservation loss를 사용하는 DreamBooth와 Custom Diffusion의 경우에는 이를 위한 이미지 합성이나 검색에 시간이 추가로 소요된다는 점을 생각하면 시간 차이는 더 커진다.

 

아래의 sample을 보면, ViCo는 object의 디테일을 재현하면서도 prompt도 충실하게 따른 결과를 생성한다. 특히 마지막 줄에서는 ViCo만이 사자 얼굴, 항아리 모양, 목재 재질을 모두 표현하는 이미지를 합성한 것을 볼 수 있다.

 

저자들은 visual conditioning과 attention map mask, initialization of $S_*$의 효과를 확인하기 위해서 Ablation study를 진행했다.

위의 표에서 $(0)$은 Textual Inversion, 즉 $S_*$ 학습만 진행한 모델, $(1)$은 visual condition을 추가한 모델, $(2)$는 masking을 추가한 모델이다. 위의 그림에서 볼 수 있듯이 visual condition은 object의 디테일을 보존하는데 도움을 준다. 표에서도 visual condition을 추가한 모델이 모델의 성능을 크게 높이는 것을 확인할 수 있다. Masking은 그림에서 볼 수 있듯이 결과물을 왜곡하는 것을 방지한다. 표에서도 object fidelity가 증가하는 것을 확인할 수 있다. 또한, mask regularization을 사용하면 mask가 object와 잘 정렬되는 것을 확인할 수 있다.

 

Textual Inversion은 embedding $S_*$의 초기값에 큰 영향을 받는다. 아래 그림과 같이 "object", "animal"과 같은 일반적인 단어로 초기화하는 경우에 왜곡이 나타난다. 반면 ViCo는 visual condition 덕분에 specific한 초기값 없이도 잘 작동하는 것을 학인할 수 있다.

아래는 ViCo를 사용가능한 작업들의 예시를 보여준다.

---

4. Limitations

Diffusion model을 frozen하는 방법은 diffusion model을 fine-tuning하는 방법에 비해 낮은 성능을 초래할 수 있다. 또한 mask binarization을 위한 otsu thresholding 과정은 sampling 과정을 지연시킨다. 단, 시간 지연에 있어서는 짧은 training과 전처리가 필요하지 않다는 장점으로 보완이 가능하다.

---

5. Summary

본 논문은 image conditioning을 통해 object fidelity를 높이는 기술인 ViCo를 제안한다. ViCo는 image cross-attention module을 통해 모델에 image condition을 주입하고, attention map으로부터 object mask를 추출한다. Masking을 통해 target object로부터 배경을 분리하여 과적합을 방지할 수 있다. ViCo는 모델을 fine-tuning하는 과정이 없어 training 시간이 짧고 pretrained된 모델에 쉽게 적용이 가능하다. 단, 모델을 fine-tuning하지 않기 때문에 성능에서 손해를 볼 수 있고, inference 시간이 길다는 한계가 존재한다.