[평범한 학부생이 하는 논문 리뷰] InstructAny2Pix : Flexible Visual Editing via Multimodal Instruction Following

https://arxiv.org/abs/2312.06738

InstructAny2Pix: Flexible Visual Editing via Multimodal Instruction Following

https://github.com/jacklishufan/InstructAny2Pix/tree/main?tab=readme-ov-file

GitHub - jacklishufan/InstructAny2Pix

아카이브에 올라온 논문과 깃헙에 있는 논문 링크가 살짝 다릅니당... 왜지....

1. Introduction

 ControlNet과 T2I와 같은 기존의 SOTA는 그들이 학습된 특정한 edit만 수행할 수 있다. 이는 이들을 일반적인 사용자들에게 접근이 어렵게 한다. 이를 완화하기 위해, InstructPix2Pix와 같은 instruct-based image edit 방식들은 자연어인 그들의 instruction을 사용할 수 있게한다. 그러나 이러한 방식들도 그들이 학습된 간단한 instruction에 제한되고, 여러 editing operation과 여러 객체를 포함한 복잡한 instruction에 일반화될 수 없다. 게다가 그들은 추가적인 audiovisual input을 취할 수 없다.

 본 논문에서는 InstructAny2Pix를 제안하는데, 이는 복잡하고 multi-modal이고, multi-object 지시를 따를 수 있는 첫 instruction-following image editing 시스템이다. Instruction의 예시로는 "add the [sound] to [image]"나 "add [object A] and remove [object B] from [image]"가 있다. 여기서 object들은 이미지, 텍스트, 오디오로 표현될 수 있다. 본 논문의 시스템은 "change [image A] to the style of [image B]"나 "fit [image] to [music]"와 같이 자유 형식 instruction들도 수용한다. 이 연구는 이미지 수정 instruction들의 범위를 상당히 넓혀준다. 정성적, 정량적 평가를 통해 본 논문의 method는 다양한 editing task에서 높은 performance를 보인다.

 본 논문의 framework는 multimodal encoder의 multi-modal perception 능력, diffusion model의 높은 수준의 생성 능력, LLM의 instruction 이해 능력을 다 가져온다. 이 시스템의 performance는 이 세 구성에 의존한다.

 첫째, 본 논문의 모델은 많은 양의 multi-modal edit instruction들로 학습된다. 학습 데이터셋은 incorporating audio into image, adding image to image, adding multiple objects, replacing object with another object (이때 object는 image, text, audio로 표현됨)를 포함한다. 이 데이터셋은 각 instruction의 카테고리에서 인간이 만든 예시들로 SOTA LLM을 prompting함으로써 생성한다. 이 다양한 데이터셋은 image-editing 시스템의 instruction-following 능력들을 확장하는데 중요하다.

 둘째, 언어 모델의 instruction tuning에 대한 최신 기술들과 여러 modality들의 input을 포함한 instruction을 효율적으로 이해하는 multi-modal representation learning을 통합한다. 

 마지막으로, 효율적인 학습을 용이하게 하고 생성 퀄리티를 강화하기 위해 refinement prior module을 제안한다. Refinement module의 필요성은 두 가지 관측에서부터 발생한다. 첫 번째, LLM은 학습하는게 느리고, 생성 task에서 사용할만한 높은 퀄리티의 representation을 생성하도록 학습시키는데 상당한 시간이 걸릴 것이다. Multi-modal alignment에 사용되는 데이터의 질은 높은 퀄리티의 생성 모델을 학습시키는데 사용되는 데이터의 질보다 안 좋다. 이러한 관측들로 원해지지 않는 bias들을 완화하고 학습을 가속시키는 방식을 도입할 필요성이 생겼다. 본 논문은 refinement prior module을 LLM의 output embedding을 refine하는 것을 학습시키느 transformer로 구현했다.

2. Related Works

2.1 Image editing

Image2Image translation Base Methods

종류 : ControlNet, T2I

한계 : source image의 상당한 정보를 잃는다. 디테일 보존에 애먹는다. 

Text Prompt Base Methods

종류 : Prompt2Prompt, Plug-and-Play, Null-text Inversion

