Mamba 아키텍처는 시퀀스 모델링 분야에서 트랜스포머(Transformer)의 한계를 극복하기 위해 개발된 새로운 신경망 구조입니다. Mamba는 긴 시퀀스를 효율적으로 처리하고, 연산 비용을 선형적으로 증가시키며, 메모리 사용과 병렬 처리 기능을 최적화하는 것을 목표로 합니다.
Mamba 아키텍처
Mamba 아키텍처는 긴 데이터 시퀀스를 처리하기 위해 구조화된 상태 공간 시퀀스 모델(Structured State Space sequence model, S4)을 통합합니다.
S4는 연속 시간, 순환(Recurrent), 컨볼루션 모델을 결합하여 긴 의존성을 효과적으로 모델링합니다. 이를 통해 불규칙하게 샘플링된 데이터와 무한한 문맥(unbounded context)을 처리할 수 있으며, 학습과 추론 과정에서도 높은 계산 효율성을 보여줍니다.
Mamba 아키텍처의 주요 개선점
1. S4의 강화
Mamba는 S4를 기반으로 하되, 특히 시간 가변(time-variant) 연산 처리에서 획기적인 개선을 이루었습니다.
고유한 선택 메커니즘을 도입하여 입력값에 따라 Structured State Space Model(SSM)의 매개변수를 동적으로 조정합니다. 이러한 메커니즘으로 Mamba는 시퀀스 내에서 관련성이 높은 정보를 선택적으로 포착하고 불필요한 데이터를 효과적으로 걸러냅니다.
더불어 시간 불변(time-invariant) 프레임워크에서 시간 가변(time-varying) 프레임워크로 발전하여 계산 효율성과 모델 성능을 모두 향상시켰습니다.
2. GPU 하드웨어 최적화 알고리즘
커널 결합(Kernel Fusion), 병렬 스캔(Parallel Scan), 재계산(Recomputation) 기법을 사용하여 성능을 최적화합니다. 메모리 집약적인 레이어에서 확장된 상태를 메모리에 물리적으로 저장하지 않아, 성능과 메모리 사용량 모두를 개선합니다.
이 결과, Mamba는 긴 시퀀스를 처리하는 데 있어 트랜스포머(Transformer)보다 훨씬 효율적입니다.
3. 간소화된 아키텍처
Mamba는 SSM 설계를 MLP 블록과 통합하여 단순하고 균일한 구조를 제공합니다. 이를 통해 언어, 오디오, Genomics 등 다양한 데이터 유형에 대해 일반적인 시퀀스 모델링을 수행할 수 있습니다. 간소화된 아키텍처는 학습과 추론 모두에서 효율성을 유지하며, 모델의 성능을 극대화하는 핵심이 됩니다.
Mamba 아키텍처에 대한 생각
코넬대학교의 Alexander Rush 교수(Hugging Face 연구원)는 SSM 관련 강의에서 "트랜스포머의 대체 가능성"에 대해 논의하며, "현재로서는 가능하나, 10년 후의 상황은 불확실하다"라는 견해를 제시했습니다.
트랜스포머는 현재 데이터 처리 및 학습 능력에서 타의 추종을 불허하는 성능을 보여주고 있으며, 이는 다른 아키텍처가 아직 도달하지 못한 수준입니다.
다수의 AI 반도체 기업들이 트랜스포머 아키텍처의 연산 최적화에 주력하고 있습니다. 그러나 향후 10년간의 기술 발전 방향은 예측하기 어렵습니다.
이러한 맥락에서 Mamba의 SSM은 주목할 만한 대안으로 부상하고 있으며, 이에 대한 심층적인 이해를 위해 관련 논문과 코드 분석을 권장드립니다.
정리
Mamba의 뼈대를 이루고 있는 아키텍처는 State Space Model입니다. State Space Model(SSM)은 제어이론에서 사용되는 상태공간방정식과 동일한 개념입니다. 제어이론의 상태공간방정식은 연속형 변수를 가정하여 사용하지만, 딥러닝에서는 모든 변수가 이산형이므로 이산화 과정이 필요합니다. 이산화 방식에는 오일러 방식, ZOH 방식, 이중선형방식의 세 가지가 있으며, Mamba는 이 중 ZOH 방식을 채택하고 있습니다.
또한 Mamba의 코드와 사전 훈련된 모델은 오픈 소스로 공개되어 있어 연구자와 개발자들이 쉽게 활용할 수 있습니다. 또한, 다양한 신경망 아키텍처에 통합하여 시퀀스 모델링의 백본으로 활용될 수 있습니다.
Mamba 아키텍처는 시퀀스 모델링의 새로운 패러다임을 제시하며, 트랜스포머의 한계를 보완하는 강력한 대안으로 주목받고 있습니다.
* 참고문헌
1. https://openreview.net/forum?id=AL1fq05o7H
2. https://modulabs.co.kr/blog/introducing-mamba
3. https://en.wikipedia.org/wiki/Mamba_(deep_learning_architecture)