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

이 논문은 대규모 언어 모델(LLM)의 효율성과 성능 향상을 위해 조건부 계산(conditional computation)을 사용하는 MoE(Mixture-of-Experts) 방식에 더해, 조건부 메모리(conditional memory)라는 새로운 희소성 축(sparsity axis)을 제안한다. 기존 Transformer 모델은 지식 검색을 위한 고유한 기본 요소를 가지고 있지 않아, 계산을 통해 비효율적으로 정보를 검색한다고 지적한다. 이러한 문제를 해결하기 위해 고전적인 N-gram 임베딩을 현대화하여 O(1) 조회(lookup)가 가능한 Engram 모듈을 소개한다.

핵심 방법론: Engram 아키텍처

Engram은 Transformer 백본(backbone)에 정적 패턴 저장과 동적 계산을 구조적으로 분리하여 증강시키는 조건부 메모리 모듈이다. 입력 시퀀스 $X = (x_1, \dots, x_T)$와 층 $\ell$의 히든 스테이트 $H^{(\ell)} \in \mathbb{R}^{T \times d}$가 주어졌을 때, 각 위치 $t$에서 '검색(retrieval)'과 '융합(fusion)'의 두 단계를 거친다.

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

• Tokenizer Compression: 일반적인 subword 토크나이저의 의미론적 중복성 문제를 해결하기 위해 어휘 압축 레이어를 구현한다. NFKC 정규화 및 소문자 변환 등을 통해 토큰 ID를 정규화된 캐노니컬 ID($x'_t = P(x_t)$)로 변환하는 사상 함수 $P: V \to V'$를 사용한다. 이는 유효 어휘 크기를 줄여 의미론적 밀도를 높인다. 이 정규화된 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 해시로 구현된다. 최종 메모리 벡터 $e_t \in \mathbb{R}^{d_{mem}}$는 검색된 모든 임베딩을 연결하여 생성된다:
$$e_t \triangleq \parallel_{n=2}^N \parallel_{k=1}^K e_{t,n,k}$$

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

• 검색된 임베딩 $e_t$는 컨텍스트에 독립적인 사전 정보를 제공하지만, 해시 충돌이나 다의성으로 인한 노이즈가 발생할 수 있다. 이를 개선하기 위해 현재 히든 스테이트 $h_t$를 동적 Query로 사용하고 $e_t$를 Key 및 Value 투영의 소스로 사용하는 컨텍스트 인식 게이팅 메커니즘을 사용한다:

$$k_t = W_K e_t, \quad v_t = W_V e_t$$
여기서 $W_K, W_V$는 학습 가능한 투영 행렬이다. 게이트 $\alpha_t \in (0, 1)$는 다음과 같이 계산된다:
$$\alpha_t = \sigma \left( \frac{\text{RMSNorm}(h_t)^\top \text{RMSNorm}(k_t)}{\sqrt{d}} \right)$$
게이팅된 출력은 $\tilde{v}_t = \alpha_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, $M=4$)와 같은 다중 분기 아키텍처에 Engram을 통합하기 위해 매개변수 공유 전략을 사용한다. 단일 희소 임베딩 테이블과 $W_V$ 투영 행렬은 모든 $M$개 분기에 공유되는 반면, $M$개의 고유한 Key 투영 행렬 $\{W_K^{(m)}\}_{m=1}^M$는 분기별 게이팅 동작을 가능하게 한다. $m$-번째 분기의 히든 스테이트 $h_t^{(m)}$에 대한 게이팅 신호는 다음과 같다:

$$\alpha_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)} = \alpha_t^{(m)} \cdot (W_V e_t)$.

1. 시스템 효율성: 계산과 메모리 분리 (System Efficiency: Decoupling Compute and Memory):

• 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의 조건부 메모리 간의 상보성을 분석하며, 희소성 할당 문제(Sparsity Allocation problem)를 정의한다. 즉, 총 매개변수 예산이 고정되어 있을 때 MoE 전문가와 Engram 임베딩 사이에 용량을 어떻게 분배해야 하는가에 대한 문제이다.

