Zamba2-VL Technical Report

문제 정의

시각-언어 모델(VLM)은 일반적으로 트랜스포머 기반의 대규모 언어 모델(LLM) 백본에 의존하며, 이는 멀티모달 추론을 위한 표준 인터페이스가 되었습니다. 그러나 트랜스포머의 셀프 어텐션(self-attention) 메커니즘은 프리필(prefill) 단계에서 이차 함수적(quadratic) 계산 비용을 발생시키며, 시퀀스 길이에 따라 선형적으로 증가하는 키-값(KV) 캐시를 필요로 합니다. 이는 단일 고해상도 이미지나 비디오 클립이 수천에서 수십만 개의 토큰을 생성할 수 있는 멀티모달 모델에서 심각한 병목 현상이 됩니다. 시각적 토큰 압축과 같은 입력 측 완화 기법이 존재하지만, 이는 백본의 근본적인 스케일링 동작을 변화시키지는 못합니다.

상태 공간 모델(SSM), 특히 Mamba 및 Mamba2는 선형 시간 시퀀스 처리와 일정한 크기의 재귀적 상태를 제공하여 이론적으로 이러한 효율성 제약을 해결합니다. 그러나 기존의 SSM 기반 VLM은 어텐션 메커니즘이 여전히 우월한 정밀한 인컨텍스트 룩업(예: 시각적 그라운딩 및 미세 검색) 작업에서 어려움을 겪어왔습니다. 또한, 기존의 강력한 SSM 기반 VLM은 종종 트랜스포머 모델로부터 증류(distilled)되었거나 더 작은 SSM 백본을 기반으로 구축되었기에, 하이브리드 SSM-트랜스포머 아키텍처가 처음부터 대규모로 학습되었을 때 경쟁력이 있을지에 대한 의문이 남아 있었습니다.

방법론

저자들은 Zamba2 하이브리드 아키텍처를 기반으로 구축된 VLM 제품군인 Zamba2-VL을 제시합니다. Zamba2는 Mamba2 상태 공간 백본과 소수의 공유 트랜스포머 블록을 교차 배치합니다. 트랜스포머 블록은 인컨텍스트 검색 능력을 제공하기 위해 LoRA 프로젝터로 보강되었으며, Mamba2 레이어는 선형 시간 복잡도와 고정된 재귀적 상태를 통해 계산의 대부분을 처리합니다.

아키텍처:

• 백본: Zamba2 (1.2B, 2.7B, 7B 파라미터 변체).
• 비전 인코더: 2D 회전 위치 임베딩(RoPE)과 네이티브 동적 해상도 처리를 사용하여 미세한 공간 구조를 보존하는 Qwen2.5-VL의 Vision Transformer(ViT)를 선택했습니다.
• 커넥터: 시각적 특징을 언어 임베딩 공간으로 투영하는 표준 2계층 MLP 어댑터입니다. 이는 패치 임베딩의 2x2 윈도우를 풀링하여 시각적 토큰을 4분의 1로 줄입니다.
• 위치 임베딩: 과도한 계산 요구량을 피하기 위해 멀티모달 RoPE 변형을 채택하는 대신, 기본 Zamba2 모델의 특정 위치 임베딩 방식(1.2B/7B는 1D RoPE, 2.7B는 회전 임베딩 없음)을 유지합니다.

학습 파이프라인:
학습은 3단계 커리큘럼을 따릅니다:

1. 정렬(Alignment): 백본과 비전 인코더는 동결된 상태에서 저해상도 이미지-캡션 데이터(LLaVA-ReCap-558K)를 사용하여 MLP 어댑터를 학습시킵니다.
2. 사전 학습(Pretraining): 전체 모델은 약 30B 토큰의 이질적인 혼합 데이터로 학습됩니다. 이 단계에서는 순수 트랜스포머에 비해 Zamba2 백본이 텍스트 중심 작업에서 보이는 초기 지연을 보완하기 위해 문서 이해 및 OCR 데이터를 의도적으로 업샘플링합니다.
3. 지시어 튜닝(Instruction Tuning): 다회차 대화, 그라운딩(포인팅 및 바운딩 박스), 추론을 강조하는 약 20B 토큰의 정제된 지시어 준수 데이터로 모델을 미세 조정합니다.

주요 기여

