기술 요약: 대규모 룩업 레이어 (Large Lookup Layers, L3)

문제 정의

현대적인 희소 언어 모델(sparse language models)은 주로 Mixture-of-Experts (MoE) 아키텍처에 의존하여 전체 파라미터 대비 활성 파라미터의 비율을 높입니다. MoE는 토큰을 특정 "전문가(expert)" MLP로 동적으로 라우팅함으로써 다운스트림 성능을 개선하지만, 이러한 문맥 의존적 라우팅은 심각한 시스템적 과제를 야기합니다. 구체적으로, 동적 라우팅은 라우터 붕괴(router collapse)를 방지하기 위한 보조 손실(예: 부하 분산)을 필요로 하며, 학습 안정성을 복잡하게 만들고, 하드웨어 효율성을 저해합니다. 토큰이 라우터에 도달하기 전까지는 어떤 전문가가 활성화될지 알 수 없기 때문에, 추론 시 파라미터를 CPU 메모리로 오프로딩(offloading)하는 것은 상당한 오버헤드를 발생시키며, 예측 불가능한 데이터 이동으로 인해 병렬화 전략이 복잡해집니다.

반면, 표준 토크나이저 임베딩 테이블은 본질적으로 희소하며 하드웨어 효율적이지만(토큰당 단일 임베딩 선택), 문맥 정보가 부족하며 정적인 룩업(lookup) 역할만 수행합니다. 본 논문은 동적 MoE 라우터의 "시스템적 짐(systems baggage)" 없이도 문맥 의존적 라우팅의 이점을 보존할 수 있는 희소 아키텍처의 필요성을 제기합니다.

방법론

저자들은 토크나이저 임베딩 테이블의 개념을 디코더 레이어 내에서 문맥 의고적으로 작동하도록 일반화한 새로운 아키텍처인 **Large Lookup Layer (L3)**를 소개합니다.

핵심 아키텍처

L3 레이어는 은닉 상태(hidden state) $x$와 토큰 ID $t$를 입력으로 받습니다. $x$를 기반으로 일부 전문가에게 동적으로 라우팅하는 대신, L3는 정적, 토큰 기반 라우팅을 사용합니다.

1. 정적 라우팅: 토큰 ID $t$가 생성되면, 모델은 즉시 어떤 임베딩 세트를 검색해야 하는지 알 수 있습니다. 이 세트는 토큰 ID로부터 대규모 임베딩 테이블의 일부로 매핑되는 사전 계산된 매핑에 의해 결정됩니다.
2. 문맥 의존적 집계 (Context-Dependent Aggregation): 토큰 ID에 기반하여 특정 임베딩(Keys $K_t$ 및 Values $V_t$)이 검색되면, 모델은 은-상태 $x$를 사용하여 이 임베딩들에 대해 어텐션 기반 집계를 수행합니다. 이를 통해 모델은 캐싱된 임베딩의 선택과 가중치 부여에 문맥 정보를 통합할 수 있습니다.
3. 수학적 공식화: 레이어는 $W_{mix} \begin{bmatrix} \text{LayerNorm}(W_{up}(V_t^T \text{Softmax}(K_t x))) \\ x \end{bmatrix}$를 출력합니다. 여기서 $K$와 $V$는 토큰 $t$에 할당된 임베딩 행렬입니다.

정보 이론적 임베딩 할당

L3의 핵심 구성 요소는 고정된 임베딩 예산($v$)을 어휘 사전($|\tau|$) 전체에 어떻게 할당할 것인가를 결정하는 것입니다. 저자들은 LZW(Lempel-Ziv-Welch) 압축 알고리즘의 변형인 Algorithm 1을 제안합니다.

• 메커니즘: 이 알고리즘은 코퍼스를 스캔하여 빈번하게 발생하는 문맥 접미사(suffixes of contexts)를 식별합니다. 임베딩 할당 문제를 무손실 텍스트 압축의 쌍대 문제(dual)로 취급하여, 자주 발생하는 문맥 접미사에 더 많은 임베딩을 할당합니다.
• 제약 조건: 할당은 모든 토큰이 최소 하나 이상의 임베딩을 받고 최대 $k$개(예: $k=512$)를 넘지 않도록 제한됩니다. 이는 토큰당 활성 파라미터의 최악의 경우 수를 제한하여, 예측 가능한 FLOPs와 메모리 사용량을 보장합니다.
• 결과: 이는 흔한 토큰(예: "then")은 수백 개의 임베딩을 받고, 희귀한 토큰은 적게 받는 헤비 테일(heavy-tailed) 분포를 생성하여, 동적 오버헤드 없이 문맥 의존적 라우터를 효과적으로 모사합니다.

시스템 친화적 설계

L3는 하드웨어 효율성을 고려하여 설계되었습니다:

