본문 바로가기

인공지능

추론 성능 중심으로 진화한 상태공간모델, Mamba-3 기술 정리

728x90
반응형
728x170

이 글에서는 최신 상태공간모델(State Space Model, SSM) 아키텍처인 Mamba-3를 중심으로, 왜 이 모델이 등장했는지, 기존 Mamba-2와 무엇이 달라졌는지, 그리고 실제 추론(inference) 환경에서 어떤 성능적 의미를 가지는지를 정리합니다.
특히 최근 LLM 환경에서 학습보다 추론과 배포가 더 중요해진 흐름 속에서, Mamba-3가 어떤 문제의식을 가지고 설계되었는지와 그 기술적 선택이 어떤 결과로 이어졌는지를 독자가 이해하기 쉽게 설명하는 것이 목적입니다.

반응형

추론이 병목이 된 LLM 환경의 변화

Mamba-2가 등장한 2024년 중반 이후, 많은 SSM 기반 모델들은 학습 효율(training efficiency) 을 최우선 목표로 설계되었습니다.
Mamba-2는 상태 전이 행렬을 단순화하고 계산 구조를 경량화하여, 이전 세대 대비 2~8배 빠른 학습 속도를 달성했고 이는 빠른 확산으로 이어졌습니다.

하지만 이후 LLM 환경은 분명히 달라졌습니다.

  • 사전학습 이후의 후처리(post-training) 비중 증가
  • RLVR 기반 학습으로 인한 대규모 생성 롤아웃
  • Codex, Claude Code 같은 에이전트형 워크플로우 확산

이 모든 흐름의 공통점은 압도적으로 많은 추론 연산을 요구한다는 점입니다.
문제는 기존 선형 모델들이 여전히 “학습 우선” 관점에서 설계되었다는 데 있습니다.


기존 SSM의 한계: 너무 단순해진 추론 단계

SSM의 가장 큰 장점은 시퀀스 길이에 선형으로 증가하는 계산량입니다. 이는 고정 크기의 상태(state)를 유지하기 때문에 가능합니다.
하지만 이 구조는 동시에 치명적인 제약을 가집니다.

  • 모든 과거 정보를 하나의 고정 상태로 압축
  • Transformer처럼 KV 캐시로 정보를 누적 저장할 수 없음

기존 Mamba 계열은 학습을 빠르게 만들기 위해

  • recurrence를 단순화하고
  • 전이 행렬을 축소

그 결과, 디코딩 시 계산량이 너무 적어 GPU가 메모리 이동에만 묶이는(memory-bound) 현상이 발생했습니다.
즉, 하드웨어를 충분히 활용하지 못하는 구조가 된 것입니다.


Mamba-3의 핵심 질문

“추론을 최우선으로 설계한 SSM은 어떤 모습이어야 할까?”

Mamba-3는 이 질문에서 출발합니다.
고정된 상태 크기를 유지하면서도, 그 상태가 더 많은 일을 하도록 만드는 것이 핵심 목표입니다.

이를 위해 연구팀은 세 가지 레버를 선택했습니다.

  1. recurrence 자체를 더 표현력 있게 만들 것
  2. 상태 전이를 더 풍부하게 만들 것
  3. 한 스텝 안에서 병렬 연산을 늘릴 것

Mamba-3의 3가지 핵심 기술 변화

1. 더 일반화된 recurrence 설계

Mamba-3는 exponential–trapezoidal discretization 기반의 새로운 recurrence를 도입했습니다.
이 방식은 기존보다 훨씬 풍부한 동역학(dynamics)을 표현할 수 있으며, 디코딩 시에도 연산 밀도를 높여 GPU 활용도를 개선합니다.

2. 복소수 기반 상태 추적 (Complex-valued SSM)

상태를 실수(real)가 아닌 복소수(complex) 로 모델링함으로써,
단일 상태 안에서 더 많은 정보를 표현할 수 있게 되었습니다.

아키텍처 상에서는 이를 RoPE 기반 회전 표현으로 구현하여, 커널을 새로 작성하지 않고도 효율적인 계산을 유지합니다.

3. SISO에서 MIMO로의 확장

기존 Mamba는 단일 입력-단일 출력(SISO) SSM 구조였습니다.
Mamba-3는 이를 MIMO (Multi-Input Multi-Output) 로 확장하여, 여러 SSM을 병렬로 실행합니다.

