CVPR 2024. [Paper] [Page] [Github]
Zhongwei Ren, Zhicheng Huang, Yunchao Wei, Yao Zhao, Dongmei Fu, Jiashi Feng, Xiaojie Jin
Beijing Jiaotong University | University of Science and Technology Beijing | ByteDance Inc
4 Dec 2023

Introduction

LLM의 성공을 바탕으로 구축된 large multimodal model (LMM)은 높은 수준의 시각적 인식 및 사용자 상호작용 경험을 크게 향상시켰다. 그러나 대부분은 물체에 대한 마스크와 같은 픽셀 수준 응답에 대한 능력이 제한되어 글로벌한 이미지 또는 영역에 대한 텍스트 설명을 생성한다. 최근 연구들은 LLM을 사용하여 실제 응용에 더 까다롭고 유연한 새로운 reasoning segmentation task에서 마스크를 생성하는 방법을 탐구하였다. 물체를 명시적으로 지정하는 기존 segmentation과 달리 reasoning segmentation에는 보다 복잡한 명령에 대한 복잡한 추론이 필요하며 이는 LMM의 능력과 잘 일치한다. 그러나 이 방법에는 두 가지 주요 단점이 있다.

  1. 여러 물체와 관련된 task를 처리할 수 없다.
  2. SAM과 같은 사전 학습된 이미지 segmentation model에 의존한다. 이러한 의존은 상당한 계산량이 필요하며 전체 모델의 성능이 segmentation model의 능력으로 제한되므로 모델의 잠재력을 방해한다.

본 논문에서는 픽셀 수준의 추론과 이해를 위한 효과적이고 효율적인 LMM인 PixelLM을 소개한다. PixelLM은 임의 개수의 타겟과 다양한 추론 복잡성을 사용한다. 이 디자인은 비용이 많이 드는 추가 segmentation model을 피하면서 LMM의 기본 구조를 유지하여 효율성과 다양한 응용 가능성을 모두 향상시킨다.

PixelLM의 핵심은 새로운 픽셀 디코더와 segmentation codebook이다. Codebook에는 다양한 시각적 척도에서 참조하는 대상과 관련된 컨텍스트 및 지식을 인코딩하는 학습 가능한 토큰이 포함되어 있다. 그런 다음 픽셀 디코더는 이미지 feature와 함께 codebook 토큰의 hidden embedding을 기반으로 마스크를 생성한다. 이 디자인 덕분에 PixelLM은 외부 segmentation model 없이 고품질 마스크를 생성하여 효율성을 크게 높일 수 있다. 또한, 저자들은 모델이 여러 타겟을 구별하는 능력을 향상시켜 마스크 품질을 더욱 향상시키기 위해 target refinement loss를 제안하였다.

저자들은 모델 학습 및 평가를 촉진하기 위해 포괄적인 multi-target reasoning segmentation 데이터셋인 MUSE를 구축했다. GPT4V를 사용한 데이터 큐레이션 파이프라인을 활용하여 90만 개의 인스턴스를 포함하는 246,000개의 질문-답변 쌍을 생성하였다.

PixelLM은 MUSE를 포함한 다양한 벤치마크에서 SOTA 성능을 달성하였다. 특히, PixelLM은 외부 segmentation model을 사용하는 모델들과 비교하였을 때 성능을 유지하거나 향상시키면서 계산 비용을 최대 50%까지 절감하였다.

Method

1. Model Design

Framework overview

PixelLM은 네 가지 주요 부분으로 구성된 간소화된 아키텍처를 사용한다.

  1. 사전 학습된 CLIP-ViT 비전 인코더 $\mathcal{I}$
  2. LLM $\mathcal{F}$
  3. 경량 픽셀 디코더 $\mathcal{D}$
  4. Segmentation codebook $C_\textrm{seg}$

PixelLM은 이미지 $x_\textrm{img}$와 쿼리 텍스트 $x_\textrm{txt}$를 처리하여 다양한 수의 타겟에 대해 텍스트 설명과 마스크를 생성한다.

비전 인코더와 LLM은 호환성을 보장하기 위해 잘 확립된 LMM 아키텍처를 준수하지만, 픽셀 디코더와 segmentation codebook은 다양한 시나리오에서 마스크 생성 능력을 갖추도록 LMM을 강화하는 데 중추적인 역할을 한다. Segmentation codebook은 타겟에 관련된 정보를 codebook token의 임베딩들로 인코딩한다. 그러면 픽셀 디코더가 이미지 feature와 함께 정확한 마스크로 변환한다.

Segmentation codebook

