맘바(Mamba)와 상태공간모델(SSM) 쉽게 이해하기: 트랜스포머의 대안이 된 원리

핵심 메시지

• 상태공간모델(SSM)은 선형 RNN처럼 토큰을 순차 처리하지만, 트랜스포머보다 빠르고 메모리를 적게 써 긴 시퀀스 처리에 유리하다.
• 기존 SSM은 모든 입력에 같은 행렬을 적용해 유연성이 부족했고, 그 때문에 속도·메모리 이점은 정확도를 희생한 대가였다.
• 맘바는 Delta·B·C 파라미터를 각 토큰 임베딩에 따라 다르게 만드는 '선택적 SSM'으로 정확도를 트랜스포머 수준까지 끌어올렸다.
• 선택성을 넣으면 합성곱 트릭을 못 쓰지만, 병렬 연관 스캔과 HBM↔SRAM을 활용한 하드웨어 맞춤 구현으로 선형 시간을 유지한다.
• 맘바는 언어뿐 아니라 DNA·오디오 등 다양한 데이터에서 트랜스포머에 필적하거나 앞서며, 시퀀스가 길수록 격차가 벌어진다.

쉽게 이해하기

맘바는 나오자마자 어디에나 쓰이는 트랜스포머의 잠재적 대체재로 불리며 큰 반향을 일으켰다. 흥미롭게도 맘바 논문은 ICLR 첫 심사에서 탈락했지만, 커뮤니티는 이미 그 위에 여러 후속 연구를 쌓아 올렸다. 맘바는 플래시 어텐션으로 알려진 연구자들이 참여했고, 핵심 기여는 이미 트랜스포머보다 빠르고 메모리를 적게 쓰던 SSM의 정확도를 끌어올린 것이다.

SSM은 큰 신경망의 한 부품으로, 높은 수준에서 보면 선형 RNN처럼 동작한다. 이전 토큰의 출력 표현과 현재 입력 토큰의 임베딩을 변환해 결합한다. SSM은 Delta·A·B·C 네 종류의 행렬을 쓰는데, Delta는 A와 B의 가중치를 조정(이산화)하고, 조정된 A는 은닉 상태가 토큰에서 토큰으로 얼마나 전달될지를, B는 입력이 은닉 상태로 얼마나 들어올지를, C는 은닉 상태가 출력으로 어떻게 바뀌는지를 결정한다.

SSM이 빠른 이유는 확장성이 선형이기 때문이다. 트랜스포머의 자기어텐션은 시퀀스가 두 배가 되면 시간·메모리가 네 배로 늘지만, SSM은 토큰이 하나 늘면 계산도 한 번 더 늘 뿐이다. 순차 처리라 느릴 것 같지만, SSM은 선형이라 학습 시에는 전체 시퀀스가 주어지므로 여러 행렬 곱을 미리 계산해 하나의 큰 행렬 K로 묶고 합성곱으로 모든 토큰을 병렬 처리할 수 있다. 오직 추론 시에만 토큰을 하나씩 처리한다.

다만 기존 SSM은 모든 입력에 똑같은 Delta·A·B·C를 적용해, 어떤 토큰을 더 기억하고 무엇을 무시할지 선택하지 못한다. 그래서 맘바는 Delta·B·C를 각 토큰 임베딩으로부터 선형 층을 통해 매번 다르게 계산하는 '선택적 SSM'을 제안한다. 이렇게 하면 트랜스포머의 어텐션이 하던 '특정 토큰에 집중하기'를 SSM도 해낼 수 있다.

선택성의 대가로 입력에 의존하는 파라미터 때문에 합성곱 트릭을 쓸 수 없게 되지만, 저자들은 병렬 연관 스캔이라는 알고리즘과 느린 HBM에서 빠른 SRAM으로 값을 옮겨 계산하는 하드웨어 맞춤 구현으로 이를 빠르게 유지했다. 그 결과 맘바는 트랜스포머와 같은 스케일링 법칙을 보이면서 시퀀스가 길수록 더 유리했고, 256,000 토큰 스캔에서 단순 구현이 메모리 부족으로 죽는 반면 맘바 스캔은 약 10밀리초, 플래시 어텐션 2는 약 1,000밀리초가 걸렸다. DNA 분류와 오디오 모델링 등 다양한 데이터에서도 강했다.

주요 인사이트

• SSM은 연속 미분방정식에서 유래했고, 그것을 행렬 형태의 이산 방정식으로 바꾸는 과정이 이산화(discretization)다. 스텝 크기 Delta가 얼마나 세밀하게 이산화할지를 정하며, 너무 크면 곡선을 벗어나 근사 오차가 생긴다.
• 합성곱 트릭은 계산이 선형이고 입력과 무관할 때만 가능하다. 맘바가 선택성(입력 의존 파라미터)을 도입하는 순간 합성곱은 불가능해지고, 대신 병렬 연관 스캔이 그 자리를 대신한다.
• 맘바의 속도는 알고리즘만이 아니라 하드웨어 인지 구현에서 나온다. Delta·A·B·C를 SRAM으로 읽어 이산화와 스캔을 그곳에서 수행하고 결과만 HBM에 되쓰는 방식이 핵심이다.
• 맘바는 첫 심사에서 탈락했지만 연구 진행은 심사 승인을 기다리지 않았다. 이미 MoE-Mamba, 이미지용 Vision Mamba, 토크나이저 없이 원시 바이트에서 학습하는 MambaByte 같은 확장이 등장했다.
• 선택적 SSM은 결국 RNN의 귀환으로 볼 수 있다. 강력하지만 느리다고 여겨지던 RNN이, 빠르게 만드는 연구가 뒷받침되자 트랜스포머와 겨룰 수준에 이르렀다.

자주 묻는 질문

관련 AI 소식