Mamba는 긴 문서를 어떻게 기억할까: 모든 것을 저장하지 않는 모델

긴 문서를 읽는 모델을 만들 때는 기억력을 늘리고 싶어진다. 더 많은 토큰을 넣고, 더 큰 캐시를 붙이고, 더 긴 문맥을 학습시킨다. 그런데 그 문서에 정작 필요한 정보는 몇 줄뿐이라면 어떨까. 기억해야 할 양보다, 버려도 되는 내용을 구분하는 능력이 더 중요한 문제일 수 있다.

Mamba를 흥미롭게 읽은 이유도 여기에 있다. 어텐션을 싸게 흉내 내겠다는 접근보다, 고정된 크기의 기억 안에 무엇을 남길 것인가라는 질문이 앞에 나온다. 논문의 핵심은 입력의 내용에 따라 기억의 갱신 방식을 바꾸는 selective state space model이다. Mamba 원문

상태 하나로 문서를 읽는다는 것

자기회귀 Transformer는 과거 토큰의 key와 value를 KV 캐시에 보관한다. 새 토큰은 이 기록을 참조해 필요한 정보를 읽는다. 반면 순환 모델은 지금까지 읽은 내용을 상태 h에 압축한다. 다음 토큰을 처리하면 상태가 바뀌고, 이전 입력을 원형 그대로 다시 꺼내 볼 수는 없다.

이 차이를 회의록으로 생각해 보면 쉽다. 한 사람은 모든 발언을 녹음하고 필요할 때 되감는다. 다른 사람은 메모 한 장을 계속 고쳐 쓴다. 두 번째 방식은 기록을 들고 다니기 가볍지만, 처음에 지운 발언이 나중에 중요한 사실로 드러나면 난처해진다. 순환 모델의 문제는 이런 압축 결정을 토큰이 들어오는 순간 내려야 한다는 데 있다.

상태공간 모델의 이산 시간 갱신은 대략 h_t = A_bar h_(t-1) + B_bar x_t, y_t = C h_t로 볼 수 있다. A_bar는 기존 기억을 어떻게 유지하고 변형할지, B_bar는 새 입력을 어떻게 넣을지, C는 기억에서 무엇을 읽을지 정한다. 실제 모델은 여러 채널과 더 높은 차원의 상태를 사용하지만, 이 그림만으로도 중요한 질문을 잡을 수 있다. 새로 들어오는 x_t가 다른데도 기억을 다루는 규칙이 항상 같아야 할까.

고정된 갱신 규칙은 구현에 유리하다. 같은 연산이 반복되니 긴 시퀀스를 합성곱 형태로 계산할 수 있다. 하지만 중요한 토큰의 간격이 매번 달라지는 텍스트에서는 일정한 시간 간격에 잘 맞는 필터만으로 충분하지 않을 수 있다.

메모에 형광펜을 칠하는 갱신 규칙

Mamba는 B와 C, 그리고 연속 시간 시스템을 이산화할 때 쓰는 간격 Δ를 입력의 함수로 만든다. 기본 행렬 A 자체가 매 토큰 새로운 신경망으로 생성되는 구조는 아니다. 입력이 바뀌면 Δ와 B가 갱신을 조절하고, C가 읽기를 조절한다. Δ를 통해 이산화된 A_bar도 입력에 따라 달라진다. 이 구분이 있어야 선택성을 설명하면서도 모든 파라미터가 토큰마다 새로 만들어진다는 오해를 피할 수 있다. 선택 메커니즘과 알고리즘

설명을 위한 예를 하나 만들어 보자. 로그에 일반 상태 메시지 백 개가 나오고, 중간에 “복구 키: violet”이 등장한 뒤, 또 일반 메시지가 이어진다. 마지막 질문은 복구 키를 묻는다. 모든 메시지를 같은 정도로 상태에 섞으면 중요한 단어가 희석되기 쉽다. 좋은 갱신 규칙은 상태 메시지에서는 기억을 거의 유지하고, 복구 키를 발견하면 관련 상태를 크게 바꾸는 방향을 배울 수 있다.

이 예에서 선택은 단순히 토큰의 중요도 점수를 출력하는 일이 아니다. 어떤 내용이 이후 상태로 전달될지에 직접 관여한다. 읽을 때 중요한 곳을 찾는 능력과, 읽는 동안 중요한 내용을 남기는 능력은 서로 닮았지만 동일하지 않다. Mamba는 후자에 구조적인 자리를 만들어 준다.

다만 “복구 키”라는 문자열을 보면 특정 슬롯에 정확히 저장한다는 식의 프로그램을 사람이 지정한 것은 아니다. 그런 동작을 유도할 수 있는 학습 가능한 통로를 마련한 것이다. 데이터에 유사한 관계가 충분히 나타나고 학습이 성공해야 유용한 선택이 생긴다. 입력 의존성만 넣었다고 자동으로 완벽한 기억 장치가 되는 것은 아니다.

