Diffusion은 신이고 GPT는 무적이다

ICML 2026. [Paper]
Albert Tseng, Christopher De Sa
Cornell University
29 Jan 2026

Introduction

본 논문에서는 tokenizer 임베딩 테이블의 개념을 확장하여 디코더 layer 내에서 컨텍스트에 따라 작동하는 Large Lookup Layer (L³)를 소개한다. L³의 핵심은 모델이 정보를 캐싱하는 데 사용할 학습된 lookup table 역할을 하는 토큰별 임베딩 모음이다. L³를 통해 모델은 디코더 layer 주변의 계산 시간을 단축할 수 있다. L³는 MoE 아키텍처와는 대체로 독립적이며, MoE layer를 대체하는 것이 아니라 모델의 sparsity를 더욱 확장하는 방법이다.

L³ 속도를 향상시키는 핵심은 토큰 ID별로 고정된 임베딩 세트를 선택하는 정적인 context-independent 라우팅을 사용하는 시스템 친화적인 아키텍처이다. MoE와 달리, 토큰이 생성되는 순간 필요한 정확한 L³ 파라미터를 알 수 있다. 이를 통해 L³ 파라미터 처리를 간소화하고, 기존의 context-dependent 라우팅에서는 불가능했던 fetching을 L³ 이전 단계의 연산과 중첩할 수 있다. 이러한 임베딩이 효율적으로 사용되도록 토큰 hidden state에서 임베딩을 관리하여 모델이 context-dependent한 조회를 수행할 수 있도록 한다.

L³ 압축 품질을 좌우하는 주요 요소는 토큰 ID에 임베딩을 할당하는 방식이다. L³는 정적인 토큰 라우터를 사용하기 때문에, LZW와 같은 무손실 압축 알고리즘을 사용하여 codeword의 빈도에 따라 임베딩을 할당할 수 있다. 단순한 균일 임베딩 할당 방식 대신 LZW 기반 할당 알고리즘을 사용하면 dense model 대비 perplexity 차이를 두 배로 늘릴 수 있다.

Method


단일 토큰 수준에서 L³ layer는 hidden state \(x \in \mathbb{R}^{d_\textrm{in}}\)과 해당 토큰 ID \(t \in \{1, \ldots, \vert \tau \vert\}\)를 입력으로 받는다. $\vert \tau \vert$는 tokenizer $\tau$의 vocabulary 크기이다. L³ layer는 mixing 행렬 \(W_\textrm{mix}\), up projection 행렬 \(W_\textrm{up}\), LayerNorm, key 행렬 시퀀스 \(K = \{K_t\}_{t=1}^{\vert \tau \vert}\), value 행렬 시퀀스 \(V = \{V_t\}_{t=1}^{\vert \tau \vert}\)로 parameterize된다.

\[\begin{aligned} \textrm{L}^3 (x, t; W_\textrm{mix}, W_\textrm{up}, K, V) = W_\textrm{mix} \begin{bmatrix} \textrm{LayerNorm}(W_\textrm{up} (V_t^\top \textrm{Softmax}(K_t x))) \\ x \end{bmatrix} \in \mathbb{R}^{d_\textrm{out}} \\ \textrm{where} \quad W_\textrm{mix} \in \mathbb{R}^{d_\textrm{out} \times (d_\textrm{in} + d_\textrm{up})}, \quad W_\textrm{up} \in \mathbb{R}^{d_\textrm{out} \times (d_\textrm{up} + d_\textrm{emb})}, \quad K_t \in \mathbb{R}^{d_t, d_\textrm{in}}, \quad V_t \in \mathbb{R}^{d_t, d_\textrm{emb}} \end{aligned}\]

품질과 속도를 조절하는 \(D \in \{d_t\}_{t=1}^{\vert \tau \vert} \in \mathbb{Z}^{+}\)를 제외한 모든 파라미터가 학습된다.

1. Embedding Allocation

\(v = \sum_i d_i\)개의 임베딩이라는 고정된 예산 내에서, 토큰에 임베딩을 할당하는 방식은 L³ layer의 성능에 상당한 영향을 미칠 수 있다. 본 논문에서는 $D$를 사용하여 MoE와 같은 context-dependent routing을 모방하고자 한다. 이러한 routing에 대한 자연스러운 조회 기반 근사는 컨텍스트에서 자주 나타나는 suffix를 캐싱하고, 임베딩을 선택할 때 주어진 컨텍스트의 가장 긴 suffix와 일치시키는 것이다. 이 근사는 $v$에 대한 할당 문제를 주어진 corpus에 대한 대표적인 prefix 집합을 찾는 문제로 축소한다.

suffix: 특정 토큰 위치를 끝점으로 하는 연속 토큰 시퀀스. (ex. 문장이 abc라는 3개의 토큰을 구성된 경우, suffix는 a, b, ab, c, bc, abc)

이는 사실상 무손실 텍스트 압축 문제이다. 연속된 구간에서 가장 긴 코드가 일치하도록 매칭할 때 기대 인코딩 길이를 최소화하는 코드 집합, 즉 토큰 튜플의 집합을 찾는 것이 목표이다. 따라서 고전적인 압축 알고리즘을 사용하여 임베딩을 할당할 codeword 집합을 구성할 수 있다.