Segmentation codebook은 타겟에 관한 정보의 인코딩을 강화하여 고품질 마스크 생성을 촉진하는 것을 목표로 고안되었다. 이 codebook에는 시각적 개념에 대한 다양한 수준의 세분성 또는 scale을 나타내는 다양한 토큰들의 그룹이 포함되어 있으며, segmentation task의 요구 사항을 충족하도록 맞춤화되었다. 좋은 segmentation을 위해서는 semantic 카테고리와 기하학적 모양을 모두 이해해야 한다. 본 논문에서는 두 요소를 하나의 codebook으로 통합하여 정확한 segmentation에 필요한 semantic 측면과 기하학적 측면을 효과적으로 포착한다.

특히, codebook은 여러 토큰들의 그룹으로 구성되며, 각각은 이미지 인코더의 visual feature의 semantic scale에 해당한다.

\[\begin{equation} C_\textrm{seg} = \{c_n^\ell \in \mathbb{R}^d \}_{n=1, \ell=1}^{N, L} \end{equation}\]

여기서 $L$과 $N$은 각각 그룹별 시각적 scale과 토큰의 수를 나타내고 $d$는 LMM의 hidden dimension이다.

먼저 $N = 1$로 두고 codebook 토큰이 LMM 내에 통합되어 타겟 마스크 생성에 필요한 정보를 인코딩하는 방법을 살펴보자. 입력 이미지 $x_\textrm{img}$에 대하여 비전 인코더 $\mathcal{I}$의 레이어에서 출력되는 $L$개의 visual feature로 구성된 멀티스케일 visual feature \(I_\textrm{img} = \{I_\textrm{img}^\ell\}_{\ell=1}^L\)을 추출한다. 마지막 레이어 \(I_\textrm{img}^L\)은 글로벌한 이미지 정보를 캡슐화하고 vision-to-language projection layer $p_{V \rightarrow T}$를 통해 언어 공간으로 변환된다. 동시에, vision-to-decoder projection layer $p_{V \rightarrow D}$는 모든 feature $I_\textrm{img}$를 변환한다.

\[\begin{equation} \{f_\textrm{img} = p_{V \rightarrow D} (I_\textrm{img}^\ell) \}_{\ell=1}^L \end{equation}\]

입력 이미지 및 텍스트와 결합된 codebook 토큰은 LLM에 의해 처리되어 autoregressive 방식으로 응답 $y_\textrm{res}$를 생성한다.

\[\begin{equation} y_\textrm{res} = \mathcal{F} (p_{V \rightarrow T} (I_\textrm{img}^L), x_\textrm{txt}, C_\textrm{seg}) \end{equation}\]

예를 들어, “Segment the apple on the left”라는 텍스트 쿼리를 생각해보자. 출력 $y_\textrm{res}$에는 $C_\textrm{seg}$의 $L$개의 토큰이 포함된다.

“The apple is $c^1, \ldots, c^L$”

$C_\textrm{seg}$의 hidden embedding, 즉 $\mathcal{F}$의 마지막 레이어의 출력 $h = {h^\ell}_{\ell=1}^L$는 마스크 생성을 위한 이미지 feature 파악과 함께 픽셀 디코더 $\mathcal{D}$에 대한 입력이다.

$N > 1$인 경우를 살펴보자. LLM이 정확한 텍스트 응답을 제공하여도 여러 타겟 또는 복잡성 증가와 관련된 시나리오에서는 하나의 토큰이 타겟 semantic의 전체 범위를 적절하게 포착하지 못할 수 있다. 복잡한 추론 시나리오를 해석하는 모델의 역량을 강화하기 위해 각 scale 그룹 내에 여러 토큰을 도입하고 토큰 융합 연산을 수행한다.

\[\begin{equation} c^\ell = \{c_n^\ell\}_{n=1}^N \end{equation}\]

디코딩 프로세스 전에 linear projection layer $\phi$는 이러한 그룹화된 토큰의 hidden state를 통일된 형태로 변환한다.

\[\begin{equation} h^\ell = \phi (h_1^\ell, \ldots, h_N^\ell) \end{equation}\]

아래 그림은 각 그룹 내에서 여러 토큰을 적용하는 모습을 보여준다. 디코딩 후의 attention map은 개별 토큰이 고유하면서도 보완적인 정보를 제공하여 단일 토큰 방식보다 더 효과적인 마스크 생성이 가능함을 보여준다.

Pixel decoder

비전 인코더의 멀티스케일 feature를 적절하게 활용하기 위해 새롭고 경량 픽셀 디코더 $\mathcal{D}$를 사용한다. 이 디코더는 $C_\textrm{seg}$의 hidden embedding과 함께 이러한 feature를 정확한 segmentation mask로 변환하는 방법을 학습한다. 이러한 설계는 추가 비용이 드는 segmentation model의 필요성을 없애므로 효율성이 크게 향상된다.

