[평범한 학부생이 하는 논문 리뷰] EmerDiff: Emerging Pixel-level Semantic Knowledge in Diffusion Models (ICLR 2024)

Paper : https://arxiv.org/abs/2401.11739

EmerDiff: Emerging Pixel-level Semantic Knowledge in Diffusion Models

0. Abstract

 Diffusion model을 이용해서 fine-grained segmentation mask를 생성하는 것은 annotated dataset으로 추가적인 학습을 필요로 하고, pre-trained diffusion model로만 생성된 이미지의 semantic relation을 이해하는 정도가 불명확하다. 이 문제를 해결하기 위해서, 본 논문은 Stable Diffusion으로부터 추출된 semantic knowledge를 활용하고 추가적인 training없이 fine-grained segmentation map 생성이 가능한 image segmentor를 개발하는 것이 목표이다. Semantically meaningful feature maps는 spatially lower-dimensional layer에 존재하고 이 때문에 이러한 feature maps에서 pixel-level semantic relation을 직접 추출하는 것이 challenge이다. 이 문제를 극복하기 위해서, 본 논문의 framework는 SD의 generation process를 활용해서 image pixel과 low-dimensional feature map의 spatial location사이의 semantic correspondence를 인지하고 이를 활용해서 image-resolution segmentation map을 구축한다.

1. Introduction

 diffusion model로 semantic relation을 capture하기 위한 공통의 unsupervised approach는 semantically meaningful feature map에 k-means를 적용하는 것이다. 이 방식은 diffusion model의 semantic awareness를 직관적으로 시각화하지만, semantically meaningful feature map은 spatially low-dimensional layer에 존재한다는 사실 때문에 얻어진 semantic relation은 종종 coarse하다. 그럼에도 불구하고, diffusion model은 low-resolution feature map에서 embedding된 semantic knowledge에 기초해서 high-resolution image를 구성하는 능력을 가진다. 이러한 점에서 영감을 받아, 본 논문은 generation process를 통해서 semantically meaningful low-resolution feature map이 output image에 어떻게 영향을 끼치는 지 분석하고, 이를 통해 본 논문은 diffusion model로부터 pixel-level semantic knowledge를 추출하는 중요한 step이 존재한다고 가정한다.

 본 논문의 key discovery는 low-resolution feature map의 sub-region의 값을 perturb시킬 때, 생성된 이미지는 sub-region에 semantically하게 관계된 pixel들만 바뀌는 방식으로 대체된다는 것이다. 그 결과, pixel level에서 변화를 측정함으로써 image pixel과 low-dimensional feature maps의 sub-region사이의 semantic correspondence를 식별할 수 있다.

 이러한 insight를 통해, 본 논문은 어떠한 추가적인 knowledge에 대한 필요 없이 fine-grained segementation maps를 생성하는 image segmentor를 제안한다.

2. Methods

 본 논문의 목표는 pre-trained diffusion model로부터 뽑은 semantic knowledge만을 사용해서 fine-grained segmentation map을 생성하는 것이다.

2.1 Preliminaries

 본 논문에서는 Stable Diffusion(SD)를 사용한다. SD의 U-Net은 residual block, self-attention layer, cross-attention layer로 구성된 modular block들이 쌓여있는 구조이다. 이 modular block은 4가지의 spatial resolution level에 속하는데, 본 논문은 upward path에서 $16 \times 16$ 해상도에 속하는 modular block에 집중한다. 이 block은 특히 semantically enriched representation을 포함한다. 특히, upward path에서 $16 \times 16$에 세 개의 연속된 modular block이 있는데, 첫 번째 cross-attention은 contents를 제작하는데 주요하게 작용하고 마지막 cross-attention layer는 appearance를 더 control하려 한다는 것이 발견됐다. 본 논문에서는 이러한 semantically meaningful layer를 사용해서 fine-grained segmentation map을 생성한다.

2.2 Constructing Low-Resolution Segmentation Maps

