[평범한 학부생이 하는 논문 리뷰] Prompt-to-Prompt Image Editing with Cross-Attention Control (ICLR 2023)

Paper | Github | Project Page

Prompt-to-Prompt

1. Introduction

 기존의 large-scale language-image (LLI) 모델들은 image editing 능력이 없었다. Text prompt를 살짝 바꾸면 output이 크게 바뀌기 때문에 text prompt를 통한 editing을 할 수 없다. 그래서 이러한 모델들로 editing을 하기 위해서는 사용자가 masking을 하는 등의 방식을 이용했다. 하지만 이러한 방식은 여러 제약을 지닌다.

 본 논문에서는 직관적이고 강력한 textual editing method인 Prompt-to-Prompt를 도입한다. 본 논문의 key idea는 diffusion process동안 cross-attention map을 주입함으로써 이미지를 editing하는 것이다.

2. Method

 본 논문에서의 목표는 text prompt $\mathcal{P}$와 random seed $s$로 text-guided diffusion model이 image $\mathcal{I}$를 생성했을 때, 수정된 prompt $\mathcal{P}^*$만 사용해서 수정된 prompt에 대응되는 부분만 제외하고 original image의 content와 structure를 유지하는 editied image $\mathcal{I}^*$를 얻는 것이다.

위 이미지에서 두번째 줄의 이미지들을 보면 random seed를 고정하고 prompt만 바꿔서 이미지를 생성하면 구조나 content를 유지하고 이미지 editing하는 결과를 얻는데 실패한다.

 본 논문의 key observation은 생성된 structure과 appearance는 random seed뿐만 아니라, diffusion process동안에 text embedding에 대한 pixel들 사이의 interaction에 의존한다는 것이다. Cross-attention layer에서 발생하는 pixel-to-text interaction을 수정함으로써, 본 논문은 Prompt-to-Prompt image editing 능력을 제공한다. 특히 input image $\mathcal{I}$의 cross-attention maps을 주입하는 것은 원래의 composition과 structure를 유지하도록 한다.

2.1 Cross-Attention in text-conditioned Diffusion Models

 본 논문에서는 Imagen을 backbone으로 사용하는데, 이 때 $64 \times 64$ 해상도에서 composition과 geometry가 결정되기 때문에 super-resolution은 그대로 사용하고 text-to-image diffusion model만 조정한다. Diffusion model에서 두 모달리티의 interaction은 noise prediction에서 발생하는데, 여기서는 각 textual token에 대한 spatial attention maps를 생성하는 Cross-attention map을 사용해 visual과 textual featur들의 embedding이 융합된다. 

 위 이미지처럼 noisy image의 deep spatial features $\phi(z_t)$는 query matrix $Q = l_Q(\phi(z_t))$에 projection되고 textual embedding은 key matrix $K = l_K(\psi(\mathcal{P}))$과 value matrix $V = l_V(\psi(\mathcal{P}))$에 projection된다. 여기서 $l_Q, l_K, l_V$는 learnable linear projection이다. Attention map은 다음과 같아진다.

여기서 $M_ij$는 $i$번째 픽셀에 대한 $j$번째 token의 값에 대한 weight를 나타낸다. 결국 cross-attention output은 $\hat{\phi}(z_t) = MV$로 정의되고, 이는 spatial features $\phi(z_t)$를 업데이트하는데 사용된다.

 직관적으로 cross-attention의 output $MV$는 value $V$의 weighted average인데 여기서 weght가 $Q$와 $V$사이의 유사도에 연관된 attention maps $M$인 것이다. 표현력을 높이기 위해 multi-head attention을 사용해서 결과들을 concat하고 학습된 linear layer를 통과해 최종 결과를 낸다.

2.2 Controlling the Cross-attention

생성된 이미지의 spatial layout과 geometry는 cross-attention maps에 의존한다. 위 이미지에서 볼 수 있듯이 픽셀들은 그것들을 묘사하는 단어에 더 attract된다. 흥미롭게도, 이미지의 구조는 이미 diffusion의 초기 process에서 결정된다는 것을 볼 수 있다.

