Elamur Research⌕
☰
home / papers / 2024 / Better & Faster Large Language Models via Multi-token Prediction

Better & Faster Large Language Models via Multi-token Prediction

Paper · 2024 · Fabian Gloeckle, Badr Youbi Idrissi, Baptiste Roziere et al.

Summary

LLMs such as GPT and Llama learn by guessing only the single next token. The authors argue this objective needs a lot of data and favours short-range patterns. Gloeckle et al. (FAIR at Meta) instead train models to forecast several upcoming tokens from every position. A shared transformer trunk feeds n separate transformer-layer heads, and all heads use one common unembedding matrix. To limit GPU memory, the heads run their forward and backward passes one after another, so only one head's logits are held at a time. All comparisons are parameter-matched: any layers added to the heads are taken out of the trunk.

The models were trained from scratch on code (300M–13B parameters) and also on natural language. The approach slightly hurts small models but helps more as size grows. The clearest gains appear on generative coding benchmarks like MBPP and HumanEval. The extra heads also allow self-speculative decoding without a separate draft model, which makes inference about 3× faster. The method therefore improves both data efficiency and serving speed without extra training cost.

Key points

  • At 13B parameters, multi-token models solve about 12% more HumanEval problems and 17% more MBPP problems than comparable next-token baselines. Below roughly 1–3B parameters, they underperform.
  • In ablations on 7B models trained on 200B code tokens with a 32k vocabulary, n = 4 was the best setting overall. It reached 33.8 MBPP pass@1 versus 30.0 for the baseline.
  • Byte-level training benefits strongly. A 7B model predicting 8 bytes scored 32.3 MBPP pass@1, compared with 19.3 for the next-byte model.
  • With a 4-token model, self-speculative decoding ran 3.0× faster on code, accepting about 2.5 of 3 proposed tokens, and 2.7× faster on text. An 8-byte model reached a 6.4× speedup.
  • Running the heads sequentially cuts peak memory from O(nV + d) to O(V + d). Gains persist but shrink at 1T tokens (4 epochs), and all experiments together used about 500K A100/H100 GPU hours.

Our summary of the paper, in our own words — not its text. The work belongs to its authors; read it at the links above.

 Institutions

 Fields, topics and concepts

Cite

Cite as Fabian Gloeckle, Badr Youbi Idrissi, Baptiste Roziere, David Lopez-Paz and Gabriel Synnaeve. 2024. Better & Faster Large Language Models via Multi-token Prediction. arXiv:2404.19737.

The paper itself, not this page. BibTeX from the library's record of the work. DOI 10.48550/arxiv.2404.19737 · arXiv 2404.19737