1. DDPM-based inversion을 통해 real image를 특정 수의 denoising step으로 invert한다.
2. Timestep $t_f$에서 upward $16 \times 16$ modular blocks의 첫 번째 cross-attention layer로부터 query vector(low-dimensional feature map)을 뽑는다.
3. 뽑은 feature map에 대해 k-means를 적용해서 low-resolution segmentation masks 역할을 하는 $K$개의 cluster를 얻는다.

2.3 Building Image-Resolution Segmentation Maps

 본 논문의 다음 목표는 image pixel들과 low-resolution mask들 사이의 semantic correspondence들을 식별함으로써 low-resolution segmentation map들로부터 image-resolution segmentation maps를 얻는 것이다. 이를 위해, 본 논문은 먼저 SD 안의 low-dimensional layers가 생성된 이미지의 pixel들에 어떻게 영향을 끼치는지 관찰한다. 본 논문은 $16 \times 16$ cross-attention layers에서 feature maps의 sub-region의 값들을 조정하고 이 local change가 생성된 이미지들의 pixel값에 어떻게 영향을 끼치는지 관찰한다.

 Cross-attention layer를 조정하기 위해서, 본 논문은 cross-attention layer에서의 attention 계산을 다음과 같이 대체한다.

$$f(\sigma(\frac{QK^T}{\sqrt{d}})\cdot V) + cM \in \mathbb{R}^{hw \times d}$$

여기서 $c$는 스칼라 값이며 $M \in \{ 0,1 \}^{hw \times 1}$은 spatial locations를 특정화하는 binary mask이다. 즉, feature maps의 sub-region에 constant offset $c$를 균등하게 더해준다. Modulated denoising process동안, timestep $t_m$에 특정 cross-attention layer에 이 modulation을 적용한다.

 변화를 관찰하기 위해서, $c = +\lambda, -\lambda$로 modulated denoising process를 진행해, 두 개의 altered images $I^-, I^+ \in \mathbb{R}^{H \times W \times 3}$를 얻는다. 그리고 나서 difference maps $d = \Vert I^- - I^+ \Vert_2 \in \mathbb{R}^{H \times W}$를 계산한다. Figure 3에서 볼 수 있듯이, modulated sub-region과 semantically하게 관련된 pixel들은 분명히 변하는 반면, 다른 부분들은 거의 그대로 남는다. 그러므로 얻어진 difference map은 image pixel들과 low-resoluion feature map의 sub-region 사이의 semantic correspondence의 강도(strength)로 해석될 수 있다.

 이를 통해서, image pixel들과 low-resolution segmentation mask들의 모든 페어들 사이의 semantic correnpondence의 강도를 계산할 수 있다. 각 low-resolution mask $M^i \in \{ 0,1 \} ^{hw \times 1}$에 대해, difference map $d^i \in \mathbb{R}^{H \times W}$를 구한다. 이 때, $d^i_{x,y} \in \mathbb{R}$는 pixel $(x,y)$와 mask $i$사이의 semantic correspondence의 강도를 나타낸다. Image-resolution segmentation maps을 만들기 위해, 본 논문은 각 픽셀 $(x,y)$를 가장 강력한 semantic correspondence를 가진 low-resolution mask $k$로 라벨링한다.( i.e., $k = argmax_i d^i_{x,y}$)

 마지막으로 pixelated artifacts를 막기위해서 Gaussian filtering을 적용한다.

3. Experiments

3.1 Implementation Details

 Low-resolution segmentation maps를 생성하기 위해, 본 논문은 timestep $t_f = 1$(mininum noise)에서 feature maps를 추출한다. 또한, timestep $t_m = 281$과 $\lambda = 10$으로 $16 \times 16$ upward blocks의 세 번째 cross-attention layer에 modulation을 적용한다.

3.2 Qualitative Analysis

3.3 Quantitative Analysis

4. Limitation and Conclusion

Hyperparameter Analysis

Number of masks per image

Selection of cross-attention layer to extract feature maps

Timestep for extracting feature maps

Modulation timestep and strength

Selection of cross-attention layer to modulate

Attention injection