합성곱 기법이 사라졌어. 그다음은?
선택성이 합성곱 관점을 깨뜨렸으니 Mamba에는 학습을 병렬화할 새 방법이 필요했어. 답은 하드웨어 인지 selective scan이야. GPU의 메모리 계층에 맞춘 병렬 prefix-scan 알고리즘이지. 중간 상태는 빠른 온칩 SRAM에 붙잡아 두고, 실행 시간을 잡아먹는 느린 HBM 왕복은 피했어.
공학 철학은 FlashAttention과 같아. 둘 다 Tri Dao 연구진에서 나왔고 질문도 같지. 복잡한 시퀀스 연산으로 현대 GPU를 실제로 바쁘게 만들려면 어떻게 해야 할까? mamba_ssm.ops.selective_scan_interface의 selective scan 커널은 그 답을 수백 줄의 CUDA로 구현해. 시퀀스를 타일로 나눠 SRAM에 올린 다음 그 안에서 순환식을 계산하고, 타일마다 마지막 상태만 써서 돌려보내. 행렬 곱이 아니라 순환식을 실행하면서도 잘 만든 어텐션 커널에 가까운 GPU 활용률을 얻어.
아키텍처도 단순해졌어
Mamba는 구조상의 부대 비용도 크게 덜어 냈어. 표준 Transformer 블록에는 어텐션 하위 층과 MLP 하위 층이 함께 들어가. Mamba 블록에는 SSM 하위 층만 있고 별도의 MLP는 없어. 내부에서는 두 갈래가 나란히 돌아. 한쪽은 Conv1D를 거쳐 선택적 SSM으로 들어가고, 다른 쪽은 SiLU 게이트 선형 투영을 맡아. 둘을 원소별로 합친 뒤 모델 차원으로 다시 투영해. 층마다 매개변수와 연산이 줄고 계산 그래프도 단순해져.
실제 실행 시간도 빨라
Mamba 논문은 같은 품질에서 Transformer보다 추론 처리량이 5배 높다고 보고했어. 시퀀스가 길수록 격차도 벌어져. Transformer 어텐션에는 제곱 비용이 있지만 Mamba가 토큰 하나를 처리하는 비용은 시퀀스 길이와 무관하기 때문이야. 3B 규모에서는 매개변수가 거의 두 배인 Transformer와 맞먹거나 앞섰어. FLOP 표에서만 좋아 보이는 결과가 아니야. 하드웨어 인지 커널이 FLOP 차이를 실제 경과 시간의 차이로 바꿨어.
효율적인 알고리즘만으로는 부족해. 하드웨어의 생김새에 맞춘 구현까지 있어야 해. SSM이 2010년대에 다시 살아난 건 수학이 갑자기 생겼기 때문이 아니야. 수십 년 묵은 수학을 GPU에서 제대로 돌릴 커널을 마침내 누군가 만들었기 때문이지.