UniControl (NeurIPS 2023, Diffusion)

 

 

Abstract

• SD 모델에서 spatial, structural, geometric control에 대한 image를 생성하는데 종종 실패
• 광범위한 control에 대해서 C2I(Control-to-Image) task를 수행하는 foundation model인 UniControl 제안
• T2I diffusion model을 augment하고 task-aware한 HyperNet을 사용
• 20M ITC triplet을 포함하는 MultiGen-20M dataset 제안

 

 

Motivation

diffusion model에서 condition을 넣어준다는 아이디어는 SD에서 시작해서 ControlNet까지 발전되었다. 하지만 ControlNet에서도, 특정한 control modality에 대해서 train된 경우에만 사용이 가능하고, 일반적인 control은 사용 가능하지 않다. 

 

UniControl에서는 language와 visual condition을 동시에 사용할 수 있다. UniControl은 MOE(Mixture of Expert)-style adapter와 task-aware HyperNet을 사용한다. MOE-style adapter는 visual condition의 feature를 추출하고 HyperNet은 natural language로 작성된 task instruction을 task-aware embedding으로 encode한다. 

 

Table 1: Architecture and Model Size: UniControl vs. Multi-ControlNet.

 

Tab. 1을 보면 이러한 unified structure를 사용하므로 task마다 다른 ControlNet을 사용하지 않아도 된다는 것을 알 수 있다. 

 

UniControl의 pretrain에는 5개의 categories의 9개의 task에 대해서 pretrain했다. 이를 위해 MultiGen-20M dataset을 construct해서 사용했다. 이는 다음 categories로 구성된다:

 

1. edges (Canny, HED, User Sketch)

2. region-wise maps (Segmentation Maps, Bounding Boxes)

3. skeletons (Human Pose Skeletons)

4. geometric maps (Depth, Surface Normal)

5. editing (Image Outpainting)

 

→ 이것만으로 모든 zero-shot condition에 대해서 잘 수행할 수 있을지는 의문. 근데 사실 condition이라고 하는게 완전히 무궁무진한 형태는 아니라서, 어느 정도만 수행된다면 실제로 사용하는 데에는 무리가 없을 듯 

 

UniControl

Training Setup

UniControl은 다른 모델과 다르게 넓은 범위의 visual condition과 language prompt를 입력으로 받는다. 이를 위해 formulation하면, $K$개의 task와 그에 따른 dataset $\mathcal {D}$가 있다고 하자:

 

$$ \mathcal{D} := \{ \mathcal{D}_1 \cup \cdots \cup \mathcal{D}_K \} $$  

 

각각의 training pair는 다음과 같이 구성된다:

 

$$\mathcal {D_i} = ([c_\text{text}, c_\text{task}], \mathcal{I}_c, \mathbb x)$$

 

이때 $c_\text{text}$는 text prompt, $c_\text{task}$는 task instruction, $\mathcal {I}_c$는 visual condition, $\mathbb x$는 target image이다. 이에 대한 training sample pair의 example은 다음과 같다:

 

Figure 2: Training image pair example.

 

→ iDeA; 만약 만화가가 데생 작업을 위해 사용한다면 그림체와 인물 묘사가 consistent하게 사용될 수 있을까? 

 

training loss는 다음과 같이 정의된다:

 

$$ l^k(θ) := \mathbb E _{z, ε,t, c_\text{task}, c_\text{text}, \mathcal{I}_c} \left [ \| ε-ε_θ(z_t, t, c_\text{task}, c_\text{text}, \mathcal{I}_c ) \|^2_2\right] \text{, with } ([c_\text{task}, c_\text{text}], \mathcal{I}_c, \mathbb x) \sim \mathcal {D}_k$$

 

$t$는 time step, $z_t$는 noise-corrupted latent tensor at time step $t$, $z_0\ = mathbb E(\mathbb x)$, $θ$는 trainable parameters of UniControl이다. 

 

classifier-free guidance에 30%의 text prompt를 random drop해서 input visual condition의 controllability를 향상한다. $K$개의 task에 대해서 $k$개의 task를 randomly select하고, $\mathcal {D}_k$에서 sample한 뒤 loss를 계산하여 optimize한다. 

 

Model Design

모델은 1) different task의 low-level feature 간 misalignment를 overcome할 수 있어야 하고, 2) 다른 task 사이에서 meta-knowledge를 배워 adapt할 수 있어야 한다.

 

→ meta-knowledge가 정확하게 무엇인지 간단한 문장으로 설명하기: 각 정보가 어떤 정보를 의미하는지에 대한 정보

 

먼저 첫 번째 property는 모든 task에서 필요한 정보를 UniControl이 학습할 수 있음을 의미한다. 예컨대 segmentation map이 visual condition으로 주어질 경우 $\mathcal I_c$에서 3d information는 ignore해야 할 것이다. 

 

두 번째 property는 model이 task 간 shared knowledge와 difference를 학습할 수 있어야 함을 의미한다.

Figure 3: UniControl Model Overview.

 

• MOE-Style Adapter

convolution module을 활용해서 low-level visual condition의 feature를 extract한다. 이는 다음과 같이 표현된다:

 

$$ \mathcal F_\text{Adapter}(\mathcal{I}_c^k ):= \sum^k_{i=1} \mathbb 1 (i==k) \cdot \mathcal {F} ^{(i)} _\text{Cov1} \circ \mathcal {F} ^{(i)} _\text{Cov2} (\mathcal I _c^k)$$

 

$\mathbb 1$는 indicator이고, $\mathcal {I}^k_c$는 task $k$에서 conditioned image이고, $\mathcal {F}^{(i)}_\text{Cov1}$와 $\mathcal {F} ^{(i)} _\text{Cov2}$는 $i$번째 adapter module의 convolution layer이다.

 

 

original MOE의 weight를 제거하고, 

 

 

 

• Task-Aware HyperNet

먼저 task instruction $c_\text{task}$를 CLIPText encoder로 projection한다. 그 후 ControlNet의 trainable copy에 task embedding을 zero-conv layer에 multiply해서 inject한다. 이는 ControlNet과 비슷한 구조이다:

→ 추가 필요 ControlNet 논문 본 이후 

 

 

Task Generation Ability

 

 

 

 

Experiments

 

Results

 

Discussion

 

 

---

References

 

Footnotes