기술 요약: 확장 가능한 룩업을 통한 조건부 메모리 (Engram)

1. 문제 정의

현재의 대규모 언어 모델(LLM), 특히 Mixture-of-Experts(MoE) 아키텍처를 활용하는 모델들은 용량을 확장하기 위해 **조건부 연산(conditional computation)**에 의존합니다. 그러나 표준 트랜스포머에는 **지식 룩업(knowledge lookup)**을 위한 고유한 프리미티브(primitive)가 부족합니다. 결과적으로, 모델은 동적 연산을 통해 검색을 비효율적으로 시뮬레이션해야 합니다.

저자들은 언어적 신호가 본질적인 이질성을 보인다고 주장합니다:

• 구성적 추론(Compositional Reasoning): 깊고 동적인 연산을 필요로 합니다.
• 지식 검색(Knowledge Retrieval): 텍스트의 상당 부분(예: 고유 명사, 정형화된 패턴)은 국소적이고, 정적이며, 매우 정형화되어 있습니다.

표준 트랜스포머는 (예: "Diana, Princess of Wales"를 해결하는 것과 같은) 정적인 룩업을 재구성하기 위해 여러 초기 레이어의 어텐션과 피드포워드 네트워크(FFN)를 소모하도록 강제하며, 이는 이론적으로 단순한 테이블 룩업을 통해 검색될 수 있는 작업에 귀중한 순차적 깊이(sequential depth)를 낭비하게 만듭니다. 이는 모델이 더 높은 수준의 추론에 자원을 할당하는 것을 방해합니다.

2. 방법론: Engram 모듈

이를 해결하기 위해, 본 논문은 조건부 메모리를 조건부 연산의 보완적인 희소성 축(sparsity axis)으로 도입합니다. 이는 고전적인 $N$-gram 임베딩을 현대화하여 $O(1)$ 상수 시간 룩업을 가능하게 하는 Engram을 통해 구현됩니다.

2.1. 아키텍처

Engram은 정적 패턴 저장과 동적 연산을 구조적으로 분리함으로써 트랜스포머 백본을 확장합니다. 이는 두 단계로 작동합니다:

1. 해시된 $N$-gram을 통한 희소 검색:
   • 토크나이저 압축: 원시 토큰 ID를 정규화된 식별자(대소문자, 구두점 등을 정규화)로 매핑하는 전사 함수(surjective function)를 사용하여 의미적 밀도를 극대화하고 어휘 크기를 줄입니다.
   • 멀티 헤드 해싱: $N$-gram의 조합적 공간을 처리하기 위해, 모듈은 $N$-gram 차수당 $K$개의 별도 해시 헤드를 사용합니다. 이들은 압축된 컨텍스트를 결정론적 함수를 통해 임베딩 테이블의 인덱스로 매핑하여 정적 벡터를 검색합니다.
   • 결합(Concatenation): 다양한 $N$-gram 차수와 해시 헤드에서 검색된 임베딩을 결합하여 메모리 벡터를 형성합니다.
2. 컨텍스트 인식 게이팅 및 퓨전(Fusion):
   • 정적 메모리는 컨텍스트가 없으므로, 게이팅 메커니즘은 현재의 은닉 상태(집계된 글로벌 컨텍스트)를 Query로 사용하여 검색된 메모리를 조절합니다.
   • 스칼라 게이트 $\alpha_t$는 정규화된 은닉 상태와 메모리 Key 사이의 내적을 통해 계산됩니다. 만약 메모리가 컨텍스트와 모순될 경우, 게이트는 메모리를 억제합니다(노이즈 감소).
   • 경량화된 depthwise causal convolution이 게이트가 적용된 메모리를 백본에 퓨전하기 전에 수용 영역(receptive field)을 확장합니다.

2.2. 시스템 효율성 및 인프라 공동 설계

동적 라우팅과 은닉 상태 통신이 필요한 MoE와 달리, Engram은 오직 입력 토큰 ID에 기반한 **결정론적 주소 지정(deterministic addressing)**을 사용합니다.

• 프리페칭(Prefetching): 인덱스는 포워드 패스 이전에 이미 알려지므로, PCIe를 통해 호스트 메모리(CPU/DRAM)로부터 임베딩을 비동기적으로 프리페칭할 수 있습니다.
• 오버랩(Overlap): 이를 통해 통신을 이전 트랜스포머 블록의 연산과 중첩시켜 발생할 수 있는 오버헤드를 무시할 수 있는 수준으로 낮출 수 있습니다.
• 오프로딩(Offloading): 거대한 임베딩 테이블(예: 100B 파라미터)을 호스트 메모리로 오프로드하여 GPU 메모리 제약을 우회할 수 있습니다.

2.3. 희소성 할당(Sparsity Allocation)

저자들은 희소성 할당 문제를 공식화합니다: 고정된 파라미터 예산이 주어졌을 때, 용량을 MoE 전문가(동적 연산)와 Engram 메모리(정적 룩업) 사이에 어떻게 나누어야 하는가?

