[논문리뷰] Object-Centric Slot Diffusion (Latent Slot Diffusion, LSD)

NeurIPS 2023 (Spotlight). [Paper] [Page]
Jindong Jiang, Fei Deng, Gautam Singh, Sungjin Ahn
Rutgers University | KAIST
20 Mar 2023

Introduction

물리적 세계의 기본 구조는 합성적(compositional)이고 모듈식(modular)이다. 이 구조는 토큰이나 단어 형태의 언어와 같은 일부 데이터 modality에서 자연스럽게 드러나지만 이미지와 같은 다른 modality에서는 이 구조를 발견할 수 있는 방법이 어렵다. 그러나 이러한 합성성과 모듈성은 높은 수준의 인지 능력을 달성하기 위해 지식 조각의 체계적인 조작이 필요한 다양한 애플리케이션에 필수적이다. 여기에는 추론, 인과 추론 및 분포 외 일반화가 포함된다.

Object-centric learning은 관련 feature들을 바인딩하는 방법을 학습하여 구조화되지 않은 관찰에서 잠재적 합성 구조를 발견하여 unsupervised 방식으로 유용한 토큰을 형성하는 것을 목표로 한다. 이미지의 경우 가장 널리 사용되는 접근 방식 중 하나는 Slot Attention 인코더를 통해 이미지를 오토인코딩하는 것이다. Slot Attention은 competitive spatial attention을 적용하여 이미지를 별도의 로컬 영역으로 분할한 다음 각 영역에서 슬롯이라고 하는 표현을 얻는다. 그런 다음 디코더는 재구성 오차를 최소화하기 위해 슬롯에서 이미지를 생성한다. 제한된 용량과 슬롯 간의 경쟁으로 인해 각 슬롯은 개체와 같은 재사용 가능하고 합성적인 엔터티를 캡처하도록 권장된다.

Unsupervised object-centric learning의 프레임워크에 남아 있는 주요 과제는 복잡한 자연 이미지에 대해 작동하도록 만드는 것이다. 최근까지 대부분의 object-centric 모델은 혼합 디코더로 알려진 특별한 유형의 디코더를 채택했다. 상대적으로 단순한 장면 이미지에서 object-centric 표현의 출현을 촉진하는 효능으로 인해 이 혼합 디코더에서 약한 슬롯별 디코더가 주로 사용되지만, 추가 연구에 따르면 이 강력한 우선순위가 복잡한 자연주의적 장면 이미지를 더 어렵게 처리할 수 있음이 밝혀졌다. 기존의 믿음과는 달리 최근 이 저용량 혼합 디코더 접근 방식에서 벗어나 object-centric learning에서 표현적인 transformer 기반 autoregressive 이미지 generator를 사용할 것이 제안되었다. 디코더 용량을 늘리는 것이 이 프레임워크에서 복잡하고 자연스러운 장면을 처리하는 데 중요하다는 것이 나타났다.

Object-centric learning에서 transformer 기반 이미지 생성 모델링의 성공은 자연스럽게 질문으로 이어진다.

표현력이 뛰어난 생성 능력으로 알려진 현대 심층 생성 모델링의 또 다른 기둥인 diffusion model도 object-centric learning에 도움이 될 수 있는가?

Diffusion model은 denoising process를 기반으로 하며 다양한 이미지 생성 task에서 인상적인 성능을 입증했으며 때로는 transformer 기반 autoregressive model을 능가한다. 또한 diffusion model은 transformer 기반 autoregressive model이 제공할 수 없는 고유한 모델링 능력을 가지고 있다. 그러나 잠재성에도 불구하고 unsupervised object-centric learning에 대한 diffusion model의 적용 가능성은 대체로 탐구되지 않은 상태로 남아 있다. 결과적으로 이 접근법의 타당성을 검토하고 관련 이점과 한계를 식별하는 것이 중요하다.

