Mamba‑3: 상태공간 원리로 한 단계 앞선 시퀀스 모델링

Mamba‑3은 기존 선형 모델의 한계를 극복하기 위해 세 가지 핵심 개선을 도입한다. (1) 지수‑사다리꼴(discretization) 방식으로 더 풍부한 재현성을 확보하고, (2) 복소수‑값 상태 업데이트를 통해 정교한 상태 추적 능력을 부여하며, (3) 다중입·다중출(MIMO) 구조로 디코딩 시 연산 효율을 높인다. 1.5 B 파라미터 규모에서 Mamba‑3은 평균 정확도를 0.6 %p 끌어올리고, MIMO 변형은 추가로 1.2 %p 상승…

저자: Aakash Lahoti, Kevin Y. Li, Berlin Chen

Mamba‑3: 상태공간 원리로 한 단계 앞선 시퀀스 모델링
본 논문은 대규모 언어 모델(Large Language Model, LLM)의 인퍼런스 비용이 성능 향상의 주요 병목이 되고 있음을 출발점으로 삼는다. Transformer 기반 모델은 뛰어난 품질을 제공하지만, KV‑캐시 메모리의 선형 증가와 자기‑어텐션의 2차 연산 복잡도 때문에 추론 시 높은 연산량과 메모리 사용을 요구한다. 이러한 한계를 극복하고자 최근에는 선형 복귀 모델, 특히 상태공간 모델(State Space Model, SSM) 기반 구조가 주목받고 있다. Mamba‑1·2와 Gated DeltaNet(GDN) 등은 상수 메모리와 선형 연산을 통해 효율성을 확보했지만, 모델 품질·특히 장기 의존성 및 상태 추적 능력에서 기존 Transformer에 비해 뒤처지는 경우가 많았다. 또한 이론적으로 선형이라 하더라도 실제 하드웨어에서의 연산 집약도가 낮아 GPU/TPU 활용도가 떨어지는 실용적 문제도 존재한다. 이에 저자들은 “인퍼런스‑퍼스트” 설계 철학을 채택해, 기존 Mamba‑2를 기반으로 세 가지 핵심 개선을 제안한다. 첫 번째는 **지수‑사다리꼴(Exponential‑Trapezoidal) 이산화**이다. 연속‑시간 SSM을 이산화할 때, 기존 Mamba‑1/2는 ZOH(Zero‑Order Hold)와 유사한 근사에 추가적인 “지수‑오일러” 방식을 적용했지만, 이는 LTV(Linear‑Time‑Varying) 시스템에 대한 엄밀한 이론적 근거가 부족했다. 논문은 연속식 ḣ(t)=A(t)h(t)+B(t)x(t) 에 대해 상태전이와 입력 적분을 각각 별도 근사화하고, 입력 적분에 트라페조이달(Trapezoidal) 규칙을 적용함으로써 αₜ, βₜ, γₜ 의 새로운 폐쇄형식을 도출한다. 이 방식은 기존의 단순 지수‑감쇠 마스크 L 을 두 개의 밴드 마스크와 곱 형태로 확장해, 암시적 컨볼루션 효과를 제공한다. 결과적으로 장기 의존성을 더 정확히 모델링하면서도 연산량은 크게 증가하지 않는다. 두 번째는 **복소수‑값 상태공간**이다. 실수형 상태만을 사용하면 비선형적인 상태 추적, 예를 들어 비트열의 짝·홀수 판단, 순환 패턴 인식 등에 한계가 있다. Mamba‑3은 Aₜ, Bₜ, Cₜ 를 복소수로 확장하고, 복소수 회전(복소수 RoPE) 연산을 도입한다. 이 연산은 입력 시퀀스에 위상 정보를 부여해, 데이터‑의존적인 회전 임베딩과 동등한 효과를 제공한다. 복소수 연산은 실수 연산과 동일한 FLOP 수준으로 구현 가능하며, 실험에서는 복소수 SSM이 기존 선형 모델이 무작위 수준에 머물렀던 산술·패턴 추론 과제를 99 % 이상의 정확도로 해결함을 보였다. 세 번째는 **다중입·다중출(MIMO) SSM**이다. 기존 SISO 구조는 상태 업데이트를 외적(outer‑product) 형태로 수행해 FLOP 효율이 낮았다. Mamba‑3은 입력 X∈ℝ^{T×D} 에 대해 B∈ℝ^{T×N}, C∈ℝ^{T×N} 을 행렬 형태로 유지하고, L⊙(C Bᵀ) 연산을 직접 매트릭스 곱으로 수행한다. 이는 신호처리 관점에서 SISO→MIMO 전이와 일치하며, 디코딩 단계에서 메모리‑바운드 상태 업데이트에 추가 연산을 삽입해 GPU/TPU의 연산 집약도를 3‑4배 상승시킨다. 중요한 점은 **지연 시간은 변하지 않는다**는 것으로, 하드웨어 활용도는 크게 개선되면서도 실시간 추론 요구사항을 만족한다. **실험**은 세 가지 축을 모두 검증한다. 1.5 B 파라미터 규모에서 Mamba‑3(SISO)은 GDN보다 +0.6 %p, Mamba‑3(MIMO)은 +1.8 %p 전체 정확도 향상을 보이며, 특히 상태 크기 64 에선 Mamba‑2(128)와 동등한 퍼플렉시티를 달성한다. 합성 상태‑추적 벤치마크에서는 복소수 SSM이 99 % 이상의 정확도를 기록, 기존 선형 모델은 50 % 이하에 머물렀다. 또한 공개된 커스텀 커널을 통해 디코딩 FLOP는 4배 증가했지만, 실제 지연 시간은 Mamba‑2와 거의 동일했다. **결론**적으로, Mamba‑3은 이산화 기법, 복소수 상태, MIMO 구조라는 세 축을 통해 기존 선형 모델의 표현력·상태 추적·하드웨어 효율성을 동시에 끌어올렸다. 이는 차세대 LLM 설계에서 “인퍼런스‑우선” 접근법이 얼마나 강력한지를 실증적으로 보여주는 중요한 사례이며, 성능‑효율 파레토 프론티어를 크게 확장한다는 점에서 향후 연구와 산업 적용에 큰 영향을 미칠 것으로 기대된다.

원본 논문

고화질 논문을 불러오는 중입니다...

댓글 및 학술 토론

Loading comments...

의견 남기기