기술 요약: 마르코프에서 라플라스까지: Mamba가 마르코프 체인을 인컨텍스트 학습하는 방법

문제 정의

Transformer 기반 모델들이 최근 AI의 발전을 주도해 왔으나, 이들의 계산 복잡도(시퀀스 길이에 따른 이차적 스케일링)는 Mamba 및 Mamba-2와 같은 선택적 SSM(Selective SSM)을 포함한 구조적 상태 공간 모델(SSM)과 같은 효율적인 대안에 대한 관심을 불러일으켰습니다. 이러한 모델들의 경험적 성공과 추론 속도 향상에도 불구하고, Mamba의 근본적인 학습 능력은 여전히 제대로 이해되지 않고 있습니다. 구체적으로, Transformer의 ICL(In-Context Learning)을 설명하는 점증하는 연구들(예: 유도 헤드/induction heads)과 비교하여 Mamba의 ICL 메커니즘에 대한 이론적 이해가 부족한 상황입니다.

본 연구가 다루는 핵심 질문은 다음과 같습니다: Mamba의 ICL 능력을 체계적으로 규명할 수 있는가? 저자들은 Transformer의 ICL 메커니즘을 밝혀내는 데 사용되었던 프레임워크인 무작위 마르코프 체인을 통해 Mamba의 성능을 조사함으로써 이 질문에 답하고자 합니다.

방법론

본 논문은 입력 토큰이 $k$차 무작위 마르코프 체인에서 추출되는 Markov-ICL 프레임워크를 채택합니다.

1. 데이터 생성: 전이 커널(Transition kernels)은 디리클레 사전(Dirichlet prior)으로부터 독립적으로 샘플링되어, 모든 테스트 시퀀스가 고유한 마르코프 분포를 따르도록 보장합니다. 모델은 명시적인 파라미터 업데이트 없이 인컨텍스트로 다음 토큰 분포를 추정해야 합니다.
2. 최적 베이스라인: 이 작업에 대한 이론적 최적 추정기는 라플라스 평활화(Laplacian smoothing) (add-$\beta$) 추정기이며, 이는 베이즈(Bayes) 및 미니맥스(minimax) 최적입니다. 이 추정기는 특정 컨텍스트 내에서의 토큰 발생 횟수를 세고 가산 평활화(additive smoothing)를 적용하는 과정을 필요로 합니다.
3. 아키텍처 분석: 저자들은 Mamba 아키텍처(특히 Mamba-2)의 재귀성(recurrence), 선택성(selectivity), 그리고 컨볼루션(convolution) 메커니즘에 초점을 맞춥니다. 이들은 핵심 학습 메커니즘을 격리하기 위해 게이팅(gating)과 비선형성을 제거하고 임베딩, 컨볼루션 기반 입력 선택성, 선형 레이어만을 유지한 단순화된 변형 모델인 MambaZero를 도입합니다.
4. 이론적 규명: 논문은 MambaZero가 최적의 라플라스 추정기를 표현할 수 있음을 보여주는 구성적 증명(constructive proof)을 제공합니다. 또한, 재귀적 아키텍처가 이러한 추정기를 표현하기 위해 필요한 은닉 차원(hidden dimension)의 하한선을 설정합니다.

주요 기여

1. Mamba는 최적의 라플라스 추정기를 학습한다

저자들은 단일 레이어 Mamba조차 모든 마르코프 차수에 대해 효율적으로 인컨텍스트 라플라스 평활화 추정기를 학습함을 입증했습니다.

• 경험적 증거: 실험 결과, 1레이어 Mamba는 이진 및 다중 상태 마르코프 체인에 대해 최적의 추정기와 일치하는 성능을 보였으며, 이는 1레이어 Transformer(실패함)보다 훨씬 뛰어나고 2레이어 Transformer와 대등한 수준입니다.
• 놀라운 발견: 유도 헤드(counting 메커니즘)를 구현하기 위해 일반적으로 여러 레이어가 필요한 Transformer와 달리, Mamba는 단일 레이어로 이러한 최적의 카운팅 능력을 달 수 있습니다.

2. 컨볼루션의 중추적 역할

아블레이션 연구(ablation studies)를 통해, 본 논문은 게이팅이나 비선형 활성화 함수보다 컨볼루션이 Mamba의 ICL 능력을 가능하게 하는 근본적인 구성 요소임을 식별했습니다.

