더 나은, 더 빠른 대형 언어 모델: 다중 토큰 예측을 통한 접근

Better & Faster Large Language Models via Multi-token Prediction

April 30, 2024

저자: Fabian Gloeckle, Badr Youbi Idrissi, Baptiste Rozière, David Lopez-Paz, Gabriel Synnaeve

cs.AI

초록

GPT와 Llama와 같은 대규모 언어 모델은 다음 토큰 예측 손실을 통해 학습됩니다. 본 연구에서는 언어 모델이 한 번에 여러 개의 미래 토큰을 예측하도록 학습시키는 것이 샘플 효율성을 높인다는 점을 제안합니다. 구체적으로, 학습 코퍼스의 각 위치에서 모델이 공유된 모델 트렁크 위에서 작동하는 n개의 독립적인 출력 헤드를 사용하여 다음 n개의 토큰을 예측하도록 요구합니다. 다중 토큰 예측을 보조 학습 작업으로 간주하여, 코드 및 자연어 모델 모두에 대해 학습 시간의 추가 부담 없이 향상된 다운스트림 능력을 측정했습니다. 이 방법은 모델 크기가 커질수록 점점 더 유용해지며, 여러 에포크 동안 학습할 때도 그 매력을 유지합니다. 특히 코딩과 같은 생성적 벤치마크에서 이점이 두드러지며, 우리의 모델은 강력한 베이스라인보다 몇 퍼센트 포인트 더 높은 성능을 보입니다. 13B 파라미터 모델은 HumanEval에서 12%, MBPP에서 17% 더 많은 문제를 해결했습니다. 소규모 알고리즘 작업에 대한 실험은 다중 토큰 예측이 귀납 헤드와 알고리즘적 추론 능력의 발달에 유리하다는 것을 보여줍니다. 추가적인 이점으로, 4-토큰 예측으로 학습된 모델은 큰 배치 크기에서도 추론 속도가 최대 3배 빨라집니다.

English

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.