[논문리뷰] Prompt-to-Prompt Image Editing with Cross Attention Control

Prompt-to-Prompt Image Editing with Cross Attention Control
Amir Hertz, Ron Mokady, Jay Tenenbaum, Kfir Aberman, Yael Pritch, Daniel Cohen-Or
[paper] [github]

---

Contents

1. Introduction

2. Method

3. Applications

4. Limitations

5. Summary

---

1. Introduction

최근 Text-To-Image synthesis 분야가 많은 관심을 받고 있다. 대용량의 Text-Image datasets으로 학습한 diffusion models은 이미지 합성 분야에서 엄청난 성능을 보이고 있다. 반면에 합성한 이미지의 일부를 원하는 대로 수정하는 image editing 분야에서는 아직 많은 어려움이 있다. Text prompt를 일부만 변경해도 아예 다른 이미지를 생성하기 때문에 이미지 편집은 더욱 쉽지 않다.

 

기존 연구는 편집할 부분을 명시적으로 나타내는 mask를 입력받는 방식으로 문제를 해결했다. 이러한 방법은 괜찮은 품질을 보여주지만 mask를 만드는 과정이 번거롭고 느리다는 문제가 있다. 게다가 masking을 이용한 방법은 이미지의 공간 정보를 잃게 만들어 이미지의 품질을 손상시킨다.

 

저자들은 이러한 과정없이 prompt만으로 이미지를 편집하는 textual editing 기술인 Prompt-to-Prompt를 제안한다. 아래 사진과 같이 추가적인 입력 없이 prompt의 수정만으로 기존 이미지의 요소를 변경하거나, 스타일을 바꾸거나, 단어의 의미를 강화하는 작업을 할 수 있다.

---

2. Method

Random seed $s$와 prompt $\mathcal{P}$로 text-guided generation된 이미지를 $\mathcal{I}$라고 하자. Textual editing의 목적은 원본 이미지 $\mathcal{I}$와 수정된 prompt $\mathcal{P}^{*}$를 이용하여 편집된 이미지 $\mathcal{I}^{*}$를 생성하는 것이다. 예를 들어 prompt "my new bicycle"로 생성된 이미지를 생각해 보자. 우리가 원하는 것은 자전거의 색이나 재질을 변경하거나, 자전거 대신 오토바이를 그리도록 하는 것이다. 가장 먼저 떠오르는 방법은 단순히 prompt를 변경하는 것이다. Random seed를 고정하고 regeneration한 결과는 다음과 같다.

기존 이미지인 레몬 케이크의 형태는 유지하면서 사용한 재료를 수정하고 싶었지만, 모델은 모양과 크기가 다른 결과들을 생성한다. 이로부터 관찰할 수 있는 것은 이미지의 구조나 형태가 random seed에만 의존하지 않고, 각 픽셀과 텍스트 임베딩 사이의 상호작용에 영향을 받는다는 것이다. 저자들은 cross-attention layer에서 pixel-to-text 상호작용을 수정하는 방식으로 Prompt-to-prompt 기술을 구현한다.

 

저자들은 논문의 backbone으로 Imagen 모델을 사용하였다. 각 denoising step $t$에서 U-net 기반의 모델은 노이즈가 낀 이미지 $z_t$와 텍스트 임베딩 $\psi(\mathcal{P})$으로부터 노이즈 $\epsilon$을 예측한다. 이러한 과정을 반복적으로 거쳐 노이즈 $z_T$로부터 이미지 $\mathcal{I}$를 얻는다. U-net 내부에서 이미지와 텍스트 임베딩은 cross-attention layers를 통해 결합하여 각 텍스트 토큰에 대해 spatial attention map을 생성한다. 

이미지의 features $\phi\left(z_t\right)$는 쿼리 행렬 $Q=\ell_Q\left(\phi\left(z_t\right)\right)$로 표현된다. 마찬가지로 텍스트 임베딩은 키 행렬 $K=\ell_K(\psi(\mathcal{P}))$과 값 행렬 $V=\ell_V(\psi(\mathcal{P}))$로 표현된다. Attention maps은 다음과 같이 계산된다.

$$M=\operatorname{Softmax}\left(\frac{Q K^T}{\sqrt{d}}\right)$$

$M_{ij}$는 픽셀 $i$와 토큰 $j$ 사이의 가중치를 나타낸다. 따라서 cross-attention의 출력 $MV$는 attention maps $M$이 가중치인 $V$의 weighted average임을 알 수 있다.