본 논문에서는 Latent Slot Diffusion (LSD)라는 새로운 모델을 도입하여 이 문제를 해결한다. LSD 모델은 두 가지 관점에서 해석할 수 있다. Object-centric learning의 관점에서 볼 때 LSD는 조건부 latent diffusion model로 기존 슬롯 디코더를 대체하는 첫 번째 모델로 볼 수 있으며 조건은 Slot Attention에서 제공하는 object-centric slot이다. Diffusion model에서 LSD는 최초의 unsupervised 합성적 조건부 diffusion model이다. 기존의 조건부 diffusion model은 합성적 생성을 위해 이미지에 대한 텍스트 설명과 같은 supervised 주석이 필요하지만, LSD는 unsupervised object-centric learning을 통해 이미지에서 추출된 시각적 개념 측면에서 이러한 합성적 설명을 합성할 수 있는 diffusion model이다.

Latent Slot Diffusion

1. Object-Centric Encoder

입력 이미지 $x \in \mathbb{R}^{H \times W \times C}$가 주어지면, object-centric encoder는 $x$를 $N$개의 벡터 또는 슬롯 $S \in \mathbb{R}^{N \times D}$의 컬렉션으로 분해하고 표현하려고 한다. 여기서 각 슬롯 $s_n \in \mathbb{R}^D$는 이미지의 합성 요소를 나타낸다. 이를 위해 SOTA object-centric learning에서도 사용되는 아키텍처인 Slot Attention을 채택한다.

Slot Attention에서 먼저 입력 이미지 $x$를 backbone network $f_\phi^\textrm{backbone}$을 통해 $M$개의 입력 feature 집합 \(E \in \mathbb{R}^{M \times D_\textrm{input}}\)로 인코딩한다. 즉, $E = f_\phi^\textrm{backbone}$이다. $f_\phi^\textrm{backbone}$는 최종 출력 feature map이 flatten되어 집합을 형성하는 CNN으로 구현된다. 다음으로, $E$의 feature는 $N$개의 공간 그룹으로 그룹화되고 각 그룹의 정보는 슬롯을 생성하기 위해 집계된다. 그룹화는 반복적인 슬롯 정제 절차를 통해 이루어진다. 정제 절차가 시작될 때 슬롯 $S$는 임의의 Gaussian noise로 채워진다. 그런 다음 $N$개의 슬롯이 query 역할을 하고 $M$개의 입력 feature가 key와 value 역할을 하는 입력 feature에 대한 competitive attention을 통해 정제된다. Query와 key는 내적을 거쳐 $N \times M$개의 attention 비율을 생성한다. 다음으로, 이러한 attention 비율에서 softmax 함수가 $N$ 축을 따라 적용되어 슬롯에 대한 각 입력 feature의 소프트 할당을 캡처하는 attention 가중치 $A$를 생성한다. 그런 다음 각 $n$에 대해 모든 입력 feature는 attention 가중치 $A_{n,1}, \cdots, A_{n, M}$으로 가중치를 합산하여 attention readout $u_n \in \mathbb{R}^D$를 생성한다.

\[\begin{equation} A = \underset{N}{\textrm{softmax}} \bigg( \frac{q (S) \cdot k (E)^\top}{\sqrt{D}} \bigg) \\ A_{n, m} = \frac{A_{n, m}}{\sum_{m=1}^M A_{n, m}} \\ u_n = \sum_{m=1}^M v (E_m) A_{n, m} \end{equation}\]

여기서 $q$, $k$, $v$는 슬롯과 입력 feature를 공통 차원 $D$에 매핑하는 linear projection이다. Readout $u_n$에 의해 캡처된 상향식 정보를 사용하여 슬롯은 RNN에 의해 업데이트된다.

\[\begin{equation} s_n = f_\phi^\textrm{RNN} (s_n, u_n) \end{equation}\]

실제로 competitive attention 및 RNN 업데이트는 여러 번 반복적으로 수행되며 마지막 iteration의 슬롯은 최종 슬롯 표현 $S$로 간주된다.

2. Latent Slot Diffusion Decoder

LSD 디코더의 설계는 diffusion 기반 생성 모델링의 최근 발전을 활용한다. 디코딩 접근 방식에 대한 개요는 아래 그림에 나와 있다.

