[논문 리뷰] Conditional Memory via Scalable Lookup: LLM의 새로운 희소성 축, Engram

TL;DR

대규모 언어 모델(LLM)은 방대한 지식을 매번 추론 과정에서 다시 계산해야 하는 비효율성 문제를 안고 있습니다. 이 논문은 LLM에 **조건부 메모리(Conditional Memory)**라는 새로운 개념을 도입하여 이 문제를 해결합니다. Engram이라는 O(1) 조회 속도의 메모리 모듈을 통해, 모델이 정적인 지식(예: "파리의 수도는 프랑스")은 빠르게 찾아 쓰고, 동적인 추론에 계산 자원을 집중하도록 설계했습니다. 실험 결과, Engram과 기존의 MoE(Mixture-of-Experts)를 결합한 모델은 순수 MoE 모델보다 지식 검색 및 일반 추론 능력에서 월등한 성능을 보였습니다. 이 연구는 정적 지식 조회와 동적 연산을 분리하는 새로운 희소성 축을 제시하며, 차세대 LLM 아키텍처의 중요한 방향을 제시합니다.

연구 배경 및 동기

LLM은 놀라운 성능을 보여주지만, 모든 입력을 처리할 때 모델의 전체 파라미터를 활성화하는 'Dense' 방식은 막대한 계산 비용을 유발합니다. 마치 우리가 "1+1"을 계산할 때마다 덧셈의 원리를 처음부터 증명하는 것처럼 비효율적입니다.

이 문제를 해결하기 위해 **MoE(Mixture-of-Experts)**가 등장했습니다. MoE는 입력 토큰에 따라 가장 적합한 '전문가(Expert)' 신경망 일부만 활성화하는 **계산적 희소성(Computational Sparsity)**을 구현하여 효율을 높였습니다. 하지만 MoE 역시 정적인 '사실'을 기억하기 위해 여전히 신경망 연산에 의존합니다.

이 논문은 여기서 한 걸음 더 나아가, 아예 **정적 지식(static knowledge)**과 **동적 추론(dynamic reasoning)**을 분리하자는 아이디어를 제안합니다. 즉, 이미 학습된 사실적 정보는 계산 없이 바로 '조회(lookup)'하고, 복잡한 추론이 필요할 때만 신경망을 사용하자는 것입니다. 이것이 바로 조건부 메모리의 핵심이며, 이를 구현한 모듈이 Engram입니다.

관련 연구

LLM의 효율성을 높이려는 시도는 꾸준히 이어져 왔습니다.

• Mixture-of-Experts (MoE): 입력에 따라 특정 전문가 네트워크를 선택적으로 활성화하여 계산량을 줄입니다. 하지만 이는 '어떻게 계산할지'를 효율화할 뿐, '무엇을 기억할지'의 문제는 해결하지 못합니다.
• Transformer: 어텐션 메커니즘으로 문맥을 파악하지만, 모든 토큰에 대해 연산을 수행하므로 정적 정보도 매번 다시 계산해야 합니다.
• Retrieval-Augmented Generation (RAG): 외부 데이터베이스에서 관련 정보를 검색하여 모델 입력에 추가합니다. 효과적이지만, 외부 DB에 의존하며 추론 시 지연이 발생할 수 있습니다.

본 논문의 Engram은 RAG처럼 정보를 '검색'하지만, 이를 모델 내부에 완전히 통합된 파라미터화된 메모리로 구현했다는 점에서 차별화됩니다.

핵심 기여

1. 조건부 메모리 개념 도입: 정적 지식 조회를 위한 새로운 희소성 축을 제안하여, 동적 추론과의 분리를 가능하게 했습니다.
2. Engram 모듈 설계: N-gram 임베딩을 현대적으로 재해석하여, 해시 기반의 $O(1)$ 시간 복잡도로 정보를 조회하는 확장 가능한 메모리 모듈을 구현했습니다.
3. U자형 스케일링 법칙 발견: 희소성 할당에 대한 중요한 발견으로, 순수 Engram이나 순수 MoE보다 둘을 적절히 혼합했을 때 최적의 성능을 보인다는 것을 실험적으로 증명했습니다.
4. 모델 효율성 증대: Engram이 모델의 초기 레이어에서 간단한 정적 패턴을 처리해주면, 후속 MoE 레이어는 더 복잡한 추론에 집중할 수 있어 전체적인 효율과 성능이 향상됩니다.

제안 방법론

핵심 아이디어: 지식과 추론의 분리

Engram의 핵심은 LLM의 역할을 두 가지로 나누는 것입니다.