@misc{gloeckle2024better,
  title = {{Better \& Faster Large Language Models via Multi-token Prediction}},
  author = {Fabian Gloeckle, Badr Youbi Idrissi, Baptiste Roziere, David Lopez-Paz and Gabriel Synnaeve},
  year = {2024},
  doi = {10.48550/arxiv.2404.19737},
  eprint = {2404.19737},
  archivePrefix = {arXiv},
  url = {https://doi.org/10.48550/arxiv.2404.19737}
}

References

57 references, 25 of them with a card of ours. ↗ leads to the work, «find» searches for one we have no id for yet.

  1. Alexander R. Fabbri, Irene Li, Tianwei She, Suyi Li, and Dragomir R. Radev. Multi-news: a large-scale multi-document summarization dataset and abstractive hierarchical model, 2019.
  2. Jacob Austin, Augustus Odena, Maxwell Nye, Maarten Bosma, Henryk Michalewski, David Dohan, Ellen Jiang, Carrie Cai, Michael Terry, Quoc Le, et al. Program synthesis with large language models. arXiv preprint arXiv:2108.07732, 2021.
  3. Mehrdad Farahani. Summarization using bert2bert model on wikisummary dataset. https://github.com/m3hrdadfi/wikisummary, 2020. ↗
  4. Gregor Bachmann and Vaishnavh Nagarajan. The pitfalls of next-token prediction, 2024. find
  5. Mehrdad Farahani, Mohammad Gharachorloo, and Mohammad Manthouri. Leveraging parsbert and pretrained mt5 for persian abstractive text summarization. In 2021 26th International Computer Conference, Computer Society of Iran (CSICC). find
  6. IEEE, March 2021. doi: 10.1109/ csicc52343.2021.9420563. URL http://dx.doi.org/10.1109/CSICC52343.2021.9420563. — [doi:10.1109/csicc52343.2021.9420563](https://doi.org/10.1109/csicc52343.2021.9420563) ↗
  7. Samy Bengio, Oriol Vinyals, Navdeep Jaitly, and Noam Shazeer. Scheduled sampling for sequence prediction with recurrent neural networks, 2015. find
  8. Yonatan Bisk, Rowan Zellers, Ronan Le Bras, Jianfeng Gao, and Yejin Choi. Piqa: Reasoning about physical commonsense in natural language, 2019.
  9. Michael C Frank. Bridging the data gap between children and large language models. Trends in Cognitive Sciences, 2023. find
  10. Tianle Cai, Yuhong Li, Zhengyang Geng, Hongwu Peng, Jason D. Lee, Deming Chen, and Tri Dao. Medusa: Simple llm inference acceleration framework with multiple decoding heads, 2024.
  11. Bogdan Gliwa, Iwona Mochol, Maciej Biesek, and Aleksander Wawer. Samsum corpus: A human-annotated dialogue dataset for abstractive summarization. In Proceedings of the 2nd Workshop on New Frontiers in Summarization. Association for Computational Linguistics, 2019. doi: 10.18653/v1/d19-5409. URL http…
  12. Rich Caruana. Multitask learning. Machine learning, 28: 41–75, 1997. find
  13. Mark Chen, Jerry Tworek, Heewoo Jun, Qiming Yuan, Henrique Ponde, Jared Kaplan, Harri Edwards, Yura Burda, Nicholas Joseph, Greg Brockman, et al. Evaluating large language models trained on code. arXiv preprint arXiv:2107.03374, 2021.
  14. Sachin Goyal, Ziwei Ji, Ankit Singh Rawat, Aditya Krishna Menon, Sanjiv Kumar, and Vaishnavh Nagarajan. Think before you speak: Training language models with pause tokens, 2023. find
  15. Nakhun Chumpolsathien. Using knowledge distillation from keyword extraction to improve the informativeness of neural cross-lingual summarization. Master’s thesis, Beijing Institute of Technology, 2020. find
  16. Dan Hendrycks, Steven Basart, Saurav Kadavath, Mantas Mazeika, Akul Arora, Ethan Guo, Collin Burns, Samir Puranik, Horace He, Dawn Song, et al. Measuring coding challenge competence with apps. arXiv preprint arXiv:2105.09938, 2021.
  17. Karl Cobbe, Vineet Kosaraju, Mohammad Bavarian, Mark Chen, Heewoo Jun, Lukasz Kaiser, Matthias Plappert, Jerry Tworek, Jacob Hilton, Reiichiro Nakano, et al. Training verifiers to solve math word problems. arXiv preprint arXiv:2110.14168, 2021.
  18. Ari Holtzman, Jan Buys, Li Du, Maxwell Forbes, and Yejin Choi. The curious case of neural text degeneration, 2020.
  19. Jianyu Zhang Leon Bottou. Multi-label classification as an auxiliary loss for language modelling. personal communication, 2024. find
  20. Li Dong, Nan Yang, Wenhui Wang, Furu Wei, Xiaodong Liu, Yu Wang, Jianfeng Gao, Ming Zhou, and Hsiao-Wuen Hon. Unified language model pre-training for natural language understanding and generation. In Proceedings of the 33rd International Conference on Neural Information Processing Systems, pages 130…
  21. Mandar Joshi, Eunsol Choi, Daniel S. Weld, and Luke Zettlemoyer. Triviaqa: A large scale distantly supervised challenge dataset for reading comprehension, 2017.
  22. Alexandre Défossez, Jade Copet, Gabriel Synnaeve, and Yossi Adi. High fidelity neural audio compression. arXiv preprint arXiv:2210.13438, 2022.
  23. Diederik Kingma and Jimmy Ba. Adam: A method for stochastic optimization. ICLR, 2015.
  24. Ryan Koo, Minhwa Lee, Vipul Raheja, Jong Inn Park, Zae Myung Kim, and Dongyeop Kang. Benchmarking cognitive biases in large language models as evaluators. arXiv preprint arXiv:2309.17012, 2023. find
  25. Moussa Kamal Eddine, Antoine J. P. Tixier, and Michalis Vazirgiannis. Barthez: a skilled pretrained french sequence-to-sequence model, 2021. find
  26. Tom Kwiatkowski, Jennimaria Palomaki, Olivia Redfield, Michael Collins, Ankur Parikh, Chris Alberti, Danielle Epstein, Illia Polosukhin, Matthew Kelcey, Jacob Devlin, Kenton Lee, Kristina N. Toutanova, Llion Jones, Ming-Wei Chang, Andrew Dai, Jakob Uszkoreit, Quoc Le, and Slav Petrov. Natural questi… find
  27. Shashi Narayan, Shay B. Cohen, and Mirella Lapata. Don’t give me the details, just the summary! topic-aware convolutional neural networks for extreme summarization, 2018. — Don't Give Me the Details, Just the Summary! Topic-Aware Convolutional Neural Networks for Extreme Summarization
  28. Guillaume Lample, Marie-Anne Lachaux, Thibaut Lavril, Xavier Martinet, Amaury Hayat, Gabriel Ebner, Aurélien Rodriguez, and Timothée Lacroix. Hypertree proof search for neural theorem proving, 2022.
  29. Yann LeCun. A path towards autonomous machine intelligence version 0.9. 2, 2022-06-27. Open Review, 62(1), 2022. find
  30. Benjamin Lefaudeux, Francisco Massa, Diana Liskovich, Wenhan Xiong, Vittorio Caggiano, Sean Naren, Min Xu, Jieru Hu, Marta Tintore, Susan Zhang, Patrick Labatut, and Daniel Haziza. xformers: A modular and hackable transformer modelling library. https://github.com/facebookresearch/xformers, 2022. ↗
  31. Yaniv Leviathan, Matan Kalman, and Yossi Matias. Fast inference from transformers via speculative decoding, 2023. find
  32. Yujia Li, David Choi, Junyoung Chung, Nate Kushman, Julian Schrittwieser, Rémi Leblond, Tom Eccles, James Keeling, Felix Gimeno, Agustin Dal Lago, et al. Competition-level code generation with alphacode. Science, 378(6624):1092–1097, 2022.
  33. Catherine Olsson, Nelson Elhage, Neel Nanda, Nicholas Joseph, Nova DasSarma, Tom Henighan, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, Tom Conerly, Dawn Drain, Deep Ganguli, Zac Hatfield-Dodds, Danny Hernandez, Scott Johnston, Andy Jones, Jackson Kernion, Liane Lovitt, Kamal Ndousse, Dario Amode… — In-context Learning and Induction Heads
  34. OpenAI. Gpt-4 technical report, 2023.
  35. Long Ouyang, Jeff Wu, Xu Jiang, Diogo Almeida, Carroll L. Wainwright, Pamela Mishkin, Chong Zhang, Sandhini Agarwal, Katarina Slama, Alex Ray, John Schulman, Jacob Hilton, Fraser Kelton, Luke Miller, Maddie Simens, Amanda Askell, Peter Welinder, Paul Christiano, Jan Leike, and Ryan Lowe. Training la… — Training language models to follow instructions with human feedback
  36. Koyena Pal, Jiuding Sun, Andrew Yuan, Byron C. Wallace, and David Bau. Future lens: Anticipating subsequent tokens from a single hidden state, 2023.
  37. Weizhen Qi, Yu Yan, Yeyun Gong, Dayiheng Liu, Nan Duan, Jiusheng Chen, Ruofei Zhang, and Ming Zhou. Prophetnet: Predicting future n-gram for sequence-to-sequence pre-training, 2020. find
  38. Chin-Yew Lin. ROUGE: A package for automatic evaluation of summaries. In Text Summarization Branches Out, pages 74–81, Barcelona, Spain, July 2004. Association for Computational Linguistics. URL https://aclanthology.org/W04-1013. ↗
  39. Jesse Read, Bernhard Pfahringer, Geoffrey Holmes, and Eibe Frank. Classifier chains: A review and perspectives. Journal of Artificial Intelligence Research, 70:683–718, 2021. find
  40. Zhenghao Lin, Zhibin Gou, Yeyun Gong, Xiao Liu, Yelong Shen, Ruochen Xu, Chen Lin, Yujiu Yang, Jian Jiao, Nan Duan, and Weizhu Chen. Rho-1: Not all tokens are what you need, 2024.
  41. Melissa Roemmele, Cosmin Adrian Bejan, and Andrew S Gordon. Choice of plausible alternatives: An evaluation of commonsense causal reasoning. In 2011 AAAI Spring Symposium Series, 2011. find
  42. Ilya Loshchilov and Frank Hutter. Sgdr: Stochastic gradient descent with warm restarts, 2017. find
  43. Maarten Sap, Hannah Rashkin, Derek Chen, Ronan LeBras, and Yejin Choi. Socialiqa: Commonsense reasoning about social interactions, 2019.
  44. Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization, 2019. find
  45. David Silver, Aja Huang, Chris J Maddison, Arthur Guez, Laurent Sifre, George Van Den Driessche, Julian Schrittwieser, Ioannis Antonoglou, Veda Panneershelvam, Marc Lanctot, et al. Mastering the game of go with deep neural networks and tree search. nature, 529(7587): 484–489, 2016. find
  46. Michael Mathieu, Camille Couprie, and Yann LeCun. Deep multi-scale video prediction beyond mean square error, 2016. find
  47. Ramesh Nallapati, Bowen Zhou, Cicero Nogueira dos santos, Caglar Gulcehre, and Bing Xiang. Abstractive text Aaditya K Singh, Stephanie CY Chan, Ted Moskovitz, Erin Grant, Andrew M Saxe, and Felix Hill. The transient nature of emergent in-context learning in transformers. arXiv preprint arXiv:2311.08… — [arXiv:2311.08360](https://arxiv.org/abs/2311.08360) find
  48. Eleftherios Spyromitros-Xioufis, Grigorios Tsoumakas, William Groves, and Ioannis Vlahavas. Multi-target regression via input space expansion: treating targets as inputs. Machine Learning, 104:55–98, 2016. find
  49. Nitish Srivastava, Elman Mansimov, and Ruslan Salakhutdinov. Unsupervised learning of video representations using lstms, 2016. find
  50. Mitchell Stern, Noam Shazeer, and Jakob Uszkoreit. Block-wise parallel decoding for deep autoregressive models, 2018. find
  51. Yi Tay, Mostafa Dehghani, Vinh Q Tran, Xavier Garcia, Jason Wei, Xuezhi Wang, Hyung Won Chung, Siamak Shakeri, Dara Bahri, Tal Schuster, et al. Ul2: Unifying language learning paradigms. arXiv preprint arXiv:2205.05131, 2022.
  52. Vladimir Vapnik and Akshay Vashist. A new learning paradigm: Learning using privileged information. Neural networks, 22(5-6):544–557, 2009. find
  53. Carl Vondrick, Hamed Pirsiavash, and Antonio Torralba. Anticipating visual representations from unlabeled video, 2016. find
  54. Willem Waegeman, Krzysztof Dembczyński, and Eyke Hüllermeier. Multi-target prediction: a unifying view on problems and methods. Data Mining and Knowledge Discovery, 33:293–324, 2019. find
  55. Vikas Yadav, Steven Bethard, and Mihai Surdeanu. Quick and (not so) dirty: Unsupervised selection of justification sentences for multi-hop question answering. arXiv preprint arXiv:1911.07176, 2019. find
  56. Zhilin Yang, Zihang Dai, Yiming Yang, Jaime Carbonell, Russ R Salakhutdinov, and Quoc V Le. Xlnet: Generalized autoregressive pretraining for language understanding. In Advances in neural information processing systems, pages 5753–5763, 2019.
  57. Rowan Zellers, Ari Holtzman, Yonatan Bisk, Ali Farhadi, and Yejin Choi. Hellaswag: Can a machine really finish your sentence?, 2019. 3.0 2.5 2.0 1.5 1.0 0.5 0.0 1.0 find

Updated

↑ back to top

This page in Markdown → https://md.research.elamur.ai/free/en/papers/gloeckle-2024-better-faster-large-language-models/