Pre-Trained Image Auto-Encoder

LSD 디코더의 핵심 디자인 요소 중 하나는 사전 학습된 오토인코더(AE)로, 인코더 $f_\phi^\textrm{AE}$를 통해 이미지 $x$를 저차원 latent 표현 $z_0$에 매핑하는 방법을 제공한다. 이를 통해 LSD는 저차원 latent $z_0$를 중간 재구성 타겟으로 사용하여 고해상도 이미지를 재구성하는 계산 부담을 줄일 수 있다. 또한 AE 디코더 $g_\theta^\textrm{AE}$를 사용하여 latent $z_0$를 디코딩하여 이미지 콘텐츠와 충실도를 손상시키지 않고 나중에 원본 해상도 이미지를 얻을 수 있다.

\[\begin{equation} z_0 = f_\phi^\textrm{AE} (x), \quad \hat{x} = g_\theta^\textrm{AE} (z_0) \\ \textrm{where} \quad z_0 \in \mathbb{R}^{H_\textrm{AE}, W_\textrm{AE}, D_\textrm{AE}}, \quad \hat{x} \in \mathbb{R}^{H \times W \times C} \end{equation}\]

본 논문의 모델에서는 OpenImages에서 사전 학습된 오토인코더를 사용한다.

Slot-Conditioned Diffusion

LSD에서 diffusion 모델링을 활용하여 슬롯 $S$에 컨디셔닝된 이미지 latent $z_0$를 재구성한다. 이 모델링 접근 방식은 Latent Diffusion Model (LDM)의 text-to-image 생성과 같은 supervised 컨텍스트에서 주로 탐색되었다. 그러나 LDM과 달리 본 논문에서는 supervised 레이블의 임베딩에 디코더를 컨디셔닝하는 대신 슬롯 자체를 얻는 프로세스, 즉 Slot Attention이 supervision 없이 디코더와 공동으로 학습되는 슬롯에 컨디셔닝한다. LDM을 따라 디코더는 슬롯 $S$가 주어진 이미지 latent $z_0$의 log-likelihood $\log p_\theta (z_0 \vert S)$를 최대화하기 위해 디코딩 분포 $p_\theta (z_0 \vert S)$를 학습하여 작동한다. 이 디코딩 분포 $p_\theta (z_0 \vert S)$는 $T$-step denoising process의 모델링과 같다.

\[\begin{equation} p_\theta (z_0 \vert S) = \int p(z_T) \prod_{t = T,\cdots,1} p_\theta (z_{t-1} \vert z_t, t, S) dz_{1:T} \end{equation}\]

여기서 $p(z_T) = \mathcal{N} (0,I)$이고, $p_\theta (z_{t-1} \vert z_t, t, S)$는 one-step denoising 분포이며, $z_T, \cdots, z_0$는 점진적으로 denoise된 latent들의 시퀀스이다. $p_\theta (z_{t-1} \vert z_t, t, S)$는 신경망 $g_\theta^\textrm{LSD}$에 의해 다음과 같이 parameterize된다.

\[\begin{equation} p_\theta (z_{t-1} \vert z_t, t, S) = \mathcal{N} \bigg( \frac{1}{\sqrt{\alpha_t}} \bigg( z_t - \frac{\beta_t}{\sqrt{1 - \bar{\alpha}_t}} \hat{\epsilon}_t \bigg), \beta_t I \bigg), \\ \textrm{where} \quad \hat{\epsilon}_t = g_\theta^\textrm{LSD} (z_t, t, S), \quad \alpha_t = 1 - \beta_t, \quad \bar{\alpha}_t = \prod_{i=1}^t \alpha_i \end{equation}\]