• CPU 오프로딩: 필요한 파라미터가 토큰 생성 즉시 알려지므로, L3 파라미터는 이전 레이어의 연산 중에 CPU에서 GPU로 프리페치(prefetch)될 수 있습니다. 이는 I/O를 연산과 중첩(overlap)시킬 수 있게 하며, 이는 MoE에서는 달성하기 어려운 작업입니다.
• 배치 정렬 (Batch Sorting): 학습 중에는 배치의 토큰들을 토큰 ID 순으로 정렬할 수 있습니다. 이는 L3 레이어에 대한 "어텐션 마스크"를 블록 대각(block-diagonal) 구조로 변환하여, 최적화된 어텐션 커널(예: FlexAttention, MegaBlocks)을 사용할 수 있게 하고 MoE 특유의 불규칙한 메모리 액세스 패턴을 피할 수 있게 합니다.

주요 기여

1. L3 아키텍처: 토큰 기반의 정적 라우팅을 통해 문맥 의존적 집계를 가능하게 함으로써 임베딩 테이블을 디코더 레이어로 일반화한 새로운 희소 레이어입니다.
2. LZW 기반 할당 알고리즘: 코퍼스 빈도에 따라 토큰에 임베딩을 할당하는 정보 이론적 방법으로, 균등 할당보다 성능이 훨씬 뛰어납니다.
3. 시스템 효율성: 정적 라우팅이 CPU 오프로딩과 효율적인 배치 정렬을 가능하게 하여, 동적 MoE 라우팅과 관련된 오버헤드를 제거함을 입증했습니다.
4. 실증적 검증: 최대 2.6B의 활성 파라미터(및 훨씬 더 큰 전체 파라미터)를 가진 모델의 사전 학습을 통해, 밀집(dense) 및 MoE 베이스라인보다 우수한 성능을 입증했습니다.

결과

저자들은 디코더 레이어 사이에 L3 레이어를 삽 삽입하여, 활성 파라미터 수의 2~4배의 희소도를 목표로 하는 Llama 기반 트랜스포머를 사전 학습했습니다.

• 밀집 모델 대비 성능: L3 모델은 Perplexity(FineWeb-Edu 기준) 및 다운스트림 태스크(Wikitext2, ARC, Hellaswag 등)에서 iso-FLOP 밀집 모델을 유의미하게 능가합니다. 예를 들어, 활성 파라미터 2.6B, 총 파라미터 7B인 L3 모델은 2.6B 밀집 모델의 Perplexity가 15.43인 것에 비해 14.51을 달 기록했습니다.
• MoE 대비 성능: L3 모델은 iso-FLOP 및 iso-sparse MoE 대응 모델보다 우수한 성능을 보입니다. 특히 800M 파라미터 규모에서, 두 개의 L3 레이어와 동일한 수준의 MoE는 단일 L3 레이어보다 성능이 낮습니다. 라우터를 학습하는 동안 초기에 밀집 베이스라인보다 성능이 떨어질 수 있는 MoE와 달리, L3의 개선은 학습 과정 내내 즉각적이고 일관적입니다.
• 스케일링: L3를 통한 희소도 증가는 밀집 모델과 유사한 데로를 따르면서도, MoE의 라우팅 불안정성 없이 동등한 밀집 모델 대비 파라미터의 선형적 증가를 보여줍니다.
• 추론 속도: 단일 B200 GPU에서, CPU 오프로딩된 파라미터를 사용하는 L3 모델은 데이터 전송이 이전 레이어의 연산에 의해 마스킹되므로, 밀집 모델과 비교했을 때 최소한의 추론 오버헤드(최악의 경우 약 10%)만을 추가합니다.

의의 및 주장

본 논문은 L3가 언어 모델을 위한 **새롭고 더 효율적인 희소 스케일링 축(axis of sparse scaling)**을 제공한다고 주장합니다. 희소성의 이점(대규모 파라미터 수)을 동적 라우팅의 시스템 비용으로부터 분리함으로써, L3는 모델이 학습된 임베딩에 정보를 캐싱하고 연산을 "지름길(shortcut)"로 처리할 수 있게 합니다.

저자들은 L3가 MoE와 크게 무관하며(orthogonal), MoE의 직접적인 대체가 아니라 희소성을 더욱 확장하기 위한 방법론임을 강조합니다. 이 연구는 동적 라우팅 아키텍처에서 흔히 상실되는 하드웨어 친화적 특성(빠른 학습, 낮은 오버헤드의 오프로딩, 예측 가능한 메모리 액세스)을 유지하면서도 대규모 희소 모델의 성능 이점을 얻는 것이 가능하다는 것을 보여줍니다. 결과적으로, MoE의 "시스템적 짐"은 동적 문맥 의존적 라우팅의 필연적인 결과가 아니라, 정적이고 정보 이론적인 할당으로 대체될 수 있는 것임을 시사합니다.