• 그들은 비활성 파라미터의 비율을 나타내는 할당 비율 $\rho$를 정의합니다.
• 실험 결과 U자형 스케일링 법칙이 나타났습니다: 순수 MoE($\rho=100\%$)와 순수 Engram($\rho=0\%$)은 모두 최적이 아닙니다. 최적의 성능은 하이브리드 할당(예: $\rho \approx 75\%–80\%$)에서 나타나며, 여기서 희소 예산의 약 20–25%가 Engram으로 재할당됩니다.

3. 주요 결과

3.1. 사전 학습 성능

저자들은 활성화된 파라미터(3.8B)와 훈련 FLOPs를 엄격히 일치시킨 상태에서 262B 토큰으로 네 가지 모델을 학습시켰습니다:

• Dense-4B: 베이스라인.
• MoE-27B: 표준 MoE 베이스라인.
• Engram-27B: MoE-27B와 동일한 파라미터 규모를 가지며, 5.7B 파라미터를 Engram으로 재할당함(전문가를 72개에서 55개로 감소).
• Engram-40B: Engram 메모리를 18.5B 파라미터로 확장(총 39.5B).

결과:

• MoE 대비 우월성: Engram-27B는 다양한 도메인에서 동일 파라미터 규모의 MoE-27B 베이스라인을 능가했습니다.
• 추론 능력 향의: 메모리가 지식 검색에만 도움이 된다는 직관과 달리, Engram은 일반 추론(BBH +5.0, ARC-Challenge +3.7) 및 코드/수학(HumanEval +3.0, MATH +2.4)에서도 유의미한 이득을 보였습니다.
• 지식 습득: 지식 집약적 작업(MMLU +3.0, CMMLU +4.0)에서 성능이 향상되었습니다.
• 스케일링: Engram-40B는 손실(loss)을 더욱 감소시켜, 메모리 축이 추가 연산 없이도 효과적으로 스케일링됨을 시사했습니다.

3.2. 긴 컨텍스트 능력

로컬 의존성을 룩업으로 오프로드함으로써, Engram은 글로벌 컨텍스트를 위한 어텐션 용량을 확보합니다.

• RULER 벤치마크: Engram-27B는 Multi-Query Needle-in-a-Haystack(NIAH)에서 **97.0%**의 정확도를 달성하여, MoE 베이스라인의 **84.2%**를 상회했습니다.
• 효율성: Engram-27B는 단 **82%**의 사전 학습 FLOPs만으로 베이스라인의 LongPPL 성능을 맞추었습니다.

3.3. 메커니즘 분석

• 유효 깊이(Effective Depth): LogitLens와 CKA(Centered Kernel Alignment)를 사용한 분석 결과, Engram의 초기 레이어는 MoE 베이스라인의 훨씬 더 깊은 레이어와 의미적으로 동등한 표현을 생성함을 발견했습니다. 이는 Engram이 정적 특징을 재구성하기 위한 초기 레이어의 필요성을 제거함으로써 네트워크를 "깊게" 만든다는 것을 시사합니다.
• 기능적 이분법: 절제 연구(Ablation studies)에 따르면, Engram을 비활성화했을 때 사실적 지식 벤치마크는 붕괴(성능이 29–44%만 유지)한 반면, 독해 능력은 견고하게 유지되었습니다(81–93%). 이는 Engram이 파라미터 지식의 주요 저장소 역할을 한다는 것을 확인시켜 줍니다.

3.4. 시스템 효율성

• 처리량(Throughput): 100B 파라미터의 Engram 테이블을 호스트 메모리로 오프로드했을 때, 처리량 저하는 <3%(8B 백본 기준 2.8%)로 미미했습니다. 이는 프리페치 및 오버랩 전략의 유효성을 입증합니다.

4. 의의 및 주장

본 논문은 조건부 메모리가 차세대 희소 모델을 위한 필수적인 모델링 프리미티브라고 주장합니다. 그 의의는 다음 세 가지 측면에 있습니다:

1. 구조적 상보성: 정적 룩업과 동적 연산은 서로 배타적인 것이 아니라 구조적으로 상보적임을 입증했습니다. 둘 사이의 할당을 최적화하면 순수 MoE보다 뛰어난 U자형 스케일링 법칙을 얻을 수 있습니다.
2. 디커플링을 통한 효율성: 결정론적 주소 지정으로 파라미터 저장과 연산을 분리함으로써, Engram은 GPU 메모리나 연산 비용에 구애받지 않고 거대한 파라미터 수(예: 40B+)로 확장할 수 있게 하여 하드웨어 병목 현상을 효과적으로 우회합니다.
3. 네트워크의 심화: 메커니즘 분석은 명시적인 지식 룩업을 제공하는 것이 초기 레이어가 정적 재구성을 수행해야 하는 부담을 덜어준다는 것을 보여줍니다. 이는 복잡한 추론을 위한 네트워크의 가용 깊이를 실질적으로 증가시키며, 추론 및 코드 작업에서의 예상치 못한 이득을 설명해 줍니다.

저자들은 Engram을 단순한 검색 증강(retrieval augmentation)이 아니라, $N$-gram 개념을 대규모 희소 모델의 시대로 현대화하여 기존의 MoE 패러다임을 보완하는 새로운 희소성 축을 제시하는 근본적인 아키텍처 변화로 규정합니다.