빨라지려면 GPU의 기억도 다뤄야 한다

입력에 따라 갱신이 달라지면 고정된 합성곱 커널을 사용하던 계산 경로가 사라진다. 여기서 Mamba의 두 번째 아이디어가 필요하다. 선택성을 얻는 대신 순차 계산 때문에 느려진다면 긴 문맥용 모델로서의 장점이 약해지기 때문이다.

논문은 병렬 scan과 커널 융합, 역전파 때의 재계산을 함께 사용한다. 핵심은 큰 중간 상태를 매 단계 GPU의 상대적으로 느린 메모리에 써 놓지 않고, 빠른 메모리에서 갱신을 수행하는 것이다. 덧셈의 누적합을 병렬화할 수 있듯이, 적절하게 구성한 상태 전이의 결합도 병렬 scan으로 계산할 수 있다. 하드웨어를 고려한 구현

이 대목은 모델 연구를 볼 때 좋은 습관을 준다. O(L)이라는 복잡도만으로 실제 속도를 예상하면 부족하다. 중간 텐서의 크기, 메모리 왕복, 커널 호출 횟수까지 실제 비용을 결정한다. 같은 수식을 Python 반복문으로 구현한 모델과 최적화된 커널로 구현한 모델은 전혀 다른 사용 경험을 줄 수 있다.

또한 학습과 생성의 계산 형태를 나눠 봐야 한다. 생성에서는 상태를 한 번씩 갱신하므로 과거 길이에 따라 KV 캐시가 커지는 방식과 다르다. 학습에서는 많은 토큰을 동시에 처리하고 역전파해야 하므로 병렬 알고리즘과 중간 상태 관리가 중요해진다. “순환 모델이니 학습도 한 토큰씩 기다려야 한다”는 직관을 구현이 어떻게 깨는지가 이 논문의 재미있는 부분이다.

긴 입력을 처리하는 것과 긴 입력을 기억하는 것

원문은 선택적 복사와 문맥 기반 연상 같은 합성 과제뿐 아니라 언어, 오디오, DNA 모델링을 평가한다. 저자들은 선택성을 갖춘 구조가 언어 모델링에서도 경쟁력 있고, 긴 시퀀스를 효율적으로 다룰 수 있음을 보고한다. 여기서 내가 받아들이는 결론은 “어텐션만이 유용한 내용 선택을 할 수 있는 구조는 아니다” 정도다. 모든 길이와 모든 추론 과제에서 우월하다는 결론으로 넓히기는 어렵다. 실험과 논의

고정 크기 상태에는 분명한 정보 병목이 있다. 긴 소설의 분위기를 유지하는 일과, 서로 다른 수천 개 주문 번호를 정확하게 회수하는 일은 요구하는 기억이 다르다. 압축된 상태가 전자에 충분하더라도 후자에서는 세부 정보 손실이 문제가 될 수 있다. 이는 위 예제에서 필요한 복구 키가 하나였다는 조건과도 연결된다. 키가 수만 개로 늘어나면 같은 설명으로 성능을 보장할 수 없다.

나는 긴 문맥 모델을 검토할 때 입력 길이만 늘리는 실험보다, 필요한 사실의 개수와 방해 정보의 밀도를 따로 바꾸는 실험이 더 보고 싶어진다. 같은 길이에서도 사실 하나를 찾는 문제와 사실 수백 개를 연결하는 문제의 난도가 다르다. 이 구분은 Mamba의 논문 주장이라기보다, 그 구조를 읽고 얻은 평가 설계의 시사점이다.

실무에서도 먼저 확인할 것은 “우리 입력은 얼마나 긴가”보다 “과거에서 어떤 정보를 얼마나 정확하게 보존해야 하는가”다. 스트리밍 센서, 로그, 음성처럼 입력이 계속 들어오는 문제에서는 상태를 유지하며 처리하는 방식이 매력적이다. 반대로 원문의 특정 구절을 정확히 인용해야 한다면 외부 검색이나 명시적 저장소가 필요할 수 있다. 모델의 압축 능력과 애플리케이션의 기록 보존 책임을 같은 것으로 취급하지 않는 편이 좋다.

Mamba가 남기는 질문은 기억을 크게 만드는 방법보다 좋은 기억의 정의에 가깝다. 모든 것을 쌓아 두는 것이 아니라, 미래의 예측에 필요한 내용을 남기는 일. 그 결정을 학습 가능한 상태 전이에 집어넣고 실제 GPU에서 빠르게 실행한 지점이 이 논문의 가장 인상적인 성취다.

논문: Albert Gu, Tri Dao, Mamba: Linear-Time Sequence Modeling with Selective State Spaces, 2023. 원문 · 공식 코드