Mamba 아키텍처 원리: 선형 RNN으로 이해하는 상태공간모델과 선택적 기억

핵심 메시지

• Mamba는 입력 길이에 대해 n log n 연산만 써서, n²이 드는 트랜스포머보다 적은 계산으로 훨씬 긴 문맥을 다룰 수 있다.
• 난해한 상태공간모델 이론 없이도, Mamba는 이해하기 쉬운 '선형 RNN의 확장'으로 설명할 수 있다.
• 선형 RNN은 병렬 스캔으로 log n 시간에 계산되고, 행렬 대각화로 연산량을 크게 줄인다.
• 가중치를 1 근처로 초기화하면 RNN의 고질병인 기울기 소실·폭발 문제를 완화할 수 있다.
• Mamba의 핵심은 입력에 따라 가중치를 바꿔(선택적으로) 정보를 언제 기억·망각할지 결정하고, 하드웨어 친화적으로 구현한 점이다.

쉽게 이해하기

Mamba는 언어 모델링에서 트랜스포머를 앞선 새 신경망 아키텍처다. 아직 수십억 파라미터 규모까지만 검증됐지만, n log n 연산으로 긴 문맥을 효율적으로 처리할 수 있어 주목받는다. 흔히 상태공간모델(SSM)의 확장으로 소개되지만, 이론이 지나치게 복잡하다. 이 영상은 대신 이해하기 쉬운 순환신경망(RNN) 관점에서 Mamba를 풀어낸다.

RNN은 이전 출력과 현재 입력을 함께 신경망에 넣어, 하나의 층만으로도 앞선 모든 입력의 정보를 담을 수 있다. 하지만 두 가지 치명적 문제가 있다. 각 단계가 이전 단계 결과에 의존해 병렬화가 불가능하고, 기울기 소실·폭발 때문에 학습이 매우 어렵다. 그래서 지난 10년간 RNN은 거의 쓰이지 않았다.

해법은 신경망을 선형 함수로 바꾼 '선형 RNN'이다. 선형 순환은 결합법칙이 성립하는 스캔(scan) 연산으로 볼 수 있어, 누적합처럼 log n 단계의 병렬 알고리즘으로 계산된다. 여기서 매 단계 d³짜리 행렬곱이 생기는 문제는 행렬 대각화로 해결한다. 대각 행렬은 원소별 곱만으로 거듭제곱되므로 연산이 d로 줄어든다.

학습 안정성 문제는 초기화로 다룬다. 선형 RNN의 기울기는 가중치가 1에 가까울 때만 안정적이므로, 가중치 크기를 0.999~1 사이로, 입력 배율은 아주 작은 값으로 초기화해 처음부터 안정 상태에서 학습을 시작한다. 흥미롭게도 상태공간모델(SSM)은 초기화 방식만 다를 뿐 사실상 선형 RNN과 같다.

선형 RNN은 장거리 벤치마크에선 강했지만 언어 모델링은 트랜스포머보다 약했다. 원인은 정보를 '선택적으로 망각'하지 못한다는 점이다. Mamba는 입력마다 다른 가중치를 생성해 언제 저장하고 언제 지울지 스스로 결정하게 했고, 출력 벡터를 16배 키우되 GPU 메모리 전송이 병목이라는 점을 이용해 사실상 공짜로 용량을 늘렸다. 한편 이 논문이 재현 성공에도 ICLR 2024에서 거절되며 동료심사 논쟁을 낳았다.

주요 인사이트

• 선형 RNN은 장거리 아레나(Long Range Arena) 벤치마크에서 트랜스포머를 평균 33% 앞서지만, 정작 언어 모델링에서는 트랜스포머보다 뒤처진다 — 선택적 망각 능력의 부재 때문이다.
• Mamba는 입력마다 다른 가중치를 만들어 정보를 언제 저장하고 지울지 학습함으로써, 선형 RNN이 못했던 언어 모델링 성능을 확보했다.
• Mamba는 출력 벡터를 16배로 키워 더 많은 과거 정보를 저장한다. 실제 병목이 연산이 아니라 GPU 고속 메모리로의 데이터 전송이라, 이 확장이 전체 속도에 거의 영향을 주지 않는다.
• 행렬 대각화 덕분에 매 단계 누적 행렬곱(d³)을 원소별 곱(d)으로 대체해, 병렬 스캔의 계산량을 현실적인 수준으로 낮췄다.
• Mamba 논문은 여러 팀이 결과를 재현했음에도 '언어모델과 무관한 벤치마크 미제출' 같은 부적절한 이유로 거절돼, 학술 동료심사 제도에 대한 논쟁을 불러일으켰다.

자주 묻는 질문

관련 AI 소식