이 논문은 대규모 언어 모델(LLM)의 성능에서 추론 시간(inference-time) 컴퓨팅의 중요성이 커짐에 따라, 모델 품질과 함께 추론 효율성을 핵심적인 설계 목표로 삼고 있습니다. 현재의 Transformer 기반 모델은 우수한 품질을 제공하지만, 이차적인 컴퓨팅 요구량과 선형적인 메모리 사용으로 인해 추론 비용이 높습니다. 이를 해결하기 위해 선형적인 컴퓨팅과 상수적인 메모리 요구사항을 가지는 sub-quadratic 모델들이 개발되었지만, 이들은 종종 모델 품질과 특정 기능(예: 상태 추적)에서 약점을 보이며, 이론적으로는 효율적이지만 실제 하드웨어에서는 비효율적인 경우가 많습니다.

이러한 한계를 극복하기 위해, 본 논문은 상태 공간 모델(SSM)의 관점에서 영감을 받은 세 가지 핵심적인 방법론적 개선 사항을 도입하여 Mamba-3 모델을 제안합니다. 이 개선 사항들은 Mamba-2를 기반으로 하며, 모델 품질, 기능, 그리고 효율성의 발전을 목표로 합니다.

1. Exponential-Trapezoidal Discretization:

State Space Model (SSM)은 연속 시간(continuous-time) 동적 시스템으로 정의되지만, 시퀀스 모델링에서는 데이터가 이산 시간(discrete time) 스텝에서 관찰되므로 SSM의 연속적인 동역학을 이산적인 재귀 관계(recurrence)로 변환하기 위해 이산화(discretization) 단계가 필요합니다. 기존의 이산화 방법들은 주로 선형 시불변(LTI) 시스템에 적용되었으며, Mamba-1과 Mamba-2는 Time-Varying SSM을 위한 휴리스틱한 방법을 사용했습니다. 본 논문은 새로운 이산화 프레임워크를 제안하며, Mamba-1/2의 이산화가 "exponential-Euler" 이산화의 한 형태임을 이론적으로 정립하고, 더욱 강력한 "exponential-trapezoidal" 이산화 방법을 도입합니다.

"exponential-trapezoidal" 방법은 상태-입력(state-input) 적분을 일반화된 사다리꼴 규칙(trapezoidal rule)을 사용하여 근사합니다. 이 방법은 Euler 규칙보다 더 높은 차수의 정확도를 제공하며, 지역 절단 오차(local truncation error)가 $O(\Delta_t^3)$로 줄어듭니다. 이로 인해 얻어지는 재귀 관계는 다음과 같습니다:
$$\mathbf{h}_t = e^{\Delta_t A_t} \mathbf{h}_{t-1} + (1 - \lambda_t)\Delta_t e^{\Delta_t A_t} \mathbf{B}_{t-1}\mathbf{x}_{t-1} + \lambda_t \Delta_t \mathbf{B}_t\mathbf{x}_t$$
이를 간단히 표현하면:
$$\mathbf{h}_t = \alpha_t \mathbf{h}_{t-1} + \beta_t \mathbf{B}_{t-1}\mathbf{x}_{t-1} + \gamma_t \mathbf{B}_t\mathbf{x}_t$$
여기서 $\alpha_t \triangleq e^{\Delta_t A_t}$, $\beta_t \triangleq (1 - \lambda_t)\Delta_t e^{\Delta_t A_t}$, $\gamma_t \triangleq \lambda_t \Delta_t$ 이며, $\lambda_t \in [0, 1]$는 데이터에 의존적인 스칼라입니다. 이 공식은 $\lambda_t = 1$일 때 Mamba-2의 Euler 규칙으로 수렴하고, $\lambda_t = \frac{1}{2}$일 때 고전적인 사다리꼴 규칙으로 수렴하는 일반화된 형태입니다.

이 새로운 이산화 방법은 SSM의 상태-입력에 대해 너비가 2인 데이터 의존적인 (data-dependent) 컨볼루션(convolution)을 적용하는 것과 동일합니다. 이는 Mamba-2의 마스크 $L$이 대각 행렬(diagonal matrix)을 가지는 반면, Mamba-3의 마스크 $L$은 2-band 행렬(two-band matrix)을 포함함으로써 더 표현력이 풍부한 Structured Masked Representation (SSD)을 형성합니다.

2. Complex-Valued State Space Model:

Mamba-2와 같은 기존의 SSM들은 효율성을 위해 상태 전이 행렬(state-transition matrix)을 단순화하여, 실수(real) 값을 가지고 비음수(non-negative) 고유값(eigenvalue)을 갖도록 제한했습니다. 이러한 제한은 모델이 parity와 같은 간단한 상태 추적(state-tracking) 작업을 수행하는 능력을 저하시킨다는 것이 밝혀졌습니다. 이는 실수 고유값이 "회전(rotational)"하는 은닉 상태 동역학을 표현할 수 없기 때문입니다.