한계 : 원하는 output 생성을 위해 long detailed description을 필요로 한다.

Instruction-guided Methods

종류 : InstructPix2Pix, MagicBrush, MGIE

한계 : 간단한 instruction에 제한되고 복잡한 instruction을 위해서는 반복적으로 모델에 적용해야한다. condition 정보로 audiovisual input을 취할 수 없다.

2.2 Multi-modal Conditioned Generation

Multi-modal language model Base Method

종류 : BLIP-Diffusion, Kosmos-G

한계 : 객체의 removal과 replacement, 일반적인 free-form instruction을 지원하지 않는다. Reference image의 공간적인 구조를 반영하지 못한다.

Audio Guided Image Generation

종류 : AAI

한계 : 3~5개의 reference 이미지를 하고 추론 시간에 각 audio input에 대해 gradient descent 최적화 스텝을 수행하는 것을 필요로 한다.

 

Method부분은 깃헙에 있는 paper랑 아카이브에 올라온 paper가 좀 다른거 같아서 둘 다 정리해봤다.

3. Method (깃헙 paper)

깃헙 논문 이미지

 전체적인 파이프라인은 위 이미지와 같다. 본 논문에서 제안한 method의 전체적인 과정은 다음과 같다.

1. Pretrain된 multi-modal encoder를 통해 다른 modality들을 통합된 하나의 embedding space로 인코딩한다.
2. Multimodal instruction을 LLM에 통과시켜 embedding sequence를 생성한다. (nonlinguistic token들의 embedding은 multimodal encoder를 통해 생성한 embedding으로 대체)
3. LLM이 [gen]과 [base] embedding을 생성한다.
4. [base]를 통해 수정할 이미지를 찾고 [gen]을 이용해 diffusion model을 통해 수정된 이미지를 생성한다.

 학습 효과를 최대화하고 수정된 output에 대한 더 좋은 control을 제공하기 위해, LLM과 diffusion을 ent-to-end 대신 각자 학습시킨다. 학습동안, LLM은 수정된 이미지의 representation을 예측하도록 학습하고 diffusion model은 동일한 embedding space에서 그 이미지의 representation을 기초로 이미지를 생성하도록 학습한다. 또한 refinement prior module을 통해 LLM의 output embedding을 refine한다.

 Multi-modal inputs $ \{ X_i \}_{i = 1:n}$를 참고하는 [image],[audio]와 같은 reference token들을 포함하는 text sequence $T$가 input으로 주어지면, 먼저 multi-modal encoder $F_{enc}$를 이용해 $\{H_i \}_{i=1:n} = F_{enc}(\{ X_i \}_{i=1:n})$의 embedding을 생성한다. 그런 다음 LLM의 input embedding layer $E$를 이용해 embedded representation $H_T = E(T)$를 얻는다. 다음 단계는 $H_T$의 대응되는 위치에 $\{H_i \}_{i=1:n}$를 넣는다. LLM의 embedding space로 $\{H_i \}_{i=1:n}$를 project시키기 위해 projection layer $P_{enc}$를 사용한다. 이를 통해 새로운 embedding sequence $H_{T}^{\prime}$를 얻는다. 이를 LLM에 넣음으로써 $H_{LLM} = F_{LLM}(H_{T}^{\prime})$를 추출하고 그런 다음 두 개의 특수 토큰 [base]와 [gen]에 해당하는 위치를 찾는다. 이렇게 하면 두 개의 임베딩 $H_{base}, H_{gen}$이 생성된다. 임베딩을 인코더의 임베딩 공간에 다시 매핑하기 위해 추가 projector $P_{out}$을 사용한다.

 $H_{base}$와 가장 높은 유사도를 지닌 embedding을 찾음으로써 base image id $k$를 찾는다. Source image $X_t$에 대해 DDIM inversion을 수행하고 이를 통해 벡터 $z_k$로 변환한다. $H_{gen}$를 컨디션으로 $z_k$에 대해 diffusion process를 수행한다. 이를 통해 edited image $X_{out}$를 얻는다. 추론 알고리즘은 아래와 같다.

