Conditional Memory via Scalable Lookup: A New Axis of Sparsity for Large Language Models

"Conditional Memory via Scalable Lookup: A New Axis of Sparsity for Large Language Models" 논문은 Transformer 모델의 비효율적인 지식 검색 문제를 해결하기 위해 MoE(Mixture-of-Experts)와 상호 보완적인 새로운 희소성 축인 conditional memory 개념을 제안하고, 이를 O(1) lookup이 가능한 Engram 모듈로 구현한다.

기존 Transformer 모델은 지식 검색을 위한 고유한 기본 요소가 없어 계산을 통해 비효율적으로 정보를 검색한다고 지적하며, 이 논문은 고전적인 N-gram 임베딩을 현대화하여 조건부 메모리 기능을 제공하는 Engram 아키텍처를 도입한다. Engram은 Transformer 백본(backbone)에 정적 패턴 저장과 동적 계산을 구조적으로 분리하여 증강시키는 모듈이다.

핵심 방법론인 Engram 아키텍처는 크게 '검색(retrieval)'과 '융합(fusion)'의 두 단계를 거친다.

1. 희소 검색 (Sparse Retrieval via Hashed N-grams):

• Tokenizer Compression: 일반적인 subword 토크나이저의 의미론적 중복성을 줄이기 위해 어휘 압축 레이어를 구현한다. 이는 NFKC 정규화 및 소문자 변환을 포함하는 사상 함수 $P: V \to V'$를 사용하여 토큰 ID $x_t$를 정규화된 캐노니컬 ID $x'_t = P(x_t)$로 변환하여 유효 어휘 크기를 줄이고 의미론적 밀도를 높인다.
• N-gram 형성: 정규화된 ID들을 이용하여 suffix N-gram $g_{t,n} = (x'_{t-n+1}, \dots, x'_t)$을 형성한다.
• Multi-Head Hashing: 모든 N-gram 조합 공간을 직접 매개변수화하는 것은 비현실적이므로, 해싱 기반 접근 방식을 채택한다. 충돌을 완화하기 위해 각 N-gram 순서 $n$에 대해 $K$개의 개별 해시 헤드를 사용한다. 각 헤드 $k$는 압축된 컨텍스트를 임베딩 테이블 $E_{n,k}$ (크기 $M_{n,k}$) 내의 인덱스로 사상한다: $z_{t,n,k} \triangleq \varphi_{n,k} (g_{t,n})$. 여기서 $\varphi_{n,k}$는 경량의 곱셈-XOR 해시(multiply-XOR hash)로 구현된다.
• 메모리 벡터 생성: 검색된 모든 임베딩을 연결하여 최종 메모리 벡터 $e_t \in \mathbb{R}^{d_{mem}}$를 생성한다: $e_t \triangleq \parallel_{n=2}^N \parallel_{k=1}^K E_{n,k}[z_{t,n,k}]$.

1. 컨텍스트 인식 게이팅 (Context-aware Gating):

• 검색된 임베딩 $e_t$는 컨텍스트에 독립적인 사전 정보를 제공하지만, 해시 충돌이나 다의성으로 인한 노이즈가 발생할 수 있다. 이를 개선하기 위해 현재 히든 스테이트 $h_t$를 동적 Query로 사용하고 $e_t$를 Key 및 Value 투영의 소스로 사용한다: $k_t = W_K e_t$, $v_t = W_V e_t$. 여기서 $W_K, W_V$는 학습 가능한 투영 행렬이다.
• 게이트 $α_t \in (0, 1)$는 다음과 같이 계산된다: $α_t = \sigma \left( \frac{\text{RMSNorm}(h_t)^\top \text{RMSNorm}(k_t)}{\sqrt{d}} \right)$.
• 게이팅된 출력은 $\tilde{v}_t = α_t \cdot v_t$로 정의된다. 이 설계는 검색된 메모리가 현재 컨텍스트와 상충할 경우 게이트가 노이즈를 억제하도록 한다.
• 수용장(receptive field)을 확장하고 비선형성을 강화하기 위해 커널 크기 $w=4$, 팽창(dilation) $\delta$ (최대 N-gram 순서), SiLU 활성화 함수를 사용하는 짧고 깊이별(depthwise) 인과 컨볼루션(causal convolution)을 적용한다. 최종 출력 $Y$는 다음과 같다: $Y = \text{SiLU} (\text{Conv1D}(\text{RMSNorm}(\tilde{V}))) + \tilde{V}$.
• Engram 모듈은 잔여 연결(residual connection)을 통해 백본에 통합된다: $H^{(\ell)} \leftarrow H^{(\ell)} + Y$. Engram은 모든 층에 적용되는 것이 아니라 특정 층에 배치된다.

1. 다중 분기 아키텍처와의 통합 (Integration with Multi-branch Architecture):