• 지식 저장소 (Engram): 세상의 정적인 사실(N-gram 패턴)을 거대한 키-값(key-value) 테이블에 저장합니다.
• 추론 엔진 (MoE, Attention): Engram에서 조회한 지식과 현재 문맥을 바탕으로 동적인 추론을 수행합니다.

모델 아키텍처 상세 설명

Engram 모듈은 Transformer 블록 내에서 Attention, MoE와 나란히 위치하며, **검색(Retrieval)**과 **융합(Fusion)**의 두 단계로 작동합니다.

1. 희소 검색 (Sparse Retrieval)

입력 토큰 시퀀스에서 N-gram(연속된 N개의 토큰 묶음)을 키(key)로 사용하여 메모리에서 값(value)을 가져옵니다.

1. 토크나이저 압축 (Tokenizer Compression): 의미적으로 동일하지만 토큰 ID가 다른 단어들을 정규화합니다. 예를 들어, Apple, apple, AAPL을 모두 같은 개념 ID로 매핑하여 메모리 효율성을 높입니다.
2. 다중 헤드 해싱 (Multi-Head Hashing): 텍스트의 N-gram(예: 2-gram, 3-gram)을 여러 개의 독립적인 해시 함수로 해싱하여 거대한 임베딩 테이블의 인덱스로 변환합니다. 예를 들어, "Large Language Models"라는 텍스트가 있다면, hash_1("Large Language"), hash_2("Language Models") 와 같이 여러 개의 키를 생성하여 메모리를 조회합니다. 이는 해시 충돌을 줄이고 더 풍부한 정보를 검색하는 효과를 낳습니다.

2. 문맥 인식 융합 (Context-aware Fusion)

검색된 N-gram 임베딩은 문맥 정보가 없는 순수한 정적 정보입니다. 이를 현재 문맥에 맞게 융합하는 과정이 필요합니다.

1. 문맥 인식 게이팅 (Context-aware Gating): 검색된 정적 임베딩 벡터($e_t$)가 현재 문맥에 얼마나 중요한지 결정합니다. 현재 레이어의 은닉 상태($h_t$)를 쿼리(Query)로, 검색된 임베딩을 키(Key)/값(Value)으로 사용하는 어텐션과 유사한 메커니즘을 통해 게이트 값($\alpha_t$)을 계산합니다. 예를 들어, "river bank"라는 문맥에서 bank를 조회했다면, 게이트는 '금융 기관'이 아닌 '강둑'의 의미에 더 높은 가중치를 부여하게 됩니다.
2. 경량 컨볼루션 (Lightweight Convolution): 게이팅을 거친 값에 깊이별 인과적 컨볼루션(depthwise causal convolution)을 적용하여, 지역적인 문맥 정보를 통합하고 비선형성을 추가합니다.
3. 잔차 연결 (Residual Connection): 최종 결과는 기존 은닉 상태($h_t$)에 더해져 다음 모듈로 전달됩니다.

핵심 수식

1. 다중 헤드 해싱: N-gram $g_{t,n}$을 $k$번째 해시 함수 $\phi_{n,k}$로 해싱하여 임베딩 테이블 $E_{n,k}$에서 임베딩 $e_{t,n,k}$를 조회합니다. $e_{t,n,k} = E_{n,k}[\phi_{n,k}(g_{t,n})]$

2. 문맥 인식 게이트 값 ($\alpha_t$): 현재 은닉 상태 $h_t$와 검색된 임베딩 $k_t$ (여기서는 $e_t$를 의미)의 유사도를 계산하여 시그모이드 함수를 통과시킨 값입니다. $\alpha_t = \sigma \left( \frac{[\text{RMSNorm}(h_t)]^T [\text{RMSNorm}(k_t)]}{\sqrt{d}} \right)$

3. 희소성 할당 비율 ($\rho$): 전체 희소 파라미터($P_{sparse}$) 중 MoE와 Engram에 얼마만큼의 파라미터를 할당할지 결정하는 비율입니다. $\rho=1$이면 순수 MoE, $\rho=0$이면 순수 Engram 모델이 됩니다. $P_{\text{MoE}}^{\text{sparse}} = \rho \times P_{\text{sparse}}, \quad P_{\text{Engram}} = (1-\rho) \times P_{\text{sparse}}$

실험 설정

• 데이터셋: C4, The Pile 등 262B 토큰을 포함하는 대규모 데이터셋
• 평가 지표: MMLU(종합 지식), BBH(추론), HumanEval(코딩), MATH(수학) 등 다양한 벤치마크
• 베이스라인: Dense-4B(40억 파라미터 밀집 모델), MoE-27B(270억 희소 파라미터 MoE 모델)
• 제안 모델: Engram-27B(Engram+MoE 혼합), Engram-40B

