Evermind & Shanda Group — MSA: Memory Sparse Attention for Efficient End-to-End Memory Model Scaling

1. 배경 및 개요

현대 대규모 언어 모델(LLM)은 수학적 추론, 프로그래밍, 역할 수행 등에서 뛰어난 능력을 보여주지만, 인간 지능의 초석인 장기적이고 미세한 기억 유지(long-term, fine-grained memory retention) 능력은 여전히 큰 과제로 남아 있습니다. 인지 과학 연구에 따르면 인간의 평생 기억 용량은 약 2억~3억 개의 토큰에 해당하며, 디지털 트윈(Digital Twin)이나 장기 에이전트 추론과 같은 복잡한 시나리오를 구현하기 위해서는 모델이 수억 개의 토큰에 달하는 컨텍스트를 효과적으로 처리할 수 있어야 합니다.

기존의 접근 방식들은 각각의 한계를 지니고 있습니다:

• 파라미터 기반 기억(Parameter-Based): 지식을 모델 파라미터에 직접 내재화하지만, 용량 확장성이 부족하고 치명적 망각(catastrophic forgetting)에 취약합니다.
• 외부 저장소 기반 기억(External Storage-Based): RAG(Retrieval-Augmented Generation)와 같은 방식은 확장성은 좋으나 엔드투엔드 미분 가능성이 없어 모델의 내부 추론 공간과 검색 메트릭 간의 최적화 격차가 존재합니다.
• 잠재 상태 기반 기억(Latent State-Based): 모델 내부의 KV 캐시 등을 활용하여 높은 정밀도를 유지하려 하지만, 1억 토큰 이상의 극한의 컨텍스트로 확장하기에는 계산 비용이 너무 높거나 정보 손실이 발생하는 딜레마가 있습니다.

MSA(Memory Sparse Attention)는 이러한 병목 현상을 해결하기 위해 제안된 엔드투엔드 학습 가능하고 효율적이며 대규모 확장이 가능한 메모리 모델 프레임워크입니다. MSA는 메모리 용량을 추론 능력과 분리함으로써 일반 목적의 모델에 평생 규모(lifetime-scale)의 내재적 기억 능력을 부여하는 기초를 제공합니다.

2. 시스템 아키텍처 및 결과

MSA는 표준 조밀 주의(dense self-attention) 메커니즘을 문서 기반 검색 희소 주의(document-based retrieval sparse attention)로 대체합니다.

• 희소 주의 메커니즘(Sparse Attention): * 백본 모델의 프로젝션 가중치를 사용하여 표준 Key($K$)와 Value($V$) 행렬을 생성하는 동시에, 전용 라우터(Router) K/Q 프로젝터를 도입하여 특수 라우팅 키($K^R$)와 쿼리($Q^R$) 행렬을 생성합니다.
• 메모리 점유율과 검색 복잡도를 줄이기 위해 문서를 고정 길이의 청크로 분할하고 청크 단위 평균 풀링(mean pooling)을 수행하여 상태를 압축합니다.
• 사용자 쿼리와 메모리 청크 간의 코사인 유사도를 계산하여 가장 관련성이 높은 Top-k 문서를 선택하고, 이들의 압축된 KV 행렬을 쿼리의 로컬 캐시와 연결하여 답변을 생성합니다.
• 문서별 RoPE(Document-wise RoPE): 각 문서에 독립적인 회전 위치 임베딩(RoPE)을 할당하여, 훈련 시보다 훨씬 많은 수의 문서가 들어오는 추론 시점에도 성능 저하 없이 컨텍스트를 확장할 수 있게 합니다.

주요 성과:

• 비전례적인 확장성: 1.6만 개에서 1억 개 토큰까지 컨텍스트를 확장하는 동안 성능 저하가 9% 미만에 불과한 뛰어난 정밀도 안정성을 보였습니다.
• 효율적 하드웨어 활용: KV 캐시 압축과 메모리 병렬(Memory Parallel) 기술을 결합하여, 단 2개의 NVIDIA A800 GPU만으로 1억 토큰에 대한 추론을 수행할 수 있습니다.
• 성능 비교: 장문 컨텍스트 벤치마크에서 프론티어 언어 모델, 최신 RAG 시스템 및 선도적인 메모리 에이전트들을 크게 능가했습니다.

3. 데이터 및 평가 방법론

MSA의 효능을 검증하기 위해 질의응답(QA) 태스크와 장문 컨텍스트 충실도를 측정하는 “Needle-In-A-Haystack(NIAH)” 태스크를 수행했습니다.