Attention은 전체적인 composition을 반영하기 때문에, original prompt $\mathcal{P}$로부터 얻은 attention map $M$을 수정된 prompt $\mathcal{P}^*$를 사용한 두번째 생성에 주입할 수 있다. 이를 통해 input image $\mathcal{I}$의 구조를 보존할 뿐만 아니라 edited prompt에 따라 조정된 edited image $\mathcal{I}^*$의 합성을 가능하게 한다. 

 $DM(z_t, \mathcal{P}, t, s)$를 diffusion process의 $t$번째 스텝이라고 하고, 그에 대한 output을 $z_{t-1}, M_t$라고 하자. 이에 따라 $DM(z_t, \mathcal{P}, t, s) \{ M \leftarrow \hat{M} \}$를 제공된 prompt로부터 values $V$를 유지하면서 attention map $M$을 $\hat{M}$으로 덮어쓰는 것으로 표현할 수 있다. 또 수정된 prompt $\mathcal{P}^*$를 사용해 생성한 attention map을 $M^*_t$로 나타낸다. 마지막으로 general edit function을 $Edit(M_t, M^*_t, t)$로 나타낸다.

 이를 가지고 전체적인 알고리즘은 다음과 같이 나타낼 수 있다.

Editing function의 예시는 다음과 같다.

Word Swap

 이 경우에는 original prompt에서 토큰들을 다른 것들로 바꾼다.("a big red bicycle" $\rightarrow$ "a big red car") Main challenge는 새로운 prompt의 content를 다루면서 original composition은 보존하는 것이다. 이를 해결하기 위해, source image의 attention map을 수정된 prompt로 진행하는 generation에 주입힌다. 그러나 제안된 attention injection은 geometry를 과도하게 제한하기에 다음과 같은 softer attention constrain을 제안한다.

여기서 $\tau$는 어떤 step까지 injection이 적용될지 정하는 파라미터이다. 

Adding a New Phrase

 이 경우에는 prompt에 새로운 토큰을 추가한다. ("a castle next to a river" $\rightarrow$ "children drawing of a castle next to river") 공통의 detail들을 보존하기 위해, 두 prompt로부터의 공동 토큰들에 대해서만 attention injection을 적용한다. 본 논문에서는 target prompt $\mathcal{P}$로부터의 token index을 받아 $\mathcal{P}$에서 상응하는 token index나 $None$을 뱉는 alignment function $A$를 사용한다.

여기서도 동일하게 $\tau$가 적용된다. Index $i$는 pixel value에 대응되고, $j$는 text token에 대응된다. Original prompt에 없는 index에 해당하는 attention map은 target prompt에서 생성된 것을 사용하고 있으면 original prompt를 통해 생성된 attention map을 사용한다.

Attention Re-weighting

 이 경우에는 결과 이미지에 영향을 끼치는 토큰의 강도를 강화하거나 약화하는 경우이다. $j^*$ token의 attention map의 scale을 $c \in [-2,2]$로 scaling해서 효과를 강화 혹은 약화시킨다. 다른 attention map은 변화하지 않는다.

3. Applications

Text-Only Localized Editing

위 결과는 본 논문의 결과와 random seed만 고정했을 때의 결과를 나타낸다.

위 결과는 "butterfly"의 attention map만 주입했을 때의 결과이다.

위 실험은 attention injection하는 diffusion step 수를 늘려서 나온 결과이다. 더 많은 수에 injection을 할 수록 original image에 대한 fidelity가 높아진다. 하지만 꼭 많은 step 수에 injection하는 것이 좋은 결과를 내는 것은 아니라고 논문에서 설명한다.

위 실험은 add a new phrase에 대한 실험 결과이다. 

Global editing

Fader Control using Attention Re-weighting

Real Image Editing

 DDIM inversion을 이용해서 real image editing이 가능하다.

하지만 다음과 같이 reconstruction이 잘 안되는 경우가 있다.

이러한 문제를 완화하기 위해 본 논문에서는 attention map으로부터 추출한 mask를 이용해 original image의 unedited region을 복원하는 것을 제안한다. 여기서 mask는 user가 직접 주는 것이 아니라는 것에 주목해야한다.

4. Conclusions

 본 논문은 text-to-image diffusion model에서 cross-attention layers의 능력을 발견했다. 본 논문은 이러한 고차원의 layer들이 합성된 이미지의 spatial layout을 text prompt안의 word들을 묶는 것에서 key role을 수행하는 해석 가능한 spatial maps의 representation을 가진다는 것을 보였다.

 본 논문의 한계점은 다음과 같다.

1. Inversion process는 몇 test images에 대해 visible distortion을 보인다.
2. Cross-attention은 network의 bottleneck에 위치하기 때문에 최근 attention maps는 low resolution이다. 이는 더 정확한  localized editing을 수행하는 능력을 제한한다.
3. 본 논문의 방식은 이미지 내에서 존재하는 object를 공간적으로 움직이는데 사용할 수 없다.