Mamba 상태공간 모델과 Transformer 어텐션 [5pt] 제어공학, 선형대수학 및 시퀀스 모델의 연결

▶ Open interactive ANITEX · equations and animations

This accessible text edition preserves the document's narrative. See the interactive edition for typeset equations, diagrams and playback.

본 문서는 Mamba의 선택적 상태공간 모델(Selective State Space Model, SSM)과 Transformer의 스케일드 닷프로덕트 어텐션을 수학적으로 해석한다. 연속시간 상태방정식의 영차유지 이산화, 입력 의존적 상태 갱신, 선택적 스캔, 어텐션의 행렬 차원, 수치 예제, 계산 복잡도 및 두 구조의 관계를 다룬다. 수식은 교육 목적의 단순화된 단일 채널 또는 단일 헤드 예제와 실제 다차원 구조를 구분한다.

기초: 시퀀스와 행렬

입력 시퀀스는 X^T d_model로 둔다. T는 토큰 수, d_model은 각 토큰 벡터의 차원이다. Transformer는 토큰 간 쌍별 유사도를 계산하고, Mamba는 입력을 받아 상태를 순차적으로 갱신한다. 둘 다 학습 가능한 매개변수와 비선형 처리를 갖는 신경망 블록이다.

Mamba: 연속시간 상태공간 모델

제어공학의 출발점

선형 시불변 연속시간 모델은 h(t)=A h(t)+B x(t), y(t)=C h(t)+D x(t) 여기서 h(t)^N은 내부 상태, x(t)^d는 입력이다. A는 상태 전이, B는 입력의 상태 유입, C는 상태의 관측, D는 직접 전달을 뜻한다. 원래의 선형 SSM에서는 A,B,C,D가 시간에 따라 일정하다.

영차유지(ZOH) 이산화의 유도

구간 [t_k,t_k+]에서 입력을 일정하게 유지하면 미분방정식의 해는 h(t_k+)&=e^A h(t_k)+_0^ e^A(-)B x_k\,d, A&=e^A, B=_0^ e^AsB\,ds. 따라서 h_k= A h_k-1+ B x_k, y_k=C h_k+D x_k A가 가역이면 B=A^-1(e^A-I)B이다. 하지만 A가 특이행렬이어도 적분식은 유효하므로 수치 구현에서는 적분식 또는 안정적인 행렬지수 계산을 사용한다.

Mamba의 선택적 상태공간

Mamba의 핵심은 일부 SSM 매개변수가 현재 입력에 의존한다는 것이다. 단순화한 형태는 _k&=softplus(W_ x_k+ b_), B_k&=f_B( x_k), C_k=f_C( x_k), A_k&=(_k A), B_k&=_0^_ke^AsB_k\,ds, h_k&= A_k h_k-1+ B_k x_k, y_k=C_k h_k. 위 식은 개념을 드러내기 위한 형태이다. 실제 구현은 채널별 , 구조화된 A, 입력 투영, 게이트, 로컬 합성곱, 잔차 연결 및 효율적인 이산화 계산을 사용한다. D x_k 직접 경로도 일반적으로 추가된다.

선택적 갱신의 직관

스칼라 상태에서 A=-a (a>0), B=1이면 h_k=e^-a_kh_k-1+1-e^-a_kax_k. _k가 작으면 이전 상태를 상대적으로 더 보존하고, 클수록 이전 상태의 영향이 더 빨리 감소한다. 이는 B_k,C_k의 선택성과 함께 토큰별 정보 보존·갱신을 조절한다. 단, 만으로 모든 기억 선택이 결정되는 것은 아니다.

선택적 스캔의 결합법칙

시간별 갱신을 h_k=a_k h_k-1+b_k로 쓰자. 두 단계의 합성은 (a_2,b_2)(a_1,b_1)=(a_2a_1,\,a_2b_1+b_2). 이 연산은 결합법칙을 만족한다. 따라서 병렬 프리픽스 스캔으로 학습 중 여러 단계의 상태를 효율적으로 계산할 수 있다. 추론 시에는 고정 크기의 상태를 순차 갱신할 수 있다.