저자들은 LZW 압축 기반 알고리즘을 도입하였다. 먼저 corpus를 스캔하여 아직 나타나지 않은 가장 긴 suffix를 찾아 codeword 집합과 각 codeword의 빈도를 구축한다. 새로운 suffix가 발견되면, 해당 suffix의 빈도를 1로 설정하고, 그 새로운 suffix에서 맨 앞 토큰 하나를 제거한 바로 다음 suffix의 빈도를 1 증가시킨다. 그다음 codeword들을 빈도가 높은 순서대로 순회하며, codeword가 토큰 $t$로 끝나는 경우 토큰 $t$에 임베딩 하나를 할당한다.

이 과정을 전체 임베딩 수가 목표값에 도달할 때까지 반복한다. 단, 각 토큰에는 최소 하나의 codeword가 할당되어야 하며, 최대 $k$개를 넘을 수 없다는 추가 조건을 둔다. 이러한 상한과 하한은 모든 토큰이 반드시 표현되도록 보장하는 동시에, 특정 토큰에 대해 활성화될 수 있는 파라미터 수의 최악의 경우를 제한한다.


위 그래프는 실험에 사용한 tokenizer에서 이 알고리즘을 실행한 결과이다.

2. Efficient Training and Inference

L³를 실제로 구현하기 위해 $K$를 단일 행렬 \(W_K \in \mathbb{R}^{v \times d_\textrm{in}}\)으로, 마찬가지로 $V$는 \(W_V \in \mathbb{R}^{v \times d_\textrm{emb}}\)로 연결한다. L³의 정적 라우팅 덕분에 토큰이나 시퀀스가 ​​생성되는 순간 $W_K$와 $W_V$의 활성 파라미터를 알 수 있다. 이러한 특성을 학습과 inference 모두에서 활용할 수 있다.


첫째, 최소한의 오버헤드로 L³ 파라미터를 offloading할 수 있다. 학습 과정에서 2.6B 파라미터를 가진 모델과 2,048의 컨텍스트 길이를 사용했을 때, 8,000개의 토큰으로 구성된 microbatch당 약 100M 파라미터만 활성화된다. 더 작은 batch를 사용하는 inference에서는 필요한 파라미터 수가 훨씬 더 적다. 이러한 파라미터는 L³ layer에 도달하기 전에 계산 중에 미리 가져올 수 있다. 또한, offloading된 L³는 dense MLP보다 컨텍스트 병렬 처리에 더 많은 복잡성을 추가하지 않으므로 컨텍스트 병렬 방식의 샤딩을 간단하게 만든다.


둘째, large-batch 설정에서 전체 batch를 정렬하여 메모리 접근 패턴을 개선할 수 있다. Hidden state \(B = \{x_1, \ldots, x_n\}\)와 이에 대응하는 토큰 \(T = \{t_1, \ldots, t_n\}\)으로 구성된 입력 batch를 생각해 보자. $B$에 L³를 적용하는 것은 $Q = B$, $K = W_K$, $V = W_V$인 일반 attention을 적용하고, attention mask를 $T$에서 허용된 토큰으로 설정하는 것과 같다. L³는 channel mixing만 수행하므로 $B$를 $T$로 정렬하여 attention mask를 block-diagonal 형태로 만들 수 있다. 이를 통해 FlexAttention과 같은 attention 커널을 사용하거나 간단하게 대각 block을 순회할 수 있다. 추가적인 오버헤드는 무시할 수 있을 정도로 작다.

반면, 이러한 최적화는 MoE에서 구현하기 어렵다. MoE에서는 관련된 router에 도달하기 전까지 토큰별 active expert들을 알 수 없다. 또한, 단일 MoE layer의 sparsity ratio가 L³ layer보다 훨씬 낮기 때문에 (4~32 vs >1000) 단일 학습 batch로 모델의 모든 expert를 거의 확실하게 처리할 수 있다. 반면 L³에서는 $\vert \tau \vert$가 합리적인 batch 크기를 훨씬 초과하므로 이는 불가능하다. MoE의 경우 이러한 특성으로 인해 offloading이 불가능해지므로 MoE를 샤딩해야 하는 장치 수가 늘어나고 병렬 처리 전략에 오버헤드가 추가된다.

Expertiments

  • 모델: 800M (디코더 400M), 1.5B (디코더 1B), 2.6B (디코더 1.9B)
  • 데이터셋: FineWeb-Edu (각 모델마다 10B, 20B, 30B 토큰)
  • 구현 디테일
    • 2~4배의 sparsity ratio를 목표로 하기 때문에 1~2개의 L³ layer만 사용
    • 전체 임베딩 테이블 크기: $v = 710,000$
    • 토큰당 최대 임베딩 수: $k = 512$

1. Pretraining Experiments

다음은 학습 perplexity를 비교한 결과이다.


다음은 다양한 다운스트림 성능을 비교한 결과이다.


다음은 동등한 FLOP과 sparsity의 MoE와 비교한 결과이다.

2. Architectural Design Choices

다음은 (왼쪽) L³ layer 수, (중간) 할당 알고리즘, (오른쪽) weight tying ablation 결과이다.


다음은 L³ layer의 위치에 따른 최종 학습 perplexity를 비교한 결과이다. (800M 모델, 355K 임베딩, L³ layer 1개, 1B 학습 토큰)

3. Training and Inference

다음은 다양한 batch 크기에 대한 inference 속도 비교 결과이다 (B200 GPU 1개). 첫 번째 L³ layer가 디코더 layer 뒤에 바로 위치하는 최악의 경우에도 L³가 추가하는 오버헤드는 10%에 불과하다.

You may also enjoy