VLM에 비전 트랜스포머가 꼭 필요한가? 비전 인코더로서의 상태 공간 모델 평가

대부분의 시각 언어 모델(VLM)이 비전 트랜스포머(ViT)에 의존하는 상황에서, 상태 공간 모델(SSM)이 더 효율적이고 정교한 시각 인코더가 될 수 있음을 입증했다. 특히 객체의 위치를 찾는 Localization 작업에서 압도적인 성능을 보여주며, 모델의 크기나 사전 학습 정확도가 VLM의 최종 성능과 항상 일치하지 않는다는 중요한 통찰을 제공한다.

비전 트랜스포머(ViT)는 이미지 패치 간의 관계를 계산할 때 Self-Attention을 사용하는데, 이는 모든 패치를 동등하게 처리하므로 공간적 구조를 유지하기 위해 위치 인코딩(Positional Encoding)이라는 추가 정보에 의존해야 한다. 이 과정에서 모델이 깊어질수록 세부적인 공간 정보가 희석되는 경향이 있다.

반면 상태 공간 모델(SSM)인 VMamba는 이미지를 여러 방향으로 스캔하며 상태를 업데이트하는 방식을 취한다. 이는 모델 구조 자체에 인접한 픽셀 간의 관계를 중시하는 '공간적 귀납적 편향'을 내장하게 만든다. 마치 사람이 눈으로 이미지를 훑으며 정보를 수집하듯, SSM은 데이터의 순서를 통해 자연스럽게 공간적 맥락을 파악한다.

결과적으로 SSM은 ViT보다 훨씬 적은 파라미터로도 이미지 내 객체의 정확한 위치와 경계를 더 잘 기억하며, 이는 VLM이 '저 사과가 어디에 있니?'와 같은 질문에 답할 때 훨씬 더 정확한 시각적 근거를 제시할 수 있게 한다.

LLaVA 스타일의 VLM 아키텍처를 기반으로 비전 인코더, 경량 커넥터, Vicuna-7B 언어 모델을 결합했다. 비전 인코더는 고정(frozen)된 상태로 유지하고 커넥터와 언어 모델만 인스트럭션 튜닝을 진행하여 인코더 아키텍처 자체의 효과를 격리했다. 비전 인코더로는 ViT, MaxViT, MambaVision, VMamba를 사용했으며, 모든 모델은 ImageNet-1K로 사전 학습된 체크포인트를 사용했다.

계층형 아키텍처인 VMamba와 MaxViT의 경우, ViT/16과 동일한 196개의 시각적 토큰을 생성하기 위해 스테이지 3(Stage 3)의 특징 맵을 추출했다. [입력 이미지 $x_{img}$를 비전 인코더 $V_{\omega}$에 입력으로] -> [특징 추출 연산을 수행해 시각적 토큰 $f_{img}$를 얻고] -> [이를 2개 층의 MLP 커넥터 $C_{\psi}$에 통과시켜] -> [언어 모델의 임베딩 공간에 매핑된 벡터 $e_{img}$를 생성한다]. 커넥터 연산은 $C_{\psi}(x) = W_2 \text{GELU}(W_1 x + b_1) + b_2$ 구조를 따르며, 시각 정보를 텍스트와 결합 가능한 형태로 변환한다.

학습은 665K개의 멀티모달 인스트럭션 데이터를 사용했으며, 4개의 NVIDIA H200 GPU에서 FSDP(Fully Sharded Data Parallel)를 통해 효율적으로 진행했다. 이미지 전처리 시에는 원본 비율을 유지하는 레터박스 리사이징을 적용하여 기하학적 왜곡을 최소화했다.

ImageNet-1K 기반의 동일 조건 비교에서 VMamba-T(30M)는 종합 성능 59.00점을 기록하여 ViT-S(22M, 51.95점)와 ViT-B(87M, 51.36점)를 압도했다. 특히 Localization 성능을 측정하는 RefCOCO 벤치마크에서 VMamba-T는 58.25점을 기록해 ViT-S(32.32점)보다 2배 가까이 높은 정확도를 보였다.

객체 탐지(COCO)나 세그멘테이션(ADE20K) 데이터로 추가 학습된 인코더를 사용했을 때 전반적인 VLM 성능이 향상되었으나, 일부 고해상도 설정에서 성능이 급락하는 'Localization Collapse'가 관찰되었다. 예를 들어 ViTDet-L은 기본 설정에서 13.05점의 낮은 Localization 점수를 기록했다.

제안된 안정화 기법인 3층 MLP 커넥터와 512x512 정사각형 입력을 적용한 결과, VMamba-B의 종합 성능은 52.72점에서 61.34점으로, ViTDet-L은 51.65점에서 61.00점으로 크게 복구되었다. 이는 VLM의 성능이 인코더 자체뿐만 아니라 인코더와 언어 모델 사이의 인터페이스 설계에 매우 민감함을 시사한다.

VMamba는 2D Selective-Scan(SS2D) 메커니즘을 통해 이미지의 공간 정보를 처리한다. 이는 이미지를 네 가지 방향(가로, 세로 및 각각의 역방향)으로 스캔하여 1차원 SSM 연산을 수행함으로써, Transformer의 전역 Self-Attention이 갖는 $O(N^2)$ 복잡도 문제를 해결하면서도 인접 픽셀 간의 강력한 상관관계를 유지한다.

연구팀은 'Localization Collapse'의 원인을 두 가지 가설로 분석했다. 첫째는 커넥터의 용량이 부족하여 공간 정보를 LLM으로 전달하지 못하는 '전송 병목(Transmission Bottleneck)'이며, 둘째는 LLM이 비정형(non-square) 입력의 공간 단서를 해석하지 못하는 '활용 병목(Utilization Bottleneck)'이다. 실험을 통해 두 요인이 복합적으로 작용함을 확인하고, 커넥터 층을 늘리고 입력 해상도를 정사각형으로 고정하는 해결책을 제시했다.

또한, 계층형 모델에서 특징을 추출할 때 가장 마지막인 Stage 4보다 Stage 3의 특징을 사용하는 것이 VLM 성능에 더 유리함을 발견했다. Stage 4는 시각적 추상화 수준은 높지만 공간적 세부 정보가 많이 소실된 상태인 반면, Stage 3는 적절한 추상화와 풍부한 공간 정보를 동시에 보유하고 있어 언어 모델과의 결합에 더 적합하기 때문이다.

본 연구는 주로 ImageNet-1K로 사전 학습된 체크포인트를 사용했으며, CLIP이나 DINOv2와 같은 대규모 대조 학습(Contrastive Learning) 또는 자가 지도 학습 기반의 SSM 모델이 부족하여 해당 영역에서의 직접적인 비교는 수행하지 못했다.