Sampling Procedure. $z_0 \sim p_\theta (z_0 \vert S)$를 샘플링하기 위해 반복적인 denoising 절차를 채택한다. 샘플링 프로세스는 임의의 Gaussian noise로 채워진 latent 표현 $z_T \sim \mathcal{N}(0,I)$으로 시작한다. 다음으로, 슬롯을 조건으로 $t = T, \cdots, 1$에 대해 one-step denoising 분포 $z_{t−1} \sim p_\theta (z_{t-1} \vert z_t, t, S)$에서 순차적으로 샘플링하여 $T$번 denoise한다. 이는 latent의 시퀀스 $z_T, \cdots, z_0$의 시퀀스를 생성하며, 점진적으로 깨끗해진다. 마지막으로 $z_0$는 재구성된 latent 표현으로 간주될 수 있다.

Training Procedure. LDM을 따라 $p_\theta (z_0 \vert S)$의 학습은 $g_\theta^\textrm{LSD}$를 학습하기 위한 간단한 절차로 변환될 수 있다. 이미지 $x$, 슬롯 표현 $S$, 이미지 latent $z_0$가 주어지면 먼저 균일 분포에서 랜덤하게 noise 레벨 \(t \in \{1, \cdots, T\}\)를 뽑는다. 주어진 $t$에서 $z_0$를 손상시키고 다음과 같이 noised latent $z_t$를 얻는다.

\[\begin{equation} z_t = \sqrt{\vphantom{1} \bar{\alpha}_t} z_0 + \sqrt{1 - \bar{\alpha}_t} \epsilon_t, \\ \textrm{where} \quad \epsilon_t \sim \mathcal{N}(0,I), \quad \bar{\alpha}_t = \prod_{i=1}^t (1 - \beta_i) \end{equation}\]

그런 다음 noise $\epsilon_t$를 예측하기 위해 noised latent $z_t$를 슬롯 $S$와 timestep $t$와 함께 $g_\theta^\textrm{LSD}$에 대한 입력으로 제공한다. 네트워크 $g_\theta^\textrm{LSD}$는 예측된 noise \(\hat{\epsilon}_t\)와 실제 noise $\epsilon_t$ 사이의 평균 제곱 오차를 최소화하여 학습된다.

\[\begin{equation} \mathcal{L} (\phi, \theta) = \| \hat{\epsilon}_t - \epsilon_t \|^2, \\ \textrm{where} \quad \hat{\epsilon}_t = g_\theta^\textrm{LSD} (z_t, t, S) \end{equation}\]

Denoising Network

슬롯 컨디셔닝을 통합하도록 조정된 기존 UNet 아키텍처의 변형으로 denoising network $g_\theta^\textrm{LSD}$를 구현한다. Denoising network는 $L$개의 레이어의 스택으로 구성되며 각 레이어 $l$은 UNet 스타일 CNN 레이어와 슬롯으로 컨디셔닝된 transformer가 뒤따른다.

\[\begin{equation} \tilde{h}_l = \textrm{CNN}_\theta^l ([h_{l-1}, h_{\textrm{skip} (l)}], t) \\ h_l = \textrm{Transformer}_\theta^l(\tilde{h}_l + p_l, \textrm{cond} = S) \end{equation}\]

여기서 $h_0 = z_t$이고 $h_1, \cdots, h_{L-1}$은 hidden state이며 \(\hat{\epsilon}_t = h_L\)은 출력이다.

CNN Layers. UNet을 따라 convolution layer \(\textrm{CNN}_\theta^1, \cdots, \textrm{CNN}_\theta^L\)는 처음 $\frac{L}{2}$개의 레이어를 통해 feature map을 다운샘플링한 다음 나머지 레이어를 통해 원래 해상도로 다시 업샘플링한다. 표준 UNet에 따라 이러한 CNN 레이어는 $\textrm{skip}(l)$로 표시된 이전 레이어에서의 skip connection을 통해 입력을 받는다. 이 네트워크 디자인은 LDM에서도 탐색되었다.

Slot-Conditioned Transformer. Transformer의 역할은 슬롯의 정보를 UNet 기반 denoising process에 통합하는 것이다. 이를 위해 각 레이어 $l$에서 CNN 레이어에 의해 생성된 중간 feature map \(\tilde{h}_l\)은 feature의 집합으로 flatten된다. 여기에 위치 임베딩 $p_l$이 추가되고 결과 feature가 transformer에 대한 입력으로 제공된다. Transformer 내에서 이러한 feature는 서로 그리고 슬롯 $S$와 상호 작용하므로 슬롯의 정보를 denoising process에 통합한다. 그런 다음 transformer 출력은 feature map $h_l$로 다시 형성된다.

