Mamba with Hierarchical Memory: Solving Representation Bottleneck in Long Sequence Modeling

당신이 백만 단어에 달하는 이야기를 기억하려고 노력하고 있다고 상상해 보십시오. 당신의 뇌는 놀랍지만, 한 번에 '활성' 사고 공간에 담을 수 있는 양에는 한계가 있습니다. 만약 그 거대한 이야기의 모든 단어를 당신의 활성 정신 속에 유지하려고 한다면, 오래된 부분들은 밀려나거나 흐릿해질 것이고, 당신은 줄거리를 놓치게 될 것입니다. 이것은 현대 인공지능(AI) 모델이 직면한 과제와 비슷합니다. 오랫동안 가장 강력한 AI 모델들은 트랜스포머(Transformer) 구조로 만들어졌는데, 이들은 문맥을 이해하는 데는 뛰어나지만 이야기가 너무 길어지면 믿을 수 없을 정도로 느려지고 비용이 많이 듭니다. 이를 해결하기 위해 과학자들은 '재귀적 선형 주의력(Recurrent Linear Attention)' 모델(Mamba와 같은)이라는 더 새롭고 빠른 유형의 모델을 발명했습니다. 이 모델들은 매우 효율적인 컨베이어 벨트와 같습니다. 정보를 한 단어씩 처리하며, 진행되는 동안 자신의 마음속에 하나의 작고 압축된 '노트'를 업데이트합니다. 이 덕분에 이 모델들은 매우 빠르며 방대한 양의 텍텍스트를 처리할 수 있습니다. 하지만 여기에는 함정이 있습니다. 그 단일한 '노트'는 고정된 크기를 가지고 있기 때문에, 마치 바다를 담으려는 작은 양동이와 같습니다. 이야기가 길어질수록 양동이는 넘쳐흐르고, 모델은 비록 다음 단어를 올바르게 예측할 수는 있을지언정 중요한 세부 사항들을 잊기 시작합니다.

큰 질문은 과학자들이 던져온 질문입니다. 어떻게 하면 이 빠른 AI 모델들이 단순히 즉각적인 단어들이 아니라 긴 이야기의 '의미'를 기억하게 만들 수 있을 것인가? 흔한 믿음은, 만약 우리가 모델이 '망각'하지 않도록(수치적 붕괴를 줄임으로써) 더 개선하기만 한다면, 모델이 긴 텍스트를 통해 추론하는 능력이 자동으로 똑똑해질 것이라는 것이었습니다. 그러나 한 새로운 논문은 이것이 속임수일 수 있다고 시사합니다. 연구진은 모델이 '양동이'가 넘치는 것을 막는 수학적으로 완벽한 상태가 되더라도, 여전히 '건초더미 속의 바늘 찾기'와 같은 복잡한 작업이나 전체 이야기를 이해해야 하는 퍼즐을 해결하는 데 실패한다는 것을 발견했습니다. 그들은 문제가 단순히 양동이가 가득 차는 것만이 아니라, 양동이가 서로 다른 복잡한 아이디어들을 하나의 흐릿한 형태 안에 억지로 밀어 넣으려 한다는 점에 있다고 주장합니다. 그들은 이를 '의미적 에일리어싱(semantic aliasing)'이라고 부르는데, 이는 두 개의 매우 다른 이야기가 모델 내부에서 정확히 똑같이 보이게 되는 현상을 말합니다. 이 문제를 해결하기 위해 팀은 인간의 기억에서 영감을 얻었습니다. 인간이 감각 기억(지금 보고 있는 것), 작업 기억(우리가 생각하고 있는 것), 그리고 장기 기억(나중을 위해 저장하는 사실들)을 가진 것처럼, 연구진은 **계층적 메모리 맘바(Hierarchical Memory Mamba, HMM)**라는 새로운 시스템을 구축했습니다. 이 시스템은 단순히 하나의 작은 양동이에 의존하지 않습니다. 대신, 텍스트를 읽는 빠른 '감각' 층, 단어들을 의미 있는 단락으로 그룹화하는 '작업' 층, 그리고 단락들의 압축된 요약본을 저장하는 '장기' 도서관을 가지고 있습니다. 모델이 문제를 해결해야 할 때, 그것은 현재의 생각만을 보는 것이 아니라 도서관으로 손을 뻗어 적절한 요약본을 가져와 자신의 정신 앞부분으로 불러옵니다.

이 논문은 이러한 접근 방식이 놀라울 정도로 잘 작동한다는 것을 입증합니다. 표준 Mamba 모델 위에 이 계층적 메모리 시스템을 추가함으로써, AI는 처음부터 다시 학습시키거나 엄청난 양의 추가 컴퓨터 자원을 사용하지 않고도 긴 시퀀스 작업에서 성능이 크게 향상되었습니다. 모델이 긴 텍스트 속에서 숨겨진 '패스키(passkey)'를 찾아야 했던 테스트에서, 이 새로운 시스템은 기존의 강력한 모델들과 비교했을 때 성공률을 **34.3%에서 37.1%**까지 높였습니다. 추론 작업(긴 문서를 바탕으로 질문에 답하는 것 등)에 있어서는 정확도가 **1.6%에서 14.2%**까지 뛰어올랐습니다. 아마도 가장 인상적인 점은, 이 모든 것이 단 **2%**의 파라미터(모델이 학습하는 내부 설정)만을 추가하고 매우 적은 추가 학습 시간만으로 달성되었다는 것입니다. 연구진은 또한 이 방법이 모델을 느리게 만들지 않는다는 것을 보여주었습니다. 즉, 원래 Mamba 아키텍처의 빠른 속도를 그대로 유지합니다.

결정적으로, 이 논문은 단순히 모델을 수학적으로 더 안정적으로 만드는 것(붕족을 줄이는 것)이 긴 문맥 문제를 해결하는 충분한 방법이라는 생각에 반박합니다. 그들은 완벽한 안정성을 갖추더라도 고정된 크기의 메모리 상태는 결국 서로 다른 역사들을 하나로 뭉뚱그려 버려, 두 개의 서로 다른 긴 이야기를 구별하는 것을 불가능하게 만든다는 것을 증명했습니다. 그들의 해결책은 양동이를 더 크게 만드는 것이 아니라, 정보가 조직되는 방식을 바꾸는 것이었습니다. '단락 수준의 의미(paragraph-level semantics, 즉 텍스트 덩어리의 핵심 아이디어)'를 추출하여 별도의 검색 가능한 메모리에 저장함으로써, 모델은 자신의 내부 상태라는 병목 현상을 우회할 수 있습니다. 결과는 이 방식이 모델이 정보를 효과적으로 검색하고 사용할 수 있도록 학습하게 하여, 매번 새로운 유형의 퍼즐에 맞춰 미세 조정(fine-tuning)될 필요 없이 다양한 작업에 걸쳐 일반화될 수 있음을 시사합니다. 이 논문은 언어 모델링과 추론에 초점을 맞추고 있지만, 기억을 계층적으로 조직하는 것이 방대한 양의 정보를 다루는 핵심이라는 기본 아이디어는, 자신의 긴 이야기 속에서 길을 잃지 않는 더 똑똑하고 효율적인 AI를 구축하기 위한 새로운 길을 제시합니다.