Diffusion은 신이고 GPT는 무적이다

CVPR 2024. [Paper] [Page]
Jiarui Xu, Xingyi Zhou, Shen Yan, Xiuye Gu, Anurag Arnab, Chen Sun, Xiaolong Wang, Cordelia Schmid
Google Research | UC San Diego
14 Dec 2023

Introduction

본 논문에서는 각 출력 단어를 픽셀 위치에 dense하게 정렬하여 세분화된 위치 파악 (localization) 능력을 갖춘 vision-language model (VLM)인 PixelLLM을 소개한다. 언어 모델의 단어 feature 위에 작은 MLP를 추가하여 각 단어의 픽셀 위치로 회귀함으로써 이를 실현한다. 언어 모델의 가중치는 고정된 상태로 유지되거나 LoRA fine-tuning을 통해 업데이트될 수 있다. 또한 PixelLLM은 위치 프롬프트나 텍스트 프롬프트를 사용하여 프롬프트에 맞는 출력을 생성할 수 있다.

아키텍처는 이미지 인코더, 프롬프트 인코더, 프롬프트 feature 추출기로 구성된다. 프롬프트 feature 추출기는 프롬프트를 조건으로 이미지 feature를 텍스트 임베딩 공간에 매핑한다. 이미지 feature와 텍스트 프롬프트는 캡션과 단어별 localization 출력을 생성하는 LLM의 접두사로 직접 제공된다. 이 아키텍처는 일반적이며 언어나 위치를 입력 또는 출력으로 조합하여 다양한 비전-언어 task에 적응할 수 있다.

학습을 위해 전용 단어-픽셀 정렬 데이터가 필요하지만 이러한 주석은 이미 대규모로 존재한다. Localized Narratives (LN) 데이터셋에는 내레이션 중에 주의를 기울이는 마우스 궤적과 함께 주어진 이미지를 설명하는 인간 주석이 포함되어 있다. 이는 내레이션 문장의 모든 단어에 대해 동기화된 위치를 제공하며, 이는 모델을 학습시키는 데 사용할 수 있다. 모든 단어-위치 쌍이 시각적으로 의미가 있거나 정확하지는 않지만 실제 인간의 관심에서 나온 것이므로 저자들은 가치가 있다고 주장한다.

PixelLLM은 referring localization (RefCOCO), location conditioned captioning (RefCOCOg, Visual Genome), dense object captioning (Visual Genome)에서 모두 SOTA 성능을 달성하였다. 또한 dense한 픽셀별 localization이 향상된 성능의 핵심이다.

Method

VLM의 프레임워크에서 출력 문장의 각 단어를 픽셀 위치에 정렬하여 localization을 공식화한다. 구체적으로, 문장 출력 $\textbf{s}$ 외에도 문장과 동일한 길이의 점들의 시퀀스 \(\textbf{p} = [p_1, \ldots, p_n]\)도 출력한다. 각 $p_i$는 문장의 각 단어 토큰에 해당한다. 모델이 시각적이지 않은 토큰을 무시하도록 강제하지 않으므로 모델이 상대적 의미를 가진 단어도 학습할 수 있다.

1. Architecture


입력은 이미지 $\textbf{I}$와 선택적 위치 프롬프트 $\textbf{b}$이다. 위치 프롬프트가 제공되지 않으면 전체 이미지의 박스 프롬프트 $b = (0, 0, H, W)$를 사용한다. 위치 프롬프트가 제공되면 모델은 프롬프트에 표시된 위치에 초점을 맞출 것으로 예상된다. 출력은 문장 $\textbf{s}$와 단어-픽셀로 정렬된 점 궤적 $\textbf{p}$이다.

이미지 인코더 $\mathcal{V}$와 위치 프롬프트 인코더 $\mathcal{P}$를 사용하여 이미지 feature $\textbf{f} = \mathcal{V}(\textbf{I})$와 위치 프롬프트 feature $\mathcal{P}(b)$를 얻는다. 여기서 $\textbf{f}$는 전체 이미지의 feature이다. 프롬프트 feature 추출기 $\mathcal{E}$를 사용하여 위치 프롬프트에 해당하는 feature \(\textbf{f}_l\)을 추출한다.

\[\begin{equation} \textbf{f}_l = \mathcal{E} (\textbf{f}, \mathcal{P} (b)) \in \mathbb{R}^{N \times C} \end{equation}\]