3.1 Multi-Modal Encoder

 ImageBind의 image와 audio encoder를 multi-modal encoder로 사용한다. 다른 modality는 다른 정규 분포를 가지는 것을 관찰했기에 LLM으로 feature들을 통과시키기 전에 $L_2$ norm을 수행한다. 학습동안 encoder는 freeze시킨다. Projection layer만 trainable하다.

3.2 Instruction Following Multi-Modal LLM

 Base language model로 Vicuna-7b를 사용한다. Input sequence는 image와 audio에 대한 다양한 referece들을 가진 text로 구성된다.("remove [audio] from [image]") Input text는 LLM의 input embedding layer를 통해 latent space로 project된다. 이를 통해 sequence $H_T \in \mathbb{R} ^{L \times R}$가 나온다. 그 후로 [image]와 [audio]와 같은 multi-modal input의 reference에 상응하는 이 sequence에서의 위치를 확인하고 이를 multi-modal feature로 교체한다. 두 embedding space의 차원이 같을 필요가 없기 때문에, 대체하기 전에 MLP projector $P_{enc}$를 통합한다. 이렇게 하면 "image"와 "audio"와 같은 reference word에 대응되는 위치를 제외하고 $H_T$와 동일한 $H_T^\prime \in \mathbb{R}^{L \times D}$를 얻는다.

[base]와 [gen]라는 두개의 토큰을 더함으로써 LLM의 output을 늘린다. 마지막 LLM layer에서의 hidden state에 MLP projection layer $P_{out}$을 적용하고 이를 통해 $H_{base}, H_{gen}$을 얻는다. $H_{base}$와 encoder를 통해 얻은 image feature들의 cosine 유사도를 측정해 source image를 찾고 이와 $H_{gen}$를 diffusion model에 condition으로 준다.

 Loss 함수는 다음과 같다.

3.3 Refinement Priors

 LLM의 output을 image embedding space에 완전히 align시키는 것이 가능한 반면, LLM의 큰 사이즈와 느린 학습때문에 적절한 훈련 시간과 데이터가 주어지면 수렴에 실패한다. 본 논문의 시스템을 효율적으로 학습시키기 위해, 본 논문은 손상된 이미지 임베딩을 원래 상태로 복원하는 방법을 학습하는 Refinement Prior를 사용한다. 본 논문에서는 Gaussian noise와 domain shift라는 두가지 종류의 corruption을 고려한다. Gaussian noise를 위해, Gaussian noise를 샘플링하고 ground truth 이미지의 embedding에 이를 더한다. Domain shift에 대해서는 이미지-텍스트 및 이미지-오디오 쌍을 사용하고 해당 텍스트 또는 오디오 임베딩이 주어지면 ImageBind 이미지 임베딩을 예측하는 방법을 학습한다.

 본 논문은 이 refinement prior로 decoder-only transformer를 사용한다. Refinement 능력을 더 강화하기위해, 추가 input으로 이미지의 LAION aesthetic score $f$를 사용한다. 이는 모델이 생성한 샘플이 사람 눈에 더 좋게 보여지게 하는데 유용하다. 이미지 $X$와 embedding에 노이즈를 추가하거나 다른 modality의 페어 embedding을 찾는 corruption operator $C(\cdot)$이 주어지면 prior module의 loss는 다음과 같다.

이 refinement scheme는 LLM과 jointly하게 학습하지 않아도 생성 퀄리티를 향상시킨다. 심지어 prior가 LLM의 output을 refine하도록 명시적으로 학습하지도 않았다.

아래의 아카이브 paper의 method는 굳이 안읽어봐도 충분히 이해 가능할 듯 하다...

---

3. Method (아카이브 paper)

 전체적인 pipeline은 다음과 같다. Pretrained multi-modal encoder를 활용해서 embeddings의 set으로 non-textual input들을 encode한다. 이러한 embedding들을 text instruction과 나란히 LLM에 넣어주고 control token들의 sequence를 생성한다. 여기서 control token들은 다음 세 가지 중 하나에 속하게 된다.