$\mathcal{D}$는 $L$개의 attention block \(\{\textrm{Attn}^\ell\}_{\ell=1}^L\)로 구성되며, 각각은 이미지 feature의 고유한 scale과 codebook에 해당한다. 각 타겟 마스크 생성에 대해 $\mathcal{D}$는 각 scale $\ell$에서 점수 마스크 맵 $m^\ell$을 순차적으로 생성한 다음 이후 scale $\ell - 1$에서 관련성이 더 높은 영역으로 모델의 attention을 보낸다. 이 전략은 모델이 $m^\ell$에서 신뢰도 점수가 높은 영역에 집중하도록 가이드하여 더 정확한 마스크가 생성된다.

\[\begin{aligned} f_\textrm{img}^{\ell^\prime} &= \begin{cases} f_\textrm{img}^L & \quad \ell = L \\ f_\textrm{img}^\ell \odot (\sigma (m^{\ell+1}) + 1) & \quad \ell < L \end{cases} \\ m^\ell &= \textrm{Attn}^\ell (h^\ell, f_\textrm{img}^{\ell^\prime}) \end{aligned}\]

여기서 $f_\textrm{img}^{\ell^\prime}$는 ℓ scale의 변조된 feature이고, $\sigma$는 시그모이드 함수, $\odot$은 element-wise multiplication이다.

마지막으로 최종 segmentation 결과를 얻기 위해 모든 scale에서 마스크 맵을 결합하는 가중치 $\gamma = [\gamma^\ell]_{\ell=1}^L$을 학습한다.

\[\begin{equation} \hat{M} = \sum_{\ell=1}^L m^\ell \quad \textrm{where} \; \vert \gamma \vert = 1 \end{equation}\]

2. Traning Objectives

Target refinement loss

타겟 수가 증가하면 모델이 혼란을 겪고 중복된 마스크를 생성할 가능성이 높아진다. 이 문제를 완화하기 위해 target refinement loss를 도입한다. 이 전략은 여러 타겟이 함께 예측되는 불분명한 픽셀에 중점을 둔다. 이는 모델이 다양한 타겟을 명확하게 식별하고 학습하는 데 도움이 된다.

마스크 예측을

\[\begin{equation} \{\hat{M}_k \in \mathbb{R}^{H \times W}\}_{k=1}^K \end{equation}\]

라 하자. 여기서 $K$는 타겟의 총 개수이고, $H$와 $W$는 마스크의 모양이다. \(\hat{M}_{k_i}\)는 각 픽셀의 바이너리 값을 나타낸다. 그런 다음 여러 타겟을 예측하는 영역에 증가된 가중치를 할당하기 위해 맵 $A$를 다음과 같이 정의한다.

\[\begin{equation} A_i = \begin{cases} \alpha & \quad \sum_k \hat{M}_{k_i} \ge 2 \\ 1 & \quad \sum_k \hat{M}_{k_i} < 2 \end{cases} \end{equation}\]

여기서 $\alpha$는 hyperparameter이다. 가중된 loss는 각 마스크의 GT $M_k$에 대해 다음과 같이 계산된다.

\[\begin{equation} \mathcal{L}_\textrm{ref} = \frac{1}{KHW} \sum_k \sum_i A_i \mathcal{L}_\textrm{BCE} (\hat{M}_{k_i}, M_{k_i}) \end{equation}\]

\(\mathcal{L}_\textrm{BCE}\)는 per-pixel binary cross-entropy loss이다.

Overall loss

모델은 텍스트 생성을 위한 autoregressive cross-entropy loss \(\mathcal{L}_\textrm{txt}\), DICE loss \(\mathcal{L}_\textrm{dice}\), target refinement loss \(\mathcal{L}_\textrm{ref}\)를 사용하여 end-to-end로 학습된다. 전체 loss $\mathcal{L}$은 다음과 같다.

\[\begin{equation} \mathcal{L} = \mathcal{L}_\textrm{txt} + \lambda_\textrm{ref} \mathcal{L}_\textrm{ref} + \lambda_\textrm{dice} \mathcal{L}_\textrm{dice} \end{equation}\]

Multi-target Reasoning Segmentation

본 논문의 목표는 임의의 수의 open-set 타겟과 다양한 추론 복잡성과 관련된 task를 처리할 수 있는 LLM을 개발하는 것이다. 가장 큰 문제는 모델 학습에 적합한 데이터셋이 없다는 것이다. 저자들은 기존 공개 데이터셋을 검토하고 중요한 제한 사항을 확인하였다.

  1. Segmentation mask의 디테일 및 물체 표현이 부적절하다.
  2. 복잡한 추론과 다양한 타겟을 갖춘 질문-답변 쌍이 부족하다.