• 아블레이션 결과: Mamba에서 컨볼루션을 제거하면 여러 레이어를 사용하더라도 마르코프 예측 작업에 실패합니다. 반대로, 컨볼루션만을 포함하는 단순화된 모델(MambaZero)은 해당 작업을 성공적으로 학습합니다.
• 메커니즘: 컨볼루션 윈도우 크기 $w$는 현재 컨텍스트를 바탕으로 전이 $i \to j$를 추적하기 위해 반드시 $w \ge k + 1$ ($k$는 마르코프 차수)을 만족해야 합니다. 컨볼루션은 모델이 적절한 컨텍스트 히스토리에 기반하여 상태(state)를 유지하고 전이 횟수를 셀 수 있게 합니다.

3. 이론적 특성 및 한계

• 표현 용량: 저자들은 MambaZero가 1차 마르코프 프로세스에 대해 라플라스 추정기를 정확히 표현할 수 있음을 보여주는 구성적 증명(Theorem 1)을 제공합니다. 이 메커니즘은 상태 $H_t$를 사용하여 전이 횟수를 누적하며, 여기서 컨볼루션은 상태 업데이트가 올바른 컨텍스트 히스토리에 의존하도록 보장합니다.
• 근본적 한계: Theorem 2는 재귀적 아키텍처의 하한선을 설정합니다. 유한한 비트 정밀도($p$)를 가진 $k$차 마르코프 프로세스의 최적 추정기를 표현하려면, 은닉 차원 $d$와 정밀도 $p$의 곱이 $\Omega(2^k)$로 스케일링되어야 합니다. 이는 3개의 레이어를 사용하여 은닉 차원이 $k$에 선형적으로 증가하는 Transformer와 대조적입니다.
• 선택성(Selectivity): 비마르코프적 "스위칭(switching)" 마르코프 프로세스에서, 본 논문은 Mamba가 스위치 토큰이 발생했을 때 카운트를 리셋하기 위해 선택성 메커도(감쇠 인자 $a_t$)를 활용함을 보여주며, 이를 통해 학습된 메커니즘의 적응성을 입증합니다.

결과

• 마르코프 작업: 단일 레이어 Mamba는 다양한 마르코프 차수와 상태 공간 크기에 대해 최적의 라플라스 추정기와의 손실 격차(loss gap)가 거의 제로에 가깝게 달성되었습니다.
• Transformer와의 비교: 2레이어 Transformer는 이 작업을 해결할 수 있지만, 1레이어 Transformer는 불가능합니다. Mamba는 1레이어로 이를 해결합니다.
• 자연어: WikiText-103 데이터셋에 대한 실험 결과, Mamba-2에서 컨볼루션을 제거하면 퍼플렉시티(perplexity)가 11% 증가하는 반면, Transformer에서 제거할 경우 퍼플렉시티는 2%만 증가했습니다. 이는 컨볼루션이 Transformer보다 Mamba의 복잡한 언어 작업 성능에 더 결정적인 역할을 한다는 것을 시사합니다.
• 스위칭 프로세스: Mamba는 "스위치" 토큰을 만났을 때 상태 감쇠 인자 $a_t$를 0으로 설정하여 전이 카운트를 리셋하는 전략을 정확히 구현하며, 이는 비정상적(non-stationary) 환경에 대한 최적의 베이즈 전략을 모방합니다.

의의 및 주장

본 논문은 Mamba와 최적 통계적 추정기 사이의 첫 번째 공식적 연결을 제공한다고 주장합니다.

• 이론적 통찰: 이는 Mamba의 ICL 능력이 단순히 창발적인 속성이 아니라, 컨볼루션과 재귀를 통해 최적의 카운팅을 구현하는 능력에 구조적으로 근거하고 있음을 밝힙니다.
• 아키텍처적 이해: 본 연구는 Mamba에게 있어 컨볼루션이 컨텍스트 의존적 통계를 학습하는 능력의 주요 동력이며, 이는 어텐션 기반 메커니즘을 사용하는 Transformer와 구별되는 지점임을 명확히 합니다.
• 일반성: 이 발견은 합성 마르코프 데이터를 넘어 자연어 모델링으로 확장되며, Mamba의 성공을 뒷받침하는 메커니즘(컨볼루션 기반 카운팅)이 광범위한 언어 작업에서도 근본적임을 시사합니다.

저자들은 Mamba의 표현 능력을 규명하는 데 집중했음을 언급하며, Mamba가 이러한 추정기로 수렴하는 학습 역학(learning dynamics)은 향후 연구 과제로 남겨두는 신중함을 보였습니다. 또한, 본 연구의 목적이 Transformer와의 직접적인 아키텍처 비교보다는 Mamba의 고유한 능력을 이해하는 데 있음을 강조하면서도, 맥락상 필요한 비교를 수행하였습니다.