• [base] : 수정 작업의 source image를 되찾기 위해 사용됨. 하나만 존재.
• [sub] : 원하는 output에서 연관된 subject의 reference image들을 되찾기 위해 사용됨. 다양한 reference image들을 설명하기 위해 두개 이상 존재.
• [gen] : diffusion model을 condition하기 위해 사용됨. refine module에 의해 refine됨.

마지막으로, diffusion model은 retrieved source image, reference image의 retrieved embedding과 [gen] token의 refine된 embedding으로부터 image output을 생성한다.

 InstructAny2Pix의 training pipeline은 다음 이미지와 같이 5단계로 구성된다.

아카이브에 있는 논문 이미지

3.1 Multi-Modal Language Alignment

 본 논문에서 multi-modal encoder의 output과 LLM의 embedding space를 mapping하기 위해 MLP projector(input projector)를 사용한다. 이 단계에서의 목표는 이 두 공간의 좋은 mapping을 학습하는 것이다. 많은 양의 image-text pair와 audio-text pair를 이용한다. Projector를 제외하고 나머지 부분은 freeze시킨다. 모델에게 "describe the [x]"라는 text instruction을 prompt한다. 여기서 [x]는 image나 audio 조각의 embedding이다. Target sequence로 ground truth caption을 사용한다. LLM의 output에 standard cross-entropy loss를 적용한다.

3.2 Multi-Modal Generation Learning

 본 논문에서 multi-modal encoder의 representation 공간에 대해 LLM의 output을 다시 mapping하기 위해 MLP projector(output projector)를 사용한다. Input projector와 output projector를 제외하고 다 freeze시킨다. 여기서는 image-text pair와 image-audio pair를 사용한다. 모델에 "generate [x]"를 prompt한다. 여기서 [x]는 ground truth caption이거나 audio piece의 embedding이다. Target으로 multi-modal encoder에 의해 encoding된 ground truth image embedding를 사용하고, L2 Loss를 output projector에 적용한다.

3.3 Refinement Learning

 Alignment step에서 사용된 이미지들에서 low quality 이미지들이 많다는 것이 관찰됐다. 이들 중 대부분은 audio-image 페어에서 왔다. 데이터셋의 LAION aesthetic score를 평가함으로써 이를 검증했다. Low quality 데이터의 영향을 완화하기 위해, refinement module을 통합한다. 이는 embedding space에서 이미지의 quality를 향상시키도록 학습하는 transformer이다. 이미지와 그의 향상된 버전의 페어로 된 데이터셋을 얻기 힘들기 때문에, 존재하는 image-text와 image-audio pair를 이용한다. Multi-modal encoder의 embedding space가 align됐기 때문에, $Pr(x|e,s)$를 직접 모델링할 수 있다. ($x$는 이미지, $e$는 embedding, $s$는 aesthetic score) 학습동안 text와 audio embedding을 사용할 수 있고, 추론 때는 image embedding으로 이를 대체할 수 있다.

 이 단계에서 refinement module만 사용된다. Image-text나 image-audio 페어가 주어지면, refinement module은 input으로 이미지의 aesthetic score와 text나 audio의 embedding이 들어오면 ground truth image의 embedding을 예측하도록 학습한다. 이때 L2 regression loss가 적용된다.

