2404.19737v1

## Better & Faster Large Language Models via Multi-token Prediction Fabian Gloeckle * 1 2 Badr Youbi Idrissi * 1 3 Baptiste Rozière 1 David Lopez-Paz + 1 Gabriel Synnaeve + 1 ## Abstract 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 × faster at inference, even with large batch sizes. ## 1. Introduction Humanity has condensed its most ingenious undertakings, surprising findings and beautiful productions into text. Large Language Models (LLMs) trained on all of these corpora are able to extract impressive amounts of world knowledge, as well as basic reasoning capabilities by implementing a simple-yet powerful-unsupervised learning task: next-token prediction. Despite the recent wave of impressive achievements (OpenAI, 2023), next-token pre- * Equal contribution + Last authors 1 FAIR at Meta 2 CERMICS Ecole des Ponts ParisTech 3 LISN Université Paris-Saclay. Correspondence to: Fabian Gloeckle <fgloeckle@meta.com>, Badr Youbi Idrissi <byoubi@meta.com>. diction remains an inefficient way of acquiring language, world knowledge and reasoning capabilities. More precisely, teacher forcing with next-token prediction latches on local patterns and overlooks 'hard' decisions. Consequently, it remains a fact that state-of-the-art next-token predictors call for orders of magnitude more data than human children to arrive at the same level of fluency (Frank, 2023). Figure 1: Overview of multi-token prediction. (Top) During training, the model predicts 4 future tokens at once, by means of a shared trunk and 4 dedicated output heads. During inference, we employ only the next-token output head. Optionally, the other three heads may be used to speed-up inference time. (Bottom) Multi-token prediction improves pass@1 on the MBPP code task, significantly so as model size increases. Error bars are confidence intervals of 90% computed with bootstrapping over dataset samples. <details> <summary>Image 1 Details</summary>  ### Visual Description ## Diagram and Chart: 4-Token Prediction Model and Performance ### Overview The image presents a two-part visualization. The top portion is a diagram illustrating a 4-token prediction model architecture, showing inputs, a shared layer, heads, and targets. The bottom portion is a horizontal bar chart displaying the MBPP Pass@1 gains achieved with 4-token prediction for various model sizes, relative to a baseline. ### Components/Axes **Top Diagram:** * **Title:** None explicitly stated, but the diagram illustrates a "4-token prediction" model. * **Elements:** * **Inputs:** Labeled "Inputs" on the bottom-left. Four input nodes are numbered 1 through 4. * **Shared Layer:** A dark blue-purple rectangular block labeled "Shared." * **Heads:** A green block labeled "Head 1 2 3 4". * **4-token targets:** Labeled "4-token targets" on the top-left. Each target has a stack of numbers, representing tokens. * **Discarded at inference:** A light pink circle with the text "Discarded at inference (or used to speed up model up to 3 times)" in the top-right. **Bottom Chart:** * **Title:** "MBPP Pass@1 gains with 4-token prediction" * **Y-axis:** "Model size" with labels: 13B, 6.7B, 3B, 1.3B, 0.6B, 0.3B. * **X-axis:** Implicitly represents the "MBPP Pass@1 gains". The baseline is marked with "-1.7" and "+4.5". * **Data Series:** Horizontal bars representing the gains for each model size. Error bars are also present. ### Detailed Analysis **Top Diagram:** * The diagram shows four inputs feeding into a shared layer. The shared layer then connects to four heads. Each head predicts a sequence of tokens (4-token targets). Some tokens are marked as "Discarded at inference". **Bottom Chart:** * **Model Sizes and Gains:** * **13B:** Green bar, gain of approximately 26.0. * **6.7B:** Green bar, gain of approximately 23.9. * **3B:** Green bar, gain of approximately 11.1. * **1.3B:** Green bar, gain of approximately 6.8. * **0.6B:** Orange bar, gain of approximately 4.7. * **0.3B:** Orange bar, gain of approximately 1.8. * **Trend:** The MBPP Pass@1 gain generally increases with model size. The 13B model has the highest gain, while the 0.3B model has the lowest. ### Key Observations * The diagram illustrates a model architecture designed for 4-token prediction. * The chart shows a positive correlation between model size and MBPP Pass@1 gain when using 4-token prediction. * The 0.6B and 0.3B models are represented with orange bars, while the larger models are green. This color difference might indicate a performance threshold or a different training regime. ### Interpretation The data suggests that using a 4-token prediction approach improves the performance (MBPP Pass@1) of models, and this improvement is more pronounced for larger models. The color difference between the smaller (0.6B and 0.3B) and larger models (1.3B to 13B) might indicate a change in the effectiveness of the 4-token prediction strategy or a difference in training methodology for smaller models. The "Discarded at inference" tokens likely represent a technique to optimize inference speed, potentially at the cost of some accuracy. </details> In this study, we argue that training LLMs to predict multiple tokens at once will drive these models toward better sample efficiency. As anticipated in Figure 1, multi-token prediction instructs the LLM to predict the n future tokens from each position in the training corpora, all at once and in parallel (Qi et al., 2020). Contributions While multi-token prediction has been studied in previous literature (Qi et al., 2020), the present work offers the following contributions: 1. We propose a simple multi-token prediction architecture with no train time or memory overhead (Section 2). 2. We provide experimental evidence that this training paradigm is beneficial at scale, with models up to 13B parameters solving around 15% more code problems on average (Section 3). 3. Multi-token prediction enables self-speculative decoding, making models up to 3 times faster at inference time across a wide range of batch-sizes (Section 3.2). While cost-free and simple, multi-token prediction is an effective modification to train stronger and faster transformer models. We hope that our work spurs interest in novel auxiliary losses for LLMs well beyond next-token prediction, as to improve the performance, coherence, and reasoning abilities of these fascinating models. ## 2. Method Standard language modeling learns about a large text corpus x 1 , . . . x T by implementing a next-token prediction task. Formally, the learning objective is to minimize the crossentropy loss <!-- formula-not-decoded --> where P θ is our large language model under training, as to maximize the probability of x t +1 as the next future token, given the history of past tokens x t :1 = x t , . . . , x 1 . In this work, we generalize the above by implementing a multi-token prediction task, where at each position of the training corpus, the model is instructed to predict n future tokens at once. This translates into the cross-entropy loss <!-- formula-not-decoded --> To make matters tractable, we assume that our large language model P θ employs a shared trunk to produce a latent representation z t :1 of the observed context x t :1 , then fed into n independent heads to predict in parallel each of the n future tokens (see Figure 1). This leads to the following factorization of the multi-token prediction cross-entropy loss: <!-- formula-not-decoded --> <!-- formula-not-decoded --> In practice, our architecture consists of a shared transformer trunk f s producing the hidden representation z t :1 from the observed context x t :1 , n independent output heads implemented in terms of transformer layers f h i , and a shared unembedding matrix f u . Therefore, to predict n future tokens, we compute: <!-- formula-not-decoded --> for i = 1 , . . . n , where, in particular, P θ ( x t +1 | x t :1 ) is our next-token prediction head. See Appendix B for other variations of multi-token prediction architectures. Memory-efficient implementation One big challenge in training multi-token predictors is reducing their GPU memory utilization. To see why this is the case, recall that in current LLMs the vocabulary size V is much larger than the dimension d of the latent representation-therefore, logit vectors become the GPU memory usage bottleneck. Naive implementations of multi-token predictors that materialize all logits and their gradients, both of shape ( n, V ) , severely limit the allowable batch-size and average GPU memory utilization. Because of these reasons, in our architecture we propose to carefully adapt the sequence of forward and backward operations, as illustrated in Figure 2. In particular, after the forward pass through the shared trunk f s , we sequentially compute the forward and backward pass of each independent output head f i , accumulating gradients at the trunk. While this creates logits (and their gradients) for the output head f i , these are freed before continuing to the next output head f i +1 , requiring the long-term storage only of the d -dimensional trunk gradient ∂L n /∂f s . In sum, we have reduced the peak GPU memory utilization from O ( nV + d ) to O ( V + d ) , at no expense in runtime (Table S5). Inference During inference time, the most basic use of the proposed architecture is vanilla next-token autoregressive prediction using the next-token prediction head P θ ( x t +1 | x t :1 ) , while discarding all others. However, the additional output heads can be leveraged to speed up decoding from the next-token prediction head with self-speculative decoding methods such as blockwise parallel decoding (Stern et al., 2018)-a variant of speculative decoding (Leviathan et al., 2023) without the need for an additional draft model-and speculative decoding with Medusa-like tree attention (Cai et al., 2024). Figure 2: Order of the forward/backward in an n -token prediction model with n = 2 heads. By performing the forward/backward on the heads in sequential order, we avoid materializing all unembedding layer gradients in memory simultaneously and reduce peak GPU memory usage. <details> <summary>Image 2 Details</summary>  ### Visual Description ## Diagram: Forward and Backward Pass in a Neural Network ### Overview The image illustrates the forward and backward pass in a neural network with two heads sharing a common layer. It includes a diagram of the network architecture and corresponding code snippet. ### Components/Axes * **Legend:** Located at the top of the image. * Forward: Represented by a blue arrow. * Backward: Represented by a red arrow. * Tensor: Represented by a gray circle. * **Network Diagram:** Located on the left side of the image. * Shared: A dark blue rounded rectangle at the bottom. * Head 1: A light green rounded rectangle above "Shared" on the left. * Head 2: A light green rounded rectangle above "Shared" on the right. * Loss 1: A light yellow rounded rectangle above "Head 1". * Loss 2: A light yellow rounded rectangle above "Head 2". * Tensor: A gray circle connecting "Shared", "Head 1", and "Head 2". * **Code Snippet:** Located on the right side of the image. ### Detailed Analysis * **Network Diagram:** * Forward Pass: * Arrow 1 (blue): From "Shared" to the gray "Tensor" circle. * Arrow 2 (blue): From the gray "Tensor" circle to "Head 1". * Arrow 4 (blue): From the gray "Tensor" circle to "Head 2". * Arrow 3 (blue): From "Head 1" to "Loss 1". * Arrow 5 (blue): From "Head 2" to "Loss 2". * Backward Pass: * Arrow 6 (red): From the gray "Tensor" circle back to "Shared". * **Code Snippet:** </details> ## 3. Experiments on real data We demonstrate the efficacy of multi-token prediction losses by seven large-scale experiments. Section 3.1 shows how multi-token prediction is increasingly useful when growing the model size. Section 3.2 shows how the additional prediction heads can speed up inference by a factor of 3 × using speculative decoding. Section 3.3 demonstrates how multi-token prediction promotes learning longer-term patterns, a fact most apparent in the extreme case of byte-level tokenization. Section 3.4 shows that 4 -token predictor leads to strong gains with a tokenizer of size 32 k. Section 3.5 illustrates that the benefits of multi-token prediction remain for training runs with multiple epochs. Section 3.6 showcases the rich representations promoted by pretraining with multi-token prediction losses by finetuning on the CodeContests dataset (Li et al., 2022). Section 3.7 shows that the benefits of multi-token prediction carry to natural language models, improving generative evaluations such as summarization, while not regressing significantly on standard benchmarks based on multiple choice questions and negative log-likelihoods. To allow fair comparisons between next-token predictors and n -token predictors, the experiments that follow always compare models with an equal amount of parameters. That is, when we add n -1 layers in future prediction heads, we remove n -1 layers from the shared model trunk. Please refer to Table S14 for the model architectures and to Table S13 for an overview of the hyperparameters we use in our experiments. ## 3.1. Benefits scale with model size To study this phenomenon, we train models of six sizes in the range 300M to 13B parameters from scratch on at least 91B tokens of code. The evaluation results in Fig- Figure 3: Results of n -token prediction models on MBPP by model size. We train models of six sizes in the range or 300M to 13B total parameters on code, and evaluate pass@1,10,100 on the MBPP (Austin et al., 2021) and HumanEval (Chen et al., 2021) benchmark with 1000 samples. Multi-token prediction models are worse than the baseline for small model sizes, but outperform the baseline at scale. Error bars are confidence intervals of 90% computed with bootstrapping over dataset samples. <details> <summary>Image 3 Details</summary>  ### Visual Description ## Bar Charts: MBPP vs. Human Eval Performance ### Overview The image presents two sets of bar charts comparing the performance of a system (MBPP) against human evaluation (Human Eval) across different model sizes (0.3B to 13B). The charts display "Pass@k" values for k=1, 10, and 100, indicating the pass rate within the top k attempts. The y-axis represents the pass rate, and the x-axis represents the model size. The bars are colored differently to indicate performance relative to a baseline, with green indicating positive performance and red indicating negative performance. Error bars are included to show the uncertainty in the measurements. ### Components/Axes * **Titles:** "MBPP" (left), "Human Eval" (right) * **Y-Axis Labels (Left):** +4.5, -1.7, +3.9, -5.4, +2.2, -9.8 * **Y-Axis Labels (Right):** +1.7, -0.6, +5.0, -1.0, +7.5, -2.3 * **X-Axis Labels (Both):** 0.3B, 0.6B, 1.3B, 3B, 6.7B, 13B * **Pass@k Labels (Right):** Pass@1 (top), Pass@10 (middle), Pass@100 (bottom) * **Bar Colors:** Green (positive performance), Red (negative performance) ### Detailed Analysis **MBPP Performance** * **Pass@1:** * 0.3B: Red bar, value approximately -1.0 * 0.6B: Red bar, value approximately -1.2 * 1.3B: Green bar, value approximately 0.2 * 3B: Green bar, value approximately 1.0 * 6.7B: Green bar, value approximately 2.0 * 13B: Green bar, value approximately 3.5 * Trend: Performance generally increases with model size. * **Pass@10:** * 0.3B: Red bar, value approximately -4.0 * 0.6B: Red bar, value approximately -4.5 * 1.3B: Green bar, value approximately 0.1 * 3B: Green bar, value approximately 1.0 * 6.7B: Green bar, value approximately 2.5 * 13B: Green bar, value approximately 3.0 * Trend: Performance generally increases with model size. * **Pass@100:** * 0.3B: Red bar, value approximately -8.0 * 0.6B: Red bar, value approximately -7.0 * 1.3B: Green bar, value approximately 0.5 * 3B: Green bar, value approximately 1.5 * 6.7B: Green bar, value approximately 2.0 * 13B: Green bar, value approximately 3.0 * Trend: Performance generally increases with model size. **Human Eval Performance** * **Pass@1:** * 0.3B: Red bar, value approximately -0.5 * 0.6B: Red bar, value approximately -0.3 * 1.3B: Green bar, value approximately 0.3 * 3B: Green bar, value approximately 0.5 * 6.7B: Green bar, value approximately 1.2 * 13B: Green bar, value approximately 1.3 * Trend: Performance generally increases with model size. * **Pass@10:** * 0.3B: Red bar, value approximately -0.8 * 0.6B: Red bar, value approximately -0.7 * 1.3B: Red bar, value approximately -0.5 * 3B: Green bar, value approximately 0.3 * 6.7B: Green bar, value approximately 2.5 * 13B: Green bar, value approximately 3.5 * Trend: Performance generally increases with model size. * **Pass@100:** * 0.3B: Red bar, value approximately -2.0 * 0.6B: Red bar, value approximately -1.5 * 1.3B: Red bar, value approximately -1.0 * 3B: Green bar, value approximately 2.0 * 6.7B: Green bar, value approximately 5.0 * 13B: Green bar, value approximately 5.5 * Trend: Performance generally increases with model size. ### Key Observations * Both MBPP and Human Eval show a clear trend of increasing performance (Pass@k) as the model size increases. * Smaller models (0.3B, 0.6B, and sometimes 1.3B) tend to have negative performance (red bars), while larger models (3B, 6.7B, 13B) consistently show positive performance (green bars). * The error bars suggest that the uncertainty in the measurements is relatively small, especially for larger models. ### Interpretation The data suggests that increasing model size significantly improves the performance of both the MBPP system and in human evaluations. The negative performance of smaller models indicates that they may not be effective for the task being evaluated. The consistent positive performance of larger models suggests that they are better suited for the task. The similarity in trends between MBPP and Human Eval indicates that MBPP is a reasonable proxy for human evaluation. The Pass@k metric shows how the model performs when given k attempts to solve a problem. A higher Pass@k value indicates better performance. </details> ure 3 for MBPP (Austin et al., 2021) and HumanEval (Chen et al., 2021) show that it is possible, with the exact same computational budget, to squeeze much more performance out of large language models given a fixed dataset using multi-token prediction. We believe this usefulness only at scale to be a likely reason why multi-token prediction has so far been largely overlooked as a promising training loss for large language model training. ## 3.2. Faster inference We implement greedy self-speculative decoding (Stern et al., 2018) with heterogeneous batch sizes using xFormers (Lefaudeux et al., 2022) and measure decoding speeds of our best 4-token prediction model with 7B parameters on completing prompts taken from a test dataset of code and natural language (Table S2) not seen during training. We observe a speedup of 3 . 0 × on code with an average of 2.5 accepted tokens out of 3 suggestions on code, and of Table 1: Multi-token prediction improves performance and unlocks efficient byte level training. We compare models with 7B parameters trained from scratch on 200B and on 314B bytes of code on the MBPP (Austin et al., 2021), HumanEval (Chen et al., 2021) and APPS (Hendrycks et al., 2021) benchmarks. Multi-token prediction largely outperforms next token prediction on these settings. All numbers were calculated using the estimator from Chen et al. (2021) based on 200 samples per problem. The temperatures were chosen optimally (based on test scores; i.e. these are oracle temperatures) for each model, dataset and pass@k and are reported in Table S12. | Training data | Vocabulary | n | MBPP | MBPP | MBPP | HumanEval | HumanEval | HumanEval | APPS/Intro | APPS/Intro | APPS/Intro | |----------------------|--------------|-----|--------|--------|--------|-------------|-------------|-------------|--------------|--------------|--------------| | Training data | Vocabulary | | @1 | @10 | @100 | @1 | @10 | @100 | @1 | @10 | @100 | | | | 1 | 19.3 | 42.4 | 64.7 | 18.1 | 28.2 | 47.8 | 0.1 | 0.5 | 2.4 | | | | 8 | 32.3 | 50.0 | 69.6 | 21.8 | 34.1 | 57.9 | 1.2 | 5.7 | 14.0 | | | | 16 | 28.6 | 47.1 | 68.0 | 20.4 | 32.7 | 54.3 | 1.0 | 5.0 | 12.9 | | | | 32 | 23.0 | 40.7 | 60.3 | 17.2 | 30.2 | 49.7 | 0.6 | 2.8 | 8.8 | | | | 1 | 30.0 | 53.8 | 73.7 | 22.8 | 36.4 | 62.0 | 2.8 | 7.8 | 17.4 | | | | 2 | 30.3 | 55.1 | 76.2 | 22.2 | 38.5 | 62.6 | 2.1 | 9.0 | 21.7 | | | | 4 | 33.8 | 55.9 | 76.9 | 24.0 | 40.1 | 66.1 | 1.6 | 7.1 | 19.9 | | | | 6 | 31.9 | 53.9 | 73.1 | 20.6 | 38.4 | 63.9 | 3.5 | 10.8 | 22.7 | | | | 8 | 30.7 | 52.2 | 73.4 | 20.0 | 36.6 | 59.6 | 3.5 | 10.4 | 22.1 | | 1T tokens (4 epochs) | | 1 | 40.7 | 65.4 | 83.4 | 31.7 | 57.6 | 83.0 | 5.4 | 17.8 | 34.1 | | 1T tokens (4 epochs) | | 4 | 43.1 | 65.9 | 83.7 | 31.6 | 57.3 | 86.2 | 4.3 | 15.6 | 33.7 | 2 . 7 × on text. On an 8-byte prediction model, the inference speedup is 6 . 4 × (Table S3). Pretraining with multi-token prediction allows the additional heads to be much more accurate than a simple finetuning of a next-token prediction model, thus allowing our models to unlock self-speculative decoding's full potential. ## 3.3. Learning global patterns with multi-byte prediction To show that the next-token prediction task latches to local patterns, we went to the extreme case of byte-level tokenization by training a 7B parameter byte-level transformer on 314B bytes, which is equivalent to around 116B tokens. The 8-byte prediction model achieves astounding improvements compared to next-byte prediction, solving 67% more problems on MBPP pass@1 and 20% more problems on HumanEval pass@1. Multi-byte prediction is therefore a very promising avenue to unlock efficient training of byte-level models. Selfspeculative decoding can achieve speedups of 6 times for the 8-byte prediction model, which would allow to fully compensate the cost of longer byte-level sequences at inference time and even be faster than a next-token prediction model by nearly two times. The 8-byte prediction model is a strong byte-based model, approaching the performance of token-based models despite having been trained on 1 . 7 × less data. ## 3.4. Searching for the optimal n To better understand the effect of the number of predicted tokens, we did comprehensive ablations on models of scale 7B trained on 200B tokens of code. We try n = 1 , 2 , 4 , 6 and 8 in this setting. Results in table 1 show that training with 4-future tokens outperforms all the other models consistently throughout HumanEval and MBPP for pass at 1, 10 and 100 metrics: +3.8%, +2.1% and +3.2% for MBPP and +1.2%, +3.7% and +4.1% for HumanEval. Interestingly, for APPS/Intro, n = 6 takes the lead with +0.7%, +3.0% and +5.3%. It is very likely that the optimal window size depends on input data distribution. As for the byte level models the optimal window size is more consistent (8 bytes) across these benchmarks. ## 3.5. Training for multiple epochs Multi-token training still maintains an edge on next-token prediction when trained on multiple epochs of the same data. The improvements diminish but we still have a +2.4% increase on pass@1 on MBPP and +3.2% increase on pass@100 on HumanEval, while having similar performance for the rest. As for APPS/Intro, a window size of 4 was already not optimal with 200B tokens of training. ## 3.6. Finetuning multi-token predictors Pretrained models with multi-token prediction loss also outperform next-token models for use in finetunings. We evaluate this by finetuning 7B parameter models from Section 3.3 on the CodeContests dataset (Li et al., 2022). We compare the 4-token prediction model with the next-token prediction baseline, and include a setting where the 4-token prediction model is stripped off its additional prediction heads and finetuned using the classical next-token prediction target. According to the results in Figure 4, both ways of finetuning the 4-token prediction model outperform the next-token prediction model on pass@k across k . This means the models are both better at understanding and solving the task and at generating diverse answers. Note that CodeContests is the most challenging coding benchmark we evaluate in this study. Next-token prediction finetuning on top of 4-token prediction pretraining appears to be the best method overall, in line with the classical paradigm of pretraining with auxiliary tasks followed by task-specific finetuning. Please refer to Appendix F for details. Figure 4: Comparison of finetuning performance on CodeContests. We finetune a 4 -token prediction model on CodeContests (Li et al., 2022) (train split) using n ′ -token prediction as training loss with n ′ = 4 or n ′ = 1 , and compare to a finetuning of the next-token prediction baseline model ( n = n ′ = 1 ). For evaluation, we generate 1000 samples per test problem for each temperature T ∈ { 0 . 5 , 0 . 6 , 0 . 7 , 0 . 8 , 0 . 9 } , and compute pass@k for each value of k and T . Shown is k ↦→ max T pass \_ at( k, T ) , i.e. we grant access to a temperature oracle. We observe that both ways of finetuning the 4-token prediction model outperform the next-token prediction baseline. Intriguingly, using next-token prediction finetuning on top of the 4-token prediction model appears to be the best method overall. <details> <summary>Image 4 Details</summary>  ### Visual Description ## Chart: Pass@k vs. k for different n and n' values ### Overview The image is a line chart showing the relationship between "pass@k (%)" on the y-axis and "k" on the x-axis for different values of "n" and "n'". There are three lines representing different combinations of n and n': (n=1, n'=1), (n=4, n'=1), and (n=4, n'=4). The chart uses a logarithmic scale for both axes. ### Components/Axes * **X-axis:** "k" (logarithmic scale), with markers at 1, 10, 100, and 1000. * **Y-axis:** "pass@k (%)" (logarithmic scale), with markers at 0.2, 0.5, 1.0, 2.0, 5.0, and 10.0. * **Legend:** Located on the right side of the chart, identifying the lines: * Solid orange line: "n=1, n'=1" * Dashed dark blue line: "n=4, n'=1" * Dotted light green line: "n=4, n'=4" ### Detailed Analysis * **n=1, n'=1 (Solid Orange Line):** * Trend: The line slopes upward, indicating that "pass@k (%)" increases as "k" increases. * Data Points: * k=1, pass@k ≈ 0.25% * k=10, pass@k ≈ 1.0% * k=100, pass@k ≈ 4.0% * k=1000, pass@k ≈ 10.0% * **n=4, n'=1 (Dashed Dark Blue Line):** * Trend: The line slopes upward, indicating that "pass@k (%)" increases as "k" increases. * Data Points: * k=1, pass@k ≈ 0.3% * k=10, pass@k ≈ 1.5% * k=100, pass@k ≈ 5.5% * k=1000, pass@k ≈ 11.0% * **n=4, n'=4 (Dotted Light Green Line):** * Trend: The line slopes upward, indicating that "pass@k (%)" increases as "k" increases. * Data Points: * k=1, pass@k ≈ 0.25% * k=10, pass@k ≈ 1.2% * k=100, pass@k ≈ 4.5% * k=1000, pass@k ≈ 10.0% ### Key Observations * For smaller values of k (e.g., k=1), the "n=4, n'=1" configuration yields the highest "pass@k (%)". * As k increases, the differences between the three configurations become smaller. * For larger values of k (e.g., k=1000), the "n=4, n'=1" configuration still yields the highest "pass@k (%)", but the "n=1, n'=1" and "n=4, n'=4" configurations are very close. ### Interpretation The chart illustrates how the "pass@k (%)" metric changes with respect to "k" for different configurations of "n" and "n'". The data suggests that increasing "k" generally improves "pass@k (%)" regardless of the "n" and "n'" values. The "n=4, n'=1" configuration appears to perform slightly better than the other configurations, especially for smaller values of "k". The logarithmic scale indicates that the rate of increase in "pass@k (%)" slows down as "k" increases. The values of n and n' seem to have a more significant impact at lower values of k. </details> ## 3.7. Multi-token prediction on natural language To evaluate multi-token prediction training on natural language, we train models of size 7B parameters on 200B tokens of natural language with a 4-token, 2-token and nexttoken prediction loss, respectively. In Figure 5, we evaluate the resulting checkpoints on 6 standard NLP benchmarks. On these benchmarks, the 2-future token prediction model performs on par with the next-token prediction baseline Figure 5: Multi-token training with 7B models doesn't improve performance on choice tasks. This figure shows the evolution of average accuracy of 6 standard NLP benchmarks. Detailed results in Appendix G for 7B models trained on 200B tokens of language data. The 2 future token model has the same performance as the baseline and the 4 future token model regresses a bit. Larger model sizes might be necessary to see improvements on these tasks. <details> <summary>Image 5 Details</summary>  ### Visual Description ## Line Chart: Average Accuracy vs. Training Step for Different 'n' Values ### Overview The image is a line chart showing the relationship between average accuracy and training step for three different values of 'n': 1, 2, and 4. The chart illustrates how the average accuracy changes as the training progresses for each 'n' value. ### Components/Axes * **X-axis:** Training step, ranging from 0 to 25000 in increments of 5000. * **Y-axis:** Average accuracy, ranging from 35.0 to 52.5 in increments of 2.5. * **Legend:** Located in the top-left corner, indicating the 'n' values: * n = 1 (solid red line) * n = 2 (dashed black line) * n = 4 (dotted green line) ### Detailed Analysis * **n = 1 (solid red line):** * Starts at approximately 35.0 accuracy at 0 training steps. * Increases rapidly to approximately 46.0 accuracy by 10000 training steps. * Continues to increase, reaching approximately 49.0 accuracy by 15000 training steps. * Reaches approximately 51.0 accuracy by 20000 training steps. * Plateaus around 52.0 accuracy by 25000 training steps. * **n = 2 (dashed black line):** * Starts at approximately 35.0 accuracy at 0 training steps. * Increases rapidly to approximately 46.0 accuracy by 10000 training steps. * Continues to increase, reaching approximately 48.5 accuracy by 15000 training steps. * Reaches approximately 50.0 accuracy by 20000 training steps. * Plateaus around 51.5 accuracy by 25000 training steps. * **n = 4 (dotted green line):** * Starts at approximately 35.0 accuracy at 0 training steps. * Increases rapidly to approximately 46.0 accuracy by 10000 training steps. * Continues to increase, reaching approximately 48.0 accuracy by 15000 training steps. * Reaches approximately 49.0 accuracy by 20000 training steps. * Plateaus around 50.5 accuracy by 25000 training steps. ### Key Observations * All three lines start at the same accuracy level and increase rapidly in the initial training steps. * The line representing n = 1 consistently achieves the highest average accuracy throughout the training process. * The line representing n = 4 consistently achieves the lowest average accuracy throughout the training process. * The lines converge as the training progresses, but the differences in accuracy remain noticeable. * All three lines show a plateauing effect towards the end of the training, indicating diminishing returns. ### Interpretation The chart suggests that the value of 'n' has a significant impact on the average accuracy achieved during training. A lower value of 'n' (n = 1) results in higher average accuracy compared to higher values (n = 2 and n = 4). This could be due to the specific algorithm or model being used, where a lower 'n' allows for better optimization or generalization. The plateauing effect indicates that increasing the training steps beyond a certain point does not significantly improve the average accuracy, suggesting that the model has reached its learning capacity for the given 'n' values. The data demonstrates that the choice of 'n' is a crucial parameter that needs to be carefully considered to optimize the model's performance. </details> throughout training. The 4-future token prediction model suffers a performance degradation. Detailed numbers are reported in Appendix G. However, we do not believe that multiple-choice and likelihood-based benchmarks are suited to effectively discern generative capabilities of language models. In order to avoid the need for human annotations of generation quality or language model judges-which comes with its own pitfalls, as pointed out by Koo et al. (2023)-we conduct evaluations on summarization and natural language mathematics benchmarks and compare pretrained models with training sets sizes of 200B and 500B tokens and with nexttoken and multi-token prediction losses, respectively. For summarization, we use eight benchmarks where ROUGE metrics (Lin, 2004) with respect to a ground-truth summary allow automatic evaluation of generated texts. We finetune each pretrained model on each benchmark's training dataset for three epochs and select the checkpoint with the highest ROUGE-L F 1 score on the validation dataset. Figure 6 shows that multi-token prediction models with both n = 2 and n = 4 improve over the next-token baseline in ROUGE-L F 1 scores for both training dataset sizes, with the performance gap shrinking with larger dataset size. All metrics can be found in Appendix H. For natural language mathematics, we evaluate the pretrained models in 8-shot mode on the GSM8K benchmark (Cobbe et al., 2021) and measure accuracy of the final answer produced after a chain-of-thought elicited by the fewshot examples. We evaluate pass@k metrics to quantify diversity and correctness of answers like in code evaluations Figure 6: Performance on abstractive text summarization. Average ROUGE-L (longest common subsequence overlap) F 1 score for 7B models trained on 200B and 500B tokens of natural language on eight summarization benchmarks. We finetune the respective models on each task's training data separately for three epochs and select the checkpoints with highest ROUGE-L F 1 validation score. Both n = 2 and n = 4 multi-token prediction models have an advantage over next-token prediction models. Individual scores per dataset and more details can be found in Appendix H. <details> <summary>Image 6 Details</summary>  ### Visual Description ## Bar Chart: Avg. ROUGE-L F1 vs. Training Tokens ### Overview The image is a bar chart comparing the average ROUGE-L F1 score for different values of 'n' (1, 2, and 4) against the number of training tokens (200 and 500). The chart shows how the ROUGE-L F1 score changes with different 'n' values and training token sizes. ### Components/Axes * **Y-axis:** Avg. ROUGE-L F1, ranging from 25.0 to 27.5 in increments of 0.5. * **X-axis:** Training tokens (B), with two values: 200 and 500. * **Legend:** Located in the top-left corner, indicating the 'n' values: * n=1 (salmon color) * n=2 (dark gray color) * n=4 (light green color) ### Detailed Analysis * **Training Tokens (200):** * n=1 (salmon): Avg. ROUGE-L F1 ≈ 26.2 * n=2 (dark gray): Avg. ROUGE-L F1 ≈ 26.7 * n=4 (light green): Avg. ROUGE-L F1 ≈ 26.7 * **Training Tokens (500):** * n=1 (salmon): Avg. ROUGE-L F1 ≈ 27.1 * n=2 (dark gray): Avg. ROUGE-L F1 ≈ 27.4 * n=4 (light green): Avg. ROUGE-L F1 ≈ 27.4 ### Key Observations * For both 200 and 500 training tokens, n=2 and n=4 have similar ROUGE-L F1 scores, which are higher than n=1. * Increasing the number of training tokens from 200 to 500 improves the ROUGE-L F1 score for all values of 'n'. * The improvement in ROUGE-L F1 score is more significant for n=1 when increasing training tokens from 200 to 500. ### Interpretation The chart suggests that increasing the number of training tokens generally improves the ROUGE-L F1 score. The values n=2 and n=4 perform similarly and better than n=1. This indicates that using higher order n-grams (n=2, 4) in the ROUGE-L evaluation metric results in better performance compared to using unigrams (n=1). The improvement observed when increasing training tokens from 200 to 500 suggests that more data leads to better model performance, as expected. The similar performance of n=2 and n=4 might indicate a saturation point where increasing 'n' beyond 2 does not significantly improve the ROUGE-L F1 score. </details> and use sampling temperatures between 0.2 and 1.4. The results are depicted in Figure S13 in Appendix I. For 200B training tokens, the n = 2 model clearly outperforms the next-token prediction baseline, while the pattern reverses after 500B tokens and n = 4 is worse throughout. ## 4. Ablations on synthetic data What drives the improvements in downstream performance of multi-token prediction models on all of the tasks we have considered? By conducting toy experiments on controlled training datasets and evaluation tasks, we demonstrate that multi-token prediction leads to qualitative changes in model capabilities and generalization behaviors . In particular, Section 4.1 shows that for small model sizes, induction capability -as discussed by Olsson et al. (2022)-either only forms when using multi-token prediction as training loss, or it is vastly improved by it. Moreover, Section 4.2 shows that multi-token prediction improves generalization on an arithmetic task, even more so than tripling model size. ## 4.1. Induction capability Induction describes a simple pattern of reasoning that completes partial patterns by their most recent continuation (Olsson et al., 2022). In other words, if a sentence contains 'AB' and later mentions 'A', induction is the prediction that the continuation is 'B'. We design a setup to measure induction Figure 7: Induction capability of n -token prediction models. Shown is accuracy on the second token of two token names that have already been mentioned previously. Shown are numbers for models trained with a next-token and a 2-token prediction loss, respectively, with two independent runs each. The lines denote per-loss averages. For small model sizes, next-token prediction models learn practically no or significantly worse induction capability than 2-token prediction models, with their disadvantage disappearing at the size of 100M nonembedding parameters. <details> <summary>Image 7 Details</summary>  ### Visual Description ## Chart: Induction Success vs. Parameters ### Overview The image is a line chart comparing the induction success rate for two different models (n=1 baseline and n=2 ours) across varying parameter sizes (in millions). The x-axis represents the number of parameters, and the y-axis represents the induction success rate. ### Components/Axes * **X-axis:** Parameters (M), with values 1, 3, 10, 30, 100, 300, 1000. The scale appears to be logarithmic. * **Y-axis:** Induction success, ranging from 0.0 to 0.5 in increments of 0.1. * **Legend:** Located on the right side of the chart. * `x` (coral): n=1 (baseline) * `x` (dark blue): n=2 (ours) ### Detailed Analysis * **n=1 (baseline) - coral line:** * The line starts at approximately (1, 0.02). * At (3, 0.15), the line increases. * At (10, 0.25), the line increases. * At (30, 0.30), the line increases. * At (100, 0.40), the line increases. * At (300, 0.45), the line increases. * At (1000, 0.47), the line increases. * The trend is generally upward, with a decreasing rate of increase as the number of parameters grows. * **n=2 (ours) - dark blue dashed line:** * The line starts at approximately (1, 0.13). * At (3, 0.27), the line increases. * At (10, 0.34), the line increases. * At (30, 0.45), the line increases. * At (100, 0.42), the line decreases slightly. * At (300, 0.48), the line increases. * At (1000, 0.46), the line decreases slightly. * The trend is generally upward, but with a slight dip around 100M parameters. ### Key Observations * Both models show an increase in induction success as the number of parameters increases. * The "n=2 (ours)" model generally outperforms the "n=1 (baseline)" model across all parameter sizes. * The "n=2 (ours)" model shows a slight dip in performance around 100M parameters. ### Interpretation The chart suggests that increasing the number of parameters in both models leads to improved induction success. The "n=2 (ours)" model consistently performs better than the "n=1 (baseline)" model, indicating that the modifications made in the "n=2" model are effective in improving induction success. The slight dip in performance for the "n=2 (ours)" model around 100M parameters could indicate a point of diminishing returns or a need for further optimization at that parameter scale. </details> capability in a controlled way. Training small models of sizes 1M to 1B nonembedding parameters on a dataset of children stories, we measure induction capability by means of an adapted test set: in 100 stories from the original test split, we replace the character names by randomly generated names that consist of two tokens with the tokenizer we employ. Predicting the first of these two tokens is linked to the semantics of the preceding text, while predicting the second token of each name's occurrence after it has been mentioned at least once can be seen as a pure induction task. In our experiments, we train for up to 90 epochs and perform early stopping with respect to the test metric (i.e. we allow an epoch oracle). Figure 7 reports induction capability as measured by accuracy on the names' second tokens in relation to model size for two runs with different seeds. We find that 2-token prediction loss leads to a vastly improved formation of induction capability for models of size 30M nonembedding parameters and below, with their advantage disappearing for sizes of 100M nonembedding parameters and above. 1 We interpret this finding as follows: multitoken prediction losses help models to learn transferring information across sequence positions, which lends itself to the formation of induction heads and other in-context learning mechanisms. However, once induction capability has been formed, these learned features transform induction 1 Note that a perfect score is not reachable in this benchmark as some of the tokens in the names in the evaluation dataset never appear in the training data, and in our architecture, embedding and unembedding parameters are not linked. Figure 8: Accuracy on a polynomial arithmetic task with varying number of operations per expression. Training with multi-token prediction losses increases accuracy across task difficulties. In particular, it also significantly improves out-of-domain generalization performance, albeit at a low absolute level. Tripling the model size, on the other hand, has a considerably smaller effect than replacing next-token prediction with multi-token prediction loss (Figure S16). Shown are two independent runs per configuration with 100M parameter models. <details> <summary>Image 8 Details</summary>  ### Visual Description ## Line Chart: Accuracy vs. Number of Operations ### Overview The image is a line chart showing the relationship between accuracy (in percentage) and the number of operations performed. There are three data series, each representing a different value of 'n' (n=1, n=2, n=4). The x-axis is divided into "in-domain" and "out-of-domain" regions, separated by a vertical dotted line at x=6. ### Components/Axes * **Y-axis:** * Label: "Accuracy (%)" * Scale: 0 to 100, with tick marks at 0, 20, 40, 60, 80, and 100. * **X-axis:** * Label: "# operations" * Scale: 1 to 10, with tick marks at each integer value. * Regions: "in-domain" (1-5) and "out-of-domain" (6-10), indicated by brackets below the axis. * **Legend (Top-Right):** * n=1 (coral color, solid line) * n=2 (dark blue-gray color, dashed line) * n=4 (light blue-gray color, dash-dot line) * **Vertical Dotted Line:** Separates the "in-domain" and "out-of-domain" regions at x=6. ### Detailed Analysis * **Data Series n=1 (coral, solid line):** * Trend: Decreases rapidly as the number of operations increases. * Data Points: * 1 operation: ~90% * 2 operations: ~60% * 3 operations: ~40% * 4 operations: ~30% * 5 operations: ~20% * 6 operations: ~10% * 7 operations: ~2% * 8 operations: ~1% * 9 operations: ~0% * 10 operations: ~0% * **Data Series n=2 (dark blue-gray, dashed line):** * Trend: Decreases rapidly as the number of operations increases. * Data Points: * 1 operation: ~90% * 2 operations: ~75% * 3 operations: ~55% * 4 operations: ~40% * 5 operations: ~30% * 6 operations: ~20% * 7 operations: ~5% * 8 operations: ~2% * 9 operations: ~1% * 10 operations: ~0% * **Data Series n=4 (light blue-gray, dash-dot line):** * Trend: Decreases rapidly as the number of operations increases. * Data Points: * 1 operation: ~100% * 2 operations: ~80% * 3 operations: ~65% * 4 operations: ~45% * 5 operations: ~35% * 6 operations: ~25% * 7 operations: ~5% * 8 operations: ~3% * 9 operations: ~1% * 10 operations: ~0% ### Key Observations * All three data series show a decreasing trend in accuracy as the number of operations increases. * The accuracy drops more sharply in the "out-of-domain" region (after 6 operations). * For a given number of operations, higher values of 'n' generally correspond to higher accuracy. * The accuracy for all values of 'n' converges to near zero as the number of operations approaches 10. ### Interpretation The chart illustrates the performance of a model or system as the complexity of the task (measured by the number of operations) increases. The "in-domain" region likely represents operations that the system is trained or designed to handle, while the "out-of-domain" region represents operations outside of its intended scope. The data suggests that the system's accuracy degrades significantly when it is applied to tasks outside of its training domain. The parameter 'n' appears to influence the system's robustness, with higher values of 'n' leading to better accuracy, especially for a lower number of operations. The rapid decline in accuracy in the "out-of-domain" region indicates that the system is not well-equipped to handle these types of operations, regardless of the value of 'n'. </details> into a task that can be solved locally at the current token and learned with next-token prediction alone. From this point on, multi-token prediction actually hurts on this restricted benchmark-but we surmise that there are higher forms of in-context reasoning to which it further contributes, as evidenced by the results in Section 3.1. In Figure S14, we provide evidence for this explanation: replacing the children stories dataset by a higher-quality 9:1 mix of a books dataset with the children stories, we enforce the formation of induction capability early in training by means of the dataset alone. By consequence, except for the two smallest model sizes, the advantage of multi-token prediction on the task disappears: feature learning of induction features has converted the task into a pure next-token prediction task. ## 4.2. Algorithmic reasoning Algorithmic reasoning tasks allow to measure more involved forms of in-context reasoning than induction alone. We train and evaluate models on a task on polynomial arithmetic in the ring F 7 [ X ] / ( X 5 ) with unary negation, addition, multiplication and composition of polynomials as operations. The coefficients of the operands and the operators are sampled uniformly. The task is to return the coefficients of the polynomials corresponding to the resulting expressions. The number m of operations contained in the expressions is selected uniformly from the range from 1 to 5 at training time, and can be used to adjust the difficulty of both in-domain ( m ≤ 5 ) and out-of-domain ( m> 5 ) generalization evaluations. The evaluations are conducted with greedy sampling on a fixed test set of 2000 samples per number of operations. We train models of two small sizes with 30M and 100M nonembedding parameters, respectively. This simulates the conditions of large language models trained on massive text corpora which are likewise under-parameterized and unable to memorize their entire training datasets. Multi-token prediction improves algorithmic reasoning capabilities as measured by this task across task difficulties (Figure 8). In particular, it leads to impressive gains in out-of-distribution generalization, despite the low absolute numbers. Increasing the model size from 30M to 100M parameters, on the other hand, does not improve evaluation accuracy as much as replacing next-token prediction by multi-token prediction does (Figure S16). In Appendix K, we furthermore show that multi-token prediction models retain their advantage over next-token prediction models on this task when trained and evaluated with pause tokens (Goyal et al., 2023). ## 5. Why does it work? Some speculation Why does multi-token prediction afford superior performance on coding evaluation benchmarks, and on small algorithmic reasoning tasks? Our intuition, developed in this section, is that multi-token prediction mitigates the distributional discrepancy between training-time teacher forcing and inference-time autoregressive generation. We support this view with an illustrative argument on the implicit weights multi-token prediction assigns to tokens depending on their relevance for the continuation of the text, as well as with an information-theoretic decomposition of multi-token prediction loss. ## 5.1. Lookahead reinforces choice points Not all token decisions are equally important for generating useful texts from language models (Bachmann and Nagarajan, 2024; Lin et al., 2024). While some tokens allow stylistic variations that do not constrain the remainder of the text, others represent choice points that are linked with higher-level semantic properties of the text and may decide whether an answer is perceived as useful or derailing . Multi-token prediction implicitly assigns weights to training tokens depending on how closely they are correlated with their successors. As an illustrative example, consider the sequence depicted in Figure 9 where one transition is a hard-to-predict choice point while the other transitions are considered 'inconsequential'. Inconsequential transitions following a choice point are likewise hard to predict in advance. By marking and counting loss terms, we find that Figure 9: Multi-token prediction loss assigns higher implicit weights to consequential tokens. Shown is a sequence in which all transitions except '5 → A' are easy to predict, alongside the corresponding prediction targets in 3-token prediction. Since the consequences of the difficult transition '5 → A' are likewise hard to predict, this transition receives a higher implicit weight in the overall loss via its correlates '3 → A', ..., '5 → C'. <details> <summary>Image 9 Details</summary>  ### Visual Description ## Model Prediction vs. Ground Truth Diagram ### Overview The image presents a diagram comparing model predictions against ground truth values. It illustrates how well a model's predictions align with the actual values across a sequence of inputs. The diagram uses numerical and alphabetical values to represent both the ground truth and the model's predictions. ### Components/Axes * **Y-Axis (Implied):** Represents the range of possible predictions for each input. * **X-Axis (Implied):** Represents the sequence of inputs, labeled as "Ground truth". * **Model predictions:** Label indicating the predictions made by the model. * **Ground truth:** Label indicating the actual values. The ground truth values are: 1, 2, 3, 4, 5, A, B. * **Arrows:** Arrows point from the ground truth values to the corresponding model predictions. * **Prediction Columns:** Each column above a ground truth value contains the model's predictions for that input. ### Detailed Analysis The diagram shows the model's predictions for each ground truth value. The predictions are arranged in columns above each ground truth value. The predictions are either correct (green) or incorrect (red). * **Ground Truth 1:** Model predicts 2, 3, 4 (all green, thus correct) * **Ground Truth 2:** Model predicts 3, 4, 5 (all green, thus correct) * **Ground Truth 3:** Model predicts 4, 5, A (A is red, thus incorrect) * **Ground Truth 4:** Model predicts 5, A, B (A and B are red, thus incorrect) * **Ground Truth 5:** Model predicts A, B, C (A, B, and C are red, thus incorrect) * **Ground Truth A:** Model predicts B, C, D (all green, thus correct) * **Ground Truth B:** Model predicts C, D, E (all green, thus correct) ### Key Observations * The model's predictions are accurate for the first two numerical inputs (1 and 2) and the last two alphabetical inputs (A and B). * The model struggles with the middle numerical inputs (3, 4, and 5), making incorrect predictions. * The incorrect predictions are highlighted in red, while the correct predictions are in green. ### Interpretation The diagram suggests that the model performs well on the extremes of the input sequence but struggles with the middle values. This could indicate a bias in the model's training data or a limitation in its ability to generalize across the entire input range. The transition from numerical to alphabetical inputs does not seem to be a factor, as the model performs well on both ends of the sequence. The red values indicate where the model's predictions deviate from the actual ground truth, highlighting areas for potential improvement. </details> n -token prediction associates a weight of n ( n +1) 2 to choice points via their correlates, and a smaller weight of n to inconsequential points. Please refer to Appendix L.3 for more details. Generally, we believe that the quality of text generations depends on picking the right decisions at choice points, and that n -token prediction losses promote those. ## 5.2. Information-theoretic argument Language models are typically trained by teacher-forcing, where the model receives the ground truth for each future token during training. However, during test time generation is unguided and autoregressive, whereby errors accumulate. Teacher-forcing, we argue, encourages models to focus on predicting well in the very short term, at the potential expense of ignoring longer-term dependencies in the overall structure of the generated sequence. To illustrate the impact of multi-token prediction, consider the following information-theoretic argument. Here, X denotes the next future token, and Y the second-next future token. The production of both of these tokens is conditioned on some observed, input context C , that we omit from our equations for simplicity. When placed before token X , vanilla next-token prediction concerns the quantity H ( X ) , while multi-token prediction with n = 2 aims at H ( X ) + H ( Y ) . We decompose these two quantities as: <!-- formula-not-decoded --> By discarding the term H ( Y | X ) -which appears again when predicting at the following position-we observe that 2-token prediction increases the importance of I ( X ; Y ) by a factor of 2 . So, multi-token predictors are more accurate at predicting tokens X that are of relevance for the remainder of the text to come. In Appendix L.2, we give a relative version of the above equations that shows the increased weight of relative mutual information in a loss decomposition of 2-token prediction loss. ## 6. Related work Language modeling losses Dong et al. (2019) and Tay et al. (2022) train on a mixture of denoising tasks with different attention masks (full, causal and prefix attention) to bridge the performance gap with next token pretraining on generative tasks. Tay et al. (2022) uses the span corruption objective, which replaces spans of tokens with special tokens for the encoder and the decoder then predicts the contents of those spans. Unlike UniLM, this allows full causal training with teacher forcing. Similarly, Yang et al. (2019) train on permuted sequences, while conserving the original positional embeddings, effectively training the model to predict various parts of the sequence given a mix of past and future information. This permuted language modeling is the closest task to ours since it allows predicting beyond the next token. However all of these language modeling tasks train on a small percentage of the input text: on average only 15% of the tokens are backwarded through. For Dong et al. (2019), where the masking is done in BERT style, it is hard to mask more than 15% since it destroys too much information. For Tay et al. (2022), it is technically possible to have a larger proportion but in practice, the settings used have between 15% and 25% of masked tokens. (Yang et al., 2019) also makes it possible to train on the whole sequence since it is only permuted, and no information is lost. Yet, in practice, since the completely random permutation is very hard to reconstruct, only 15% are predicted for training stability reasons. Multi-token prediction in language modelling Qi et al. (2020) argue that multi-token prediction encourages planning, improves representations and prevents the overfitting on local patterns that can result from teacher-forced training. However, their technical approach replicates the residual stream n -fold while ours allows for compute-matched comparisons and makes the residual representations participate more directly in the auxiliary loss terms. Stern et al. (2018) and Cai et al. (2024) propose model finetunings with multitoken prediction for faster inference but do not study the effects of such a loss during pretraining. Pal et al. (2023) use probing methods to show that next-token prediction models are able to predict additional consecutive tokens to a certain extent, but less so than our models which are specifically trained for this task. Jianyu Zhang (2024) observe improvements in language modelling tasks with multi-label binary classification over the occurrence of vocabulary words in the future as an auxiliary learning task. Self-speculative decoding Stern et al. (2018) are, to the best of our knowledge, the first to suggest a speculative decoding scheme for faster inference. Our architecture replaces their linear prediction heads by transformer layers, but is otherwise similar. By reorganizing the order of the forward/backward, we can use all loss terms instead of stochastically picking one head for loss computation. Cai et al. (2024) present a more elaborate self-speculative decoding scheme that uses the topk predictions of each head instead of the best one only. It can be used with the multi-token prediction models we train. Multi-target prediction Multi-task learning is the paradigm of training neural networks jointly on several tasks to improve performance on the tasks of interest (Caruana, 1997). Learning with such auxiliary tasks allows models to exploit dependencies between target variables and can even be preferable in the case of independent targets (Waegeman et al., 2019). While more specifically tailored architectures for multi-target prediction are conceivable (SpyromitrosXioufis et al., 2016; Read et al., 2021), modern deep learning approaches usually rely on large shared model trunks with separate prediction heads for the respective tasks (Caruana, 1997; Silver et al., 2016; Lample et al., 2022) like we do. Multi-target prediction has been shown to be a successful strategy in various domains, e.g. for learning time series prediction with more distant time steps in the future as auxiliary targets (Vapnik and Vashist, 2009) or for learning from videos with several future frames (Mathieu et al., 2016; Srivastava et al., 2016) or representations of future frames (Vondrick et al., 2016) as auxiliary targets. ## 7. Conclusion Wehave proposed multi-token prediction as an improvement over next-token prediction in training language models for generative or reasoning tasks. Our experiments (up to 7B parameters and 1T tokens) show that this is increasingly useful for larger models and in particular show strong improvements for code tasks. We posit that our method reduces distribution mismatch between teacher-forced training and autoregressive generation. When used with speculative decoding, exact inference gets 3 times faster. In future work we would like to better understand how to automatically choose n in multi-token prediction losses. One possibility to do so is to use loss scales and loss balancing (Défossez et al., 2022). Also, optimal vocabulary sizes for multi-token prediction are likely different from those for next-token prediction, and tuning them could lead to better results, as well as improved trade-offs between compressed sequence length and compute-per-byte expenses. Finally, we would like to develop improved auxiliary prediction losses that operate in embedding spaces (LeCun, 2022). ## Impact statement The goal of this paper is to make language models more compute and data efficient. While this may in principle reduce the ecological impact of training LLMs, we shall be careful about rebound effects . All societal advantages, as well as risks, of LLMs should be considered while using this work. ## Environmental impact In aggregate, training all models reported in the paper required around 500K GPU hours of computation on hardware of type A100-80GB and H100. Estimated total emissions were around 50 tCO2eq, 100% of which were offset by Meta's sustainability program. ## Acknowledgements We thank Jianyu Zhang, Léon Bottou, Emmanuel Dupoux, Pierre-Emmanuel Mazaré, Yann LeCun, Quentin Garrido, Megi Dervishi, Mathurin Videau and Timothée Darcet and other FAIR PhD students and CodeGen team members for helpful discussions. We thank Jonas Gehring for his technical expertise and the original Llama team and xFormers team for enabling this kind of research. ## References - 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. - Gregor Bachmann and Vaishnavh Nagarajan. The pitfalls of next-token prediction, 2024. - Samy Bengio, Oriol Vinyals, Navdeep Jaitly, and Noam Shazeer. Scheduled sampling for sequence prediction with recurrent neural networks, 2015. Yonatan Bisk, Rowan Zellers, Ronan Le Bras, Jianfeng Gao, and Yejin Choi. Piqa: Reasoning about physical commonsense in natural language, 2019. - 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. Rich Caruana. Multitask learning. Machine learning , 28: 41-75, 1997. - 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. - 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. - 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. - 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 13063-13075, 2019. - Alexandre Défossez, Jade Copet, Gabriel Synnaeve, and Yossi Adi. High fidelity neural audio compression. arXiv preprint arXiv:2210.13438 , 2022. - Moussa Kamal Eddine, Antoine J. P. Tixier, and Michalis Vazirgiannis. Barthez: a skilled pretrained french sequence-to-sequence model, 2021. - Alexander R. Fabbri, Irene Li, Tianwei She, Suyi Li, and Dragomir R. Radev. Multi-news: a large-scale multidocument summarization dataset and abstractive hierarchical model, 2019. Mehrdad Farahani. Summarization using bert2bert model on wikisummary dataset. https://github.com/m3hrdadfi/wikisummary, 2020. 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) . IEEE, March 2021. doi: 10.1109/ csicc52343.2021.9420563. URL http://dx.doi. org/10.1109/CSICC52343.2021.9420563 . - Michael C Frank. Bridging the data gap between children and large language models. Trends in Cognitive Sciences , 2023. 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: //dx.doi.org/10.18653/v1/D19-5409 . - 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. - 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. - Ari Holtzman, Jan Buys, Li Du, Maxwell Forbes, and Yejin Choi. The curious case of neural text degeneration, 2020. - Jianyu Zhang Leon Bottou. Multi-label classification as an auxiliary loss for language modelling. personal communication, 2024. - Mandar Joshi, Eunsol Choi, Daniel S. Weld, and Luke Zettlemoyer. Triviaqa: A large scale distantly supervised challenge dataset for reading comprehension, 2017. Diederik Kingma and Jimmy Ba. Adam: A method for stochastic optimization. ICLR , 2015. 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. 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, MingWei Chang, Andrew Dai, Jakob Uszkoreit, Quoc Le, and Slav Petrov. Natural questions: a benchmark for question answering research. Transactions of the Association of Computational Linguistics , 2019. - 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. - Yann LeCun. A path towards autonomous machine intelligence version 0.9. 2, 2022-06-27. Open Review , 62(1), 2022. - 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. - Yaniv Leviathan, Matan Kalman, and Yossi Matias. Fast inference from transformers via speculative decoding, 2023. - 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. - 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 . - 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. - Ilya Loshchilov and Frank Hutter. Sgdr: Stochastic gradient descent with warm restarts, 2017. - Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization, 2019. - Michael Mathieu, Camille Couprie, and Yann LeCun. Deep multi-scale video prediction beyond mean square error, 2016. - Ramesh Nallapati, Bowen Zhou, Cicero Nogueira dos santos, Caglar Gulcehre, and Bing Xiang. Abstractive text - summarization using sequence-to-sequence rnns and beyond, 2016. 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. 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 Amodei, Tom Brown, Jack Clark, Jared Kaplan, Sam McCandlish, and Chris Olah. In-context learning and induction heads. Transformer Circuits Thread , 2022. https://transformer-circuits.pub/2022/in-contextlearning-and-induction-heads/index.html. OpenAI. Gpt-4 technical report, 2023. 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 language models to follow instructions with human feedback, 2022. - Koyena Pal, Jiuding Sun, Andrew Yuan, Byron C. Wallace, and David Bau. Future lens: Anticipating subsequent tokens from a single hidden state, 2023. - Weizhen Qi, Yu Yan, Yeyun Gong, Dayiheng Liu, Nan Duan, Jiusheng Chen, Ruofei Zhang, and Ming Zhou. Prophetnet: Predicting future n-gram for sequence-tosequence pre-training, 2020. - Jesse Read, Bernhard Pfahringer, Geoffrey Holmes, and Eibe Frank. Classifier chains: A review and perspectives. Journal of Artificial Intelligence Research , 70:683-718, 2021. - 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. - Maarten Sap, Hannah Rashkin, Derek Chen, Ronan LeBras, and Yejin Choi. Socialiqa: Commonsense reasoning about social interactions, 2019. - 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. - 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.08360 , 2023. - 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. - Nitish Srivastava, Elman Mansimov, and Ruslan Salakhutdinov. Unsupervised learning of video representations using lstms, 2016. - Mitchell Stern, Noam Shazeer, and Jakob Uszkoreit. Blockwise parallel decoding for deep autoregressive models, 2018. - 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. - Vladimir Vapnik and Akshay Vashist. A new learning paradigm: Learning using privileged information. Neural networks , 22(5-6):544-557, 2009. - Carl Vondrick, Hamed Pirsiavash, and Antonio Torralba. Anticipating visual representations from unlabeled video, 2016. - Willem Waegeman, Krzysztof Dembczy´ nski, and Eyke Hüllermeier. Multi-target prediction: a unifying view on problems and methods. Data Mining and Knowledge Discovery , 33:293-324, 2019. - 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. - 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. - Rowan Zellers, Ari Holtzman, Yonatan Bisk, Ali Farhadi, and Yejin Choi. Hellaswag: Can a machine really finish your sentence?, 2019. ## A. Additional results on self-speculative decoding <details> <summary>Image 10 Details</summary>  ### Visual Description ## Line Chart: Throughput vs. Batch Size for Different 'k' Values ### Overview The image is a line chart that illustrates the relationship between throughput (relative) and batch size for four different values of 'k' (k=1, k=2, k=3, k=4). The chart shows how throughput changes as batch size increases, with each line representing a different 'k' value. ### Components/Axes * **X-axis:** Batch size, with markers at 1, 8, 16, 24, 32, and 40. * **Y-axis:** Throughput (relative), with markers at 0.0, 0.5, 1.0, 1.5, 2.0, 2.5, and 3.0. * **Legend (bottom-right):** * k=1 (coral color, solid line with 'x' markers) * k=2 (dark blue color, dashed line with 'x' markers) * k=3 (light orange color, dotted line with 'x' markers) * k=4 (light green color, dotted line with 'x' markers) ### Detailed Analysis * **k=1 (coral, solid line):** The throughput remains relatively constant at approximately 1.0 across all batch sizes. * Batch size 1: Throughput ~1.0 * Batch size 40: Throughput ~1.0 * **k=2 (dark blue, dashed line):** The throughput starts at approximately 1.75 at batch size 1, increases slightly to approximately 1.85 by batch size 8, and then remains relatively constant around 1.85 for larger batch sizes. * Batch size 1: Throughput ~1.75 * Batch size 8: Throughput ~1.85 * Batch size 40: Throughput ~1.85 * **k=3 (light orange, dotted line):** The throughput starts at approximately 2.45 at batch size 1, increases slightly to approximately 2.55 by batch size 8, and then remains relatively constant around 2.55 for larger batch sizes. * Batch size 1: Throughput ~2.45 * Batch size 8: Throughput ~2.55 * Batch size 40: Throughput ~2.55 * **k=4 (light green, dotted line):** The throughput starts at approximately 3.05 at batch size 1, increases slightly to approximately 3.15 by batch size 8, and then remains relatively constant around 3.15 for larger batch sizes. * Batch size 1: Throughput ~3.05 * Batch size 8: Throughput ~3.15 * Batch size 40: Throughput ~3.15 ### Key Observations * Throughput generally increases with increasing 'k' value. * For each 'k' value, throughput is relatively stable across different batch sizes, with a slight increase from batch size 1 to batch size 8. * The lines for k=2, k=3, and k=4 show a slight increase in throughput between batch sizes 1 and 8, after which they plateau. ### Interpretation The chart suggests that increasing the 'k' value leads to a higher throughput. However, the batch size has a minimal impact on throughput, especially for batch sizes greater than 8. The 'k' parameter appears to be a more significant factor in determining throughput than batch size. The data implies that optimizing 'k' is more crucial for improving throughput than adjusting the batch size, at least within the range of batch sizes considered in the chart. </details> Batch size Figure S10: Decoding speeds and latencies with self-speculative decoding relative to standard autoregressive decoding. We use k heads of a 4-token prediction model and evaluate decoding speeds of a code model as explained in Table S2. All numbers are relative to the autoregressive ( k = 1 ) baseline with the same batch size. Table S2: Relative speedups with self-speculative decoding. For wikipedia and books we prompt a 7B parameter model trained on 500B tokens, and for code we prompt a 7B parameter model trained on 1T tokens of code on 4200 sequences of 512 tokens from a test dataset not seen during training, and generate completions consisting of 512 tokens using greedy self-speculative decoding (Stern et al., 2018) using the indicated number of heads from a 4-token prediction model. Note that the maximal speedup that can be obtained with self-speculative decoding using k heads is k . The last column shows the average number of tokens retrieved from a forward containing this sequence (both verification and prediction). The speedup was evaluated at the maximal batch size of 42, but is constant across batch sizes (Figure S10). | | Wikipedia | Wikipedia | Books | Books | Code | Code | |--------------|--------------|------------------|--------------|------------------|--------------|------------------| | # Heads used | Rel. speedup | Tokens / forward | Rel. speedup | Tokens / forward | Rel. speedup | Tokens / forward | | 1 | 1.00 | 1.00 | 1.00 | 1.00 | 1.00 | 1.00 | | 2 | 1.79 | 1.88 | 1.77 | 1.87 | 1.85 | 1.94 | | 3 | 2.35 | 2.57 | 2.32 | 2.56 | 2.54 | 2.78 | | 4 | 2.74 | 3.12 | 2.67 | 3.09 | 3.05 | 3.50 | Table S3: Relative speedups with self-speculative decoding with byte-level models on code. We prompt the 7B parameter models from Section 3.3 on 4096 sequences of 1024 bytes of code not seen during training, and generate completions consisting of 1024 bytes using greedy self-speculative decoding (Stern et al., 2018) as in Table S2. The speedup was evaluated at a batch size of 16. | | n = 8 | n = 8 | n = 16 | n = 16 | n = 32 | n = 32 | |--------------|--------------|------------------|--------------|------------------|--------------|------------------| | # Heads used | Rel. speedup | Tokens / forward | Rel. speedup | Tokens / forward | Rel. speedup | tokens / forward | | 1 | 1.00 | 1.00 | 1.00 | 1.00 | 1.00 | 1.00 | | 2 | 1.94 | 1.98 | 1.94 | 1.98 | 1.93 | 1.97 | | 4 | 3.67 | 3.84 | 3.63 | 3.81 | 3.62 | 3.80 | | 8 | 6.39 | 7.04 | 6.25 | 6.92 | 6.22 | 6.89 | | 12 | - | - | 8.07 | 9.36 | 8.01 | 9.30 | | 16 | - | - | 9.24 | 11.20 | 9.15 | 11.15 | | 20 | - | - | - | - | 9.83 | 12.61 | | 24 | - | - | - | - | 10.34 | 13.67 | | 28 | - | - | - | - | 10.55 | 14.58 | | 32 | - | - | - | - | 10.84 | 15.35 | Latency (relative) 1.0 0.8 0.6 0.4 0.2 0.0 k=1 k=2 k=3 k=4 ## B. Alternative architectures Table S4: Alternative architectures improve on baseline but not as consistently. Alternative architectures for multi-token prediction are worth exploring to improve efficiency. Here we tried Anticausal, causal and linear and showed no significant improvement with respect to Parallel architecture. | | | | | MBPP | MBPP | MBPP | HumanEval | HumanEval | HumanEval | APPS/Intro | APPS/Intro | APPS/Intro | |----|-------------|--------------|---------|--------|--------|--------|-------------|-------------|-------------|--------------|--------------|--------------| | n | Head type | Architecture | +Layers | @1 | @10 | @100 | @1 | @10 | @100 | @1 | @10 | @100 | | 1 | transformer | parallel | 0 | 30.0 | 53.8 | 73.7 | 22.8 | 36.4 | 62.0 | 2.8 | 7.8 | 17.4 | | | linear | parallel | 0 | 33.6 | 55.0 | 76.2 | 21.9 | 38.5 | 63.7 | 3.1 | 10.1 | 23.0 | | | | anticausal | 0 | 30.8 | 54.8 | 75.3 | 20.9 | 38.4 | 64.5 | 2.0 | 8.7 | 21.6 | | 4 | transformer | causal | 0 | 31.9 | 54.9 | 74.9 | 20.9 | 38.1 | 67.3 | 4.0 | 11.6 | 22.8 | | | | | 0 | 33.8 | 55.9 | 76.9 | 24.0 | 40.1 | 66.1 | 1.6 | 7.1 | 19.9 | | | | parallel | 3 | 33.3 | 55.7 | 77.3 | 22.4 | 39.4 | 66.7 | 2.6 | 9.5 | 22.1 | The architecture described in Section 2 is not the only sensible option, but proved technically viable and well-performing in our experiments. We describe and compare alternative architectures in this section. Replicated unembeddings Replicating the unembedding matrix n times is a simple method for implementing multi-token prediction architectures. However, it requires matrices with shapes ( d, nV ) in the notation of Section 2, which is prohibitive for large-scale trainings. Linear heads Apart from using a single transformer layer for the heads H i , other architectures are conceivable. We experimented with a single linear layer without any nonlinearity as heads, amounting to linear probing of the model's residual representation z . Architectures with more than one layer per head are also possible, but we did not pursue this direction further. Causal and anticausal variant Instead of making the prediction heads P i ( x t + i | z t :1 ) architecturally independent of each other, we can also allow them to rely on other heads' (pre-unembedding) outputs. In a causal variant, later prediction heads are applied on top of the previous ones, i.e. the i -th prediction head P i is given by <!-- formula-not-decoded --> In another anticausal variant, the network starts by predicting the most distant tokens before gradually refining up to the following token: <!-- formula-not-decoded --> These architectures likewise allow a sequential forward/backward order as the parallel architecture from Section 2. This is described in Figure S11. ## D. Finetuning Table S6: Finetuning LLama 2 with multi-token prediction does not significantly improve performance. We tried to finetune LLama 2 with 4-token prediction but this did not yield significant improvements compared to the baseline. We suppose that this new loss changes the initialization too brutally and never really recovers. We still some improvements for example on MBPP Pass@1. All runs use 200B tokens of code. | | | | MBPP | MBPP | MBPP | HumanEval | HumanEval | HumanEval | APPS/Intro | APPS/Intro | APPS/Intro | |----|-------------|---------|--------|--------|--------|-------------|-------------|-------------|--------------|--------------|--------------| | n | Head type | +Layers | @1 | @10 | @100 | @1 | @10 | @100 | @1 | @10 | @100 | | 1 | transformer | 0 | 39.6 | 65.1 | 82.4 | 31.4 | 57.7 | 84.7 | 10.0 | 21.6 | 36.7 | | 4 | linear | 0 | 39.3 | 63.7 | 81.3 | 29.0 | 53.4 | 82.2 | 6.9 | 20.0 | 34.0 | | 4 | transformer | 0 | 38.3 | 62.2 | 80.1 | 27.9 | 53.6 | 82.4 | 5.8 | 18.2 | 34.3 | | 4 | transformer | 3 | 42.5 | 64.4 | 81.3 | 28.7 | 56.9 | 82.4 | 7.8 | 21.2 | 37.3 | Figure S11: Order of the forward/backward in a causal n -token prediction model with n = 2 heads. Like in the forward/backward depicted for parallel prediction heads in Figure 2, we avoid materializing all unembedding layer gradients in memory simultaneously and reduce peak GPU memory usage significantly. The iteration over the heads starts with the one furthest to the trunk. At each head, a gradient from the succeeding prediction heads and from the head's own loss are accumulated for both the head's output and its weights. <details> <summary>Image 11 Details</summary>  ### Visual Description ## Diagram: Multi-Head Attention Flow ### Overview The image is a diagram illustrating the flow of information in a multi-head attention mechanism. It shows the relationships between the input, trunk, head, and loss components, with arrows indicating the direction of information flow. ### Components/Axes * **Boxes:** Represent different components of the attention mechanism. * Input * Trunk * Head 1 * Head 2 * Loss 1 * Loss 2 * **Arrows:** Indicate the direction of information flow between components. * **Numbers:** Label the arrows, possibly indicating the sequence or type of operation. ### Detailed Analysis The diagram shows the following flow of information: 1. **Input** to **Trunk**: Arrow labeled "1" indicates information flows from the Input to the Trunk. 2. **Trunk** to **Head 1**: Arrow labeled "2" indicates information flows from the Trunk to Head 1. 3. **Head 1** to **Head 2**: Arrow labeled "3" indicates information flows from Head 1 to Head 2. 4. **Head 2** to **Loss 2**: Arrow labeled "4" indicates information flows from Head 2 to Loss 2. 5. **Loss 2** to **Head 2**: Arrow labeled "5" indicates information flows from Loss 2 to Head 2. 6. **Head 2** to **Head 1**: Arrow labeled "6" indicates information flows from Head 2 to Head 1. 7. **Head 1** to **Loss 1**: Arrow labeled "7" indicates information flows from Head 1 to Loss 1. 8. **Loss 1** to **Head 1**: Arrow labeled "8" indicates information flows from Loss 1 to Head 1. 9. **Head 1** to **Trunk**: Arrow labeled "9" indicates information flows from Head 1 to Trunk. 10. **Trunk** to **Input**: Arrow labeled "10" indicates information flows from Trunk to Input. ### Key Observations * The diagram illustrates a cyclical flow of information, with feedback loops between the heads and losses. * The Trunk acts as a central processing unit, receiving input and distributing information to the heads. * There are two heads (Head 1 and Head 2) and two corresponding losses (Loss 1 and Loss 2), suggesting a multi-head attention mechanism. ### Interpretation The diagram represents a multi-head attention mechanism, where the input is processed by a trunk and then distributed to multiple heads. Each head computes attention weights and produces a loss, which is then fed back into the head. This cyclical flow of information allows the model to learn complex relationships between the input and output. The presence of multiple heads allows the model to attend to different parts of the input simultaneously, improving its ability to capture long-range dependencies. The feedback loops between the heads and losses enable the model to refine its attention weights and improve its performance. </details> ## C. Training speeds Table S5: Training time relative to next-token prediction training. The slight overhead when using multi-token prediction here is explained by a suboptimal use of Fully Sharded Data Parallel. In our implementation, when doing separate backward passes for each head, we lose the overlap of layer weight communication and computation, therefore it incurs a very slight overhead that can be removed if reimplemented correctly. | Model | n=1 | n=2 | n=4 | |---------|-------|-------|-------| | 0.3B | 1 | 1.07 | 1.22 | | 0.6B | 1 | 1.05 | 1.13 | | 1.3B | 1 | 1.04 | 1.12 | | 3B | 1 | 1.02 | 1.07 | | 6.7B | 1 | 1.02 | 1.07 | | 13B | 1 | 1.04 | 1.09 | ## E. Additional results on model scaling behavior Table S7: Scaling model size Full results of scaling model size with n=1,2 and 4. | | | MBPP | MBPP | MBPP | HumanEval | HumanEval | HumanEval | |------------|-----|--------|--------|--------|-------------|-------------|-------------| | Model Size | Fut | @1 | @10 | @100 | @1 | @10 | @100 | | 0.3B | 1 | 1.8 | 10.4 | 29.9 | 1.9 | 5.0 | 10.9 | | 0.3B | 2 | 1.7 | 10.1 | 27.2 | 1.5 | 4.4 | 10.3 | | 0.3B | 4 | 1.0 | 6.3 | 20.1 | 1.2 | 4.0 | 8.6 | | 0.6B | 1 | 4.7 | 21.0 | 45.2 | 2.9 | 8.5 | 16.7 | | 0.6B | 2 | 4.6 | 21.0 | 44.7 | 3.2 | 8.9 | 16.2 | | 0.6B | 4 | 3.0 | 15.6 | 38.0 | 2.7 | 7.7 | 15.5 | | 1.3B | 1 | 6.8 | 27.0 | 51.0 | 4.6 | 13.1 | 24.3 | | 1.3B | 2 | 7.3 | 27.5 | 51.7 | 5.4 | 13.6 | 23.3 | | 1.3B | 4 | 7.4 | 27.6 | 50.1 | 4.8 | 12.3 | 22.5 | | 3B | 1 | 11.1 | 36.4 | 60.4 | 7.2 | 17.2 | 29.8 | | 3B | 2 | 11.8 | 37.2 | 60.5 | 8.0 | 18.2 | 31.2 | | 3B | 4 | 12.7 | 37.6 | 61.1 | 7.2 | 18.5 | 33.3 | | 6.7B | 1 | 23.9 | 54.2 | 74.7 | 12.8 | 29.3 | 51.7 | | 6.7B | 2 | 24.7 | 54.8 | 76.4 | 13.2 | 32.2 | 53.9 | | 6.7B | 4 | 26.0 | 55.8 | 76.0 | 13.8 | 33.2 | 58.5 | | 13B | 1 | 26.0 | 57.1 | 77.0 | 14.1 | 33.6 | 56.0 | | 13B | 2 | 30.5 | 60.5 | 79.4 | 15.2 | 36.9 | 60.0 | | 13B | 4 | 30.5 | 61.0 | 79.2 | 15.8 | 38.6 | 63.5 | ## F. Details on CodeContests finetuning We use the Python subset of the CodeContests (Li et al., 2022) train split with reward annotations ('correct' / 'incorrect') and condition on correct solutions at evaluation time. For evaluation, we generate 1000 samples per problem from the test split for each temperature T ∈ { 0 . 5 , 0 . 6 , 0 . 7 , 0 . 8 , 0 . 9 } , and compute the unbiased estimator for pass@k from Chen et al. (2021) for each value of k and T . It is possible that models that were pretrained with different losses have different respective optimal temperatures for pass@k, so we compute and show k ↦→ max T pass \_ at( k, T ) in Figure 4. In other words, we grant pass@k access to a temperature oracle. For small values of k , pass@k measures the capability of understanding and solving tasks while for large k , it additionally favors diversity in outputs. According to the results in Figure 4, multi-token prediction pretraining leads to finetuned models that are better on both axes. ## G. Additional results on natural language benchmarks We evaluate the models from Section 3.7 on standard natural language processing benchmarks: ARC Challenge (Yadav et al., 2019), COPA (Roemmele et al., 2011), Hellaswag (Zellers et al., 2019), Natural Questions (Kwiatkowski et al., 2019), PIQA (Bisk et al., 2019), SIQA (Sap et al., 2019) and TriviaQA (Joshi et al., 2017). Figure S12: Multiple token training with 7B models doesn't improve performance on choice tasks. This figure shows the evolution of average accuracy of some standard NLP benchmarks (ARC Challenge COPA Hellaswag MMLU Natural Questions PIQA SIQA and TriviaQA. For the 7B models trained on 200B tokens of language data, the 2 future token model has the same performance as the baseline and the 4 future token model regresses a bit. Larger model sizes might be necessary to see improvements on these tasks. <details> <summary>Image 12 Details</summary>  ### Visual Description ## Line Charts: Performance Comparison on Various Tasks ### Overview The image presents a series of line charts comparing the performance of a model across different tasks ("arc_challenge", "copa", "hellaswag", "nq", "piqa", "siqa", and "tqa") with varying configurations denoted by 'n' (1, 2, and 4). The charts display the 'value' (likely representing a performance metric) against the 'global_step' (training progress). ### Components/Axes * **X-axis:** 'global_step', ranging from approximately 0 to 20000, with a marked value at 10000. * **Y-axis:** 'value', with varying ranges depending on the task: * arc_challenge: 25 to 35 * copa: 70 to 80 * hellaswag: 40 to 60 * nq: 5 to 15 * piqa: 65 to 75 * siqa: 42 to 46 * tqa: 10 to 40 * **Legend (bottom-right):** * Solid Red Line: n = 1 * Dashed Black Line: n = 2 * Dotted Teal Line: n = 4 ### Detailed Analysis **1. arc_challenge:** * **n = 1 (Solid Red):** Starts at approximately 25, increases to around 37 by global_step 10000, then plateaus and slightly decreases to approximately 36 by global_step 20000. * **n = 2 (Dashed Black):** Starts at approximately 25, increases to around 35 by global_step 10000, then plateaus and slightly decreases to approximately 34 by global_step 20000. * **n = 4 (Dotted Teal):** Starts at approximately 25, increases to around 34 by global_step 10000, then plateaus and slightly increases to approximately 37 by global_step 20000. **2. copa:** * **n = 1 (Solid Red):** Starts at approximately 70, increases to around 82 by global_step 10000, then fluctuates and ends at approximately 79 by global_step 20000. * **n = 2 (Dashed Black):** Starts at approximately 70, increases to around 78 by global_step 10000, then fluctuates and ends at approximately 77 by global_step 20000. * **n = 4 (Dotted Teal):** Starts at approximately 70, increases to around 78 by global_step 10000, then fluctuates and ends at approximately 76 by global_step 20000. **3. hellaswag:** * **n = 1 (Solid Red):** Starts at approximately 40, increases to around 62 by global_step 20000. * **n = 2 (Dashed Black):** Starts at approximately 40, increases to around 61 by global_step 20000. * **n = 4 (Dotted Teal):** Starts at approximately 40, increases to around 60 by global_step 20000. **4. nq:** * **n = 1 (Solid Red):** Starts at approximately 2, increases to around 15 by global_step 20000. * **n = 2 (Dashed Black):** Starts at approximately 2, increases to around 14 by global_step 20000. * **n = 4 (Dotted Teal):** Starts at approximately 2, increases to around 13 by global_step 20000. **5. piqa:** * **n = 1 (Solid Red):** Starts at approximately 67, increases to around 76 by global_step 20000. * **n = 2 (Dashed Black):** Starts at approximately 67, increases to around 75 by global_step 20000. * **n = 4 (Dotted Teal):** Starts at approximately 67, increases to around 76 by global_step 20000. **6. siqa:** * **n = 1 (Solid Red):** Starts at approximately 42, increases to around 47 by global_step 20000. * **n = 2 (Dashed Black):** Starts at approximately 42, increases to around 46 by global_step 20000. * **n = 4 (Dotted Teal):** Starts at approximately 42, increases to around 47 by global_step 20000. **7. tqa:** * **n = 1 (Solid Red):** Starts at approximately 10, increases to around 40 by global_step 20000. * **n = 2 (Dashed Black):** Starts at approximately 10, increases to around 39 by global_step 20000. * **n = 4 (Dotted Teal):** Starts at approximately 10, increases to around 38 by global_step 20000. ### Key Observations * Across all tasks, the 'value' generally increases with the 'global_step', indicating learning or improvement over time. * The performance differences between n=1, n=2, and n=4 are task-dependent. In some tasks (e.g., 'arc_challenge', 'copa'), the performance fluctuates after a certain point. * The 'copa' task shows the most fluctuation in performance after the initial increase. ### Interpretation The charts illustrate the impact of different configurations ('n') on the performance of a model across various tasks. The 'global_step' represents the training progress, and the 'value' likely represents a performance metric such as accuracy or score. The trends suggest that the model generally improves with training, but the optimal configuration ('n') may vary depending on the specific task. The fluctuations in performance for some tasks after a certain point could indicate overfitting or the need for further optimization. The parameter 'n' could represent the number of layers in a neural network, or the number of attention heads. </details> ## H. Additional results on abstractive text summarization In this section, we report comprehensive evaluation results on summarization tasks for the 7B parameter models trained on 200B and 500B tokens of natural language from Section 3.7. Table S8: Comprehensive evaluation on abstractive text summarization. ROUGE-n (n-gram overlap) and ROUGE-L (longest common subsequence overlap) F 1 scores for 7B models trained on 200B and 500B tokens of natural language, respectively. The last three columns correspond to models trained on 500B tokens, the previous three to models trained on 200B tokens. Shown are numbers of the n = 1 baseline and the absolute difference of n = 2 and n = 4 models trained on the same number of tokens. Summary-level ROUGE-L ('ROUGE-Lsum') is reported where it differs from ROUGE-L. Model checkpoints with maximal validation ROUGE-L F 1 are selected separately for each model dataset and model type and reported in the first row corresponding to each dataset. Boldface for numbers within 0.05 difference to the best one for each dataset size separately. | Task | Metric | Baseline 200B | ∆ n =2 | ∆ n =4 | Baseline 500B | ∆ n =2 | ∆ n =4 | |----------------------------------------|------------------|-----------------|-------------|-------------|-----------------|-------------|-------------| | CNN/Dailymail (Nallapati et al., 2016) | evaluation epoch | 2 | 2 | 2 | 2 | 2 | 2 | | | ROUGE-1 | 42.88 | +0.74 | +0.74 | 43.77 | +0.55 | +0.50 | | | ROUGE-2 | 19.56 | +0.52 | +0.53 | 20.34 | +0.52 | +0.34 | | | ROUGE-3 | 11.11 | +0.39 | +0.35 | 11.69 | +0.36 | +0.19 | | | ROUGE-L | 29.72 | +0.66 | +0.49 | 30.51 | +0.48 | +0.37 | | | ROUGE-Lsum | 40.18 | +0.72 | +0.68 | 41.02 | +0.56 | +0.52 | | Multi-News (Fabbri et al., 2019) | evaluation epoch | 1 | 3 | 3 | 2 | 3 | 2 | | Multi-News (Fabbri et al., 2019) | ROUGE-1 | 44.48 | +1.70 | +1.72 | 45.87 | +1.05 | +0.69 | | Multi-News (Fabbri et al., 2019) | ROUGE-2 | 16.88 | +0.44 | +0.70 | 17.56 | +0.42 | +0.40 | | Multi-News (Fabbri et al., 2019) | ROUGE-3 | 9.63 | -0.06 | +0.17 | 9.91 | +0.22 | +0.18 | | Multi-News (Fabbri et al., 2019) | ROUGE-L | 23.82 | +0.17 | +0.40 | 24.22 | +0.20 | +0.26 | | OrangeSum (Eddine et al., 2021) | evaluation epoch | 2 | 2 | 3 | 2 | 1 | 3 | | OrangeSum (Eddine et al., 2021) | ROUGE-1 | 32.95 | +0.41 | +0.35 | 33.37 | +0.32 | +0.78 | | OrangeSum (Eddine et al., 2021) | ROUGE-2 | 13.90 | +0.31 | +0.36 | 14.22 | +0.25 | +0.53 | | OrangeSum (Eddine et al., 2021) | ROUGE-3 | 8.01 | +0.19 | +0.21 | 8.12 | +0.22 | +0.48 | | OrangeSum (Eddine et al., 2021) | ROUGE-L | 23.62 | +0.36 | +0.51 | 23.91 | +0.23 | +0.66 | | pn-summary (Farahani et al., 2021) | evaluation epoch | 1 | 1 | 1 | 1 | 2 | 3 | | pn-summary (Farahani et al., 2021) | ROUGE-1 | 1.03 | +0.02 | 0.00 | 0.92 | +0.09 | +0.05 | | pn-summary (Farahani et al., 2021) | ROUGE-2 | 0.13 | +0.02 | +0.03 | 0.15 | 0.00 | 0.00 | | pn-summary (Farahani et al., 2021) | ROUGE-3 | 0.02 | 0.00 | +0.02 | 0.02 | 0.00 | +0.02 | | pn-summary (Farahani et al., 2021) | ROUGE-L | 1.02 | +0.03 | +0.01 | 0.91 | +0.09 | +0.05 | | SAMSum (Gliwa et al., 2019) | evaluation epoch | 3 | 3 | 3 | 3 | 3 | 3 | | SAMSum (Gliwa et al., 2019) | ROUGE-1 | 51.39 | +0.70 | +0.63 | 52.54 | -0.24 | +0.69 | | SAMSum (Gliwa et al., 2019) | ROUGE-2 | 26.46 | +0.76 | +0.30 | 27.74 | -0.20 | +0.82 | | SAMSum (Gliwa et al., 2019) | ROUGE-3 | 16.40 | +0.91 | +0.28 | 17.56 | -0.30 | +0.71 | | SAMSum (Gliwa et al., 2019) | ROUGE-L | 42.59 | +0.90 | +0.51 | 43.92 | -0.10 | +0.63 | | ThaiSum (Chumpolsathien, 2020) | evaluation epoch | 2 | 3 | 3 | 3 | 3 | 3 | | ThaiSum (Chumpolsathien, 2020) | ROUGE-1 | 45.08 | +0.63 | +1.12 | 45.48 | +0.77 | +0.91 | | ThaiSum (Chumpolsathien, 2020) | ROUGE-2 | 27.85 | +0.30 | +0.73 | 28.07 | +0.74 | +0.64 | | ThaiSum (Chumpolsathien, 2020) | ROUGE-3 | 15.73 | +0.04 | +0.43 | 15.82 | +0.50 | +0.30 | | ThaiSum (Chumpolsathien, 2020) | ROUGE-L | 44.92 | +0.64 | +1.12 | 45.31 | +0.76 | +0.89 | | WikiSummary (Farahani, 2020) | evaluation epoch | 3 | 3 | 3 | 3 | 3 | 3 | | WikiSummary (Farahani, 2020) | ROUGE-1 | 10.16 | +0.67 | -0.23 | 12.80 | -0.17 | -0.99 | | WikiSummary (Farahani, 2020) | ROUGE-2 | 4.46 | -0.03 | -0.09 | 6.17 | -0.11 | -0.69 | | WikiSummary (Farahani, 2020) | ROUGE-3 | 1.31 | +0.21 | +0.13 | 1.98 | -0.08 | -0.33 | | WikiSummary (Farahani, 2020) | ROUGE-L | 10.11 | +0.65 | -0.28 | 12.69 | -0.17 | -0.99 | | XSum (Narayan et al., 2018) | evaluation epoch | 2 | 2 | 3 | 2 | 2 | 3 | | XSum (Narayan et al., 2018) | ROUGE-1 | 42.16 | +0.71 | +1.07 | 43.42 | +0.78 | +0.67 | | XSum (Narayan et al., 2018) | ROUGE-2 | 19.19 | +0.54 | +0.55 | 20.32 | +0.68 | +0.34 | | XSum (Narayan et al., 2018) | ROUGE-3 ROUGE-L | 10.43 34.03 | +0.38 +0.67 | +0.28 +0.92 | 11.23 35.18 | +0.48 +0.79 | +0.20 +0.63 | Table S9: Performance on abstractive text summarization. ROUGE-L (longest common subsequence overlap) F 1 score for 7B models trained on 200B and 500B tokens of natural language. We finetune the respective models on each task's training data separately for a given number of epochs and select the checkpoints with maximal ROUGE-L F 1 on the validation dataset. The second and fifth column report the numbers for a next-token prediction model, while the third, fourth, sixth and seventh one report the absolute improvements for 2-token and 4-token prediction models trained on the same amount of data, respectively. Boldface for numbers within 0.05 difference to the best one for each dataset size separately. | Dataset | Baseline 200B | ∆ n =2 | ∆ n =4 | Baseline 500B | ∆ n =2 | ∆ n =4 | |---------------|-----------------|----------|----------|-----------------|----------|----------| | CNN/Dailymail | 29.72 | 0.66 | 0.49 | 30.51 | 0.48 | 0.37 | | Multi-News | 23.82 | 0.17 | 0.4 | 24.22 | 0.2 | 0.26 | | OrangeSum | 23.62 | 0.36 | 0.51 | 23.91 | 0.23 | 0.66 | | pn-summary | 1.02 | 0.03 | 0.01 | 0.91 | 0.09 | 0.05 | | SAMSum | 42.59 | 0.9 | 0.51 | 43.92 | -0.1 | 0.63 | | ThaiSum | 44.92 | 0.64 | 1.12 | 45.31 | 0.76 | 0.89 | | WikiSummary | 10.11 | 0.65 | -0.28 | 12.69 | -0.17 | -0.99 | | XSum | 34.03 | 0.67 | 0.92 | 35.18 | 0.79 | 0.63 | | Average | 26.23 | 0.51 | 0.46 | 27.08 | 0.28 | 0.31 | Table S10: Summary statistics for abstractive text summarization evaluations. Reported are averages for ROUGE-n and ROUGE-L metrics across all datasets from Table S8, separately for precision, recall and F 1 score. Both 2-token and 4-token prediction models outperform the next-token prediction baseline. Trained on 500B tokens, 4-token prediction models appear better at recall metrics while 2-token prediction models appear better at precision metrics. Model checkpoints are selected as described in Table S8. Boldface for numbers within 0.05 difference to the best one for each dataset size separately. | Metric | Aspect | Baseline 200B | ∆ n =2 | ∆ n =4 | Baseline 500B | ∆ n =2 | ∆ n =4 | |-------------|-----------|-----------------|----------|----------|-----------------|----------|----------| | ROUGE-1 | F 1 | 33.77 | 0.7 | 0.68 | 34.77 | 0.39 | 0.41 | | ROUGE-1 | precision | 35.76 | 0.88 | 0.83 | 37.03 | 0.42 | -0.04 | | ROUGE-1 | recall | 34.37 | 0.45 | 0.45 | 35.14 | 0.35 | 0.68 | | ROUGE-2 | F 1 | 16.06 | 0.36 | 0.39 | 16.82 | 0.29 | 0.3 | | ROUGE-2 | precision | 16.97 | 0.4 | 0.43 | 17.91 | 0.29 | 0.03 | | ROUGE-2 | recall | 16.34 | 0.28 | 0.35 | 16.99 | 0.32 | 0.48 | | ROUGE-3 | F 1 | 9.08 | 0.26 | 0.23 | 9.54 | 0.18 | 0.22 | | ROUGE-3 | precision | 9.59 | 0.29 | 0.28 | 10.17 | 0.18 | 0.05 | | ROUGE-3 | recall | 9.26 | 0.21 | 0.2 | 9.65 | 0.21 | 0.35 | | ROUGE-L | F 1 | 26.23 | 0.51 | 0.46 | 27.08 | 0.28 | 0.31 | | ROUGE-L | precision | 27.79 | 0.62 | 0.55 | 28.85 | 0.28 | -0.09 | | ROUGE-L | recall | 26.71 | 0.37 | 0.32 | 27.4 | 0.28 | 0.57 | | ROUGE-L sum | F 1 | 27.53 | 0.52 | 0.48 | 28.4 | 0.29 | 0.33 | | ROUGE-L sum | precision | 29.07 | 0.64 | 0.58 | 30.15 | 0.29 | -0.08 | | ROUGE-L sum | recall | 28.13 | 0.35 | 0.33 | 28.81 | 0.29 | 0.6 | ## I. Additional results on mathematical reasoning in natural language Figure S13: Performance on the mathematical reasoning benchmark GSM8K (Cobbe et al., 2021). We evaluate pretrained next-token and multi-token prediction models trained on 200B and 500B tokens of natural language in 8-shot mode using nucleus sampling (Holtzman et al., 2020) with probability mass 0.95 and various sampling temperatures. Reported are the frequencies of the correct final answer to appear among k samples, for k = 1 , 10 , 100 , estimated from 200 samples like in code generation benchmarks (Chen et al., 2021). After 200B tokens, the 2-token prediction model has a clear advantage over the next-token baseline but the order reverses after 500B tokens. The 4-token prediction model is worse throughout. We interpret this similarly to the findings in Section 4.1: the follow-your-nose chains-of-thought required for GSM8K may be difficult to learn from a limited amount of data, attesting to the data efficiency of multi-token prediction training. Once the correct circuits for correct autoregressive chains-of-thought in this domain have formed, however, multi-token prediction comes at a cost. <details> <summary>Image 13 Details</summary>  ### Visual Description ## Chart: Pass Rate vs. Temperature for Different Token Sizes ### Overview The image presents six line charts arranged in a 3x2 grid. Each chart plots the "pass" rate (pass@1, pass@10, pass@100) as a function of "Temperature" for two different token sizes (200B and 500B). Each chart contains three lines representing different values of 'n' (n=1, n=2, n=4). The x-axis represents temperature, and the y-axis represents the pass rate in percentage. ### Components/Axes * **Titles:** * Top-left chart title: "200B tokens" * Top-right chart title: "500B tokens" * **X-axis:** * Label: "Temperature" * Scale: 0.2, 0.4, 0.6, 0.8, 1.0, 1.2, 1.4 * **Y-axis (Left Column):** * Top chart label: "pass@1 (%)" * Middle chart label: "pass@10 (%)" * Bottom chart label: "pass@100 (%)" * Scale (pass@1): 2, 3 * Scale (pass@10): 15, 20 * Scale (pass@100): 40, 50, 60 * **Y-axis (Right Column):** * Top chart label: "pass@1 (%)" * Middle chart label: "pass@10 (%)" * Bottom chart label: "pass@100 (%)" * Scale (pass@1): 2, 4, 6, 8 * Scale (pass@10): 20, 30 * Scale (pass@100): 40, 50, 60 * **Legend (Located at the bottom-right of the image):** * n = 1 (coral/light orange) * n = 2 (dark blue) * n = 4 (light green) ### Detailed Analysis **200B tokens** * **pass@1 (%)** * n = 1 (coral): Starts at approximately 3.6%, decreases to approximately 1.6% as temperature increases. * n = 2 (dark blue): Starts at approximately 3.2%, decreases to approximately 1.5% as temperature increases. * n = 4 (light green): Starts at approximately 3.0%, decreases to approximately 1.2% as temperature increases. * **pass@10 (%)** * n = 1 (coral): Starts at approximately 16%, increases to approximately 21% at temperature 0.6, then decreases to approximately 12% as temperature increases. * n = 2 (dark blue): Starts at approximately 16%, increases to approximately 22% at temperature 0.6, then decreases to approximately 12% as temperature increases. * n = 4 (light green): Starts at approximately 12%, increases to approximately 18% at temperature 0.6, then decreases to approximately 10% as temperature increases. * **pass@100 (%)** * n = 1 (coral): Starts at approximately 40%, increases to approximately 60% at temperature 0.8, then decreases to approximately 35% as temperature increases. * n = 2 (dark blue): Starts at approximately 40%, increases to approximately 62% at temperature 0.8, then decreases to approximately 35% as temperature increases. * n = 4 (light green): Starts at approximately 38%, increases to approximately 58% at temperature 0.8, then decreases to approximately 32% as temperature increases. **500B tokens** * **pass@1 (%)** * n = 1 (coral): Starts at approximately 8%, decreases to approximately 2% as temperature increases. * n = 2 (dark blue): Starts at approximately 7%, decreases to approximately 2% as temperature increases. * n = 4 (light green): Starts at approximately 5%, decreases to approximately 1.8% as temperature increases. * **pass@10 (%)** * n = 1 (coral): Starts at approximately 27%, increases to approximately 32% at temperature 0.6, then decreases to approximately 12% as temperature increases. * n = 2 (dark blue): Starts at approximately 25%, increases to approximately 32% at temperature 0.6, then decreases to approximately 12% as temperature increases. * n = 4 (light green): Starts at approximately 18%, increases to approximately 24% at temperature 0.6, then decreases to approximately 10% as temperature increases. * **pass@100 (%)** * n = 1 (coral): Starts at approximately 50%, increases to approximately 70% at temperature 0.6, then decreases to approximately 35% as temperature increases. * n = 2 (dark blue): Starts at approximately 50%, increases to approximately 70% at temperature 0.6, then decreases to approximately 35% as temperature increases. * n = 4 (light green): Starts at approximately 48%, increases to approximately 62% at temperature 0.6, then decreases to approximately 35% as temperature increases. ### Key Observations * For pass@1, the pass rate decreases as temperature increases for all values of 'n' and both token sizes. * For pass@10 and pass@100, the pass rate initially increases with temperature, reaches a peak around temperature 0.6-0.8, and then decreases as temperature further increases. * Generally, higher values of 'n' result in lower pass rates for all temperature values. * The 500B token models generally have higher pass rates than the 200B token models. ### Interpretation The charts illustrate the relationship between temperature, the parameter 'n', and the pass rate of a language model for different token sizes. The pass rate measures the model's ability to generate correct code. * **Temperature:** The temperature parameter controls the randomness of the model's output. Lower temperatures lead to more deterministic and predictable outputs, while higher temperatures introduce more randomness. The observed trends suggest that an intermediate temperature (around 0.6-0.8) is optimal for pass@10 and pass@100, balancing exploration and exploitation in the generation process. For pass@1, lower temperatures are better. * **Parameter 'n':** The parameter 'n' likely represents the number of samples generated by the model. Lower values of 'n' (e.g., n=1) generally lead to higher pass rates, suggesting that focusing on generating a smaller number of high-quality samples is more effective than generating a larger number of samples with potentially lower quality. * **Token Size:** The 500B token models consistently outperform the 200B token models, indicating that models trained on larger datasets have a better ability to generate correct code. </details> ## J. Additional results on induction learning Figure S14: Induction capability of n -token prediction models trained on higher-quality data. Shown is accuracy on the second token of two token names that have already been mentioned previously. Training on a 9:1 mix of a books dataset and the children storiy dataset, we observe that induction capability forms significantly earlier in training (not shown here) and to a higher degree. We believe that this is explained both because our evaluation dataset no longer contains out-of-distribution tokens (Section 4.1) and because the higher-quality data contained in the books dataset makes induction necessary earlier on (especially for small models, cf. Singh et al. (2023)). In particular, by enforcing the formation of induction capability in the model by means of the dataset - instead of the loss - the advantage of 2-token prediction models on this task disappears except for the smallest models: feature learning converts the task into a pure next-token prediction task. <details> <summary>Image 14 Details</summary>  ### Visual Description ## Line Chart: Induction Success vs. Parameters ### Overview The image is a line chart comparing the induction success of two models (n=1 baseline and n=2 ours) across varying parameter sizes (in millions). The x-axis represents the number of parameters, and the y-axis represents the induction success rate. ### Components/Axes * **X-axis:** Parameters (M). Logarithmic scale with markers at 1, 3, 10, 30, 100, 300, and 1000. * **Y-axis:** Induction success. Linear scale with markers at 0.850, 0.875, 0.900, 0.925, 0.950, 0.975, and 1.000. * **Legend:** Located in the bottom-right corner. * **Orange Line:** n=1 (baseline) * **Dark Blue Dashed Line:** n=2 (ours) ### Detailed Analysis * **n=1 (baseline) - Orange Line:** * Trend: The line starts at approximately 0.85 for 1M parameters, rises sharply to approximately 0.99 for 10M parameters, and then plateaus around 0.99-1.00 for higher parameter values. * Data Points: * 1M Parameters: ~0.85 * 3M Parameters: ~0.925 * 10M Parameters: ~0.99 * 30M Parameters: ~0.99 * 100M Parameters: ~1.00 * 300M Parameters: ~0.99 * 1000M Parameters: ~1.00 * **n=2 (ours) - Dark Blue Dashed Line:** * Trend: The line starts at approximately 0.875 for 1M parameters, rises to approximately 0.975 for 3M parameters, and then plateaus around 0.99-1.00 for higher parameter values. * Data Points: * 1M Parameters: ~0.875 * 3M Parameters: ~0.975 * 10M Parameters: ~0.99 * 30M Parameters: ~0.99 * 100M Parameters: ~1.00 * 300M Parameters: ~0.99 * 1000M Parameters: ~1.00 ### Key Observations * For lower parameter values (1M to 10M), the n=2 model (ours) shows a higher induction success rate compared to the n=1 (baseline) model. * As the number of parameters increases beyond 10M, both models converge to a similar induction success rate, plateauing around 0.99-1.00. * The n=1 model experiences a more significant jump in induction success between 1M and 10M parameters compared to the n=2 model. ### Interpretation The chart suggests that the "n=2 (ours)" model achieves better induction success with fewer parameters compared to the "n=1 (baseline)" model, especially in the range of 1M to 10M parameters. This indicates that the "n=2" model is more efficient in utilizing parameters to achieve a higher success rate. However, beyond 10M parameters, the performance of both models becomes comparable, suggesting a saturation point where increasing parameters does not significantly improve induction success. The "n=1" model requires more parameters to reach a similar level of induction success as the "n=2" model. </details> ## K. Additional results on algorithmic reasoning We investigate the following computation-sharing hypothesis for explaining the efficacy of multi-token prediction as training loss. The prediction difficulty of different tokens in natural text varies greatly. Some tokens may be the continuations of partial words that are uniquely determined from their preceding context without any effort, while others may require to predict theorem names in difficult mathematical proofs or the correct answer to an exam question. Language models with residual connections have been shown to refine their output token distribution with each successive layer, and can be trained with early exit strategies that spend variable amounts of computational resources per token position. Multi-token prediction losses explicitly encourage information-sharing between adjacent token positions and can thus be viewed as a method to learn allocating computational resources in language models more efficiently to the tokens that benefit most of it. To check the truth of this hypothesis, we augment the polynomial arithmetic task from Section 4.2 with a varying number of pause tokens (Goyal et al., 2023) inserted between the question and a token that denotes the beginning of the answer. Pause tokens introduce additional computational resources that can be expended for computations that are expected to be useful later on in the sequence, in other words: to start thinking about the answer. According to the computation-sharing hypothesis , multi-token prediction models learn information-sharing and thus computation-sharing between token positions more easily, and may be better at making use of these additional computational resources than next-token prediction models are. In Figure S15, we show the evaluation results on the polynomial arithmetic task with a fixed number of pause tokens inserted both at training and evaluation time. Multi-token prediction models likewise outperform next-token prediction models on these task variants across task difficulties and model sizes. However, we do not see strong evidence of a widening or shrinking of this gap i.e. we cannot conclude from these experiments on the veracity of the computation-sharing hypothesis. In Table S11, we report results from another experiment in the same spirit: by adding spaces and newlines to HumanEval and MBPP prompts, we add 'pause tokens' in a somewhat natural way. According to these results, multi-token prediction models have a slight advantage at using this additionally provided compute, but the effect is marginal. Figure S15: Accuracy on a polynomial arithmetic task with varying number of operations per expression and pause tokens. We train and evaluate models on the polynomial arithmetic task described in Section 4.2, modified by the addition of pause tokens (Goyal et al., 2023): between the question and the equality sign that indicates the beginning of the answer, we add a constant number of pause tokens both in training and evaluation. For both a variant with five and with ten pause tokens, respectively, we observe comparable improvements from using multi-token prediction to the ones obtained in the case without pause tokens (Figure 8). <details> <summary>Image 15 Details</summary>  ### Visual Description ## Chart: Accuracy vs. Number of Operations ### Overview The image presents two line charts comparing the accuracy of a system against the number of operations performed. The charts are differentiated by the number of pause tokens used: 5 in the left chart (a) and 10 in the right chart (b). Each chart displays three data series, representing different values of 'n' (n=1, n=2, n=4). The x-axis represents the number of operations, divided into "in-domain" (1-5 operations) and "out-of-domain" (6-10 operations) categories. The y-axis represents accuracy in percentage. ### Components/Axes * **Title (Left Chart):** (a) 5 pause tokens * **Title (Right Chart):** (b) 10 pause tokens * **X-axis Title:** # operations * **X-axis Categories:** in-domain (operations 1-5), out-of-domain (operations 6-10) * **X-axis Markers:** 1, 2, 3, 4, 5, 6, 7, 8, 9, 10 * **Y-axis Title:** Accuracy (%) * **Y-axis Markers:** 0, 20, 40, 60, 80, 100 * **Legend (Top-Right of each chart):** * n=1 (coral color, solid line with 'x' markers) * n=2 (dark blue-gray color, dashed line with 'x' markers) * n=4 (light green color, dotted line with 'x' markers) ### Detailed Analysis **Left Chart (5 pause tokens):** * **n=1 (coral):** The accuracy decreases sharply from approximately 85% at 1 operation to about 10% at 6 operations. It then remains near 0% for operations 7-10. * (1, 85%), (2, 70%), (3, 40%), (4, 30%), (5, 15%), (6, 10%), (7, 2%), (8, 1%), (9, 0%), (10, 0%) * **n=2 (dark blue-gray):** The accuracy decreases from approximately 95% at 1 operation to about 5% at 6 operations. It then remains near 0% for operations 7-10. * (1, 95%), (2, 75%), (3, 60%), (4, 35%), (5, 20%), (6, 5%), (7, 2%), (8, 1%), (9, 0%), (10, 0%) * **n=4 (light green):** The accuracy decreases from approximately 90% at 1 operation to about 5% at 6 operations. It then remains near 0% for operations 7-10. * (1, 90%), (2, 80%), (3, 55%), (4, 40%), (5, 25%), (6, 5%), (7, 3%), (8, 2%), (9, 1%), (10, 1%) **Right Chart (10 pause tokens):** * **n=1 (coral):** The accuracy decreases sharply from approximately 85% at 1 operation to about 5% at 6 operations. It then remains near 0% for operations 7-10. * (1, 85%), (2, 60%), (3, 35%), (4, 25%), (5, 15%), (6, 5%), (7, 1%), (8, 0%), (9, 0%), (10, 0%) * **n=2 (dark blue-gray):** The accuracy decreases from approximately 95% at 1 operation to about 5% at 6 operations. It then remains near 0% for operations 7-10. * (1, 95%), (2, 75%), (3, 50%), (4, 30%), (5, 20%), (6, 5%), (7, 2%), (8, 1%), (9, 0%), (10, 0%) * **n=4 (light green):** The accuracy decreases from approximately 80% at 1 operation to about 5% at 6 operations. It then remains near 0% for operations 7-10. * (1, 80%), (2, 70%), (3, 55%), (4, 40%), (5, 20%), (6, 5%), (7, 3%), (8, 2%), (9, 1%), (10, 1%) ### Key Observations * All three data series (n=1, n=2, n=4) in both charts exhibit a similar trend: a rapid decrease in accuracy as the number of operations increases. * The accuracy drops significantly when transitioning from "in-domain" to "out-of-domain" operations (after 5 operations). * The initial accuracy (at 1 operation) is highest for n=2 in both charts. * The number of pause tokens (5 vs. 10) does not appear to drastically alter the overall trend or the relative performance of the different 'n' values. ### Interpretation The charts suggest that the system's accuracy is highly sensitive to the number of operations performed, particularly when moving to "out-of-domain" scenarios. The rapid decline in accuracy indicates that the system struggles to maintain performance as the complexity or novelty of the operations increases. The 'n' value seems to have a minor impact on the initial accuracy, but the overall trend of decreasing accuracy with more operations is consistent across all 'n' values. The number of pause tokens (5 vs. 10) has a minimal effect on the overall performance, suggesting that this parameter is not a primary driver of accuracy in this context. The data implies that the system needs improvement in handling a larger number of operations and generalizing to out-of-domain scenarios. </details> Table S11: Utilization of additional whitespace tokens in code benchmarks. | Task | Whitespace | n = 1 | n = 4 | |-----------------------------------------------------|----------------------------------------------------------------------------|-------------------------------------|-------------------------------------| | APPS/Intro APPS/Intro HumanEval HumanEval MBPP MBPP | spaces + newline newline spaces + newline newline spaces + newline newline | +0.21 +0.79 -0.72 -0.26 -0.10 +0.03 | +0.34 +0.69 -0.16 +0.10 -0.06 -0.08 | | Average | | -0.01 | +0.14 | Figure S16: Accuracy on a polynomial arithmetic task for two model sizes. We train and evaluate models with 30M and 100M parameters on the polynomial arithmetic task described in Section 4.2. Tripling the model size has a smaller effect on performance than replacing next-token prediction loss by multi-token prediction. Shown are two independent runs per configuration and their means, the 100M parameter models being identical to the ones in Figure 8. <details> <summary>Image 16 Details</summary>  ### Visual Description ## Line Chart: Accuracy vs. Number of Operations ### Overview The image is a line chart that plots the accuracy (in percentage) against the number of operations. The chart compares the performance of two models (30M and 100M) with different configurations (n=1, n=2, and n=4) across a range of operations, distinguishing between "in-domain" and "out-of-domain" operations. ### Components/Axes * **X-axis:** "# operations" ranging from 1 to 10. The x-axis is divided into two regions: "in-domain" (operations 1-5) and "out-of-domain" (operations 6-10). * **Y-axis:** "Accuracy (%)" ranging from 0 to 100. * **Legend:** Located in the top-right corner, the legend identifies the different lines based on model size (30M, 100M) and configuration (n=1, n=2, n=4). * 30M, n=1: Dotted red-orange line * 30M, n=2: Dotted dark blue line * 30M, n=4: Dotted teal line * 100M, n=1: Solid red-orange line with "x" markers * 100M, n=2: Solid dark blue line with "x" markers * 100M, n=4: Solid teal line with "x" markers * A vertical dotted line separates the "in-domain" and "out-of-domain" regions at x=5.5. ### Detailed Analysis * **30M, n=1 (Dotted red-orange line):** * Trend: Accuracy decreases rapidly. * Values: Approximately 85% at operation 1, 60% at operation 2, 38% at operation 3, 25% at operation 4, 15% at operation 5, 5% at operation 6, and near 0% from operation 7 onwards. * **30M, n=2 (Dotted dark blue line):** * Trend: Accuracy decreases rapidly. * Values: Approximately 90% at operation 1, 65% at operation 2, 45% at operation 3, 30% at operation 4, 20% at operation 5, 10% at operation 6, and near 0% from operation 7 onwards. * **30M, n=4 (Dotted teal line):** * Trend: Accuracy decreases rapidly. * Values: Approximately 88% at operation 1, 68% at operation 2, 48% at operation 3, 33% at operation 4, 22% at operation 5, 12% at operation 6, and near 0% from operation 7 onwards. * **100M, n=1 (Solid red-orange line with "x" markers):** * Trend: Accuracy decreases rapidly. * Values: Approximately 87% at operation 1, 62% at operation 2, 40% at operation 3, 27% at operation 4, 17% at operation 5, 7% at operation 6, and near 0% from operation 7 onwards. * **100M, n=2 (Solid dark blue line with "x" markers):** * Trend: Accuracy decreases rapidly. * Values: Approximately 92% at operation 1, 70% at operation 2, 50% at operation 3, 35% at operation 4, 23% at operation 5, 13% at operation 6, and near 0% from operation 7 onwards. * **100M, n=4 (Solid teal line with "x" markers):** * Trend: Accuracy decreases rapidly. * Values: Approximately 100% at operation 1, 78% at operation 2, 58% at operation 3, 40% at operation 4, 25% at operation 5, 15% at operation 6, and near 0% from operation 7 onwards. ### Key Observations * All models and configurations exhibit a rapid decline in accuracy as the number of operations increases. * The 100M models generally outperform the 30M models across all configurations. * Increasing the 'n' value (likely representing the number of parallel processing units or similar) generally improves accuracy for both 30M and 100M models. * The transition from "in-domain" to "out-of-domain" operations (after operation 5) does not appear to cause a significant change in the rate of accuracy decline. The decline is consistent across both domains. * The 100M, n=4 configuration starts with the highest accuracy (100% at operation 1) and maintains a higher accuracy throughout the initial operations. ### Interpretation The data suggests that the accuracy of the models is highly sensitive to the number of operations performed. The larger model (100M) and higher configuration (n=4) consistently achieve better accuracy, indicating that model size and configuration play a crucial role in performance. The rapid decline in accuracy with increasing operations suggests that the models may struggle with longer or more complex sequences of operations. The "out-of-domain" operations do not appear to have a disproportionate impact on accuracy compared to "in-domain" operations, suggesting that the models' performance degradation is primarily driven by the number of operations rather than the specific nature of those operations. The 100M n=4 model is the best performing model. </details> Table S12: Optimal temperatures for all numbers in table 1 | Training data | Vocabulary | n | MBPP | MBPP | MBPP | HumanEval | HumanEval | HumanEval | APPS/Intro | APPS/Intro | APPS/Intro | |----------------------|--------------|-----|--------|--------|--------|-------------|-------------|-------------|--------------|--------------|--------------| | Training data | Vocabulary | n | @1 | @10 | @100 | @1 | @10 | @100 | @1 | @10 | @100 | | | | 1 | 0.2 | 0.8 | 0.8 | 0.1 | 0.8 | 0.8 | 0.8 | 0.8 | 0.8 | | | | 8 | 0.1 | 0.8 | 0.8 | 0.1 | 0.8 | 0.8 | 0.4 | 0.4 | 0.4 | | | | 16 | 0.1 | 0.8 | 0.8 | 0.1 | 0.8 | 0.8 | 0.4 | 0.4 | 0.4 | | | | 32 | 0.1 | 0.4 | 0.8 | 0.1 | 0.4 | 0.8 | 0.1 | 0.4 | 0.4 | | | | 1 | 0.1 | 0.8 | 0.8 | 0.1 | 0.8 | 0.8 | 0.1 | 0.4 | 0.8 | | | | 2 | 0.1 | 0.8 | 0.8 | 0.2 | 0.8 | 0.8 | 0.4 | 0.4 | 0.8 | | | | 4 | 0.1 | 0.8 | 0.8 | 0.1 | 0.8 | 0.8 | 0.2 | 0.8 | 0.8 | | | | 6 | 0.1 | 0.8 | 0.8 | 0.2 | 0.8 | 0.8 | 0.4 | 0.4 | 0.8 | | | | 8 | 0.1 | 0.8 | 0.8 | 0.1 | 0.8 | 0.8 | 0.2 | 0.4 | 0.8 | | 1T tokens (4 epochs) | | 1 | 0.1 | 0.8 | 0.8 | 0.1 | 0.8 | 0.8 | 0.1 | 0.4 | 0.8 | | 1T tokens (4 epochs) | | 4 | 0.1 | 0.8 | 0.8 | 0.2 | 0.8 | 0.8 | 0.4 | 0.8 | 0.8 | ## L. Additional intuitions on multi-token prediction ## L.1. Comparison to scheduled sampling In Section 5.2, we argued that multi-token prediction reduces the distribution mismatch between teacher-forced training and autoregressive evaluation of language models. Scheduled sampling (Bengio et al., 2015) is a curriculum learning method that likewise aims to bridge this gap in sequence prediction tasks by gradually replacing more and more input tokens with model-generated ones. While effective in areas such as time series forecasting, scheduled sampling is, in our opinion, inapplicable to language modelling due to the discrete nature of text. Replacing ground truth input sequences by interleavings of ground truth and model-generated tokens frequently results in ungrammatical, factually wrong or otherwise incoherent text, which should be avoided at all cost. Moreover, unlike multi-token prediction, the technique originally developed for recurrent neural networks cannot easily be adapted for parallel training setups like the ones of transformer models. ## L.2. Information-theoretic argument We give details on the information-theoretic terms appearing in the decomposition in Section 5.2 and derive a relative version that similarly allows to decompose multi-token prediction losses. As in Section 5.2, denote by X the next token and by Y the second-next one, and omit conditioning on the preceding context C for ease of notation. In Section 5.2, we decomposed H ( X ) + H ( Y ) -the quantity of interest for 2-token prediction models-as follows: <!-- formula-not-decoded --> Let us explain each of the terms. The entropy terms denote the uncertainty contained in the ground-truth random variables X and Y . 2 The term H ( Y | X ) is a classical next-token entropy for the prefix ( C, X ) . The conditional entropy H ( X | Y ) is a more theoretical entity not modelled by causal models. It describes the uncertainty about X given the prefix C and suffix Y , and therefore captures the local variations of X that do not affect the continuation of the text Y . The mutual information I ( X ; Y ) on the other hand describes the information about Y contained in X (and vice versa) and therefore captures the variations of X which constrain the continuation of the text. However, the argument given in Section 5.2 relies on the assumption that multi-token prediction losses obey a similar decomposition as the sum of the ground-truth entropies themselves. Let us make this rigorous. Denote by p ( x, y ) the joint distribution of X and Y , by p ( x ) (short for p X ( x ) ) the marginal distribution of X and by p ( y ) the one of Y . Denote the densities of the model's predictions by q ( x, y ) , q ( x ) and q ( y ) , respectively, conditional distributions by p ( x | y ) and Kullback-Leibler divergence from q to p by D ( p ∥ q ) and cross-entropy from q to p by H ( p, q ) . Definition L.1. The conditional cross-entropy H ( p X | Y , q X | Y ) of X conditioned on Y from q to p is defined as the 2 In particular, they do not refer to model predictions. expectation under y of the cross-entropy between the distributions p X and q X conditioned on y , in formulas: <!-- formula-not-decoded --> Definition L.2. The relative mutual information I p ∥ q ( X ; Y ) of X and Y from q relative to p is defined by <!-- formula-not-decoded --> We have I p ∥ q ( X ; Y ) = H ( p X , q X ) + H ( p Y , q Y ) -H ( p, q ) , I p ∥ p ( X ; Y ) = I p ( X ; Y ) reduces to standard mutual information under the distribution p and I p ∥ q ( X ; Y ) is symmetric in X and Y but can be negative. We have the following relative version of the decomposition H ( X ) = H ( X | Y ) + I ( X ; Y ) . <!-- formula-not-decoded --> Proof. We calculate <!-- formula-not-decoded --> Symmetrizing, we get the desired relative version of H ( X ) + H ( Y ) = H ( X | Y ) + 2 I ( X ; Y ) + H ( Y | X ) : <!-- formula-not-decoded --> Setting p to be the empirical distribution of the training data, the left-hand side describes the cross-entropy loss used to train 2-token prediction models. The right-hand side gives the decomposition into a local cross-entropy term, a mutual information term with weight two and a shifted next-token cross-entropy term. We interpret this as follows: by adding the term H ( p Y , q Y ) to the loss, 2-token prediction incentivizes models to precompute features which will become useful for predicting Y in the next step and increases the weight of the relative mutual information term in the loss. What does relative mutual information actually mean? By interpreting Kullback-Leibler divergence D ( p ∥ q ) as the average number of bits needed in addition to send data from p with a code optimized for q instead of p , we see that minimizing <!-- formula-not-decoded --> means minimizing the average number of additional bits needed to send data from p with a code optimized for q that treats X and Y as independent compared to one that does not. If this number is small, q managed to exploit the mutual information of X and Y under p . ## L.3. Lookahead reinforces choice points Training with multi-head prediction increases the importance of choice points in the loss in comparison to inconsequential decisions. To make this argument, we present a simplified model of language modelling. Consider a sequential decision task and a model M that is trained in a teacher-forced way on optimal trajectories. We distinguish choice points -transitions that lead to different outcomes - and inconsequential decisions which do not (Figure S17 (a) and (b)). Figure S17: Example of a sequential prediction task with derailing. The goal is to go from the arrow to the trophy. Turning around is not allowed. Most transitions are unique, but there are two turns to be taken correctly, the consequential decisions (a) and (c). Turn (b) is an inconsequential decision : the paths join right after it. Next to transitions (a) and (b), we sketch how a 4-step prediction loss can place more emphasis on consequential transitions than inconsequential ones during teacher-forced training. Next to transition (c), we sketch how a 4-step lookahead can prevent models from taking irreversible suboptimal decisions during autoregressive decoding. <details> <summary>Image 17 Details</summary>  ### Visual Description ## Diagram: Pathfinding with Obstacles and Rewards ### Overview The image depicts a pathfinding scenario represented by a grid of squares. The path starts at the bottom-left, proceeds through a series of connected squares, and aims to reach a trophy at the top of a vertical stack of squares on the right side of the diagram. There are obstacles, represented by skulls, and a potential loop or detour in the path. The path is labeled with sections (a), (b), and (c). ### Components/Axes * **Grid:** The environment is represented by a grid of squares, indicating possible locations or steps in the path. * **Start:** The path begins at the bottom of a vertical stack of squares on the left side, indicated by an upward-pointing arrow. * **End/Goal:** The goal is represented by a golden trophy at the top of a vertical stack of squares on the right side. * **Obstacles:** Two skull icons represent obstacles. One is at the top of a vertical stack of squares on the left side, and the other is at the end of a horizontal path on the right side. * **Path Sections:** The path is divided into three sections labeled (a), (b), and (c). * **Loop/Detour:** A small square labeled (b) represents a potential loop or detour in the path. Arrows indicate the flow into and out of this loop. * **Dashed Lines:** Dashed lines connect the top of the vertical stack of squares on the left to the horizontal path, indicating a potential hazard or constraint. A dotted line represents a section of the path on the right. ### Detailed Analysis * **Section (a):** This section consists of a vertical stack of squares on the left side, with an arrow indicating the starting point at the bottom. The path then extends horizontally to the right. There are 6 squares in the vertical stack. * **Section (b):** This section represents a loop or detour. The path enters the square, and arrows indicate that it can either continue forward or loop back to the same square. * **Section (c):** This section is a horizontal path leading to the skull obstacle. The path is represented by a dotted line. There are 4 squares in this section. * **Obstacles:** The skull at the top of the vertical stack on the left is connected to the horizontal path by dashed lines, suggesting a potential hazard or constraint if the path deviates upwards. The skull on the right indicates a dead end or another obstacle. * **Goal:** The trophy is located at the top of a vertical stack of squares on the right side. There are 4 squares in this vertical stack. ### Key Observations * The path is not a straight line; it involves both vertical and horizontal movements. * The loop in section (b) introduces a decision point in the path. * The dashed lines and skulls represent potential hazards or constraints that the path must avoid. * The dotted line in section (c) might indicate a less desirable or more dangerous path. ### Interpretation The diagram illustrates a pathfinding problem where an agent needs to navigate a grid to reach a goal (trophy) while avoiding obstacles (skulls) and considering potential hazards (dashed lines). The loop in section (b) introduces complexity, as the agent must decide whether to take the detour or continue forward. The dotted line in section (c) might represent a riskier path that the agent should avoid if possible. The diagram highlights the challenges of pathfinding, including obstacle avoidance, decision-making, and risk assessment. </details> More formally, assume that the language model is deployed in a reinforcement learning setting like in reinforcement learning from human feedback (Ouyang et al., 2022) (states are prompts followed by the partial sequence of tokens x t :1 generated so far, actions are single tokens x t +1 to generate, rewards are external R ( x t :1 ) ). The quantity <!-- formula-not-decoded --> is the value of the state x t :1 following the policy π , while <!-- formula-not-decoded --> quantifies the importance of the decision x t +1 on the value thereafter. Choice points can formally be viewed as steps t for which σ π ( x t :1 ) is large, while inconsequential points are steps where it is low. Note that for completion models, there is no explicit reward, and our argument is merely meant to illustrate what we mean by choice points . Derailing denotes a situation where autoregressive generation of trajectories from M at inference time results in bad outcomes after M made a mistake on a choice point. Even if subsequently, M acts optimally given this choice, the final outcome can be significantly worse than the outcome of the optimal trajectory. Staying in the teacher-forced setting, we ask: What is the impact of training M with n -step prediction instead of nextstep prediction on this task? Say x t → x t +1 is a choice point in an optimal trajectory with the suboptimal choice being x t → ˜ x t +1 (Figure S17 (a)). Assume that the trajectories preceding x t and succeeding x t +1 and ˜ x t +1 consist of inconsequential transitions, the latter denoted by ˜ x t + j → ˜ x t + j +1 . We will compare the losses of a teacher-forced next-step prediction model and a teacher-forced n -step prediction model on the partial trajectory ( x t -n +1 , . . . x t ) . For the next-step prediction model, the predictions are ( x t -n +2 , . . . , x t , ˜ x t +1 ) with a single wrong prediction. The predictions of an n -step prediction model at time t -n + i , i = 1 , . . . , n are ( x t -n + i +1 , . . . , x t , ˜ x t +1 , . . . , ˜ x t + i ) with i wrong predictions. In other words, an n -step prediction model receives 1 + . . . + n = n ( n +1) 2 loss terms pertaining to such a choice point and its consequences, while each inconsequential transition (Figure S17 (b)) is only reinforced n times as often as in a next-step prediction model. In other words, choice points receive on average n +1 2 times more importance in the loss of n -step prediction models than in next-step prediction models. As argued in Section 5.1, we believe that this model captures important features of training and inference with language models: choice points are semantically important turning points in the generated texts, such as the final answer to a question or a specific line of code, while inconsequential decisions can be a choice among synonyms or of variable names in code. Apart from this training dynamics point of view, we hypothesize that n -step prediction also allows the formation of circuits that specifically spot inconsistencies between predictions for earlier and later steps. For instance, if in an early layer of the model, it can be predicted that a decision x t → ˜ x t +1 leads to suboptimal outcomes ˜ x t + n (Figure S17 (c)), subsequent layers can reduce the probability of x t → ˜ x t +1 in the model's next-step prediction. Such behaviors also happen in next-step prediction models given enough capacity, but our experiments in Section 4.2 point to the fact that circuits of this kind are formed more easily in multi-step architectures that enforce the required information ˜ x t + n to be available to the model when predicting ˜ x t +1 . We believe that this situation appears frequently in natural language and code modelling, for instance where an initial answer to a question contradicts the results of the chain of thought brought forward with the intention to justify it. In more general terms, this situation arises whenever predicting first ˜ x n + i for some 1 < i ≤ n and then ˜ x n +1 based on ˜ x n + i is easier than predicting ˜ x n +1 directly. We discuss this phenomenon of factorization orders in the next section and present a specific instance of it that frequently appears in modelling natural language. ## L.4. Factorization orders Causal language modelling factorizes probabilities over text sequences x t · · · x 1 classically as <!-- formula-not-decoded --> While moving forward in time is certainly the most natural choice of factorization order, there exist cases where it is suboptimal. In inflectional languages, for instance, agreement between related sentence parts is a frequent pattern with one word directing the grammatical forms of others. Consider the German sentence Wie konnten auch Worte meiner durstenden Seele genügen? 3 Friedrich Hölderlin, Fragment von Hyperion (1793) where "genügen" requires a dative case object and then "Seele" requires the possessive pronoun "mein" to be in female singular dative form "meiner" and the participle "durstend" to be in female singular dative form in weak declination "durstenden" because it follows "meiner". In other words, the factorization order Wie konnten auch Worte → genügen → Seele → meiner → durstenden? is arguably an easier one for constructing the above sentence. Humans as well as language models therefore have to perform this factorization (which deviates from the causal order in which predictions take place!) within their latent activations, and a 4 -token prediction loss makes this easier as it explicitly encourages models to have all information about the successive 4 tokens in its latent representations. 3 roughly: How could words be enough for my thirsty soul? ## M. Training hyperparameters Table S13: Overview of all training hyperparameters used. We schedule all learning rates with a linear warmup and cosine decay (Loshchilov and Hutter, 2017) to a fraction of the peak learning rate which is depicted in the last column ('decay ratio'). All experiments use the Adam (Kingma and Ba, 2015) optimizer with β 1 = 0 . 9 , β 2 = 0 . 95 and decoupled L 2 weight decay (Loshchilov and Hutter, 2019) coefficient 0 . 1 . We clip gradients to a maximal Euclidean norm of 1 . 0 in all experiments except CodeContests finetunings, where we use 0 . 1 instead. Summarization finetunings correspond to three epochs on all datasets except BigPatent (1 epoch). Byte-level models use the architecture with replicated unembeddings from Appendix B. | Model | Batch size ( 2 20 ) | Steps | Tokens (B) | Warmup steps | Peak LR | Context length | Decay ratio | |---------------------------------|-----------------------|---------|--------------|----------------|------------|------------------|---------------| | Model scaling (Section 3.1) | | | | | | | | | 0.3B | 8 | 10,850 | 91.0 | 1000 | 3 × 10 - 4 | 4096 | 0.03 | | 0.6B | 8 | 10,850 | 91.0 | 1000 | 3 × 10 - 4 | 4096 | 0.03 | | 1.3B | 8 | 10,850 | 91.0 | 1000 | 3 × 10 - 4 | 4096 | 0.03 | | 3B | 8 | 10,850 | 91.0 | 1000 | 3 × 10 - 4 | 4096 | 0.03 | | 7B | 8 | 25,000 | 209.7 | 2000 | 3 × 10 - 4 | 4096 | 0.03 | | 13B | 8 | 25,000 | 209.7 | 1000 | 3 × 10 - 4 | 4096 | 0.03 | | Code models (Section 3) | | | | | | | | | 7B 200B | 8 | 25,000 | 209.7 | 2000 | 3 × 10 - 4 | 4096 | 0.03 | | 7B 500B | 7 | 68,570 | 503.3 | 2000 | 3 × 10 - 4 | 4096 | 0.03 | | 7B 1T | 7 | 136,240 | 1000.0 | 2000 | 3 × 10 - 4 | 4096 | 0.03 | | Byte-level models (Section 3.3) | | | | | | | | | 7B 314GB | 12 | 25,000 | 314.6 | 2000 | 3 × 10 - 4 | 8192 | 0.03 | | Language models (Section 3.7) | | | | | | | | | 7B 200B | 8 | 25,000 | 209.7 | 2000 | 3 × 10 - 4 | 4096 | 0.10 | | 7B 500B | 8 | 60,000 | 503.3 | 2000 | 3 × 10 - 4 | 4096 | 0.10 | | Induction task (Section 4.1) | | | | | | | | | 1M - 1B | 0.25 | 100,000 | 26.2 | 2000 | 10 - 4 | 2048 | 0.03 | | 1M - 1B (Appendix J) | 0.5 | 50000 | 26.2 | 2000 | 10 - 4 | 2048 | 0.03 | | Arithmetic task (Section 4.2) | | | | | | | | | 30M | 0.25 | 100,000 | 26.2 | 2000 | 10 - 4 | 1024 | 0.03 | | 100M | 0.25 | 100,000 | 26.2 | 2000 | 10 - 4 | 2048 | 0.03 | | Summarization (Section 3.7) | | | | | | | | | BigPatent | 0.125 | 76,680 | 10.1 | 100 | 3 × 10 - 5 | 4096 | 0.03 | | CNN/Dailymail | 0.125 | 7,140 | 0.9 | 100 | 3 × 10 - 5 | 4096 | 0.03 | | Multi-News | 0.125 | 3,330 | 0.4 | 100 | 3 × 10 - 5 | 4096 | 0.03 | | OrangeSum | 0.125 | 360 | 0.0 | 100 | 3 × 10 - 5 | 4096 | 0.03 | | pn-summary | 0.125 | 3,450 | 0.5 | 100 | 3 × 10 - 5 | 4096 | 0.03 | | SAMSum | 0.125 | 60 | 0.0 | 100 | 3 × 10 - 5 | 4096 | 0.03 | | ThaiSum | 0.125 | 23,640 | 3.1 | 100 | 3 × 10 - 5 | 4096 | 0.03 | | WikiSummary | 0.125 | 2,550 | 0.3 | 100 | 3 × 10 - 5 | 4096 | 0.03 | | XSum | 0.125 | 2,760 | 0.4 | 100 | 3 × 10 - 5 | 4096 | 0.03 | | CodeContests (Section 3.6) | | | | | | | | | 7B | 0.25 | 13,000 | 3.6 | 400 | 5 × 10 - 5 | 4096 | 0.004 | Table S14: Overview of model architectures used for scaling analyses. | Name | Dimension | Layers | Heads | |-------------|-------------|----------|---------| | 1M | 128 | 5 | 4 | | 3M | 256 | 4 | 8 | | 10M | 384 | 6 | 8 | | 30M | 512 | 10 | 8 | | 100M | 768 | 14 | 12 | | 300M | 1024 | 25 | 16 | | 1B | 1536 | 36 | 24 | | 0.3B | 1024 | 18 | 16 | | 0.6B | 1280 | 27 | 20 | | 1.3B | 2048 | 24 | 16 | | 3B | 2560 | 36 | 20 | | 6.7B ('7B') | 4096 | 32 | 32 | | 13B | 5120 | 40 | 40 |