중요한 점은 다음입니다.

  • 학습 비용은 증가
  • 추론(latency)은 거의 증가하지 않음
  • 정확도는 유의미하게 상승

이는 추론 단계가 계산보다 메모리 병목이 크다는 점을 적극 활용한 설계입니다.


아키텍처 관점에서 달라진 점

QKNorm(BCNorm) 도입

Mamba-3는 QKNorm을 추가해 학습 안정성을 개선했습니다.
이는 Transformer 및 Gated DeltaNet 계열과 아키텍처 정렬을 맞추는 역할도 합니다.

Short Convolution 제거

기존 Mamba-1/2의 핵심 요소였던 짧은 causal convolution 을 제거했습니다.

  • 새로운 recurrence 자체가 convolution과 유사한 효과를 내며
  • 성능 개선 효과는 없고 오히려 약간 저하되는 경우가 있었기 때문입니다

실험 결과, 실제 환경의 retrieval 성능에는 영향을 주지 않는 것으로 나타났습니다.

MLP와 Transformer 스타일 구조 채택

전체 아키텍처는 이제 Transformer와 유사하게

  • SSM 레이어와 MLP 레이어를 교차(interleaved) 배치하는 구조를 사용합니다.

언어 모델링 성능 결과

실험 결과는 명확합니다.

  • Mamba-3 SISO
    • 동일한 구조 크기에서 Mamba-2, Gated DeltaNet 대비 더 높은 정확도
  • Mamba-3 MIMO
    • 1B 스케일에서 1%p 이상의 추가 성능 향상
    • 디코딩 속도는 거의 동일

즉, 추론 속도를 유지하면서 품질을 끌어올린 구조라는 점이 핵심입니다.


추론 지연 시간(latency)에서의 의미

1.5B 모델, H100 GPU 기준 실험에서

  • Prefill + Decode 전체 구간
    • Mamba-3 SISO는 모든 시퀀스 길이에서 가장 빠른 성능
    • Mamba-2, Gated DeltaNet, Transformer(vLLM) 모두 상회

특히 Transformer는 긴 시퀀스에서 KV 캐시로 인해 지연 시간이 급증하는 반면,
Mamba-3는 선형 스케일 특성을 그대로 유지합니다.


커널 구현과 하드웨어 친화적 설계

Mamba-3는 커널까지 함께 오픈소스로 공개되었습니다.

  • Triton: 기본 SSM 커널, 플랫폼 독립성과 성능 균형
  • TileLang: MIMO prefill 단계에서 메모리 재사용 최적화
  • CuTe DSL: decode 단계에서 CUDA 수준의 세밀한 제어

이 조합은 단순히 구현 기술의 선택이 아니라,
Mamba-3의 알고리즘 설계가 충분히 단순하고 명확했기 때문에 가능했습니다.


728x90

Mamba-3는 단순한 후속 모델이 아닙니다.
이는 “추론 중심 시대의 SSM은 어떻게 설계되어야 하는가” 에 대한 하나의 답변입니다.

  • 학습 속도를 일부 포기하더라도
  • 추론 효율과 모델 품질의 경계를 넓히는 방향

앞으로는 순수 Transformer, 순수 선형 모델보다는
선형 레이어와 self-attention을 결합한 하이브리드 구조가 주류가 될 가능성이 큽니다.

Mamba-3는 그 전환점에서,
“더 좋은 모델을 더 빠르게 실행한다”는 목표에 가장 가까이 다가간 사례라고 볼 수 있습니다.

300x250

https://www.together.ai/blog/mamba-3?fbclid=IwY2xjawQ3CEFleHRuA2FlbQIxMABicmlkETE2VTFEa0ROOVp6eTZkQ0Vkc3J0YwZhcHBfaWQQMjIyMDM5MTc4ODIwMDg5MgABHoiaGDkLkkkc_ftr2bRCZsQPl9a2w1-lqR8f-TkAY1iN62f7o-rApeJqVAJp_aem_M80oiJ9JcoiTYYW0gjfHfg

 

Mamba-3

Meet Mamba-3: the SSM built for inference. Faster than Transformers at decode, stronger than Mamba-2, and open-source from day one.

www.together.ai

728x90
반응형
그리드형