3.4 Realigning Diffusion Models

 이번 스테이지에서의 목표는 multi-modal encoder의 representation 공간에서 conditional embedding을 수용하도록 text-to-image diffusion model을 realign하는 것이다. 특히 이는 몇 global한 의미를 encode하는 refined [gen] embedding과 [sub] 토큰으로부터 되찾은 reference image embedding 둘 다로부터 이미지를 생성하도록 학습되어야한다. 본 논문에서는 diffusion model의 representation 공간에 embedding들을 mapping하기 위해 MLP projector를 사용했다. 이 스터이지에서, diffusion model의 U-Net과 MLP projector를 학습시킨다. Aesthetic score를 통해 걸러진 고퀄리티 이미지만 사용하고 self-supervied 방식을 통해 모델을 학습시킨다. [gen] input으로 ground truth image embedding을, [sub] input으로 잘려진 region의 image embedding을 사용했다.

 여기에 몇 가지 문제가 존재한다. 첫째, ground truth image embedding과 refinement module을 통해 생성된 것들 사이의 분포 shift가 존재할 수 있다. 둘째, 학습 프로세스에서 [gen]과 [sub] input은 ground truth image로부터 오기에 모델이 오버피팅하는 것을 막는 것과 region embedding을 무시하는 것이 중요하다. 첫 번째 이유를 해결하기 위해, robustness를 향상시키기 위해 [gen] embedding에 random gaussian noise를 더한다. 두 번째 이유를 해결하기 위해, global과 region embedding을 0.5의 확률로 독립적으로 drop한다. 또한 잘린 부분에 대한 unconditioned inpainting을 랜덤하게 수행하고 이를 통해 [gen] embedding이 [sub] repainted와 상응하는 구역을 가진 이미지로부터 오게 한다. Inpainting은 비싸기 때문에, subset을 미리 선택하고 이러한 이미지를 offline으로 생성한다.

 이 스테이지에서, stable diffusion 논문에서 제안한 standart diffusion loss를 사용한다. 이미지를 자르기위해 사용된 region은 object detector에 의해 생성되거나 1:1 비율로 랜덤하게 샘플링된다. Inpainting을 위한 region은 detector에 의해 생성된다.

3.5 Instruction Guided Finetuning

 이번 스테이지에서는 multi-modal edit instruction들과 source 이미지와 수정 결과 페어를 통해 end-to-end로 LLM과 projector들을 finetune한다. Target output sequence를 "[base] [sub] ... [sub] [gen] ... [gen]"로 형식화 하고, 이 때 각 [sub]는 reference image에 대응된다. "[sub]"토큰은 원하는 결과에 reference object가 보이고 싶을 때만 추가한다. 예를 들면, “add [a]", “remove [b]", “replace [c] with [d]"에서 [a]와 [d]만 target sequence에서 대응되는 [sub] 토큰을 가진다. LLM의 마지막 linear head의 output에 cross entropy classification loss를 적용한다. 게다가 LLM의 마지막 hidden layer로부터 target sequence의 대응되는 embedding을 가지고 오고 output projector를 통해 multi-modal encoder의 representation 공간에 그들을 project한다. [base] embedding과 source image의 ground truth embedding 사이, [sub] embedding과 reference image의 ground truth embedding 사이, [gen] embedding과 target image의 ground truth embedding사이에 L2 regression loss를 적용한다. Diffusion Model은 이 프로세스에서 사용되지 않는다. Embedding space에서 loss가 직접 적용되기에 이미지를 실제로 생성할 필요가 없다.

3.6 Implementation Details

 인코더로 ImageBind를, 디코더로 SDXL를 사용했다. Base LLM으로 Vicuna-7B를 사용했다. Input과 output projector는 두 개의 층으로 이루어지고 활성화 함수로 GELU를 사용한 MLP이다.

---

4. Experiments

4.1 Instruction guided Image Editing

• $CLIP_{dir}$ : caption에서의 변화와 image에서의 변화 사이의 일치를 측정
• $CLIP_{im}$ : source와 target 이미지 사이의 유사성 측정
• $CLIP_{out}$ : 수정된 이미지와 targeted caption사이의 유사성 측정

$CLIP_{dir}$와 $CLIP_{our}$의 결과가 더 좋아 instruction을 더 잘따르는 결과를 생성한다. 다양한 범위의 instruction을 통합하고 보지 못한 instruction pattern에 일반화가 가능하기에 우월성을 보인다.

4.2 Image Conditioned Generation

 Single Image에 대해서는 DINO score로 image 유사도를 측정했고, multi image에 대해서는 잘려진 subject와 referece subject image 사이의 DINO 유사도을 측정하는 $DINO_{sub}$와 전체 결과와 background input 이미지 사이의 DINO 유사도를 측정하는 $DINO_{ref}$를 사용해서 측정했다.