본 논문의 baseline인 Imagen은 두 종류의 attention layers를 사용한다. 하나는 기본적인 cross-attention layers이고, 다른 하나는 self-attention과 cross-attention의 key-value를 결합하여 사용하는 hybrid-attention layers이다. 이 경우에는 hybrid-attention의 cross-attention 부분에만 기술을 적용한다.

 

위 사진에서 attention maps의 가중치는 각 픽셀을 설명하는 단어의 token에 집중되어 있다는 것을 확인할 수 있다. 우리는 기존 prompt $\mathcal{P}$로부터 얻은 attention maps을 새로운 prompt $\mathcal{P}^{*}$에 주입할 수 있다. 이는 편집된 이미지 $\mathcal{I}^*$가 기존 이미지 $\mathcal{I}$의 구조를 보존할 수 있도록 한다. 위 사진에서 흥미로운 점은 초기 denoising steps에서 이미지의 구조가 대부분 결정된다는 것이다.

 

이러한 과정을 다음과 같이 일반화할 수 있다. $D M\left(z_t, \mathcal{P}, t, s\right)$를 single denoising step $t$라고 하자. 각 step에서 less noisy image $z_{t-1}$과 attention map $M_t$를 얻을 수 있다. Attention map $M$을 새로운 map $\widehat{M}$으로 덮어씌운 denoising step을 $D M\left(z_t, \mathcal{P}, t, s\right)\{M \leftarrow \widehat{M}\}$이라고 하자. 이때 key matrix $V$는 그대로 사용한다. 수정된 prompt $\mathcal{P}^{*}$로부터 얻은 attention maps은 $M_{t}^{*}$라 하자. 이제 edit function을 $\operatorname{Edit}\left(M_t, M_t^*, t\right)$로 일반화할 수 있다.

 

Prompt-to-Prompt는 두 개의 prompt에 대해 denoising step을 동시에 계산하며 진행된다. 이때 두 모델의 무작위성을 결정하는 random seed $s$는 고정한다. 아래의 pseudocode로 알고리즘을 설명할 수 있다.

기존 denoising step에 비해 계산이 많아 보이지만, 두 이미지를 병렬로 계산하고, Edit 함수를 두 번째 denoising step 내부에서 바로 적용함으로써 계산을 줄일 수 있다.

 

이제 각 작업에 대해 Edit 함수의 기능을 정의할 것이다.

 

단어 교체는 prompt에서 특정 토큰을 교체하는 작업이다. 예를 들어 prompt를 "a big red bicycle"에서 "a big red car"로 바꾸는 작업이 해당된다. 해결해야 할 문제는 기존 이미지의 구성을 보존하면서 새로운 subject를 도입하는 것이다. 이를 위해 기존 이미지의 attention maps을 주입하는 방법을 사용한다. 단, attention maps을 전부 주입하는 것은 자동차를 자전거로 바꾸는 작업과 같이 큰 수정이 필요한 경우의 표현을 제한할 수 있다. 따라서 다음과 같은 방식을 사용한다.

$$\operatorname{Edit}\left(M_t, M_t^*, t\right):= \begin{cases}M_t^* & \text { if } t<\tau \\ M_t & \text { otherwise}\end{cases}$$

초기 denoising steps에서 이미지의 구조가 결정되기 때문에, timestamp parameter $tau$를 도입하여 주입 단계를 제한한다. 이는 기존 이미지의 구조를 보존하면서 새로운 subject를 표현할 수 있도록 한다.

 

구문 추가는 prompt에 새로운 토큰을 추가하는 작업이다. 예를 들어 prompt "a castle next to a river"를 "children drawing of a castle next to a river"로 확장하는 작업이 해당된다. 기존의 정보를 보존하기 위해 두 프롬프트에서 모두 나타나는 토큰에 대해서만 attention maps을 주입한다. 수정된 prompt $\mathcal{P}^{*}$로부터 토큰 인덱스를 받아 prompt $\mathcal{P}$에서 대응하는 토큰 인덱스를 출력하는 함수 $A$를 사용한다. 따라서 아래 식과 같이 표현할 수 있다.$$\left(\operatorname{Edit}\left(M_t, M_t^*, t\right)\right)_{i, j}:= \begin{cases}\left(M_t^*\right)_{i, j} & \text { if } A(j)=\text { None } \\ \left(M_t\right)_{i, A(j)} & \text { otherwise}\end{cases}$$

단어 교체와 마찬가지로 timestamp $\tau$를 도입하여 주입 단계를 제한할 수 있다.

 

가중치 조정은 prompt에서 특정 토큰의 의미를 강화하거나 약화시키는 작업이다. 예를 들어 prompt "a fluffy red ball"에서 공이 fluffly한 정도를 조절하는 작업이 해당된다. 이를 위해 특정 토큰 $j^{*}$에 해당하는 attention map을 파라미터 $c \in[-2,2]$로 스케일링하는 방법을 사용한다.

