Mamba: Linear-Time Sequence Modeling with Selective State Spaces
Information
| 일자 | 2026년 06월 01일 |
|---|---|
| 발표자 | 이건 |
Video
Overview
이 발표는 어텐션 메커니즘 없이 선택적 상태 공간 모델(Selective SSM)을 도입하여 기존 트랜스포머의 연산 복잡도 문제를 해결하고, 시퀀스 처리 속도와 성능을 대폭 향상시킨 Mamba 모델의 핵심 설계와 실험 결과를 제시합니다.
기존 모델의 한계 및 문제 정의
- 트랜스포머의 한계: 시퀀스 길이에 대해 $O(N^2)$의 연산 복잡도를 가져 메모리 및 연산 비용이 기하급수적으로 증가
- 기존 S4(SSM)의 한계: 선형 확장성은 뛰어나지만, 고정된 파라미터(LTI 특성)로 인해 텍스트 등 정보 밀도가 높은 이산적 데이터 처리 성능 저하
- 동적 선택 능력 부족: 입력 내용에 따라 필요한 정보를 취사선택하는 능력(Selective Copying, Induction Heads) 구현 불가
Mamba 핵심 설계 (Selective SSM)
- 선택적 메커니즘(Selective Mechanism) 도입: 고정되었던 파라미터($\Delta, B, C$)를 입력($X$)에 의존하는 함수로 변환
- 동적 토큰 제어: 들어오는 각 토큰의 맥락을 파악하여 관련 정보는 기억하고 불필요한 노이즈는 무시하도록 설계
- 간소화된 아키텍처: H3 블록과 Gated MLP를 단순화하여 선택적 SSM(S6) 기반의 단일 블록 구조 적용
하드웨어 인지 알고리즘 (Hardware-aware Algorithm)
- 연산 병목 해결: 파라미터의 동적 변화로 인해 기존 방식의 컨볼루션 병렬화가 불가능해진 문제를 하드웨어 최적화로 극복
- 커널 퓨전(Kernel Fusion): 여러 연산을 하나의 GPU 커널로 합쳐 HBM과 SRAM 간의 메모리 I/O 최소화
- 병렬 연관 스캔(Parallel Associative Scan): 순차적 연산을 스캔 알고리즘을 통해 병렬 처리로 전환
- 재계산(Recomputation): 역전파 시 필요한 중간 상태를 저장하지 않고 재계산하여 메모리 사용량 대폭 절감
모델 성능 및 실험 결과
- 합성 태스크: 훈련 길이보다 훨씬 긴 시퀀스에서도 문맥을 파악하는 과제(Selective Copying 등) 완벽 수행
- 언어 모델링 성능: 어텐션 프리 모델임에도 최신 트랜스포머(Transformer++)와 대등하거나 우수한 퍼플렉시티 달성
- 추론 속도 및 효율성: 유사 크기 트랜스포머 대비 4배 이상 빠른 추론 속도 및 512K 시퀀스에서도 OOM(메모리 부족) 미발생
- 어블레이션 결과: 모델 구조 자체보다 선택적 SSM(S6) 레이어 적용 유무가 성능 향상의 핵심 요인임 입증
한계점 및 향후 활용 과제
- 검증 규모의 한계: 주로 중소형 모델 위주로 실험되어 대규모 분산 학습 환경에 대한 검증 부족
- 하드웨어 의존성: 고성능 유지를 위해 복잡한 하드웨어 최적화가 필수적이라 타 환경으로의 이식성 저하
- 데이터 특성 딜레마: 이산 데이터에 최적화되면서 기존 SSM이 갖는 연속 신호 처리의 장점 손실 가능성 존재
- 응용 확장: 향후 트랜스포머 구조와의 결합(수면 파운데이션 모델의 인코더 활용 등)을 통한 다중 모달 시너지 탐색 필요