# Multi-Token Prediction, Simply Explained #104 Jun 1 LLMs https://www.aibutsimple.com/p/multi-token-prediction-simply-explained When humans read text, we don't only process the words one at a time, pausing to guess the next word before continuing. We read ahead, absorbing chunks of words at a time, anticipating where the thought is heading before it arrives, forming a broader sense of direction as we read. This "reading ahead" is more than a speed trick, as it helps you understand the deeper structure of language and its presentation over a span of language. If you’re familiar with [large language models (LLMs)](https://newsletter.aibutsimple.com/p/reasoning-llms-simply-explained), you’ll know that the dominant training paradigm has actually been the opposite: perform next-token prediction (NTP) by generating a single token at a time until completion. In April 2024, researchers at Meta's FAIR lab published a paper, ["Better & Faster Large Language Models via Multi-token Prediction,"](https://arxiv.org/abs/2404.19737?utm_source=newsletter.aibutsimple.com&utm_medium=referral&utm_campaign=multi-token-prediction-simply-explained) proposing an inference approach where LLMs adopt this human-like reading technique. By training models to simultaneously predict multiple future tokens at every position in a sequence, they unlocked meaningful gains in overall performance, efficiency, and inference speed. ![](https://media.beehiiv.com/cdn-cgi/image/fit=scale-down,width=1200,format=auto,onerror=redirect,quality=80/uploads/asset/file/0acbdbbc-cf1c-4985-8da6-ecf91370e4c7/Multi_Token_Prediction__4_.png) Incredibly, a deeper semantic and contextual understanding came without increasing the training cost. --- ### What You’ll Learn 1. The issue with next-token prediction 2. The mechanism behind multi-token prediction 3. Multi-token prediction’s memory problem 4. Self-speculative decoding, a new technique 5. Empirical results of multi-token prediction ### What’s Helpful to Know - Tokens - Portions of an input fed into a transformer model. - Vocabulary - The total number of unique tokens a model recognizes. - Next-Token Prediction (NTP) - The standard autoregressive training objective, calculating the conditional probability for the next token. - Logits - Raw, unnormalized network outputs. Models must generate a logit for every token in the vocabulary. - Latent Representation - A compressed vector representation produced by the transformer's hidden layers that encodes the context of a sequence. ![](https://media.beehiiv.com/cdn-cgi/image/fit=scale-down,width=1200,format=auto,onerror=redirect,quality=80/uploads/asset/file/9c828e76-3d7d-4001-b36d-a7cbb12e5707/Copy_of_Transformer_Tokens__1_.png) --- ### Stop re-prompting. Say it right the first time. [![](https://media.beehiiv.com/cdn-cgi/image/fit=scale-down,width=1200,format=auto,onerror=redirect,quality=80/uploads/asset/file/23f507aa-eb75-4803-b2db-1c7d7ce91ebb/flow-messages-correction-demo.png)](https://ref.wisprflow.ai/beehiiv-ai/?utm_campaign=XIAEVJHTGJ&utm_source=beehiiv&utm_term=ai_s2_q2&_bhiiv=opp_56750746-7a85-4821-bb16-d71c5f75505a_4de8c0ec) Voice-first prompts preserve the nuance you cut when typing. Speak once, paste into any AI tool, get results that don't need a follow-up. 89% of messages sent with zero edits. [Start flowing free](https://ref.wisprflow.ai/beehiiv-ai/?utm_campaign=XIAEVJHTGJ&utm_source=beehiiv&utm_term=ai_s2_q2&_bhiiv=opp_56750746-7a85-4821-bb16-d71c5f75505a_4de8c0ec) --- ### The Problem With Next-Token Prediction Let’s take a second to consider how next-token prediction is limiting. During inference, the model receives a sequence of tokens and is asked to predict only the very next one. This process repeats for every new token, requiring many sequential decoding steps, which can be costly. During training, the model is given the ground-truth previous tokens instead of the model’s prediction. This is called teacher forcing, and it ensures the model is only conditioned on the correct tokens. Since the model is already given this ground truth, supervising on one token at a time might be unnecessarily restrictive. Teacher forcing also has another a structural flaw: it encourages the model to fixate on local patterns. - Think about being taught to read by practicing reading a sentence where you can only see one word at a time. - You'd get quite good at predicting common word pairs and short phrases, but you wouldn’t be able to plan ahead far into future words. Similarly, a next-token model would tend toward predictable transitions, including our grammar rules, a reason LLMs rarely make grammatical errors, and not directly toward the decisions that require long-range planning. ![](https://media.beehiiv.com/cdn-cgi/image/fit=scale-down,width=1200,format=auto,onerror=redirect,quality=80/uploads/asset/file/efe8a2ad-1bc9-4ded-ab12-4937274c50d6/Multi_Token_Prediction__1_.png) That’s why state-of-the-art language models still require orders of magnitude more “data” and material than a human child to reach comparable written language fluency. --- ### How Multi-Token Prediction Works In Multi-Token Prediction (MTP), at every position in the training corpus, instead of predicting only the next token, the model is trained to predict the next n tokens simultaneously. - For example, if n = 4, every training step asks the model: "What are the four tokens that follow this position?" The architecture to support this has two key components layered on top of a standard transformer: - The Shared Trunk: This is just our original transformer architecture, up until the linear output projection. It processes the input context and produces a latent summary of the relevant semantic context the model has seen. This representation is shared by all the output heads. - n Independent Output Heads: Each head is an individual small set of linear layers that takes the trunk's latent output and produces a probability distribution over the vocabulary for one of the n future token positions. ![](https://media.beehiiv.com/cdn-cgi/image/fit=scale-down,width=1200,format=auto,onerror=redirect,quality=80/uploads/asset/file/4de28d95-5f0c-4a78-bf82-7da4c1dd218a/Multi_Token_Prediction__6_.png) Each head is assigned to one token. The first head predicts the immediate next token while the second predicts the token after that, and so on. A multi-headed approach to divide work is not exactly new for transformers. In fact, in earlier transformer development, researchers introduced multiple parallel attention heads within the transformer layer, which we know as [multi-headed attention (MHA)](https://newsletter.aibutsimple.com/p/transformers-and-multi-head-attention-mathematically-explained). It allowed the model to attend to multiple relationships and contextual patterns in parallel from different learned projections of the same input. In MTP, we use the same "parallel multi-head” idea in our output projection to predict multiple tokens at the same time. For efficiency, these heads output tokens in parallel using the same shared trunk output but don’t see other heads’ outputs. They also draw from a single vocabulary matrix (lookup table to map internal vectors to word probabilities), keeping total parameter growth minimal. For performance, since the latent representation from the trunk is shared across heads, for the model to score well in training, it must incorporate broader contextual information into the trunk to support multi-token predictions. --- If you enjoy our content and custom visuals, consider sharing this newsletter with others or [upgrading](http://www.aibutsimple.com/pro?utm_source=newsletter.aibutsimple.com&utm_medium=referral&utm_campaign=multi-token-prediction-simply-explained) so we can keep doing what we do. ### [Ace the Interview With AI, But Simple PRO](https://www.aibutsimple.com/pro?utm_source=newsletter.aibutsimple.com&utm_medium=referral&utm_campaign=multi-token-prediction-simply-explained) ![](https://media.beehiiv.com/cdn-cgi/image/fit=scale-down,width=1200,format=auto,onerror=redirect,quality=80/uploads/asset/file/963eba00-ab31-4912-8c1f-1470a52f9ecc/aibutsimple_academy__2_.png) Supporting us by purchasing a PRO membership will get you: - Exclusive Hands-on Tutorials With Code - Real-World Interview Questions and Cases (From Real Companies) - Full Math Explanations and Beautiful Visualizations - Ad-Free Reading for an Uninterrupted Learning Experience [Stay Ahead in AI With PRO](https://www.aibutsimple.com/pro?utm_source=newsletter.aibutsimple.com&utm_medium=referral&utm_campaign=multi-token-prediction-simply-explained) --- ### The Memory Challenge Adding n output heads introduces a memory problem. Every head added adds additional memory overhead from model parameters and intermediate activations. In modern LLMs, the vocabulary size V (often 32,000+ tokens) is much larger than the model's weight dimension d. When you create the logit vectors (the raw probability scores over every token in the vocabulary), they’re enormous, taking up much more memory than the parameters. Naively running all n heads simultaneously would require holding n tensors of all these logits in GPU memory, severely limiting batch sizes due to resource bandwidth, making MTP impractical. The solution was to process the output heads sequentially during the backward pass. Essentially, first compute the forward and backward pass for a given output head. Then, accumulate gradients at the trunk (this memory hit will grow over time), which means we can free the head’s memory, and then repeat for all of the preceding heads. ![](https://media.beehiiv.com/cdn-cgi/image/fit=scale-down,width=1200,format=auto,onerror=redirect,quality=80/uploads/asset/file/d66fd1b9-a1be-4179-8f46-4a003e3a02b9/Multi_Token_Prediction__2_.png) The result is that the most that the GPU memory ever holds is one head's logits at a time, reducing memory usage from O(nV) to O(V + d), the same as an NTP model. --- ### Self-Speculative Decoding One of the most practical benefits of multi-token prediction reveals itself at inference time. During generation, the model only uses one output head to produce the output, one token at a time. The other n-1 heads will serve instead as a built-in drafter models for a technique called speculative decoding. Speculative decoding is a way to speed up inference without compromising output quality. It is typically a two-model trick: 1. A smaller, fast model (called a drafter) generates several tokens ahead. 2. A larger model verifies them in parallel. The catch is that you need a separate, well-matched draft model, which may pose external challenges. ![](https://media.beehiiv.com/cdn-cgi/image/fit=scale-down,width=1200,format=auto,onerror=redirect,quality=80/uploads/asset/file/be112331-13b3-4d4a-8736-01dee74ffcd1/Multi_Token_Prediction__11_.png) Multi-token prediction eliminates that burden. All output heads act as built-in drafters, proposing the next several tokens. - For instance, if we had the input sequence “the cat sat” and the model had 3 output heads, they might draft the 3 tokens “on the mat.” From the draft, the model will perform standard autoregressive generation using only the first output head to verify the proposed tokens one-by-one across successive forward passes. - From the previous example, if we feed in the same sequence, “the cat sat,” and the model autoregressively generates the token “on,” it matches the drafted token “on” and is accepted. - The model then generates the next token, and if it matches the token “the,” it is accepted, and so on. If an NTP-generated token does not match the drafted token, all drafted tokens from the mismatch point onward are discarded and are replaced by the NTP ones. ![](https://media.beehiiv.com/cdn-cgi/image/fit=scale-down,width=1200,format=auto,onerror=redirect,quality=80/uploads/asset/file/7b86f0b3-d130-4fcd-84c8-cd44f4bbe81e/Multi_Token_Prediction__9_.png) The process of drafting multiple tokens, then verifying them autoregressively, is repeated until the end-of-sequence criteria (EOS token, max length, etc.) is reached. The researchers call this self-speculative decoding. In practice, a 4-token prediction model achieved a 3× inference speedup on code, with an average of 2.5/3 proposed tokens being accepted. ![](https://media.beehiiv.com/cdn-cgi/image/fit=scale-down,width=1200,format=auto,onerror=redirect,quality=80/uploads/asset/file/b278d35e-02be-4c96-80b0-2c1e7c66aa3a/image.png) Above, we observe that increasing the number of heads (represented by k) increases the relative throughput of tokens while also driving down latency using MTP. --- ### The Results On code generation, the results are consistent and meaningful. The 13B parameter multi-token model solved 12% more problems on [HumanEval](https://github.com/openai/human-eval?utm_source=newsletter.aibutsimple.com&utm_medium=referral&utm_campaign=multi-token-prediction-simply-explained) and 17% more on [MBPP](https://github.com/google-research/google-research/tree/master/mbpp?utm_source=newsletter.aibutsimple.com&utm_medium=referral&utm_campaign=multi-token-prediction-simply-explained) compared to a next-token baseline of identical size. However, these gains scale with model size. Smaller models (less than 300M parameters) can actually perform worse under MTP, requiring a model with enough representational capacity to benefit from the richer meaning encoded in the trunk. ![](https://media.beehiiv.com/cdn-cgi/image/fit=scale-down,width=1200,format=auto,onerror=redirect,quality=80/uploads/asset/file/a7844fdb-f273-468c-87d2-caef0d6c31a4/image.png) At granular levels of tokenization, where each token is a single byte rather than a word fragment, the results are striking: An 8-byte prediction model solved 67% more problems on MBPP with one inference attempt (pass@1) than the byte-level baseline. However, on natural language tasks, the story shows some cracks. While multi-token models are stellar on generative benchmarks like text summarization, where scores improved consistently across eight datasets, they show slight regressions on multiple-choice question (MCQ) benchmarks. The researchers’ hypothesis is that these tests may not effectively capture generative quality and that larger training datasets help close the gap. --- ### The Bottom Line Multi-token prediction may seem obvious in hindsight: train the model to generate multiple tokens at once. But the real challenges were making it memory-efficient and determining the scale of the gains it unlocks on code benchmarks and inference speed. Next-token prediction remains the dominant paradigm, but multi-token objectives offer a promising direction for improving inference efficiency and enhancing performance in certain settings. --- Here’s a special thanks to our [biggest supporters](https://www.aibutsimple.com/upgrade?utm_source=newsletter.aibutsimple.com&utm_medium=referral&utm_campaign=multi-token-prediction-simply-explained): - Sushant Waidande - Sai Krishna Pashapu (AI Developer, BTech, [LinkedIn](https://www.linkedin.com/in/sai-krishna-pashapu-997358283/?utm_source=newsletter.aibutsimple.com&utm_medium=referral&utm_campaign=multi-token-prediction-simply-explained), [Github](https://github.com/saikrishna-pashapu?utm_source=newsletter.aibutsimple.com&utm_medium=referral&utm_campaign=multi-token-prediction-simply-explained)) - Bob Boncella - Demet Yazici - Imran Khan --- If you enjoy our content, consider supporting us so we can keep doing what we do. Please share this with a friend! Want to reach 8000+ ML engineers? Advertising Inquiries: [WORK WITH US 📈](https://tally.so/r/A7Dgdk?utm_source=newsletter.aibutsimple.com&utm_medium=referral&utm_campaign=multi-token-prediction-simply-explained) Feedback or inquiries? Send us an email at [edwin@aibutsimple.com](mailto:edwin@aibutsimple.com). --- That’s it for this week’s issue of AI, but simple. See you next week! —AI, but simple team