• Manifold-Constrained Hyper-Connections (mHC)와 같은 다중 분기 아키텍처에 Engram을 통합하기 위해 매개변수 공유 전략을 사용한다. 단일 희소 임베딩 테이블과 $W_V$ 투영 행렬은 모든 $M$개 분기에 공유되는 반면, $M$개의 고유한 Key 투영 행렬 $\{W_K^{(m)}\}_{m=1}^M$는 분기별 게이팅 동작을 가능하게 한다.
• $m$-번째 분기의 히든 스테이트 $h_t^{(m)}$에 대한 게이팅 신호는 다음과 같다: $α_t^{(m)} = \sigma \left( \frac{\text{RMSNorm}(h_t^{(m)})^\top \text{RMSNorm}(W_K^{(m)} e_t)}{\sqrt{d}} \right)$.
• 검색된 메모리는 공유된 값 벡터 $W_V e_t$에 이 독립적인 게이트들이 적용되어 변조된다: $u_t^{(m)} = α_t^{(m)} \cdot (W_V e_t)$.

시스템 효율성 (System Efficiency) 측면에서, Engram의 결정론적 검색 메커니즘은 매개변수 저장을 계산 자원과 분리할 수 있게 한다. MoE와 달리 Engram의 검색 인덱스는 입력 토큰 시퀀스에만 의존하므로, 추론 시 비동기적인 프리페칭(prefetching)이 가능하다. Engram 모듈은 백본의 특정 층에 배치되어 선행 층의 계산을 버퍼로 활용하여 GPU stall을 방지한다. 학습 시에는 임베딩 테이블을 여러 GPU에 샤딩하여 모델 병렬화를 구현하고 All-to-All 통신을 사용한다. N-gram의 Zipfian 분포 특성을 활용하여 다단계 캐시 계층(Multi-Level Cache Hierarchy)을 구축함으로써, 자주 접근하는 임베딩은 GPU HBM이나 Host DRAM에 캐싱하고 드문 패턴은 NVMe SSD에 저장하여 효율성을 높인다.

확장 법칙 및 희소성 할당 (Scaling Laws and Sparsity Allocation)에 대한 분석에서, 이 논문은 MoE의 조건부 계산과 Engram의 조건부 메모리 간의 상보성을 분석하며, 총 매개변수 예산이 고정되어 있을 때 MoE 전문가와 Engram 임베딩 사이에 용량을 어떻게 분배해야 하는가에 대한 '희소성 할당 문제(Sparsity Allocation problem)'를 정의한다. 실험 결과, MoE 전문가 용량에 할당되는 비활성(inactive) 매개변수 예산의 비율 $\rho$와 검증 손실(validation loss) 사이에 일관된 U자형 관계를 발견했다. 이는 순수 MoE 모델이 최적이 아니며, 전체 희소 매개변수 예산의 20-25%를 Engram에 재할당할 때 최상의 성능을 달성함을 보여준다. 또한, 메모리 예산이 확장될 때 Engram의 확장 동작을 탐구한 결과, 메모리 슬롯 수가 증가함에 따라 검증 손실이 일관되게 개선되며 로그 공간에서 선형적인 관계(power law)를 따른다.

대규모 사전 학습 (Large Scale Pre-training)을 통해 제안된 아키텍처와 할당 법칙을 검증했다. 2620억 개의 토큰으로 Dense-4B, MoE-27B, Engram-27B, Engram-40B 모델을 사전 학습시켰으며, 모든 모델은 활성화된 매개변수(activated parameters) 수가 동일하게 맞춰졌다. 특히 Engram-27B는 MoE-27B와 엄격히 동일한 총 매개변수 수를 가지며, MoE의 라우팅된 전문가 수를 줄이고 해당 매개변수를 5.7B Engram 메모리에 재할당($72 \to 55$명의 전문가, $\rho=74.3\%$)했다. 실험 결과, Engram-27B는 MoE-27B보다 MMLU, BBH, HumanEval, MATH 등 광범위한 벤치마크에서 우수한 성능을 보였다. Engram-40B는 Engram-27B와 동일한 백본 및 계산 예산을 유지하면서 Engram 메모리만 18.5B 매개변수로 확장하여 지속적인 성능 향상을 달성하며 Engram의 확장성을 입증했다.

메커니즘 분석 (Mechanism Analysis)에 따르면, Engram은 백본의 초기 층이 정적 지식을 재구성하는 부담을 덜어주어 복잡한 추론을 위한 효과적인 깊이(effective depth)를 증가시킨다. 또한, 지역적 의존성(local dependencies)을 lookup에 위임함으로써 어텐션(attention) 용량을 전역 컨텍스트(global context)에 집중할 수 있게 하여 장문 컨텍스트 시나리오에서 뛰어난 성능을 가능하게 한다. 인프라 효율성 (Infrastructure Efficiency) 측면에서, Engram의 결정론적 주소 지정은 호스트 메모리로부터 런타임 프리페칭(runtime prefetching)을 가능하게 하여, 100B 매개변수 테이블을 오프로드하더라도 무시할 만한 오버헤드(<3%)만을 발생시킨다. 이는 GPU 메모리 제약을 효과적으로 우회하고 공격적인 매개변수 확장을 가능하게 한다.