선형 복잡도와 경쟁력 있는 정확도 2Mamba

초록

본 논문은 최신 선형 어텐션 모델인 Mamba‑2를 분석·단순화하고, 핵심 구성요소를 추출한 뒤 A‑mask 개선과 2차 차원의 히든 스테이트 도입으로 정확도를 크게 향상시킨 2Mamba를 제안한다. 실험 결과 2Mamba는 긴 시퀀스에서도 메모리 효율성을 유지하면서 소프트맥스 어텐션에 근접한 성능을 달성한다.

상세 분석

본 연구는 선형 어텐션이 소프트맥스 어텐션에 비해 표현력이 부족한 문제를 Mamba‑2라는 최신 모델을 기반으로 해결하고자 한다. 먼저 저자들은 Mamba‑2의 복잡한 구현을 세부 요소별로 분해하고, 각 요소가 정확도에 미치는 영향을 체계적인 Ablation 실험을 통해 검증한다. 주요 분석 포인트는 다음과 같다.

1. 핵심 구성요소 식별: Q·K·V 생성 시 사용되는 활성함수(SiLU vs ReLU vs None), A‑mask의 구현 방식(원본 음수‑exp vs softplus), 입력 컨볼루션 윈도우 크기, 잔차 연결(D residual) 및 Z‑gate와 같은 게이트 메커니즘, 정규화 방식(출력 RMSNorm vs 소프트맥스 정규화), 값 이산화 여부 등 총 7가지 요소를 독립적으로 제거·변경하여 성능 변화를 측정하였다.

2. A‑mask 개선: 원본 Mamba‑2는 A‑mask를 음수‑exp 형태로 구현하고 dt와 결합해 시간 감쇠를 적용한다. 저자들은 이를 부드러운 softplus 함수로 대체하고, dt와의 결합을 해제함으로써 마스크가 보다 연속적이고 안정적인 값 범위를 갖게 하였다. 이 변경은 특히 긴 시퀀스에서 메모리 사용량을 감소시키면서도 손실 감소에 기여한다.

3. 고차 히든 스테이트 도입: 기존 선형 어텐션은 1차 RNN 형태로, 히든 스테이트 차원이 (d_H, d_H) 수준이다. 저자들은 Taylor(맥클라우린) 전개를 이용해 소프트맥스의 무한급수를 근사함으로써, 2차 히든 스테이트(차원 (d_H², d_H))를 도입한다. 이는 쿼리·키 내적의 제곱항을 포함해 양의 값만을 생성하므로 정규화가 용이하고, 메모리 요구량은 KV‑cache(2·N·d_H)보다 낮게 유지된다.

4. 정규화 전략: 소프트맥스 정규화는 양의 내적값이 전제되어야 하지만, 원본 Mamba‑2는 RMSNorm을 사용한다. 고차 히든 스테이트와 결합된 경우, ReLU와 같은 비선형을 적용해 내적값을 양수화하고, 이후 소프트맥스 정규화를 적용함으로써 안정적인 학습과 빠른 수렴을 얻었다.

5. 실험 설계 및 결과: 300M 파라미터 규모의 작은 모델과 700M 파라미터 규모의 중간 모델을 각각 Llama‑2 기반으로 재구성하고, 동일한 CC‑MAIN‑2024‑51 데이터셋으로 90K100K 스텝 학습하였다. 테스트 손실 기준으로 2Mamba는 기존 Mamba‑2보다 약 57% 개선되었으며, 소프트맥스 어텐션과의 격차를 0.1~0.2 수준으로 축소했다. 특히 시퀀스 길이 2048에서 메모리 사용량은 KV‑cache 대비 30% 이하로 감소하였다.

6. 한계와 향후 과제: 현재 구현은 KV‑cache 없이도 동작하지만, 실제 추론 시 캐시 활용이 가능한 경우 추가적인 속도 향상이 기대된다. 또한 2차 히든 스테이트의 차원 폭이 d_H²이므로, 매우 큰 헤드 차원을 사용할 경우 메모리 부하가 급증할 수 있다. 저자들은 차원 축소 기법(예: 저랭크 근사)과 더 높은 차수(p>2) 확장의 효율성을 탐구할 필요성을 제시한다.

이와 같이 본 논문은 선형 어텐션의 구조적 약점을 체계적으로 분석하고, A‑mask와 고차 히든 스테이트라는 두 가지 핵심 개선을 통해 소프트맥스 어텐션에 근접한 정확도를 달성하면서도 선형 복잡도와 메모리 효율성을 유지한다는 점에서 의미가 크다.