맘바(Mamba)와 상태공간모델(SSM) 쉽게 이해하기: 트랜스포머의 한계와 선택적 스캔

핵심 메시지

• 트랜스포머의 셀프 어텐션은 길이 L에 대해 L^2 계산이 필요해 학습은 병렬화로 빠르지만 추론이 이차적으로 느려지는데, 상태공간모델(SSM)은 그 대안으로 제시된다.
• SSM은 상태 방정식과 출력 방정식(행렬 A·B·C·D)으로 이뤄지며, 영차 유지(zero-order hold)와 학습된 스텝 크기로 이산 텍스트 신호를 다룰 수 있게 이산화한다.
• SSM은 연속·순환(recurrent)·합성곱(convolution) 세 가지 표현을 가지며, 학습에는 병렬화가 되는 합성곱을, 추론에는 선형 확장되는 순환 표현을 골라 쓸 수 있다.
• 기존 S4까지의 SSM은 시간 불변(LTI)이라 입력에 따라 주의를 바꾸지 못하는데, 맘바는 행렬 B·C와 스텝 크기를 입력에 의존시켜 무엇을 기억하고 버릴지 '선택'하게 만든다.
• 맘바는 선택적 스캔(병렬 스캔으로 학습 가능한 RNN)과 플래시 어텐션식 커널 퓨전의 하드웨어 인식 알고리즘을 합친 S6가 핵심이며, 잠바(Jamba)·잠바(Zamba)처럼 트랜스포머와 결합한 하이브리드도 등장했다.

쉽게 이해하기

영상은 강력한 LLM이 반드시 트랜스포머여야 하는 것은 아니라는 문제 제기에서 출발한다. 트랜스포머의 셀프 어텐션은 시퀀스의 각 토큰을 서로 비교하기 때문에 길이 L에 대해 L^2의 계산이 든다. 학습 때는 어텐션 계산이 병렬화되어 빠르지만, 추론 단계에서는 이 이차적 성질 때문에 느려진다. 여기서 상태공간모델(SSM)과 맘바가 대안으로 떠오른다.

SSM은 미로 속 위치처럼 시스템이 가질 수 있는 상태들로 문제를 수학적으로 표현한다. 핵심은 두 방정식이다. 상태 방정식은 학습된 행렬 A(상태가 시간에 따라 어떻게 변하는가)와 B(입력이 상태에 어떻게 영향을 주는가)로 상태를 갱신하고, 출력 방정식은 행렬 C(상태가 출력으로 어떻게 번역되는가)와 D(입력이 출력에 직접 미치는 영향, 스킵 연결)로 예측을 만든다. 연속 신호를 다루기 위해 영차 유지 기법과 학습된 스텝 크기로 A·B를 이산화한다.

SSM은 하나의 모델을 세 가지로 표현할 수 있다. 이산화가 필요한 원래의 연속 모델, 추론이 빠르지만 병렬화가 어려운 순환(RNN) 표현, 그리고 학습 시 병렬화가 되지만 효율적 추론이 어려운 합성곱 표현이다. 이 둘의 장점을 합쳐 학습에는 합성곱을, 추론에는 순환 표현을 오가는 모델이 선형 상태공간 계층(LSSL)이다. 다만 이들 표현은 시간 불변(LTI)이라 모든 시간 단계에서 행렬이 고정된, 내용을 인식하지 못하는 정적 표현이라는 한계가 있다.

긴 거리 의존성을 다루기 위해 HiPPO(고차 다항식 사영 연산자)를 A 행렬에 결합하면, 오래된 입력일수록 크게 감쇠시키면서도 시퀀스 앞부분을 잊지 않는다. 연속 모델·세 표현·HiPPO를 모두 합친 결과가 긴 시퀀스를 효율적으로 다루는 S4 모델이다. 그렇다면 맘바는 무엇을 더했는가? 바로 불필요한 정보를 걸러 내는 선택적 스캔 알고리즘과, 결과를 효율적으로 저장하는 하드웨어 인식 알고리즘이다.

S4의 약점은 선택적 복사나 유도 헤드 같은 과제에서 드러난다. 시간 불변이라 모든 토큰을 똑같이 다루기 때문에 입력에 따라 주의를 바꾸지 못한다. 맘바는 행렬 B·C와 스텝 크기를 시퀀스 길이·배치에 의존시켜 무엇을 은닉 상태에 남기고 무엇을 무시할지 선택하게 한다. 여기에 순차적이라 병렬화가 어려운 RNN의 문제를 병렬 스캔(프리픽스 합)으로 풀고, 빠른 SRAM과 느린 DRAM 사이 복사를 최소화하는 커널 퓨전을 더한 S6가 맘바 블록의 토대다. 잠바(Jamba)·잠바(Zamba)는 맘바 블록과 트랜스포머 블록(일부는 전문가 혼합)을 번갈아 배치한 하이브리드다.

주요 인사이트

• 맘바의 통찰은 '압축의 균형'에 있다. 트랜스포머는 입력을 전혀 압축하지 않고(그래서 이차 비용), S4는 크게 압축하지만 내용을 구분하지 못한다. 맘바는 입력에 따라 선택적으로 압축해 그 사이에서 균형을 잡는다.
• 선택적 스캔의 묘미는 '병렬화 가능한 RNN'을 만든 데 있다. 원래 RNN은 이전 계산이 끝나야 다음이 시작되어 병렬화가 어려운데, 병렬 스캔의 하향 스윕과 상향 스윕으로 중간 단계를 동시에 계산해 이 제약을 우회한다.
• 하드웨어 인식 설계는 알고리즘만큼 중요하다. 맘바는 플래시 어텐션처럼 주요 파라미터는 DRAM에, 상태와 관련 계산은 SRAM에 두어 느린 메모리와의 데이터 복사를 줄이는 커널 퓨전으로 실제 속도를 확보한다.
• 트랜스포머가 당장 사라지지는 않겠지만, 맘바 같은 SSM 기반 구조와 트랜스포머를 섞은 하이브리드(잠바·잠바)가 효율적 추론과 강력한 어텐션이라는 두 장점을 모두 취하려 한다는 점에서 아키텍처 경쟁의 방향을 엿볼 수 있다.

자주 묻는 질문

관련 AI 소식