1. 최초의 오픈 하이브리드 SSM-트랜스포머 VLM: Zamba2-VL은 강력한 하이브리드 SSM-트랜스포머 백본을 기반으로 구축된 최초의 오픈 VLM 제품군으로 제시되며, 이러한 아키텍처가 1.2B에서 7B 규모의 멀티모달 모델링을 위한 실행 가능한 경쟁력 있는 기질임을 입증합니다.
2. 아키텍처 및 데이터 설계: 유능한 VLM으로 하이브리드 백본을 확장하기 위해 필요한 구체적인 아키텍처 선택(커넥터 설계, 비전 토큰 통합) 및 데이터 구성 전략(특히 문서/OCR 데이터의 업샘플링)을 상세히 설명합니다.
3. 효율성-정확도 트레이드오프: 본 연구는 이 모델들의 성능-효율성 프로필을 규명하며, 이들이 Zamba2 백본의 추론 속도 이점을 상속받으면서도 선도적인 트랜스포머 기반 VLM들과 경쟁력 있는 정확도를 달성함을 보여줍니다.

결과

저자들은 이미지 이해, 추론, OCR, 그라운딩 및 카운팅 벤치마크 전반에 걸쳐 Zamba2-VL(1.2B, 2.7B, 7B)을 선도적인 오픈 웨이트 트랜스포머 기반 VLM(Molmo2, Qwen3-VL, InternVL3.5, PerceptionLM)과 평가했습니다.

• 정확도: Zamba2-VL은 유사한 규모의 트랜스포머 베이스라인들과 전반적으로 경쟁 가능한 수준입니다.
  • 카운팅(Counting): 모델들은 시각적 카운팅에서 탁월한 성능을 보입니다. Zamba2-VL-1.2B는 PixMoCount에서 62.5를 기록하여 InternVL3.5-1B의 점수인 32.8보다 거의 두 배 높았습니다. 7B 모델은 85.3을 달성하여 해당 구간에서 최상위권을 차지했습니다.
  • 문서/OCR: 모델들은 DocVQA, ChartQA, TextVQA에서 경쟁력을 갖추었으며, 이는 사전 학습 중 문서 데이터의 의도적인 업샘플링 덕분인 것으로 분석됩니다.
  • 추론: 모델들은 무거운 추론 벤치마크(예: MMMU, MathVista)에서 강력한 트랜스포머 베이스라인보다 약간 뒤처지는데, 저자들은 이를 RL 스타일의 추론 특화보다는 능력 습득에 집중된 학습 파이프라인 때문이라고 설명합니다.
• 효율성: 가장 중요한 발견은 추론 속도입니다. Zamba2-VL은 동일한 파라미터 규모의 트랜스포머 베이스라인과 비교했을 때 약 **한 자릿수 낮은 첫 토큰 시간(TTFT)**을 제공합니다. 이 효율성 격차는 온디바이스 및 엣지 배포에 중요한 1.2B 및 2.7B 규모에서 가장 두드러집니다.
• 절제 연구(Ablation Studies):
  • 비전 인코더: Qwen2.5-VL ViT는 2D RoPE와 동적 해상도 처리 덕분에 CLIP 및 SigLIP-2보다 크게 우수한 성능을 보였습니다.
  • 백본: 동일한 학습 조건 하의 Llama-3.x 참조 모델과 비교했을 때, 하이브리드 Zamba2 백본은 대부분의 이미지 이해 작업에서 트랜스포머와 대등하거나 능가했으나, 텍스트 밀도가 높은 문서 작업에서는 뒤처졌습니다(이는 이후 데이터 업샘플링을 통해 완화되었습니다). Zamba2와 Llama 간의 성능 격차는 규모가 커짐에 따라 크게 좁혀졌으며, 7B 모델은 Llama-3.1-8B 참조 모델과 거의 구별할 수 없을 정도였습니다.

의의

본 논문은 Zamba2-VL이 특히 지연 시간 및 메모리 제약이 있는 환경에서 VLM 백본으로서 하이브리드 SSM-트랜스포머 LLM이 실질적인 대안이 될 수 있음을 보여준다고 주장합니다. 결과는 "하이브리드 레시피"—즉, 계산의 대부분에는 부차적(sub-quadratic) 재귀를 사용하고 소량의 어텐션을 사용하여 회상(recall)을 수행하는 방식—가 대규모에서 밀집(dense) 트랜스포머의 능력을 따라잡으면서도 상당한 추론 효율성 이점을 유지할 수 있음을 시사합니다. 저자들은 서빙 비용과 지연 시간이 제약 조건이 되는 시나리오를 구체적으로 겨냥하여, 연구자와 실무자들이 활용할 수 있는 토대를 제공하기 위해 모델과 코드를 공개합니다.