Mamba의 계산 블록도

[>=Latex,node distance=9mm,box/.style=draw,rounded corners,align=center,minimum width=29mm,minimum height=10mm,fill=blue!5,small/.style=draw,rounded corners,align=center,minimum width=26mm,minimum height=9mm,fill=green!6] (x) 입력 x_k; (params) _k,B_k,C_k 입력 의존 투영; (conv) 입력 투영 로컬 합성곱; (state) 선택적 상태 갱신 h_k= A_k h_k-1+ B_k x_k; (out) 출력 y_k 및 게이팅; (x.east) -- ++(4mm,0) |- (params.west); (x.east) -- ++(4mm,0) |- (conv.west); (params.east) -| (state.north); (conv.east) -- (state.west); (state.east) -- (out.west); ((state.south)+(0,-7mm)) -- node[right,font=] 이전 상태 h_k-1 (state.south); 참고: 위 그림은 수학적 핵심을 강조한 개념도이며 실제 구현의 모든 투영·게이트 경로를 표시하지 않는다.

상태 감쇠 그래프

[width=.86,height=6cm,xlabel=,ylabel=e^-a,xmin=0,xmax=4,ymin=0,ymax=1.05,grid=both,legend pos=north east] exp(-0.5*x);a=0.5 exp(-x);a=1 exp(-2*x);a=2

Transformer: 어텐션 행렬

Query, Key, Value의 선형 투영

입력 X^T d_model에 대해 Q=XW_Q, K=XW_K, V=XW_V 여기서 Q,K^T d_k, V^T d_v이다. 각 토큰의 Query는 다른 토큰의 Key와 내적되어 관련도 점수를 생성한다.

스케일드 닷프로덕트 어텐션

S&=QK^d_k^T T, P_ij&=(S_ij)_j'=1^T(S_ij'), O&=\!(QK^d_k+M)V. 마스크 M은 허용되지 않는 위치에 매우 큰 음수를 더하는 역할을 한다. 자기회귀 언어모델의 인과 마스크는 미래 토큰을 참조하지 못하게 한다. 수치적으로 안정적인 Softmax는 행별 최댓값을 빼고 지수함수를 계산한다.

왜 d_k

로 나누는가? Query와 Key의 성분이 독립이고 평균 0, 분산 1이라고 가정하면 내적 q^ k의 분산은 대략 d_k이다. d_k로 나누면 점수의 분산이 약 1이 되어 차원이 증가할 때 Softmax가 지나치게 포화되는 현상을 줄인다. 실제 학습된 성분이 반드시 독립인 것은 아니므로 이는 설계 동기를 설명하는 근사이다.

2개 토큰의 계산 예제

Q=K=1&0 0&1, V=2&0 0&4, d_k=2로 놓자. S=121&0 0&10.707&0 0&0.707. 행별 Softmax를 적용하면 p=e^0.707/(e^0.707+1)0.670이며 P0.670&0.330 0.330&0.670. 출력은 O=PV1.340&1.320 0.660&2.680. 첫 번째 출력 토큰은 첫 번째 Value에 0.670, 두 번째 Value에 0.330의 가중치를 부여한다.

어텐션 행렬의 구조도

[>=Latex,node distance=9mm,box/.style=draw,rounded corners,align=center,minimum height=10mm,minimum width=21mm,fill=blue!5] (x) X; (q) Q=XW_Q; (k) K=XW_K; (v) V=XW_V; (s) QK^/d_k; (p) 행별 Softmax; (o) PV; (x.east) -- ++(4mm,0) |- (q.west); (x.east) -- (k.west); (x.east) -- ++(4mm,0) |- (v.west); (q.east) -| (s.north); (k.east) -- (s.west); (s.east) -- (p.west); (p.east) -- (o.west); (v.east) -- ++(13mm,0) -| (o.south);

예제 어텐션 행렬의 히트맵