• 최적 할당 비율 ($\rho$): MoE 전문가 용량에 할당되는 비활성(inactive) 매개변수 예산의 비율 $\rho \in [0, 1]$을 정의한다. 실험 결과, 검증 손실(validation loss)과 할당 비율 $\rho$ 사이에서 일관된 U자형 관계를 발견했다. 이는 순수 MoE 모델이 최적이 아니며, 전체 희소 매개변수 예산의 20-25%를 Engram에 재할당할 때 최상의 성능을 달성함을 보여준다. 10B 모델($C = 6 \times 10^{20}$ FLOPs)에서 검증 손실은 순수 MoE의 1.7248에서 $\rho \approx 80\%$ 근처의 1.7109로 개선되었다. 이는 두 모듈이 구조적으로 상보적임을 입증한다.
• 무한 메모리 영역 (Infinite Memory Regime): 메모리 예산이 완화되거나 공격적으로 확장될 때 Engram의 확장 동작을 탐구한다. 실험 결과, 메모리 슬롯 수가 증가함에 따라 검증 손실이 일관되게 개선되며, 로그 공간에서 선형적인 관계(power law)를 따른다. 이는 Engram이 예측 가능한 확장성을 제공하며, 추가적인 계산 없이도 더 큰 메모리가 계속해서 이점을 가져온다는 것을 나타낸다.

대규모 사전 학습 (Large Scale Pre-training)

제안된 Engram 아키텍처와 경험적으로 도출된 할당 법칙을 검증하기 위해, 2620억 개의 토큰으로 4개의 모델을 사전 학습시켰다: Dense-4B, MoE-27B, Engram-27B, Engram-40B. 모든 모델은 활성화된 매개변수(activated parameters) 수가 동일하게 맞춰졌다. Engram-27B는 MoE-27B와 엄격히 동일한 총 매개변수 수를 가지며, MoE의 라우팅된 전문가 수를 줄이고 해당 매개변수를 5.7B Engram 메모리에 재할당(72 $\to$ 55명의 전문가, $\rho=74.3\%$)했다.

실험 결과:

• 전반적 성능: 모든 희소 모델(MoE-27B, Engram-27B/40B)은 동일한 훈련 계산 예산 하에서 Dense-4B 모델을 모든 벤치마크에서 크게 능가했다.
• Engram의 우위: Engram-27B는 엄격히 동일한 활성화된 매개변수 및 총 매개변수 수를 가진 MoE-27B보다 광범위한 영역에서 우수한 성능을 보였다. 지식 집약적 작업(예: MMLU +3.4, CMMLU +4.0)뿐만 아니라, 일반 추론(예: BBH +5.0, ARC-Challenge +3.7) 및 코드/수학 도메인(HumanEval +3.0, MATH +2.4)에서도 훨씬 큰 개선을 보였다.
• 확장성: Engram-40B는 Engram-27B와 동일한 백본 및 계산 예산을 유지하면서 Engram 메모리만 18.5B 매개변수로 확장하여 (총 39.5B 매개변수), 지속적인 성능 향상을 달성하며 Engram의 확장성을 입증했다.
• 메커니즘 분석: LogitLens 및 CKA를 통한 분석에 따르면, Engram은 백본의 초기 층이 정적 지식을 재구성하는 부담을 덜어주어 복잡한 추론을 위한 효과적인 깊이(effective depth)를 증가시킨다. 또한, 지역적 의존성을 조회(lookup)에 위임함으로써 어텐션 용량을 전역 컨텍스트에 집중할 수 있게 하여 장문 컨텍스트 시나리오에서 뛰어난 성능을 가능하게 한다 (예: Multi-Query NIAH: 84.2 $\to$ 97.0).
• 인프라 효율성: Engram의 결정론적 주소 지정은 호스트 메모리로부터 런타임 프리페칭(runtime prefetching)을 가능하게 하여, 100B 매개변수 테이블을 오프로드하더라도 무시할 만한 오버헤드(<3%)만을 발생시킨다. 이는 GPU 메모리 제약을 효과적으로 우회하고 공격적인 매개변수 확장을 가능하게 함을 보여준다.

결론적으로, 이 논문은 조건부 메모리가 차세대 희소 모델을 위한 필수적인 모델링 기본 요소가 될 것이라고 제안한다.