Mamba와 S4 완전 정리: 상태공간모델부터 선택적 스캔, 병렬 스캔, 커널 융합까지

핵심 메시지

• 순환신경망은 토큰당 추론 비용이 일정한 대신 학습을 병렬화하지 못하고, 트랜스포머는 학습은 병렬이지만 길이가 늘수록 계산이 제곱으로 불어난다. 상태공간모델은 이 둘의 장점만 취하려는 시도다.
• 상태공간모델은 미분방정식을 이산화해 순환식으로 바꾸며, 파라미터가 고정이라 학습할 때는 합성곱으로 병렬 계산하고 추론할 때는 순환식으로 한 토큰씩 뽑는 두 얼굴을 갖는다.
• A 행렬은 과거 정보를 다음 상태로 얼마나 실어 나를지 정하는 부품이고, HiPPO 방식으로 초기화하면 최근 토큰은 잘 복원하고 먼 과거는 서서히 잊는 형태가 된다.
• 기존 상태공간모델은 모든 토큰에 같은 파라미터를 적용해 내용에 따라 골라내는 과제에 실패했고, Mamba는 B·C·Δ를 입력마다 다르게 만들어 이 한계를 넘었다.
• 파라미터가 매 단계 바뀌면 합성곱 경로가 사라져 순환 계산만 남지만, 병렬 스캔과 커널 융합, 역전파용 활성값 재계산으로 속도를 되찾았다.

쉽게 이해하기

강사는 먼저 시퀀스 모델 세 가지를 나란히 놓고 각각의 장단점을 정리한다. 순환신경망은 이론적으로 문맥 길이에 제한이 없고 토큰 하나를 만드는 비용이 항상 같지만, 이전 토큰이 있어야 다음을 계산할 수 있어 학습을 병렬화할 수 없다. 반대로 트랜스포머는 학습을 병렬화할 수 있지만 100번째 토큰을 만들 때 100번의 내적이 필요하고 학습 비용도 길이의 제곱으로 늘어난다.

상태공간모델을 이해하려면 미분방정식이 필요해서, 강사는 토끼 개체수가 일정 비율로 늘어나는 예제로 도입부를 채운다. 해석적으로 풀기 어렵고 실제 입력도 이산적이므로 오일러 방법으로 근사하듯 시스템을 이산화하는데, 논문에서는 더 나은 근사를 위해 영차 유지 방식을 쓴다. 이산화 간격 Δ는 사람이 정하지 않고 모델이 학습으로 찾는 파라미터다.

핵심은 같은 계산을 두 가지 방식으로 할 수 있다는 점이다. 순환식을 펼쳐 보면 출력이 CB, CAB, CA²B 같은 항으로 이뤄진 커널과 입력의 합성곱으로 표현되고, 합성곱은 각 출력이 서로를 기다리지 않아 병렬화된다. 그래서 학습은 합성곱으로, 추론은 상수 비용의 순환식으로 처리한다. 입력이 벡터일 때는 차원마다 독립된 상태공간모델을 두는데, 강사는 이것이 트랜스포머의 멀티헤드 어텐션이 차원을 나눠 처리하는 방식과 다르지 않다고 설명한다.

Mamba가 등장한 이유는 두 과제 때문이다. 단순 복사는 기존 모델도 해내지만, 흰색 토큰을 무시하고 색깔 토큰만 골라 내보내는 선택적 복사와 앞서 본 패턴을 떠올려 답하는 인덕션 헤드는 실패한다. 파라미터가 모든 입력에 똑같이 적용되니 토큰을 구분할 수 없기 때문이다. Mamba는 B와 C, Δ를 입력 토큰마다 다르게 만들어 내용에 반응하게 했고, 그 대가로 시간 불변성이 깨져 합성곱을 쓸 수 없게 됐다.

남은 순환 계산은 접두사 합과 같은 스캔 연산으로 볼 수 있고, 결합법칙이 성립하면 스캔은 병렬화가 가능해 복잡도가 스레드 수만큼 줄어든다. 여기에 GPU의 느린 고대역폭 메모리 대신 빠른 SRAM에서 이산화와 스캔을 끝내는 커널 융합, 역전파에 필요한 중간 활성값을 저장하는 대신 다시 계산하는 기법이 더해진다. 결과적으로 선택적 복사 정확도는 99.8%, 인덕션 헤드는 시퀀스 길이 10의 6제곱까지 거의 100%를 유지했고, 퍼플렉시티는 최신 트랜스포머와 견줄 만한 수준으로 나왔다.

주요 인사이트

• 트랜스포머 추론의 비대칭성이 문제의 출발점이다. 첫 토큰을 만드는 비용과 100번째 토큰을 만드는 비용이 다르기 때문에 아주 긴 입력으로 확장하기 어렵고, 상태공간모델의 상수 비용은 바로 이 지점을 겨냥한다.
• HiPPO로 초기화한 A 행렬이 하는 일은 지수이동평균과 비슷하다. 최근 구간은 원래 신호에 가깝게 복원되고 오래된 구간은 흐릿해지는데, 다음 토큰 예측에는 국소 문맥이 더 중요하다는 판단이 깔려 있다.
• 선택성을 얻는 대가로 병렬 합성곱을 잃는 맞교환이 인상적이다. 커널이 입력마다 달라지면 하나의 커널을 재사용할 수 없어, 모델 구조가 아니라 알고리즘 최적화로 손해를 메운 셈이다.
• GPU에서 느린 것은 계산이 아니라 데이터 복사라는 점을 강사는 A100 사양으로 못박는다. 초당 연산량과 메모리 대역폭의 격차가 수십 배라서, 중간 결과를 큰 메모리에 쓰지 않는 커널 융합이 곧 속도가 된다.
• 역전파용 활성값을 저장했다가 다시 읽는 것보다 그때그때 다시 계산하는 편이 빠를 수 있다는 발상은 메모리와 연산의 가격표가 뒤바뀐 현재 하드웨어를 잘 보여 준다. 강사 본인도 Mamba가 여전히 순환 구조이며 초대규모 데이터에서의 성능은 아직 알 수 없다는 한계를 덧붙인다.

자주 묻는 질문

관련 AI 소식