MoME는 같은 토큰을 문맥별 메모리 슬롯으로 보내는 조건부 메모리 기법이다.
MoME: Mixture-of-Memory Embeddings for Context-Aware Sparse Lookup
무엇인가
이 논문은 LLM의 희소 용량 확장에서 조건부 메모리(conditional memory)가 가진 근본적 한계를 다룬다. Mixture-of-Experts가 토큰마다 일부 전문가만 활성화해 연산을 아끼듯, 조건부 메모리는 토큰 인덱스 임베딩 테이블을 조회해 백본에 저렴한 파라미터 프라이어를 주입한다. 그런데 Per-Layer Embedding, Value Embedding, STEM, Engram, Bigram 같은 기존 방법은 토큰 ID나 고정 n-gram 해시처럼 표면형의 결정적 함수로 메모리를 검색한다. 그래서 cats와 cat처럼 의미가 겹치는 토큰은 각자 행을 차지해 용량이 중복되고, bank나 python처럼 문맥에 따라 의미가 달라지는 토큰은 하나의 고정 벡터에 모든 의미가 뭉개진다. MoME는 이 검색을 문맥 인지적으로 만들되 토큰 인덱스 조회의 저렴함을 유지하려 한다.
어떻게 동작하나
구조적으로 MoME는 메모리 테이블을 E_mem ∈ R^{N×M×d_value}로 정의한다. N은 1단계 인덱서 f(x_t)가 고르는 행 수이고, M은 행마다 둔 메모리 슬롯 수, d_value는 주입 지점의 헤드별 값 차원이다. 먼저 f(x_t)=n_t가 토큰 행을 정하고, 2단계 문맥 인덱서 g_θ(h_t)가 그 행 안에서 읽을 슬롯을 정한다. 각 값 헤드 i에 대해 은닉 상태 h_t를 선형 사상해 슬롯 로짓 ℓ_{t,i}=W_g^{(i)}h_t+b_g^{(i)} ∈ R^M을 만들고, TopK로 K개 슬롯 A_{t,i}를 고른다. K>1이면 시그모이드 정규화로 선택된 슬롯에만 가중치 α를 배분하고, K=1이면 소프트맥스를 쓴다. 최종 메모리 m̃_{t,i}=Σ_{a∈A_{t,i}} α_{t,a}^{(i)} m_{n_t,a}를 값 스트림에 주입한다. 주입은 헤드별 값 잔차 게이트 γ_θ(h_t)=2σ(W_γ h_t+b_γ) ∈ R^H를 통해 ṽ_{t,i}=v_{t,i}+γ_{θ,t,i}m̃_{t,i}로 이뤄진다. 선택적으로 토큰 임베딩 공간에서 kNN으로 의미가 비슷한 토큰을 묶는 f를 써서 행 수를 줄이고, 줄어든 용량만큼 슬롯 M을 늘릴 수 있다. 기본 설정은 M=H, 활성 슬롯 K=2, 그룹화 없음이며 메모리 블록과 일반 트랜스포머 블록을 번갈아 둔다. 모든 값 헤드가 하나의 메모리 테이블을 공유해 파라미터를 N M d_value로 제한하되, 슬롯 라우터는 헤드별로 둔다. 라우팅과 검색을 값 사영과 병렬로 배치해 qwen3_4b 두 메모리 레이어 벤치마크에서 0.509ms 오버헤드를 냈고, 은닉 상태 주입 변형의 0.996ms보다 거의 절반이다.
무엇과 다른가
실험은 nanochat, Llama-3/MobileLLM, Qwen3 계열 백본에서 통제된 사전학습으로 진행됐다. 주요 표는 FineWeb-Edu를 524,288토큰 배치로 학습했고, 100B 토큰 실험과 d24 통제 실험은 ClimbMix를 1,048,576토큰 배치로 썼다. 모든 실험은 2048토큰 시퀀스와 32k BPE 토크나이저를 쓰고, 행렬 파라미터는 Muon, 임베딩과 메모리 테이블 등은 AdamW로 최적화했다. 비교 대상은 Base, Value Embedding, STEM, Bigram, Dense, VEmbedding이고, 외부 참조로 Qwen3-0.6B와 Llama 3.2-1B를 둔다. 지표는 bits per byte(bpb), CORE, CORE-22, 그리고 다의어 라우팅 분석을 위한 Word-in-Context(WiC)다.
어떻게 쓰나
nanochat 12층 135M 백본의 iso-FLOP 실험에서 MoME는 3×10^18 FLOPs, 약 3.3B 토큰을 학습했고 홀수 층마다 메모리 블록 하나를 넣었다. Bigram은 N_hash=6V로 151M 메모리 파라미터 예산을 맞췄다. 이 예산에서 MoME는 검증 bpb와 CORE 모두 Bigram과 VEmbedding보다 좋았다. 메모리를 302M로 두 배 늘리면 세 그룹화 인자 모두 검증 bpb가 더 낮아졌고, CORE는 c_grp=1과 2에서 개선됐다. 파라미터 수를 고정하고 c_grp∈{1,2,4}를 스윕했을 때 c_grp=2가 두 예산 모두에서 가장 높은 CORE를 냈고 302M에서는 가장 낮은 학습·검증 bpb를 냈다. 151M에서는 c_grp=4가 학습·검증 bpb를 약간 더 낮췄다. 라우팅 오버헤드는 iso-parameter Bigram 처리량의 2% 미만이었다.
전제와 한계
Llama/MobileLLM과 Qwen3 계열에서는 Llama3 스타일 MobileLLM 125M을 5B 토큰, MobileLLM 350M과 Qwen3 0.6B를 20B 토큰으로 학습했다. STEM은 원 논문의 1/2 설정처럼 절반의 층에 메모리를 넣었다. MoME는 메모리 없는 Base 대비 이득이 125M, 350M, 0.6B 세 스케일과 두 계열에서 안정적으로 전이됐고, 세 스케일 모두에서 STEM보다 검증 bpb와 CORE가 좋으면서 메모리 파라미터는 더 적었다. 학습 벽시계 시간은 메모리 없는 Base의 1.04~1.08배였다. nanochat d12에서 메모리 크기를 키우는 iso-FLOP 스케일링에서는 작은 메모리에서 Bigram과 비슷하게 시작하지만, 테스트한 모든 메모리 크기에서 동일 메모리 Bigram보다 검증 bpb가 낮았고 학습·검증 bpb의 실행 간 분산도 더 작았다. Bigram과의 복합 스케일링도 가능하다. Bigram 인덱서를 1단계 토큰 인덱서로 쓰고 메모리 테이블을 공유하도록 조정하면 같은 파라미터 예산에서 Bigram보다 좋아졌고, 더 큰 예산에서는 Bigram 테이블과 MoME 슬롯을 함께 키운 변형이 가장 좋았다. 추론 지연은 백본이 커질수록 상대 오버헤드가 줄었고, 100B ClimbMix 학습에서는 MoME가 Dense와 VEmbedding보다 세 가지 도메인 외 평가(FineWeb-Edu, enwik9, Shakespeare) 모두에서 낮은 bpb와 높은 CORE-22를 냈다. 다만 도메인 내 ClimbMix bpb는 VEmbedding보다 약간 높았고, Qwen3-0.6B의 CORE-22에 근접하면서도 사전학습 토큰은 훨씬 적었으며 전체 시퀀스 스코어링 처리량은 외부 참조와 비슷한 범위였다.
모델 분석은 메모리가 깊이에 따라 균일하게 쓰이지 않음을 보여준다. 각 메모리 층에서 헤드와 토큰 위치를 평균한 주입 게이트와 최종 메모리 잔차 크기를 재면, 메모리 사용은 모델 초반과 특히 후반에서 강하고 중간층에서는 상당히 낮다. 저자들은 이를 초반 문맥화와 후반 과제별 정교화에 메모리가 기여하고 중간층은 백본 표현 변환에 더 의존하는 패턴으로 해석한다. 라우팅의 의미 해석 가능성도 확인된다. bank, bug, drive 같은 다의어에서 같은 의미의 프롬프트는 같은 메모리 슬롯으로, 의미가 바뀐 프롬프트는 다른 슬롯으로 라우팅됐다. WiC 데이터셋의 엄격히 필터링된 670쌍을 d24 체크포인트의 144개 메모리 층·헤드 지점에서 평가했을 때, 다른 의미 쌍의 라우팅이 같은 의미 쌍보다 더 다른지를 나타내는 ΔJSD가 117개 지점에서 양수였고, 선택 슬롯 중첩을 뒤집은 ΔÕ가 110개 지점에서 양수였다. 즉 대부분의 메모리 라우팅 헤드가 같은 의미 문맥에서는 비슷하게, 의미가 바뀌면 다르게 라우팅한다.
저자들이 밝힌 가장 큰 한계는 계산 규모다. 확장된 100B 토큰 실험은 총 파라미터 약 1.4B에 도달하지만 실제 dense 백본은 약 0.8B이고 단일 시드 실험이다. CORE 활성화와 WiC 분석은 선택적 메모리 사용과 의미 민감 라우팅의 서술적 증거일 뿐, 이런 행동이 다운스트림 정확도나 도메인 이동 견고성을 실제로 개선하는지 인과적으로 입증하려면 개입 실험이 필요하다고 적는다. 향후 작업으로는 효율적 토큰 라우팅과 메모리 혼합을 결합하는 복합 스케일링 전략, 그리고 학습된 메모리 경로를 제어해 모델 행동을 조종하는 방향을 제안한다. 개발자 입장에서는 서브빌리언 규모에서 토큰 인덱스 메모리의 문맥 맹점을 줄이고 싶을 때 MoME를 검토할 만하다. 다만 메모리 예산, 활성 슬롯 수 K와 전체 슬롯 M, 그룹화 인자, 주입 위치, 추론 지연, 학습 안정성을 자신의 백본과 데이터로 다시 확인해야 하며, 특히 대규모·다중 시드 검증과 실제 태스크에서의 인과 효과 확인이 필요하다.