Ternary Mamba: W1.58A16 상태 공간 모델 (SSM)의 그룹화된 양자화 인식 학습 (QAT)

Mamba-2와 같은 상태 공간 모델 (State Space Models (SSMs))은 선형 시간 추론 (linear-time inference)을 제공하지만, 메모리 점유율로 인해 엣지 배포 (edge deployment)에 한계가 있습니다. 이전의 삼진법 SSM 연구 (Slender-Mamba)는 1,500억 개의 토큰으로 처음부터 학습(from scratch)을 진행했습니다. 본 연구에서는 사전 학습된 체크포인트 (pretrained checkpoint)만으로도 충분하다는 것을 보여주며, 이를 통해 한계 토큰 예산을 1,000배 줄였습니다. 동결된 FP16 교사 모델 (teacher model)로부터의 지식 증류 (knowledge distillation)를 결합한 그룹화된 양자화 인식 학습 (grouped quantization-aware training (QAT))을 사용하여, Mamba-2 1.3B를 3.61배 압축 (2,687MB에서 744MB로)하였으며, 단 1억 2백만 개의 토큰 (4 GPU-시간, 단일 H100)만으로 48.1%의 제로샷 정확도 (7개 태스크 평균)를 달성했습니다. 이는 Bi-Mamba의 48.4%에 근접하는 수치입니다 (+/-0.9pp 신뢰 구간 내). 이러한 사전 학습 모델 기반의 QAT 설정은 학습 가능한 양자화 스케일 (learnable quantization scales)로 인해 발생하는 새로운 불안정성인 제로 비율 붕괴 (zero-ratio collapse)를 드러내며, 이는 처음부터 학습하는 방식에서는 발생하지 않습니다. 나아가 우리는 Transformer에는 효과적이었던 사후 교정 전략 (post-hoc correction strategies)이 재귀 (recurrence)를 통한 오차 누적으로 인해 SSM에서는 실패한다는 것을 보여줍니다. 이러한 결과는 삼진법 SSM이 값비싼 처음부터의 학습을 요구하지 않음을 입증합니다. 즉, 지식 증류 (KD)를 결합한 사전 학습 체크포인트로부터의 QAT가 데이터 효율적인 대안임을 보여줍니다.