MSA: Memory Sparse Attention은 LLM의 효과적인 컨텍스트 길이 제약을 해결하고 100M 토큰에 이르는 극도로 긴 컨텍스트를 처리하기 위해 제안된 end-to-end trainable, scalable sparse latent-state memory 프레임워크입니다. 기존 LLM은 풀 어텐션(full attention)의 계산 비용 때문에 컨텍스트 길이가 128K~1M 토큰으로 제한되며, 선형 어텐션(linear attention), 고정 크기 상태 메모리(fixed-size state memory), 외부 저장소(external storage)를 사용하는 RAG/Agents 같은 시도들은 극단적인 스케일에서 정밀도 저하 및 지연 시간 증가, end-to-end 미분 불가능성, 또는 복잡한 파이프라인과 같은 단점을 가집니다.

MSA는 이러한 한계를 극복하기 위해 세 가지 핵심 아이디어를 통합합니다:

1. Memory-Sparse Attention (MSA) Layer:
   • Scalable Sparse Attention: 문서(document)의 latent state인 $K$, $V$, 그리고 라우터 키(router key) $K_r$을 chunk-mean pooling하여 압축합니다. 쿼리(query)는 router projector를 통해 $Q_r$로 변환되며, 이 $Q_r$은 $K_r$과의 코사인 유사도(cosine similarity)를 통해 문서들의 관련성(relevance)을 계산합니다. 유사도 계산은 헤드(head)별 평균 후 토큰별 최대값(token-wise max)을 사용합니다. 이를 통해 가장 관련성이 높은 Top-k 문서를 선택하고, 선택된 문서들의 압축된 K/V를 쿼리의 로컬 K/V와 연결(concatenation)하여 autoregressive decoding에 사용합니다. 이러한 Sparse Attention 메커니즘은 트레이닝과 인퍼런스 모두에서 $O(L)$의 거의 선형적인(near-linear) 복잡도를 달성합니다.
   • Document-wise RoPE (Rotary Positional Embedding): 각 문서는 위치 임베딩(positional embedding)을 0부터 다시 시작하여 적용하며, 이는 "train-short, infer-long" 문제, 즉 짧은 컨텍스트로 학습하고 긴 컨텍스트로 추론할 때 발생하는 위치 표류(position drift)를 방지하여 64K 토큰 학습이 100M 토큰으로 외삽(extrapolate)될 수 있도록 합니다.
   • Global RoPE: 쿼리의 시작 인덱스에 Top-k로 검색된 블록(blocks)의 개수 $k$를 offset으로 추가하여 배경(background) → 쿼리(query) → 생성(generation)의 인과적 순서(causal ordering)를 보존합니다.
   • 이러한 라우팅(routing)은 모델의 상위 계층(upper layers)에만 적용되며, 하위 계층(lower layers)은 계층적 정렬(hierarchical alignment)을 위해 독립적인 문서 처리를 유지합니다.

2. KV Cache Compression + Memory Parallel Inference Engine:
   • KV Cache Compression: 메모리 계층 구조(tiered storage)를 활용합니다. GPU에 상주하는 라우팅 키($K_r$)와 CPU의 DRAM에 저장되는 실제 콘텐츠 K/V를 분리하여 저장합니다.
   • Memory Parallel: 분산된 GPU 환경에서 작동하는 인퍼런스 엔진으로, 라우팅 키 $K_r$을 여러 GPU에 샤딩(sharding)합니다. 쿼리는 모든 GPU에 브로드캐스트(broadcast)되고, 각 GPU는 로컬에서 스코어링(scoring)을 수행한 후 전역적으로 Top-k를 선택하기 위한 리듀스(reduce) 연산을 수행합니다. 선택된 콘텐츠 K/V는 CPU DRAM에서 비동기적으로 페치(fetch)되어 GPU로 전송되며, 이를 통해 제한된 VRAM 내에서 100M 토큰 규모의 처리량(throughput)을 가능하게 합니다.

3. Memory Interleave:
   • "generative retrieval → context expansion → generation" 단계를 적응적으로 반복하는 메커니즘입니다. 이를 통해 여러 문서에 걸쳐 다중 홉 추론(multi-hop reasoning) 능력을 크게 향상시킵니다.

MSA의 인퍼런스 파이프라인은 세 단계로 구성됩니다:

4. Global Memory Encoding (오프라인): 전체 코퍼스(corpus)에 대해 포워드 패스(forward pass)를 수행하여 chunk-pooled된 $(\bar{K}, \bar{V}, \bar{K_r})$를 캐시합니다.
5. Online Routing & Context Assembly: 쿼리를 $Q_r$로 투영하고, $\bar{K_r}$와 매칭하여 Top-k 문서를 선택한 후, 선택된 $\bar{K}/\bar{V}$만 로드하여 로컬 컨텍스트와 연결합니다.
6. Sparse Generation: 최종적으로 스파스(sparse) 컨텍스트에 대해 autoregressive하게 생성합니다.

평가 결과, MSA는 16K에서 100M 토큰 범위에서 9% 미만의 성능 저하를 보이며, 이는 메모리 용량과 추론 능력 간의 디커플링(decoupling)이 가능하다는 것을 시사합니다. 다양한 QA 데이터셋(MS MARCO v1, NQ, TriviaQA(10M) 등)과 NIAH(Needle-in-a-Haystack) 벤치마크(RULER)에서 MSA는 동일한 백본(backbone)을 사용하는 RAG, 최상위 RAG 스택(best-of-breed RAG stacks), 그리고 선도적인 장문 컨텍스트 모델들을 능가하는 성능과 안정성을 입증했습니다. 특히, RULER NIAH 벤치마크에서는 1M 토큰에서 94.84%의 정확도를 유지하며, 이는 128K 토큰 이상에서 성능이 급격히 저하되는 기존 모델들과 대조됩니다.

MSA는 158.95B 토큰의 연속 사전 학습(continuous pretraining)과 보조 라우팅 손실(auxiliary routing loss)을 포함한 트레이닝 과정, 그리고 8K에서 64K로 확장되는 2단계 SFT(Supervised Fine-Tuning) 커리큘럼(curriculum)을 통해 학습됩니다. MSA는 memory capacity와 reasoning을 분리하는 실용적인 경로를 제시하며, 초장문 컨텍스트 AI 시스템 개발에 기여합니다.