이러한 문제를 해결하기 위해 multi-target reasoning segmentationn (MUSE) 데이터를 구성하기 위한 주석 파이프라인을 도입한다. MUSE는 open-set 개념, 자세한 물체 설명, 여러 타겟에 대한 복잡한 질문-답변 쌍과 인스턴스 수준의 마스크 주석을 가지고 있다.

1. MUSE Dataset

LVIS 데이터셋에서 이미지를 기반으로 한 자세한 텍스트 설명과 함께 총 91만 개의 고품질 instance segmentation mask가 선택되었다. 이러한 인스턴스를 활용하여 246,000개의 질문-답변 쌍을 구성하고 답변당 평균 3.7개의 타겟을 구성한다. 그런 다음 이 데이터셋은 각각 23.9만 개, 2,800개, 4,3000개의 질문-답변 쌍을 포함하는 train, val, test의 세 가지 split으로 나뉜다. Test split은 질문에 포함된 타겟의 수가 3개보다 적은 것과 많은 것으로 구분된다.

2. Dataset Generation Pipeline


먼저 이미지 캡션에 LLaVA를 사용한 다음 여러 이미지 영역에 대한 질문을 생성하는 데 GPT-4를 사용한다. 주석 비용을 줄이기 위해 기존 마스크 주석이 있는 이미지를 활용한다. 이미지 캡션, 수동으로 선택한 물체 이름, 이미지의 bounding box 좌표가 GPT-4에 입력되어 답변 선택 및 질문 구성이 용이해진다. 그러나 이미지 내용을 직접적으로 인식할 수 없기 때문에 이 방법으로 생성된 질문-답변 쌍의 내용은 캡션 설명에만 국한되어 데이터의 다양성이 크게 제한되는 경우가 많다.

이러한 문제점을 해결하기 위해 파이프라인은 두 가지 주요 방식으로 개선되었다. 먼저, 시각적 콘텐츠 이해에 강력한 능력을 보여주는 GPT-4V로 전환한다. GPT-4V는 보다 미묘하고 자연스러운 질문을 생성하는 데 중요한 역할을 한다. 또한 답변 생성을 위해 보다 동적인 접근 방식을 구현한다. 특히 이미지의 모든 인스턴스 카테고리 이름과 bounding box 좌표를 GPT-4V에 제공한다. 신중하게 제작된 프롬프트를 통해 GPT-4V는 인스턴스를 자동으로 선택하여 이미지와 관련된 질문-답변 쌍을 구성한다.

3. Evaluation

평가를 위해 다음과 같은 세 가지 측면, 즉 마스크와 정렬된 텍스트 설명 생성, 마스크와 텍스트 설명 간의 일치 정확도, 마스크 품질에 중점을 두었다. 각 타겟의 마스크 품질에 중점을 두기 때문에 이미지 레벨 캡션을 평가하지 않는다.

각 질문에 대한 평가 프로세스는 4단계로 구성된다.

  1. DETR과 유사한 이분 매칭을 사용하여 마스크 IoU 점수를 기반으로 예측 마스크를 실제 마스크와 일치시킨다. 할당되지 않은 예측이나 GT 정보에는 빈 마스크가 할당된다.
  2. 생성된 텍스트의 마스크 위치를 해당 GT 물체 설명으로 바꾼다.
  3. 이 수정된 텍스트를 GPT-3.5에 입력하여 각 예측에 1부터 10까지 점수를 매기도록 한다. 점수가 높을수록 품질이 더 좋고 할당되지 않은 예측은 0점을 받는다.
  4. 각 예측의 최종 점수는 GPT와 IoU의 곱이다. gIoU와 cIoU를 사용한다.

Experiments

  • 구현 디테일
    • 멀티모달 모델: LLaVA-7B, LLaVA-Llama2-13B (LoRA 적용)
    • 비전 인코더: CLIP-ViT-L/14-336
    • 학습되는 부분: pixel decoder, LoRA 파라미터, segmentation codebook, vision-to-language & vision-to-decoder projection layers

1. Results on MUSE

다음은 MUSE 벤치마크에서의 결과를 비교한 표이다.

2. Results on Referring Segmentation

다음은 multi-referring segmentation 벤치마크에서의 결과를 비교한 표이다.


다음은 referring segmentation 벤치마크에서의 결과를 비교한 표이다.

3. Ablation Study

다음은 scale의 수, 토큰 수, 토큰 융합 메커니즘, target refinement loss, 데이터 생성 방법, 질문-답변 쌍의 개수에 대한 ablation 결과이다.


다음은 PixelLM과 PixelLM† (토큰 융합 메커니즘과 target refinement loss이 없는 PixelLM)의 결과를 비교한 것이다.