Mamba-3는 이러한 한계를 극복하기 위해 복소수(complex-valued) SSM을 도입합니다. 복소수 SSM은 회전 동역학을 표현할 수 있어 상태 추적 능력을 회복합니다. 복소수 SSM은 이산화되었을 때, 상태 차원이 두 배가 된 실수 값 SSM과 동일하다는 것이 증명됩니다 (Proposition 2). 이 등가성은 상태 전이 행렬이 2x2 회전 행렬(rotation matrices)들의 블록-대각 행렬(block-diagonal matrix)로 구성되어 있음을 보여줍니다.

구체적으로, $\mathbf{h}(t) \in \mathbb{C}^{N/2}$인 복소수 SSM은 다음과 같습니다:
$$\dot{\mathbf{h}}(t) = \text{Diag}(A(t) + i\theta(t)) \mathbf{h}(t) + (B(t) + i\hat{B}(t)) x(t)$$
$$y(t) = \text{Re}\left( (C(t) + i\hat{C}(t))^\top \mathbf{h}(t) \right)$$
이 시스템은 exponential-Euler 이산화를 통해 $N$ 차원의 실수 SSM으로 등가 변환될 수 있습니다:
$$\mathbf{h}_t = e^{\Delta_t A_t} R_t \mathbf{h}_{t-1} + \Delta_t \mathbf{B}_t x_t$$
$$y_t = \mathbf{C}_t^\top \mathbf{h}_t$$
여기서 $R_t \triangleq \text{Block}(\{R(\Delta_t \theta_t[i])\}_{i=1}^{N/2}) \in \mathbb{R}^{N \times N}$ 이고, $R(\theta) \triangleq \begin{pmatrix} \cos(\theta) & -\sin(\theta) \\ \sin(\theta) & \cos(\theta) \end{pmatrix}$ 입니다. $\mathbf{B}_t$와 $\mathbf{C}_t$는 복소수 B, C 행렬의 실수 및 허수부를 stack하여 구성됩니다.

이러한 복소수 SSM은 데이터에 의존적인 Rotary Positional Embedding (RoPE)을 SSM의 B, C 구성 요소에 적용하는 것과 동일합니다 (Proposition 3). RoPE는 효율적인 계산 방법을 가지고 있으므로, 복소수 SSM의 구현은 실수 SSM에 비해 계산 오버헤드가 적습니다. 이는 "RoPE trick"이라고 불립니다.

3. Multi-Input, Multi-Output (MIMO) SSM:

기존의 SSM은 주로 단일 입력 단일 출력(SISO) 동역학을 기반으로 합니다. Mamba-3는 디코딩 시 FLOPs 효율성을 개선하기 위해, 외적(outer-product) 기반 상태 업데이트에서 행렬 곱셈(matrix-multiplication) 기반 상태 업데이트로 전환합니다. 이는 신호 처리 관점에서 SISO에서 다중 입력 다중 출력(MIMO)으로의 일반화에 해당합니다.

MIMO SSM은 상태 공간 모델의 입력 차원(D_in)과 출력 차원(D_out)을 증가시켜, 하나의 스칼라 입력 $x(t)$와 스칼라 출력 $y(t)$ 대신 벡터 $x(t) \in \mathbb{R}^{\text{D}_{\text{in}}}$와 $y(t) \in \mathbb{R}^{\text{D}_{\text{out}}}$을 처리할 수 있도록 합니다. 이는 상태 전이 행렬 A를 블록 대각 행렬로 확장하고, B, C 행렬을 각각 $N \times \text{D}_{\text{in}}$ 및 $\text{D}_{\text{out}} \times N$ 크기로 확장함으로써 이루어집니다. 이로 인해 모델링 능력이 향상되며, 고정된 상태 크기(state size)에서 디코딩 시 메모리 바운드(memory-bound) 상태 업데이트에서 더 많은 연산이 가능해져 하드웨어 활용률이 증가합니다. MIMO는 특히 추론(inference)에 적합하며, 상태 크기를 늘리지 않고도 더 많은 계산을 가능하게 하여 속도 저하 없이 성능을 향상시킵니다.

종합적으로, Mamba-3는 이러한 세 가지 핵심 개선 사항을 통해 retrieval, state-tracking, 그리고 다운스트림 언어 모델링(downstream language modeling) 작업 전반에서 상당한 성능 향상을 달성합니다. 1.5B 스케일에서 Mamba-3 (MIMO)는 Mamba-2 대비 1.9%, Gated DeltaNet (GDN) 대비 1.8%의 평균 다운스트림 정확도를 향상시켰습니다. 또한, Mamba-3 (MIMO)는 Mamba-2의 절반 상태 크기(64 vs 128)로도 유사한 perplexity를 달성하여, 동일한 언어 모델링 성능을 더 낮은 지연 시간(latency)으로 제공합니다. Mamba-3의 복소수화된 SSM 상태는 Mamba-2가 해결할 수 없었던 합성 상태 추적 작업을 해결하는 새로운 기능을 제공합니다. Mamba-3 (MIMO)는 고정된 상태 크기에서 Mamba-2 대비 최대 4배의 디코딩 FLOPs를 증가시키면서도 유사한 wall-clock 디코딩 지연 시간을 유지하여 하드웨어 활용도를 높입니다. 이러한 발전은 Mamba-3가 성능-효율성 Pareto frontier를 크게 진전시켰음을 입증합니다.