하이퍼파라미터

실험 결과 분석

주요 결과: Engram의 압도적 성능

동일한 희소 파라미터 수(27B)를 가진 모델을 비교했을 때, Engram과 MoE를 혼합한 Engram-27B 모델이 순수 MoE-27B 모델을 모든 벤치마크에서 크게 앞섰습니다.

이는 Engram이 정적 지식을 효율적으로 처리해준 덕분에, 모델의 나머지 부분이 더 복잡한 추론에 집중할 수 있었음을 시사합니다.

U자형 스케일링 법칙

이 논문의 가장 흥미로운 발견 중 하나는 U자형 스케일링 법칙입니다. 희소 파라미터를 Engram(지식)과 MoE(추론)에 어떻게 배분하는지에 따라 모델 성능이 달라졌습니다.

• $\rho=0$ (순수 Engram): 지식은 많지만 추론 능력이 부족하여 성능이 낮음.
• $\rho=1$ (순수 MoE): 추론 능력은 있지만 모든 지식을 연산해야 해 비효율적이고 성능이 낮음.
• $0 < \rho < 1$ (혼합 모델): Engram과 MoE가 시너지를 내며 최적의 성능을 보임.

이는 마치 암기 과목과 수리 과목을 균형 있게 공부해야 시험을 잘 보는 것과 같습니다. 지식과 추론 능력, 둘 중 하나만으로는 최고의 성능을 낼 수 없다는 것을 보여줍니다.

비판적 평가

강점

1. 새로운 희소성 축 제시: MoE의 계산적 희소성을 넘어, 지식 조회를 위한 메모리 희소성이라는 새로운 차원을 열었습니다.
2. 효율적인 지식 검색: Engram 모듈은 $O(1)$ 시간 복잡도로 정적 지식을 조회하여 LLM의 고질적인 연산 비효율 문제를 해결합니다.
3. 뛰어난 성능과 확장성: 기존 MoE 구조와 자연스럽게 결합하여 시너지를 내며, 실험을 통해 그 효과를 명확히 입증했습니다.

한계점과 개선 방향

1. 구현의 복잡성: 다중 헤드 해싱, 문맥 인식 게이팅 등 Engram 모듈의 구현은 기존 모델보다 복잡하여 실무 적용에 허들이 될 수 있습니다.
2. 메모리 요구량: N-gram 임베딩을 저장하기 위한 거대한 테이블은 상당한 메모리 공간을 차지합니다.
3. 정적 메모리의 한계: Engram에 저장된 지식은 사전 학습 데이터에 고정되므로, 새로운 정보를 반영하기 어렵고 데이터의 편향을 그대로 학습할 위험이 있습니다.

향후 연구 방향

1. 동적 Engram: 학습 후에도 Engram의 내용을 수정하거나 추가할 수 있는 동적 메모리 메커니즘에 대한 연구가 필요합니다.
2. 의미 기반 검색: N-gram 같은 표면적 형태가 아닌, 의미적 유사도에 기반한 검색 방식으로 Engram을 개선할 수 있습니다.
3. 최적의 혼합 비율 탐색: 다양한 모델 크기와 태스크에 대해 Engram과 MoE의 최적 혼합 비율($\rho$)을 찾는 연구가 이어질 수 있습니다.

실무 적용 가이드

• 모듈 배치: Engram은 간단한 패턴을 먼저 처리하는 것이 효과적이므로, Transformer 블록의 초기 레이어에 배치하는 것이 좋습니다.
• 하드웨어 고려: Engram의 거대한 임베딩 테이블을 수용하기 위해 충분한 RAM 또는 HBM(고대역폭 메모리)이 필수적입니다.
• 사전 계산: 학습 데이터에 대한 N-gram 해시 값은 미리 계산하여 저장해두면 학습 속도를 높일 수 있습니다.

결론

"Conditional Memory via Scalable Lookup"은 LLM의 비효율성 문제를 해결하기 위해 '지식 조회'와 '동적 추론'을 분리하는 혁신적인 접근법을 제시합니다. Engram 모듈은 이 아이디어를 성공적으로 구현했으며, MoE와의 결합을 통해 차세대 희소 모델의 청사진을 보여주었습니다. 이 연구는 단순히 모델의 성능을 높이는 것을 넘어, LLM이 어떻게 지식을 저장하고 활용해야 하는지에 대한 근본적인 질문에 중요한 단서를 제공합니다.

참고 자료

• 논문 링크: arXiv:2401.07372
• 코드 저장소: (논문 공개 시 추가될 수 있음)
• 관련 자료: (관련 블로그나 발표 자료 링크)