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

[Paper] [GitHub]

📝 Summary

Image Editing을 더이상 마스킹 같은 것이 아니라, 텍스트를 이용해 의미적으로 편집할 수 있는 기법을 소개한다.

• 확산 과정 중 cross-attention map을 주입하여, 어떤 픽셀이 어떤 프롬프트 토큰에 주목할지를 제어함으로써 이미지를 편집

Introduction

LLI (Large-scale Language Image) 기반의 기존 방법들은 사용자가 이미지의 일부를 명시적으로 마스킹하고, 마스킹된 영역에서만 이미지가 편집되도록 유도한다.

기존의 한계점

• 마스킹 과정이 번거롭기 때문에 직관적이고 빠른 텍스트 기반 편집에는 제약이 있다.
• 마스킹은 편집하고 싶은 영역만 남기는 기법으로, 가려진 부분에서의 질감, 형태 등의 핵심 정보가 모두 제거된다. 따라서, 마스킹된 영역의 주변 픽셀만 보고 추측하기 때문에 원래 객체의 세밀한 질감이나 형태를 그대로 살리기 어렵다.

제안하는 방법

• 의미적으로 이미지를 편집할 수 있는 텍스트 기반 편집 기법인 Prompt-to-Prompt를 제안한다.
• 확산 과정 중 cross-attention map을 주입하여, 어떤 픽셀이 어떤 프롬프트 토큰에 주목할지를 제어함으로써 이미지를 편집할 수 있다.

Methods

텍스트 프롬프트 $\mathcal{P}$와 랜덤 시드 $s$를 사용해 생성된 이미지를 $\mathcal{I}$라고 한다.

입력 이미지를 수정된 프롬프트 $\mathcal{P}^$만을 사용해 편집된 이미지 $\mathcal{I}^$를 생성하는 것이다.

1. Cross-attention in text-conditioned Diffusion Models

Imagen 모델을 백본으로 사용하며, 이미지의 구성과 기하학적 구조는 주로 64x64 해상도에서 결정되기 때문에 T2I 모델만 수정하고 SR 과정은 그대로 둔다.

이미지 임베딩과 텍스트 임베딩은 cross-attention layer에서 결합되며, 각 텍스트 토큰에 대한 spatial attention map을 생성한다.

\[M=\text{Softmax}\bigg(\frac{QK^\top}{\sqrt{d}}\bigg) \\ Q=l_Q(\phi(z_t))~,~K=l_K(\psi(\mathcal{P}))~,~V=l_V(\psi(\mathcal{P}))\]

Attention map은 위와 같이 정의되며, $M_{ij}$는 픽셀 $i$에 대한 $j$번째 토큰의 가중치를 의미한다.

\[\hat{\phi}(z_t)=MV\]

cross-attention 출력은 위와 같이 정의되며, 원래의 spatial feature $\phi(z_t)$를 업데이트할 때 사용된다.

2. Controlling the Cross-attention

위의 그림은 픽셀과 텍스트 간의 상호작용을 attention map으로 시각화한 것이다.

이미지의 구조는 확산 과정의 초반 단계에서 이미 결정되어 있다.

어텐션이 전반적인 구성을 반영하므로, 원래 프롬프트 $\mathcal{P}$로 생성할 때 얻은 어텐션 맵 $M$을 수정된 프롬프트 $\mathcal{P}^*$의 두 번째 생성 과정에 주입할 수 있다. 이를 통해 편집된 이미지 $\mathcal{I}^*$는 단순히 수정된 프롬프트에 맞춰 조작될 뿐만 아니라, 원본 이미지 $\mathcal{I}$의 구조도 유지된다.

알고리즘은 두 개의 프롬프트에 대해 동시에 확산 과정을 수행하며, 원하는 편집 task에 따라 각 단계에서 attention 기반 조작을 적용하는 방식을 일반적으로 사용한다.

