논문: Better & Faster Large Language Models via Multi-token Prediction
2026-10-25 원문 재검토: 2024년 v1 의 §2–3, 부록 A·C·D·E를 기준으로 기존 GPT 대화 요약을 수정했습니다. 근거 없는 Llama 2 메모리 용량과 DDR 가속 추정치를 제거하고, 학습 목표·검증을 포함한 추론·실측 조건을 구분했습니다.
무엇을 바꾸는가
일반적인 next-token prediction은 각 문맥에서 바로 다음 토큰을 맞히도록 학습합니다. 이 논문은 같은 문맥으로부터 여러 미래 위치를 예측하는 보조 과제를 추가합니다. 학습된 첫 번째 헤드만 사용하면 기존처럼 한 토큰씩 생성할 수 있고, 추가 헤드를 사용하면 자기 추측 디코딩(self-speculative decoding)으로 추론을 가속할 수 있습니다.
따라서 학습 목표 개선과 추론 가속은 별도로 평가해야 합니다. 여러 토큰을 예측하도록 학습했다는 사실만으로 생성 속도가 자동으로 빨라지지는 않습니다.
공유 트렁크와 미래 토큰별 헤드
문맥을 \(x_{1:t}\), 공유 트렁크의 출력을 \(z_t=f_s(x_{1:t})\)라고 쓰겠습니다. 기본 구조에서는 각 미래 위치 \(i\)에 독립적인 Transformer 헤드 \(f_{h_i}\)를 두고, 어휘 공간으로 투영하는 unembedding \(f_u\)는 공유합니다.
\[ p_i(\cdot\mid x_{1:t}) =\operatorname{softmax}\bigl(f_u(f_{h_i}(z_t))\bigr). \]경계를 넘어서는 타깃을 제외한 학습 손실은 다음처럼 쓸 수 있습니다.
\[ \mathcal L_n =-\sum_{i=1}^{n}\sum_{t=1}^{T-i} \log p_i(x_{t+i}\mid x_{1:t}). \]각 헤드는 같은 관측 문맥을 사용합니다. 두 번째 헤드가 첫 번째 헤드의 예측 토큰을 입력으로 받는 구조가 기본 설계는 아닙니다. 예를 들어 문맥이 The quick brown fox이면 두 헤드는 각각 다음 위치와 그다음 위치의 토큰을 예측합니다. 여기서 토큰은 반드시 단어 하나와 일치하지 않습니다.
독립 헤드의 예측을 그대로 이어 붙이면 일반적인 autoregressive 모델의 조건부 예측과 일치한다고 보장할 수 없습니다. 이 차이 때문에 추론 시 검증이 필요합니다.
학습 비용은 왜 전체 모델의 4배가 아닌가
공유 트렁크는 문맥마다 한 번 계산합니다. 미래 토큰 수를 4로 늘려도 전체 트렁크를 네 번 계산하는 방식이 아닙니다. 또한 논문의 주된 비교는 총 파라미터 수를 맞춥니다. 추가 헤드에 Transformer 층을 배치하는 대신 공유 트렁크의 층 수를 줄입니다(§3). 기존 Llama 2 전체 층 위에 헤드를 그대로 추가하는 경우와 비용 조건이 다릅니다.
로짓을 모든 헤드에 대해 동시에 유지하는 순진한 구현은 메모리가 커집니다. 논문은 헤드 하나씩 forward/backward를 수행하고 트렁크에 대한 gradient를 누적한 뒤 로짓을 해제합니다. 이로써 해당 로짓·트렁크 gradient 저장의 차수를 \(O(nV+d)\)에서 \(O(V+d)\)로 줄입니다. \(V\)는 어휘 크기, \(d\)는 은닉 차원이며, 배치·시퀀스 축은 생략한 표기입니다. 이는 모델 가중치·옵티마이저·모든 활성화를 합친 전체 VRAM 공식은 아닙니다.
부록 C, Table S5에는 실제 구현의 시간 오버헤드도 있습니다.
| 모델 크기 | next-token 학습 | 4-token 학습의 상대 시간 |
|---|---|---|
| 0.3B | 1.00 | 1.22 |
| 6.7B | 1.00 | 1.07 |
| 13B | 1.00 | 1.09 |
저자들은 이를 헤드별 backward와 FSDP 통신·계산 중첩의 구현 문제로 설명하며 개선 가능하다고 봅니다. 초록의 “추가 학습 시간 없음”을 모든 구현에서 오버헤드 0%라는 뜻으로 읽어서는 안 됩니다. 반대로 4-token이라는 이유로 연산량·활성화 메모리가 모두 4배라는 계산도 맞지 않습니다.
Llama 2 7B의 정확한 VRAM이나 메모리 트래픽은 배치, 길이, 정밀도, 체크포인팅, 옵티마이저, 헤드 구성 없이는 산출할 수 없습니다. 저장 용량(GB)과 전송률(GB/s)도 구별해야 합니다.
추론: 후보 제안 뒤 검증한다
첫 번째 헤드만 사용하면 보통의 autoregressive 디코딩입니다. 가속 경로에서는 추가 헤드가 미래 토큰 후보를 제안하고, 첫 번째 헤드의 순차 생성 기준과 일치하는 접두 구간을 검증하여 채택합니다. 중간 후보가 틀리면 그 뒤 후보도 그대로 확정하지 않습니다.
논문 §3.2와 부록 A의 속도 실험은 greedy self-speculative decoding입니다. 같은 학습 모델의 첫 번째 헤드로 수행하는 greedy 생성이 기준이며, 임의의 sampling 설정까지 동일한 보장을 확장해서는 안 됩니다. 별도 draft 모델은 필요하지 않지만 후보 예측과 검증에 드는 계산은 존재합니다.
Table S2는 7B 4-token 모델로 512-token 문맥에서 512-token 완성을 생성합니다. 코드 모델은 1T 코드 토큰, Wikipedia·Books 평가 모델은 500B 토큰으로 학습했습니다. 같은 배치 크기의 일반 autoregressive 실행 대비 결과이며, 표의 속도는 최대 배치 42에서 측정했습니다.
| 평가 입력 | 4개 헤드 사용 시 상대 속도 | forward당 평균 확정 토큰 |
|---|---|---|
| Wikipedia | 2.74배 | 3.12 |
| Books | 2.67배 | 3.09 |
| Code | 3.05배 | 3.50 |
코드에서 추가 후보 3개 중 평균 2.5개가 채택됩니다. 평균 3.5토큰을 진행해도 실제 가속은 3.05배이므로, 채택 개수와 속도 배율을 동일시하면 안 됩니다. 최대 약 3배라는 결과는 이 실험 조건의 관측값입니다.
코드 생성 성능의 12%·17%는 상대 개선
부록 E, Table S7의 13B 모델 결과를 보면 다음과 같습니다. 모델 규모 비교 실험은 코드로 처음부터 학습한 모델을 사용하며, 최소 91B 토큰을 학습합니다(§3.1). 아래 점수는 pass@1입니다.
| 평가 | next-token | 4-token | 절대 변화 | 상대 변화 |
|---|---|---|---|---|
| HumanEval | 14.1% | 15.8% | +1.7%p | 약 +12.1% |
| MBPP | 26.0% | 30.5% | +4.5%p | 약 +17.3% |
상대 변화는 예를 들어 \((15.8-14.1)/14.1\)로 계산합니다. 초록의 “12%·17% 더 많은 문제”를 점수 자체가 12·17%p 올랐다는 의미로 읽으면 안 됩니다.
작은 코드 모델에서는 성능 저하도 관찰됩니다. 하지만 기존 글에서 제시했던 “작은 모델의 학습 불안정·과적합이 원인”은 입증된 결론이 아닙니다. §5의 작동 원리 논의도 저자들이 명시한 가설입니다. 예측 범위나 모델 크기에 따른 실험 결과와 원인에 대한 해석을 구별해야 합니다.
또한 부록 D의 기존 Llama 2에 대한 multi-token finetuning은 일관된 개선을 보이지 않습니다. 사전학습부터 적용한 결과를 기존 모델에 헤드만 붙여 쉽게 재현할 수 있다고 기대해서는 안 됩니다.
DDR 장치에서는 무엇을 측정해야 하나
제 개인적 관심은 메모리 대역폭에 제한된 장치에서 이 방법이 유효한지입니다. 여러 토큰의 검증을 묶으면 가중치 접근을 분담할 여지가 있지만, 논문이 DDR 기반 Llama 2의 2배 가속이나 KV cache 트래픽의 정확한 4분의 1 감소를 측정한 것은 아닙니다. 후보 헤드, 검증, KV 읽기·쓰기, 배치와 문맥 길이까지 포함해 측정해야 합니다.
설명용 비용 모델로, 일반 디코딩의 토큰당 시간을 \(T_1\), 후보 생성·검증 한 주기의 시간을 \(T_c\), 그 주기에서 평균 확정되는 토큰 수를 \(a\)라고 두면 다음과 같습니다.
\[ S=\frac{aT_1}{T_c}. \]가령 \(a=3\), \(T_c=1.5T_1\)라는 가정이면 2배이지만, 같은 채택 수에서 \(T_c=3T_1\)이면 가속은 없습니다. 이는 하드웨어 실측이나 논문의 DDR 예측이 아닌 계산 예시입니다. 예측 헤드 수만으로 대역폭 요구나 속도를 결정할 수 없습니다.
댓글