Better & Faster Large Language Models via Multi-token Prediction

Large language models such as GPT and Llama are trained with a next-token prediction loss. In this work, we suggest that training language models to predict multiple future tokens at once results in higher sample efficiency. More specifically, at each position in the training corpus, we ask the model to predict the following n tokens using n independent output heads, operating on top of a shared model trunk. Considering multi-token prediction as an auxiliary training task, we measure improved downstream capabilities with no overhead in training time for both code and natural language models. The method is increasingly useful for larger model sizes, and keeps its appeal when training for multiple epochs. Gains are especially pronounced on generative benchmarks like coding, where our models consistently outperform strong baselines by several percentage points. Our 13B parameter models solves 12 % more problems on HumanEval and 17 % more on MBPP than comparable next-token models. Experiments on small algorithmic tasks demonstrate that multi-token prediction is favorable for the development of induction heads and algorithmic reasoning capabilities. As an additional benefit, models trained with 4-token prediction are up to 3 times faster at inference, even with large batch sizes.

한 줄 요약

다중 토큰 예측을 통해 13B 모델이 HumanEval에서 12%, MBPP에서 17% 더 많은 문제를 해결하는 LLM 훈련 방법을 제안한다.

핵심 기여도

13B 파라미터 모델에서 HumanEval 12%, MBPP 17% 성능 향상
4-token 예측 모델은 추론 속도가 기존 모델 대비 최대 3배 빠름
n개의 독립 출력 헤드를 사용한 다중 토큰 예측을 제안, 공유 트렁크 위에서 병렬 예측
코드 생성 및 알고리즘적 추론 능력 향상, 인덕션 헤드 개발 촉진

핵심 아이디어

기존 LLM은 단일 토큰 예측(Next-token prediction)을 통해 훈련되지만, 이는 지역적 패턴에만 의존하고 "복잡한 결정"을 무시할 수 있다. 본 연구는 각 토큰 위치에서 **n개의 토큰을 동시에 예측**하도록 모델을 훈련시키는 **다중 토큰 예측**(Multi-token prediction)을 제안한다. 이는 **공유 트렁크**(shared model trunk) 위에 **n개의 독립 출력 헤드**(independent output heads)를 배치하여 병렬 예측을 수행하는 방식이다. 이 접근법은 훈련 시간 증가 없이 샘플 효율성을 높이고, 특히 **코드 생성**과 같은 생성적 작업에서 강력한 성능 향상을 보인다.

기술적 접근법

**모델 구조**: 공유 트렁크 `f_s`로 잠재 표현 `z_t:1` 생성 → n개의 독립 출력 헤드 `f_h_i`로 병렬 예측
**예측 과정**:

$$
P_\theta(x_{t+1}, \dots, x_{t+n} \mid x_{t:1}) = \prod_{i=1}^n P_\theta(x_{t+i} \mid x_{t:1})
$$

**토크나이저**: 32k 크기의 토크나이저 사용
**추론 가속**: 4-token 예측 모델은 **speculative decoding**을 통해 최대 3배 빠른 추론 속도 달성
**하이퍼파라미터**: 모델 파라미터 수 일치 조건 하에, 트렁크에서 `n-1` 레이어 제거 후 헤드에 추가

주요 결과

**HumanEval**: 13B 모델에서 12% 더 많은 문제 해결
**MBPP**: 13B 모델에서 17% 더 많은 문제 해결
**추론 속도**: 4-token 예측 모델은 기존 모델 대비 최대 3배 빠름
**알고리즘 작업**: 인덕션 헤드 개발 촉진, 알고리즘적 추론 능력 향상
**코드 생성**: CodeContests 데이터셋에서 미세조정 후 성능 향상

의의 및 한계

**의의**: 샘플 효율성 향상, 생성 작업 성능 개선, 추론 속도 향상
**한계**: `n` 값의 자동 선택 방법 미비, 최적 토크나이저 크기 탐색 필요
**학술적 가치**: 인덕션 헤드 및 알고리즘적 추론 능력 개발 메커니즘 탐색 가능
**실용적 가치**: 코드 생성 및 대규모 훈련에서 효율성 증대

실용적 활용

**코드 생성**: 코드 훈련 및 생성 작업에서 높은 성능 향상
**대규모 LLM 훈련**: 샘플 효율성과 추론 속도 개선으로 산업적 활용 가능
**알고리즘 작업**: 인덕션 헤드 개발이 필요한 알고리즘 학습에 유용