맘바 3와 상태공간 모델(SSM): KV 캐시 없는 LLM 구조는 어떻게 발전했나

핵심 메시지

• 상태공간 모델은 과거 토큰을 모두 보관하는 대신 고정 크기의 은닉 상태 하나로 압축해, 추론 중 메모리 사용량이 일정하게 유지된다.
• 같은 모델을 학습할 때는 합성곱 형태로, 추론할 때는 순환 형태로 계산할 수 있다는 점이 다른 구조에 없는 유연성이다.
• 맘바 1의 핵심은 입력에 따라 파라미터가 변하는 선택 메커니즘으로, 중요한 토큰만 상태에 기록하게 만든 것이다.
• 맘바 2는 상태공간 모델과 선형 어텐션이 같은 행렬 계열임을 보이며 학습 처리량을 크게 끌어올렸다.
• 맘바 3는 이산화 방식 개선, 회전 위치 임베딩으로 구현한 복소수 동역학, 그리고 MIMO로 디코딩 병목을 겨냥했다.
• 순수 상태공간 모델은 정확한 검색에 약해, 어텐션 층을 일부 섞는 하이브리드가 현실적인 답으로 자리 잡았다.

쉽게 이해하기

트랜스포머는 토큰을 생성할 때마다 과거 키·값을 캐시에 쌓기 때문에 문맥이 길어질수록 메모리가 계속 늘어난다. 상태공간 모델은 반대로 지금까지 본 모든 것을 요약한 고정 크기 상태 하나만 들고 간다. 새 토큰이 들어오면 상태가 갱신될 뿐 커지지는 않는다. 추론 중 메모리가 일정하고 계산량이 입력 길이에 비례한다는 것이 이 계열의 출발 명제다.

이 모델의 뿌리는 제어 이론의 연속 시간 시스템이다. 입력을 상태에 써 넣는 행렬, 상태를 시간에 따라 진화시키는 행렬, 상태에서 출력을 읽는 행렬 세 가지로 구성되며, 특히 상태를 진화시키는 행렬이 정보가 얼마나 오래 남는지를 결정하는 심장부다. 언어처럼 이산적인 토큰 열에 쓰려면 이 연속 방정식을 단계별 갱신 규칙으로 바꾸는 이산화가 필요한데, 그 결과 나오는 식은 사실상 순환 신경망의 갱신식과 같다.

초기 모델의 결정적 한계는 파라미터가 학습이 끝나면 고정된다는 점이었다. 같은 필터를 모든 위치에 똑같이 적용하므로 '개가 공을 쫓았다'의 개와 '어디선가 개가 짖었다'의 개를 구분하지 못한다. 의미가 아니라 위치로만 가중치를 주는 셈이다. 맘바는 입력을 쓰고 읽는 행렬과 갱신 폭을 결정하는 값을 모두 현재 토큰의 함수로 바꿔 이 문제를 풀었다. 갱신 폭이 크면 '이 토큰은 중요하니 기록하라', 작으면 '건너뛰라'는 의미가 되어 장단기 기억망의 망각 게이트와 비슷한 역할을 한다.

파라미터가 매 시점 달라지면 합성곱 가속이 불가능해진다. 맘바는 이산화·상태 갱신·출력 읽기를 하나의 GPU 커널로 묶어 빠른 온칩 메모리 안에서 처리하고 역전파 때 중간 상태를 다시 계산하는 하드웨어 친화적 병렬 스캔으로 이를 해결했다. 이후 맘바 2는 상태공간 모델의 계산이 선형 어텐션과 같은 구조화된 삼각 행렬 계열에 속한다는 사실을 밝혀, 트랜스포머 최적화 기법을 그대로 끌어 쓸 수 있게 만들었고 학습 처리량을 크게 개선했다.

맘바 3의 개선은 세 갈래다. 입력이 구간 내내 일정하다고 가정하던 이산화를 앞뒤 입력을 평균 내는 2차 방식으로 바꿔 오차를 줄였고, 1과 0의 개수가 홀수인지 짝수인지 맞히는 단순한 문제조차 이전 모델이 못 풀던 원인인 실수 스칼라의 한계를 복소수 회전으로 해결하되 그 회전이 회전 위치 임베딩과 수학적으로 같다는 점을 이용해 기존 최적화 인프라를 그대로 썼다. 또 토큰을 하나씩 만들 때 상태 갱신이 얇은 외적 연산이라 연산 장치가 놀게 되는 문제를, 입력과 출력을 여러 개 동시에 처리해 두꺼운 행렬 곱으로 바꿔 완화했다. 다만 고정 크기 상태는 본질적으로 압축이라 문맥 깊숙한 곳의 특정 문장을 정확히 되짚는 데는 약해, 현재 업계는 대부분의 층을 상태공간 모델로 두고 일부만 어텐션 층으로 섞는 하이브리드로 수렴하고 있다.

주요 인사이트

• 맘바 계열의 발전은 성능을 더하는 방식이 아니라, 상태 진화 행렬을 계속 단순화하면서 다른 곳을 강화하는 일관된 패턴을 보인다.
• 홀짝 판별 실험은 표현력의 한계를 드러내는 좋은 도구다. 실수 스칼라는 값을 키우거나 줄일 수는 있어도 방향을 뒤집지 못하고, 그 뒤집기에는 복소수 회전이 필요하다.
• 디코딩이 느린 이유가 연산량이 아니라 산술 강도 때문이라는 진단이 중요하다. MIMO는 연산을 더 쓰는 것이 아니라 이미 놀고 있던 연산 자원을 쓰는 방법이다.
• 복소수 동역학을 회전 위치 임베딩으로 우회한 것처럼, 새로운 수학을 기존의 잘 최적화된 연산으로 환원하는 것이 실전에서는 성능을 좌우한다.
• 어텐션 층을 전부 제거하면 검색 정확도가 떨어지는 정도가 아니라 완전히 실패한다는 점이, 순수 대체 구조가 아니라 하이브리드로 정착한 이유다.

자주 묻는 질문

관련 AI 소식