1. 원래 텍스트 $\mathcal{P}$를 사용해 디노이징하여 $z_{t-1}$과 attemtion map $M_t$를 출력한다.
2. 수정된 프롬프트 $\mathcal{P}^*$를 사용해 attention map $M_t^*$를 출력한다.
3. 원래 attention map과 수정된 attention map을 이용해 새로운 attention map $\hat{M_t}$를 출력한다.
4. $V$는 유지한 채, 원래 attention map을 $\hat{M_t}$으로 덮은 후 디노이징을 수행한다.

Word Swap

사용자가 원본 프롬프트의 토큰을 다른 토큰으로 교체하는 경우이다.

\[\mathcal{P}=``\text{a big red bicycle}"\to \mathcal{P^*}=``\text{a big red car}"\]

새로운 프롬프트의 내용을 반영하면서도 원본 이미지의 구성을 유지하는 것이 주요 목표이며, 이를 위해 수정된 프롬프트로 새 이미지를 생성할 때, 원본 이미지의 attention map을 주입한다.

하지만 “bicycle”에서 “car”처럼 큰 구조적 변경이 있을 경우, attention map의 주입이 과도하게 형상을 제약할 수 있기 때문에 아래와 같이 완화된 제약을 둔다.

\[\text{Edit}(M_t, M_t^*, t) := \begin{cases}M_t^* & \text{if } t < \tau \\M_t & \text{otherwise}\end{cases}\]

$\tau$는 어느 시점에서 어떤 attention map을 주입할 지를 결정하는 파라미터이다.

• 확산 과정의 초반 단계에서는 이미지 구성이 결정되기 때문에 원본 attention map을 주입하여 구조를 유지시킨다.
• 확산 과정의 후반 단계에는 수정된 프롬프트에 대한 attention map을 주입하여 새로운 콘텐츠가 반영되도록 한다.

Adding a New Phrase

사용자가 원본 프롬프트에 새로운 토큰을 추가하는 경우이다.

\[\mathcal{P}=``\text{a castle next to a river}"\to \mathcal{P^*}=``\text{children drawing of a castle next to a river}"\]

공통된 부분을 유지하기 위해, 양쪽 프롬프트에 모두 존재하는 토큰에 대해서만 attention map을 주입한다.

\[\left( \text{Edit}(M_t, M_t^*, t) \right)_{i,j} := \begin{cases}(M_t^*)_{i,j} & \text{if } A(j) = \text{None} \\(M_t)_{i, A(j)} & \text{otherwise}\end{cases}\]

Alignment 함수 $A$는 $\mathcal{P}^*$의 토큰 인덱스 $j$를 입력 받아 $\mathcal{P}$에서 그와 일치하는 토큰의 인덱스를 출력한다. (없다면 None 출력)

• 원본 프롬프트에 없는 토큰이라면 수정된 프롬프트에 대한 attention map을 주입하여 새로운 콘텐츠가 반영되도록 한다.
• 원본 프롬프트에 있는 토큰이라면 원본 attention map을 주입하여 원래의 구조가 유지되도록 한다.

Attention Re-weighting

사용자가 특정 토큰이 이미지에 미치는 영향을 강화하거나 약화시키고자하는 경우이다. (아래 프롬프트에서 공이 더 복슬복슬해지거나 덜 복슬복슬해지도록 조정하고 싶은 경우)

\[\mathcal{P}=``\text{a fluffy red ball}"\]

조정하고 싶은 토큰 $j^*$에 해당하는 attention map만 스케일링한다.

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

Experiments

Attention injection이 적을수록 원본 이미지의 구조를 보존하지 못하고, 많을수록 원본 이미지의 구조에 제약을 많이 받는다.

---

원래 프롬프트에 단어를 추가함으로써 local edit과 global modification을 수행할 수 있다.

---

특정 단어에 대한 cross-attention을 줄이거나 늘림으로써, 이미지 생성에서 해당 단어의 영향력을 제어할 수 있다.