[VLTP] Vision-Language Guided Token Pruning for Task-Oriented Segmentation

WACV 2025

2024   ·   encoder-side   ·   efficient-vlm

Encoder WACV 2025

Hanning Chen, Yang Ni, … Mohsen Imani · UC Irvine 외

arXiv GitHub

한 줄 요약. 분할(segmentation) ViT를 가속하는 토큰 가지치기인데, 대상이 task-oriented segmentation(TOS)이다 — 같은 이미지라도 "무슨 task냐"에 따라 정답 마스크가 달라지는 추론형 분할. 기존 토큰 가지치기는 고정 클래스 기준이라 TOS에 안 맞는다. VLTP는 prune decoder를 ViT 여러 층에 끼워, MLLM의 vision-language 가이드로 각 토큰이 그 task에 관련 있는지를 점수화하고 관련 토큰만 깊은 층으로 보낸다(버린 토큰은 다음 stage에서 재활성). MLLM·ViT는 freeze, prune·mask decoder만 학습. ViT FLOPs ~25%↓ 무손실(40%↓에 1% 하락), 게다가 SOTA 대비 mIoU +2.5%.

배경

ViT는 분할의 강력한 백본이지만 연산이 비싸다. 토큰 가지치기가 답이 될 수 있는데, TOS에선 기존 방법이 무너진다.

  • TOS의 특수성 — “이 사진에서 못을 박을 물건은?”처럼 task에 따라 각 픽셀의 정답 클래스가 바뀐다. 고정 semantic class 기준으로 토큰을 자르면 task에 필요한 정보를 잃는다.
  • MLLM 호환 필요 — 최신 TOS는 MLLM의 추론으로 가이드되는데(예: LISA식 reasoning segmentation), 기존 가지치기는 이런 외부 추론 단서를 못 쓴다.

ViT가 모든 토큰을 끝까지 처리할 필요가 있을까? task와 관련된 토큰만 깊은 층으로 보내면 되지 않을까 — 단, 무엇이 관련인지는 MLLM이 알려준다.

Figure 2. VLTP 구조. ViT 여러 층에 prune decoder를 삽입 — 이미지 토큰 + MLLM의 vision-language 가이드를 받아 각 토큰의 task 관련도를 예측하고, 관련도 높은 토큰만 다음 층으로 보낸다(다단계). 버려진 토큰은 다음 stage에서 재평가로 되살린다.

핵심 아이디어

① MLLM 가이드 prune decoder

ViT의 선택된 여러 층에 prune decoder를 끼운다. 입력은 이미지 토큰 + MLLM의 vision-language 가이드, 출력은 각 토큰의 task 관련도 점수. 점수 상위(1−r)만 다음 층으로, 나머지는 그 stage 동안 self-attention에서 제외(freeze).

② 다단계 + 재활성

한 번에 안 자르고 여러 stage로. 초기 stage의 오판(잘못 버림)을 보완하려, 버린 토큰을 다음 pruning stage에서 재평가·재활성하고 마스크 생성에도 다시 쓴다 — 정보 손실 최소화.

  • 학습 — MLLM과 ViT는 freeze, prune decoder + mask decoder만 학습(reasoning 마스크 손실 CE+DICE + pruning 마스크 손실). 즉 training 기반(training-free 아님).
  • 2-step 프레임워크 — MLLM의 추론 + vision foundation model(분할)의 품질을 결합. 기존 two-stage(LLM이 명시적 rationale 생성)와 달리 가이드를 prune decoder가 바로 소비.

적용·평가

항목 내용
적용 모델 SAM의 ViT-H 분할 인코더를 가속 · 가이드 MLLM = LLaVA (LISA식 reasoning segmentation 파이프라인)
데이터셋 RIO(100+ tasks, common/uncommon) · COCO-Tasks(14 tasks) — 둘 다 MS COCO 2014 기반
Task Task-Oriented Segmentation(TOS) — task/의도에 따라 마스크가 바뀌는 추론형 분할
대표 결과 ViT FLOPs ~25%↓ 무손실 · ~40%↓에 1% 하락 · SOTA 대비 mIoU +2.5%(MLLM 가이드)

결과

Table 1. 기존 방법과의 비교 및 MLLM 가이드 효과. VLTP는 ViT 연산을 줄이면서 mIoU를 오히려 높인다.
  • 효율+정확도 동시 — ViT FLOPs ~25%↓ 무손실, ~40%↓에 1% 하락. 게다가 SOTA 대비 mIoU +2.5% — 가지치기가 정확도까지 끌어올린다.
  • RIO·COCO-Tasks(TOS)에서 vision-only 가지치기(EViT·ToMe)보다 우수.

한 줄 정리 & 의의

  • “task에 관련된 토큰만 깊은 ViT 층에” — MLLM이 관련도를 알려주는 분할용 토큰 가지치기. prune decoder가 vision-language 가이드로 관련도를 예측, 관련 토큰만 통과(+버린 토큰 재활성).
  • 차별점. 고정 semantic class 기준의 일반 토큰 가지치기와 달리, task-dependent(TOS)하며 MLLM 추론 가이드를 직접 소비. MADTP·CrossGET이 VLT 분류·검색을 노린다면, VLTP는 분할(segmentation) ViT가 무대.
  • 위치. Encoder — 분할 ViT 인코더 내부에서 task 무관 토큰을 줄인다. → Efficient VLM 개요