[scale=1.35] // in 0/1/0.670,1/1/0.330,0/0/0.330,1/0/0.670100* (,) rectangle ++(1,1); (,) rectangle ++(1,1); at (+.5,+.5) ; at (.5,0) k_1; at (1.5,0) k_2; at (0,1.5) q_1; at (0,.5) q_2;

멀티헤드 어텐션

H개 헤드에 서로 다른 투영행렬을 사용한다. head_h&=Attention(XW_Q^(h),XW_K^(h),XW_V^(h)), MHA(X)&=Concat(head_1,,head_H)W_O. 일반적인 Transformer 블록은 멀티헤드 어텐션, 잔차 연결, 정규화, 위치별 MLP(FFN)로 구성된다. 위치 정보는 위치 임베딩 또는 RoPE 등의 방식으로 제공한다.

어텐션의 미분: 역전파의 핵심

O=PV이고 손실 L의 출력 기울기를 G_O= L/ O라 하자. 그러면 G_V&=P^ G_O, G_P=G_OV^, (G_S)_ij&=P_ij, G_Q&=G_SK/d_k, G_K=G_S^ Q/d_k. 이는 마스크된 위치를 제외한 행별 Softmax 미분과 행렬 곱의 연쇄법칙에서 나온다.

Mamba와 Transformer의 수학적 비교

p31mmp57mmp57mm 항목&Mamba 선택적 SSM&Transformer Self-Attention 핵심 수식&h_t= A_t h_t-1+ B_t x_t&O=(QK^/d_k)V 정보 결합&고정 차원 상태에 압축&토큰 간 쌍별 가중합 기본 수학&미분방정식, 이산화, 상태공간&선형대수, 확률적 정규화 시퀀스 길이 T&고정된 상태 차원 가정 시 대체로 선형 스케일&일반적인 완전 어텐션은 점수행렬 T T 추론 메모리&상태 캐시 중심&일반적인 자기회귀 구현은 KV 캐시 중심 병렬화&결합적인 선택적 스캔&대규모 행렬 곱 복잡도는 구현·배치·헤드·상태 차원에 따라 달라진다. 표는 대표적인 전체 어텐션과 선택적 SSM의 구조적 비교이며, FlashAttention은 메모리 입출력을 개선하고 희소/선형 어텐션은 다른 복잡도를 갖는다.

길이에 따른 개념적 연산량 비교

아래 그래프는 T에 따른 T와 T^2의 정규화된 이론적 증가율이며 실제 모델 실행시간 측정값이 아니다. [width=.85,height=6cm,xlabel=시퀀스 길이 T (상대 단위),ylabel=정규화된 연산량,xmin=0,xmax=10,ymin=0,ymax=100,grid=major,legend pos=north west] 5*x;선형 5T x^2;이차 T^2

전자공학과의 연결 및 연구 문제

제어공학: 연속시간 상태방정식의 행렬지수 e^A와 안정성(고유값의 실수부)을 분석한다. 신호처리: SSM을 IIR 필터와 비교하고, 어텐션을 데이터 의존적 가중 FIR형 집계와 비교한다. 단, 어텐션은 일반적인 고정 계수 FIR 필터와 동일하지 않다. 반도체·가속기: Mamba의 상태 갱신과 Transformer의 행렬 곱에서 메모리 대역폭 및 병렬화 특성을 비교한다. 연습: A=-2, =0.1일 때 A=e^-0.2와 B=(1-e^-0.2)/2를 계산한다. 연습: Q,K^12864일 때 QK^의 차원과 인과 마스크의 효과를 설명한다.

참고문헌

A. Gu and T. Dao, ``Mamba: Linear-Time Sequence Modeling with Selective State Spaces,'' 2023/2024. https://arxiv.org/abs/2312.00752 A. Vaswani et al., ``Attention Is All You Need,'' 2017. https://arxiv.org/abs/1706.03762 T. Dao et al., ``FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness,'' 2022. https://arxiv.org/abs/2205.14135