프롬프트 feature 추출기 $\mathcal{E}$는 학습 가능한 토큰 세트 $\textbf{q}$가 있는 양방향 transformer이다. 특히 양방향 transformer는 각 레이어에서 $[\mathcal{P}(b), \textbf{q}]$와 $\textbf{f}$를 query 또는 key/value로 사용하고 마지막 레이어에서 학습 가능한 토큰 feature를 사용한다. 출력 feature \(\textbf{f}_l\)은 위치 프롬프트 $b$에 해당하는 feature를 전달한다. 프롬프트 feature 추출기는 ROIAlign과 유사한 함수를 가지고 있지만 학습이 가능하고 feature interpolation이나 샘플링이 필요하지 않다.

언어 모델의 dense한 위치 출력. \(\textbf{f}_l\)이 주어지면 autoregressive 디코딩을 사용하여 captioning을 위해 언어 모델에 이를 제공할 수 있다.

\[\begin{equation} w_i = \mathcal{L} (\textbf{f}_l, \textbf{w}_{1:i−1}) \end{equation}\]

언어 모델의 마지막 linear layer는 언어 feature space에서 one-hot vocabulary 인덱스로 매핑되는 vocabulary 매핑 레이어이다. $\mathcal{L}^{-}$를 마지막 vocabulary 매핑 레이어가 없는 언어 모델이라 하면, 디코딩 프로세스는 다음과 같이 다시 쓸 수 있다.

\[\begin{equation} w_i = \arg \max \; (\textbf{v} \cdot \mathcal{L}^{-} (\textbf{f}_l, \textbf{w}_{1:i-1})) \end{equation}\]

여기서 \(\textbf{v} \in \mathbb{R}^{\vert V \vert \times C}\)는 vocabulary 매핑 레이어의 가중치이다.

Localization에 동일한 언어 feature를 사용하려면 언어 feature를 2차원 위치 출력에 매핑하는 작은 MLP를 추가하기만 하면 된다.

\[\begin{equation} p_i = \textrm{MLP} (\mathcal{L}^{-} (\textbf{f}_l, \textbf{w}_{1:i-1})) \end{equation}\]

원본 텍스트 디코딩 프로세스에 영향을 주지 않기 위해 autoregressive 디코딩에 대한 localization 출력을 피드백하지 않는다. 위치 예측은 언어 디코딩과 함께 즉시 실행되며 약간의 계산 오버헤드만 추가된다. 이 디자인은 언어 모델에 구애받지 않으며 원래 언어 생성 능력을 방해하지 않고 모든 언어 모델에 적용될 수 있다. 텍스트 프롬프트를 입력으로 사용하려면 텍스트 프롬프트의 단어 임베딩을 \(\textbf{f}_l\)과 직접 concatenate한다.

2. Training

사람이 주석을 단 Localized Narratives (LN) 데이터셋을 사용하여 모델을 학습시킨다. LN 데이터셋은 설명하는 영역 위로 마우스를 이동하는 동시에 주석 작성자에게 주어진 이미지에 대한 설명을 요청한 데이터셋이다. 내레이션과 마우스 추적이 동기화되어 내레이션의 각 단어 위치를 제공한다. 마우스 궤적에 잡음이 있을 수 있지만 dense한 위치 supervision을 얻을 수 있는 저렴하고 효과적인 방법이다. LN 데이터셋에는 이미지 $\textbf{I}$, 캡션 문장 $\textbf{s}$, 위치 궤적 $\textbf{p}$ 등 필요한 모든 주석 $(\textbf{I}, \textbf{s}, \textbf{p})$가 포함되어 있다.

Label-smoothed cross-entropy loss를 사용하여 캡션 출력을 학습시키고 L1 regression loss를 사용하여 localization 출력을 학습시킨다.

\[\begin{equation} L = \frac{1}{n} \sum_{i=1}^n (\textrm{CE} (\mathcal{L} (\textbf{f}_l, \textbf{w}_{1:i-1}), w_i) + \lambda \| \hat{p}_i - p_i \|) \end{equation}\]

\(\hat{p}_i\)는 $i$번째 단어의 예측 위치이고, $n$은 캡션 길이이다.

3. Adapting to downstream vision tasks.

PixelLLM의 아키텍처는 텍스트/위치의 모든 조합을 입력 또는 출력으로 사용할 수 있으므로 다양한 위치 관련 비전 task에 적용될 수 있다.

Referring localization and segmentation


Referring localization과 referring segmentation은 이미지 $\textbf{I}$와 문장 쿼리 $\textbf{t}$를 입력으로 취하고 쿼리에 해당하는 bounding box $b \in \mathbb{R}^4$를 생성하는 것을 목표로 한다. 이 task에 프레임워크를 적용하기 위해 위치 프롬프트를 글로벌 박스 $b = (0, 0, W, H)$로 설정하고 쿼리 문장을 $\mathcal{L}^{-}$의 조건으로 사용한다.