4.3 Multi-Object Multi-Modal Instruction Editing

 본 논문의 method는 다른 modality의 많은 input을 포합하는 복잡한 수정 작업을 수행할 수 있다. 본 논문의 method는 instruction-guided image editing의 경계를 극적으로 확장시켰다.

4.4 Ablations

Refinement

$CLIP_{im}$와 $CLIP_{out}$에서의 향상을 보였고 시각적 품질에서의 향상도 보였다.

Control Strength

Image editing에서 중요한 목적은 수정된 이미지가 기존 이미지는 그대로 가져가면서 의도된 변화를 반영하는 것이다. 이러한 두 목적 사이에 trade-off가 있고 이와 관련된 하이퍼파라미터는 CFG와 $\alpha$이다. 이에 대한 실험 결과는 위와 같다. 아래의 이미지는 알고리즘 1에서의 $\alpha, \beta$에 대한 실험 결과이다.

5. Discussion

5.1 Quality of Generation

 본 논문의 method가 더 좋은 퀄리티의 결과를 생성해낸다. 이는 본 논문의 data isolation때문이라고 추정된다. 본 논문의 method는 높은 퀄리티의 이미지에 대해서만 학습된다. 그렇기에 모델은 SDXL을 pretrain할 때 사용한 데이터셋 보다 평균 퀄리티가 낮은 editing 데이터셋으로부터의 bias를 상속받지 않는다.

5.2 Music Guided Generation

6. Conclusion

 본 논문에서는 multi-modal, multi-object insturction들을 베이스로 image를 editing하는 유연한 시스템인 InstructAny2Pix를 제안한다. 본 논문은 instruction의 범위를 확장하고 반복적인 chain-of-editing에 의존하지 않고 복잡한 instruction을 수행할 수 있다. 또한 end-to-end finetuning을 제거함으로써 training 효율성을 향상시킨 새로운 training scheme을 제안한다.

Appendix

7. Additional Experiment Results

7.1 Qualitative Comparison with Text-based Methods

 InstructPix2Pix와 MagicBrush, Instruct Diffusion과 정성적으로 비교한 결과는 위와 같다. 본 논문의 method만 instructPix2Pix 데이터셋으로 학습하지 않았다는 점에서 위 결과는 매우 인상적이다.

7.2 Qualitative Comparison with Image-conditioned Generation Methods

 위 결과를 통해 InstructAny2Pix가 생성 퀄리티와 이미지 일관성에서 다른 두 모델들 보다 좋은 것을 볼 수 있다.

7.3 Additional Ablation Study

Image Quality

InstructAny2Pix는 다양한 modality의 Input을 포함한 다양한 범위의 insturction을 효율적으로 다룰 수 있다는 것을 보인다.

Image Consistency

$\alpha$에 따라 생성된 결과는 위와 같다. $\alpha = 1$이면, inversed latent에 noise가 추가 되지 않는다. 이는 생성이 공간적으로 source image를 많이 따르는 결과를 띈다.

7.4 Image Quality Ablation

Image quality에 대해 Aesthetic score와 PickScore를 통해 비교한다. Refinement module이 있을 경우 정량적으로도 image quality가 상승하는 것을 볼 수 있다.

8. Limitations

8.1 Inherited Biases

 본 논문은 pretrain된 diffusion model과 LLM을 사용한다. 이런 이유로, 이러한 모델의 학습 프로세스로부터 biases를 상속받을 수 있다.

8.2 Generation Style

 본 논문의 모델은 사실과 같은 결과 대신 artistic/painting 결과에 편향되어있다.  다음과 같은 요소때문에 발생한다.

1. LAION-Aesthetic-3M 데이터셋은 많은 art와 painting을 포함한다.
2. LAION Aesthetic score는 high saturation과 artistic 결과에 편향되어있다.
3. MM-Inst 데이터셋의 이미지들을 SDXL을 통해 생성하는데, SDXL generation은 artistic output에 편향되어있다.

8.3 Fine-grained editing

 본 논문의 method는 counting이나 객체의 작은 부분 editing과 같은 fine-grained image edit에 실패한다. 또한 customization을 하지 못한다.

8.4 Reconstruction v.s. instructability tradeoff

9. Faliure Cases