$$\left(\operatorname{Edit}\left(M_t, M_t^*, t\right)\right)_{i, j}:= \begin{cases}c \cdot\left(M_t\right)_{i, j} & \text { if } j=j^* \\ \left(M_t\right)_{i, j} & \text { otherwise }\end{cases}$$

---

3. Applications

Attention maps을 고정하고 단어 교체를 수행한 결과는 위와 같다. Attention maps을 고정하지 않고 수행한 결과와 대비되게, 레몬 케이크의 형태를 보존하면서도 새로운 재료로 만든 케이크 이미지를 잘 생성한다.

형태를 보존하고 재질을 바꾸는 작업뿐만 아니라, 물체의 형태를 변경하는 작업도 수행 가능하다. attention maps 주입의 영향을 분석하기 위해 왼쪽에서 오른쪽으로 갈수록 점차 많은 step에 attention maps을 주입한 결과를 보여준다. 주입하는 단계가 많아질수록 원본 이미지를 잘 반영한다. 위의 이미지에서 확인할 수 있듯이 모든 단계에 attention maps을 주입하는 것이 무조건 좋은 결과를 나타내지는 않는다. 

단어 교체 말고도 이미지에 새로운 subject를 추가하는 것이 가능하다. 위의 이미지 상단에서 "car"에 "crushed"라는 단어를 추가하여 배경을 보존하면서도 차의 디테일을 수정할 수 있다. 위의 이미지 하단에서는 특정 개체를 변경하는 것이 아닌 배경을 변경하는 작업도 기존 내용을 보존하면서 수행하는 것을 확인할 수 있다. 아래와 같이 그림을 현실적으로 재생성하거나 사진을 다양한 스타일의 그림으로 그리는 작업도 가능하다.

 

Prompt를 이용한 이미지 편집은 효과적이지만 완전한 제어는 힘들다. 예를 들어 "snowy mountain"이라는 prompt를 활용해 이미지를 생성할 때 텍스트로 쌓인 눈의 양을 표현하는 것은 어렵다. 따라서 저자들은 특정 단어에 해당하는 attention map을 스케일링하여 단어의 의미를 조절하는 fader control을 제안한다. 아래 그림에서 결과를 확인할 수 있다.

 

Diffusion model로 생성한 이미지가 아닌 실제 이미지를 편집하고 싶다면, 그 이미지를 생성할 수 있는 초기 노이즈 이미지를 찾아야 한다. 최근 이를 위한 기술인 inversion이 연구되고 있다. Inversion을 기반으로 하여 실제 이미지를 편집한 결과는 괜찮은 결과를 보여준다. 

다만 마지막 행과 같이 inversion이 재생성에 실패하는 경우가 존재한다. 이는 distortion-editability tradeoff 때문으로, guidance parameter를 낮추면 재생성 품질은 좋아지지만 편집 능력이 떨어지게 된다.

이러한 문제를 해결하기 위해 저자들은 attention maps을 mask로 사용하는 mask-based editing을 제안한다. 중요한 것은 추가 입력 없이 prompt만으로 mask를 만들어 사용한다는 점이다.

---

4. Limitations

Prompt-to-Prompt 기술에는 몇 가지 한계가 존재한다. 첫 번째는 실제 이미지 편집이 inversion에 의존한다는 점이다. Inversion의 문제점인 이미지 왜곡 문제를 해결해야하며, inversion을 위한 prompt를 사용자가 만들어야 한다는 문제점이 존재한다. 두 번째는 attention maps의 해상도가 낮다는 것이다. Cross-attention은 모델의 병목에 해당하는 위치에 존재하기 때문에 낮은 해상도를 가져 지역적인 편집 작업에 한계가 있다. 마지막으로 Prompt-to-Prompt를 사용한 편집은 기존 이미지에서 물체의 위치를 조정할 수 없다.

---

5. Summary

본 논문은 text prompt만으로 이미지 편집을 수행하는 기술인 Prompt-to-Prompt를 제안하였다. Prompt-to-Prompt 기술의 핵심은 U-net의 cross-attention layers 내부의 attention maps을 조작하는 것이다. Attention maps을 어떻게 조작하느냐에 따라 단어 교체, 구문 추가 등의 다양한 작업을 수행할 수 있다. 또한 inversion을 이용해 실제 이미지도 편집할 수 있다. 다만 attention maps의 낮은 해상도로 인한 품질 문제와 기존 이미지에서 물체의 위치를 조정하는 편집은 힘들다는 한계가 존재한다.