기본적으로 PixelLLM은 하나의 bounding box가 아닌 궤적을 출력한다. 궤적의 경계를 취하여 bounding box를 만들 수 있지만 부정확하다. 따라서 동일한 MLP 레이어를 사용하여 <EOS> 토큰에서 정확한 bounding box를 출력하도록 모델을 학습시킨다.

\[\begin{equation} \hat{b} = \textrm{MLP} \bigg( \mathcal{L}^{-} (\mathcal{E} (\mathcal{V} (\textbf{I}), \mathcal{P} (\hat{b})), [\textbf{t}, \langle \textrm{EOS} \rangle]) \bigg) \end{equation}\]

모델에는 이미 SAM의 이미지 backbone과 프롬프트 인코더가 포함되어 있으므로 SAM의 마스크 디코더를 연결하기만 하면 segmentation mask를 추가로 얻을 수 있다. 따라서 예측된 박스 $b$ 위에 마스크를 생성하여 referring segmentation에도 사용할 수 있다.

Location-conditioned captioning & Dense object captioning


Location-conditioned captioning은 이미지 $\textbf{I}$와 bounding box $b$를 위치 프롬프트로 취하고 박스 쿼리에서 표시된 물체에 해당하는 캡션 문장 $\textbf{s}^b$를 생성한다. 단어별 위치 출력을 무시하면서 프롬프트 인코더와 autoregressive 언어 모델을 사용하여 직접 적용할 수 있다.

\[\begin{equation} \textbf{s}_i^b = \mathcal{L} (\mathcal{E} (\mathcal{V} (\textbf{I}), \mathcal{P} (b)), \textbf{s}_{1:i-1}^b) \end{equation}\]

Dense object captioning은 먼저 주어진 이미지에서 모든 물체를 감지한 다음 captioning하는 것을 목표로 한다. Bounding box 후보를 얻기 위해 이미지 인코더 뒤에 proposal head를 추가한다. 그런 다음 bounding box를 위치 프롬프트에 개별적으로 공급하여 각각에 대해 location-conditioned captioning을 수행한다. 구체적으로 Simple Feature Pyramid를 사용하여 visual feature $\textbf{f}$를 업샘플링하고 detection을 위해 CenterNet head를 사용한다. Detection loss와 캡션 loss를 사용하여 end-to-end로 모델을 fine-tuning한다.

Experiments

  • 아키텍처 디테일
    • 비전 인코더 $\mathcal{V}$
      • SAM으로 초기화된 ViT-H와 EVA02로 초기화된 ViT-L을 동시에 사용
      • 고정된 SAM에서 localization feature를 얻고 EVA02에서 semantic feature를 학습
      • 두 ViT의 feature를 채널 차원으로 concatenate하고 프롬프트 feature 추출기에 입력
    • 언어 모델 $\mathcal{L}$
      • instruction fine-tuning된 T5-XL
      • self-attention block의 query 및 value projection layer에 LoRA를 적용 (rank 32)
    • 프롬프트 feature 추출기 $\mathcal{E}$
      • 2-layer transformer (학습 가능한 토큰 32개)
      • 추출기의 출력을 접두사 텍스트의 텍스트 임베딩과 concatenate하고 언어 모델 $\mathcal{L}$에 공급
  • 학습 디테일
    • 먼저 WebLI 데이터셋으로 captioning 목적 함수만 사전 학습 (10 epochs, 224$\times$224)
    • 그런 다음 Localized Narrative 데이터셋으로 captioning과 localization 목적 함수를 동시에 학습 (5 epochs, 384$\times$384)
    • localization 가중치: $\lambda = 0.1$
    • label-smooth factor: 0.1

1. Joint Captioning and trace generation

다음은 pixel-aligned captioning과 referring segmentation, dense object captioning에 대한 예시들이다.

2. SOTA Comparison on Downstream Tasks

다음은 RefCOCO에서 referring localization과 referring segmentation을 기존 방법들과 비교한 표이다.


다음은 Visual Genome에서 dense object captioning을 기존 방법들과 비교한 표이다.


다음은 RefCOCOg와 Visual Genome에서 location-conditioned captioning을 기존 방법들과 비교한 표이다.

3. Ablation study

다음은 localization supervision에 대한 ablation 결과이다.


다음은 언어 모델 크기와 LoRA에 대한 ablation 결과이다.

You may also enjoy