• QA 벤치마크: MS MARCO, Natural Questions, TriviaQA(10M 토큰 규모) 등 9개의 다양한 벤치마크를 사용했습니다. 동일한 백본 모델(Qwen3–4B-Instruct)을 사용하는 RAG 시스템 대비 평균 11.5%~16.0%의 성능 향상을 기록했습니다.
• NIAH 평가: RULER 데이터셋을 사용하여 3.2만 개에서 100만 개 토큰까지의 안정성을 평가했습니다. MSA는 100만 토큰 규모에서도 94.84%의 높은 검색 정확도를 유지한 반면, 수정되지 않은 백본 모델은 12.8만 토큰을 넘어서면 성능이 급격히 붕괴되었습니다.

4. 기술적 세부 분석 (Methodology Deep Dive)

MSA의 핵심 혁신은 학습, 추론, 그리고 복합 추론 단계의 최적화에 있습니다.

A. 2단계 학습 전략 (Training Strategy)

1. 지속적 사전 학습(Continuous Pre-training): 1,589.5억 개의 토큰으로 구성된 말뭉치에서 생성적 검색(Generative Retrieval) 능력을 강화합니다. 특히 보조 라우팅 손실($\mathcal{L}_{aux}$)을 도입하여 내부 희소 주의 기제가 올바른 증거에 집중하도록 감독합니다.
2. 포스트 트레이닝(Post-Training): 8k 컨텍스트에서 기본 지시 이행 능력을 확립한 후, 64k 컨텍스트로 확장하는 커리큘럼 학습(Curriculum Learning)을 통해 극한의 메모리 규모에 대한 외삽 성능을 높입니다.

B. 메모리 병렬 및 계층형 저장 (Memory Parallel & Tiered Storage)

100만 토큰 규모의 캐시(약 169GB)는 표준 노드의 VRAM 용량을 초과합니다. 이를 해결하기 위해:

• GPU 상주 라우팅 키: 검색에 필수적인 라우팅 키($\overline{K}^R$)만 GPU VRAM에 분산 저장하여 지연 시간을 최소화합니다.
• CPU 오프로드 콘텐츠 KV: 실제 데이터인 콘텐츠 KV는 호스트 DRAM(CPU 메모리)에 저장하고, 라우팅을 통해 선택된 Top-k 청크만 비동기적으로 GPU로 가져와 연산합니다.

C. 메모리 인터리브 (Memory Interleave)

멀티홉(multi-hop) 추론과 같은 복잡한 쿼리를 위해 도입된 메커니즘입니다. 한 번의 검색으로 끝내지 않고, 모델이 자율적으로 결정한 횟수만큼 생성적 검색과 컨텍스트 확장을 반복합니다. 이전 단계에서 검색된 문서가 다음 단계의 쿼리 일부가 되어 증거 체인을 구성할 수 있게 합니다.

5. 분석: 효율성 및 컨텍스트 저하

• 계산 복잡도: MSA는 학습과 추론 모두에서 메모리 크기 $L$에 대해 선형 복잡도 $\mathcal{O}(L)$를 달성합니다. 특히 오프라인 전처리 비용이 모든 쿼리에 대해 분할 상환(amortized)되므로, 쿼리가 많은 환경에서 기존 방식보다 상당한 효율성 이득을 제공합니다.
• 컨텍스트 저하 억제: 대규모의 무관한 컨텍스트가 모델의 추론력을 약화시키는 현상을 효과적으로 방지합니다. MS MARCO 벤치마크에서 1.6만 토큰 시 4.023이었던 점수가 1억 토큰 확장 시에도 3.669를 유지하며 8.8% 수준의 완만한 저하만을 보였습니다.

6. 시사점: 한계 및 향후 과제

MSA는 장문의 텍스트 컨텍스트에 대한 내재적 메모리를 강화했지만, 여러 문서에 걸쳐 강력하고 긴밀하게 결합된 의존성을 모델링하는 데는 여전히 한계가 있습니다. 증거가 여러 소스에 분산되어 고도로 상호 연결된 시나리오에서 순수 내재적 기억만으로 정확한 구조적 정렬을 유지하는 것은 어렵습니다. 메모리 인터리브가 이를 완화할 유망한 방향이지만, 문서 간 관계를 더 잘 보존할 수 있는 원칙적인 설계가 향후 필요합니다.

References:

1. MSA: Memory Sparse Attention for Efficient End-to-End Memory Model Scaling to 100M Tokens (https://github.com/EverMind-AI/MSA?tab=readme-ov-file)

Tags:

LLM #MemoryModel #SparseAttention #LongContext #ScalableAI #RetrievalAugmentedGeneration #언어모델 #메모리모델 #희소주의 #장문컨텍스트 #확장가능한AI #RAG