비전 트랜스포머를 활용한 마스크 분류 혁신, ViT‑P

초록

ViT‑P는 마스크 생성과 분류를 분리하는 2단계 프레임워크로, 클래스‑agnostic 마스크 제안기와 Vision Transformer 기반 포인트 분류기를 결합한다. 최고값 포인트를 이용해 마스크를 빠르고 정확하게 분류하며, 사전 학습 없이 다양한 사전 학습된 ViT 백본에 적용 가능하다. 거친 라벨·바운딩 박스만으로도 분류 성능을 크게 끌어올려 어노테이션 비용을 절감하고, ADE20K, Cityscapes, COCO 등에서 최첨단 성능을 달성한다.

상세 분석

본 논문은 현재 마스크 기반 세그멘테이션 모델이 마스크 품질은 뛰어나지만, 마스크에 대한 클래스 레이블 예측이 약한 점을 정확히 지적한다. 특히 Mask2Former·InternImage와 같은 최신 모델은 마스크 자체는 80% 이상의 PQ를 기록하지만, 동일 마스크에 대해 완벽히 라벨링했을 때의 ‘upper bound’와 실제 성능 사이에 20% 이상의 격차가 존재한다는 실험 결과를 제시한다. 이러한 격차는 “마스크 분류가 전체 세그멘테이션 정확도의 병목”이라는 핵심 가설을 뒷받침한다.

ViT‑P는 이 문제를 해결하기 위해 두 단계 구조를 도입한다. 1단계는 기존의 마스크 제안기(예: OneFormer, Mask2Former)를 그대로 사용해 클래스‑agnostic 마스크를 생성한다. 여기서는 마스크 품질을 높이는 데 집중하고, 라벨링은 전혀 수행하지 않는다. 2단계에서는 각 마스크의 가장 높은 값(픽셀) 위치, 즉 마스크 중심에 가까운 포인트를 추출하고, 이를 Vision Transformer에 입력한다. 포인트는 별도의 포인트 임베딩 레이어를 통해 패치 임베딩과 동일 차원으로 변환되며, 포지셔널 임베딩을 공유한다. 이렇게 구성된 시퀀스는 표준 ViT 인코더를 통과해 전역 컨텍스트를 학습하고, 최종 MLP 헤드가 각 포인트에 대한 클래스 확률을 출력한다.

핵심 기술적 기여는 다음과 같다.

1. 포인트 기반 분류: 마스크 전체를 일일이 분류하는 대신, 가장 신뢰도 높은 포인트 하나만을 사용함으로써 연산량을 크게 감소시키면서도 높은 정확도를 유지한다. 포인트가 마스크 내부에 위치하기 때문에 경계 모호성의 영향을 최소화한다.
2. 프리트레인 프리 어댑터: ViT‑P는 사전 학습된 ViT(DeiT, BEiT, CLIP 등)의 가중치를 그대로 재사용한다. 포인트 임베딩 레이어와 분류 헤드만 새로 학습하면 되므로, 대규모 라벨링 데이터가 부족한 상황에서도 빠르게 적용 가능하다.
3. 다중 어노테이션 전략: Fine annotation(정밀 마스크), Coarse annotation(거친 마스크), Box annotation(바운딩 박스) 세 종류를 조합한다. 특히 박스는 사전 학습 단계에서만 사용해 포인트 샘플링 없이도 객체 존재 여부를 학습시켜, 전체 파이프라인의 어노테이션 비용을 크게 낮춘다.
4. 범용성: 동일 모델이 semantic, instance, panoptic 세 가지 세그멘테이션 태스크에 동일하게 적용된다. 마스크 제안기만 교체하면 다른 데이터셋·태스크에 바로 전이 가능하다.

실험 결과는 주장을 설득력 있게 뒷받침한다. ADE20K 파노프틱 세그멘테이션에서 54.0 PQ, Cityscapes semantic segmentation에서 87.4 mIoU, COCO-Stuff-164K에서 53.7 mIoU를 기록했으며, 이는 기존 최첨단 모델 대비 각각 0.71.3%p, 0.4%p, 0.5%p 정도의 향상이다. 특히 “upper bound” 실험에서 마스크를 완벽히 라벨링했을 때의 성능과 실제 ViT‑P 성능 차이가 매우 좁아, 포인트 기반 분류가 병목을 거의 제거했음을 보여준다. 또한, 박스와 거친 라벨만 사용했을 때도 fine annotation 대비 12%p 수준의 성능 저하에 그쳐, 비용 효율적인 라벨링 전략의 가능성을 입증한다.

한계점으로는 포인트 선택이 마스크 중심에 국한된다는 가정이 있다. 매우 비대칭하거나 얇은 객체에서는 최고값 포인트가 경계에 가까워질 수 있어 분류 오류가 발생할 가능성이 있다. 또한, 마스크 제안기의 품질에 크게 의존하므로, 제안기가 낮은 Recall을 보일 경우 전체 성능이 제한된다. 향후 연구에서는 다중 포인트 샘플링, 포인트 위치 예측을 위한 별도 모듈, 그리고 제안기와 분류기 간의 공동 최적화를 탐색할 여지가 있다.