Compositional Image Synthesis

새로운 이미지를 합성하는 기존의 접근 방식은 일반적으로 supervise되고 텍스트 프롬프트에 의존한다. 여기에는 vocabulary의 단어를 사용하여 문장을 만든 다음 text-to-image 모델에 제공하여 원하는 이미지를 합성하는 task가 포함된다. 그러나 본 논문과 같이 완전한 unsupervised 환경에서는 먼저 레이블이 지정되지 않은 많은 이미지셋을 관찰하여 시각적 개념 라이브러리를 구축해야 한다. 그런 다음 단어를 사용하여 문장 프롬프트를 합성하는 것과 유사하게 시각적 개념 라이브러리에서 개념을 선택하여 개념 프롬프트를 합성한다. 이 개념 프롬프트를 LSD 디코더에 제공함으로써 원하는 새로운 이미지를 합성할 수 있다.

Unsupervised Visual Concept Library

레이블이 지정되지 않은 이미지에서 시각적 개념 라이브러리를 구축하기 위해 먼저 $B$개의 이미지의 대규모 배치 $x_1, \cdots, x_B$를 가져온다. 그런 다음 이러한 이미지에 대한 슬롯 $S_1, \cdots, S_B$을 얻기 위해 Slot Attention을 적용한다. 다음으로, 이 모든 슬롯을 단일 집합 $S$로 모으고 $K$-mean을 수행한다. $K$-mean 절차는 $S$의 각 슬롯을 $K$개의 클러스터 중 하나에 할당한다. $k$번째 클러스터에 할당된 슬롯 집합을 시각적 개념 라이브러리 \(\mathcal{V}_k\)로 간주한다. $K$를 클러스터 수라 하면 이 방법은 $K$개의 시각적 개념 라이브러리 \(\mathcal{V}_1, \cdots, \mathcal{V}_K\)가 된다.

Novel Image Synthesis

라이브러리 \(\mathcal{V}_1, \cdots, \mathcal{V}_K\)가 주어지면 $k$번째 라이브러리에서 각각 $K$개의 슬롯을 선택하고 함께 쌓아 개념 프롬프트 $S_\textrm{compose}$를 합성할 수 있다.

\[\begin{equation} S_\textrm{compose} (s_1, \cdots, s_K), \quad \textrm{where} \; s_k \sim \textrm{Uniform} (\mathcal{V}_k) \end{equation}\]

그런 다음 구성된 프롬프트 $S_\textrm{compose}$를 LSD 디코더에 제공하여 latent 이미지

\[\begin{equation} z_\textrm{compose} \sim p_\theta (z_0 \vert S_\textrm{compose}) \end{equation}\]

를 생성한다. $z_\textrm{compose}$에 이미지 디코더를 적용하면 원하는 새로운 장면 이미지

\[\begin{equation} x_\textrm{compose} = g_\theta^\textrm{AE} (z_\textrm{compose}) \end{equation}\]

가 생성된다. 예를 들어 FFHQ 데이터셋에서 개념 프롬프트는 선택한 머리 스타일, 얼굴, 옷, 배경의 모음일 수 있다. 이 프롬프트를 디코딩하면 이 프롬프트를 준수하는 얼굴 이미지가 생성된다.

Experiments

• 데이터셋: CLEVR, CLEVRTex, MOVi-C, MOVi-E, FFHQ

1. Object-Centric Representation Learning

다음은 segmentation 성능과 표현 품질을 비교한 표이다.

다음은 unsupervised segmentation을 시각화한 것이다.

2. Compositional Generation with Visual Concept Library

다음은 합성적 이미지 생성 결과를 비교한 표이다.

다음은 합성적 생성 샘플들이다.

3. Slot-Based Image Editing

다음은 슬롯 기반 이미지 편집 결과이다.