Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Pre-Training Curriculum for Multi-Token Prediction in Language Models

Ansar Aynetdinov, Alan Akbik

TL;DR

This work tackles the challenge of leveraging multi-token prediction (MTP) for small language models (SLMs) by introducing curriculum learning strategies to gradually adjust the number of prediction heads during pre-training. The authors propose forward and reverse MTP curricula and rigorously evaluate them across 1.3B and 3B decoder-only models, with subword and byte-level tokenizations, on datasets including MiniPile. Key findings show that the forward curriculum improves downstream NTP performance and generative quality while preserving self-speculative decoding, whereas the reverse curriculum achieves strong NTP performance and output quality but does not provide inference speedups. The results highlight a practical path to harness MTP for SLMs and point to byte-level models and forward curricula as the most effective combination for balancing accuracy, speed, and quality in pre-training regimes.

Abstract

Multi-token prediction (MTP) is a recently proposed pre-training objective for language models. Rather than predicting only the next token (NTP), MTP predicts the next $k$ tokens at each prediction step, using multiple prediction heads. MTP has shown promise in improving downstream performance, inference speed, and training efficiency, particularly for large models. However, prior work has shown that smaller language models (SLMs) struggle with the MTP objective. To address this, we propose a curriculum learning strategy for MTP training, exploring two variants: a forward curriculum, which gradually increases the complexity of the pre-training objective from NTP to MTP, and a reverse curriculum, which does the opposite. Our experiments show that the forward curriculum enables SLMs to better leverage the MTP objective during pre-training, improving downstream NTP performance and generative output quality, while retaining the benefits of self-speculative decoding. The reverse curriculum achieves stronger NTP performance and output quality, but fails to provide any self-speculative decoding benefits.

Pre-Training Curriculum for Multi-Token Prediction in Language Models

TL;DR

This work tackles the challenge of leveraging multi-token prediction (MTP) for small language models (SLMs) by introducing curriculum learning strategies to gradually adjust the number of prediction heads during pre-training. The authors propose forward and reverse MTP curricula and rigorously evaluate them across 1.3B and 3B decoder-only models, with subword and byte-level tokenizations, on datasets including MiniPile. Key findings show that the forward curriculum improves downstream NTP performance and generative quality while preserving self-speculative decoding, whereas the reverse curriculum achieves strong NTP performance and output quality but does not provide inference speedups. The results highlight a practical path to harness MTP for SLMs and point to byte-level models and forward curricula as the most effective combination for balancing accuracy, speed, and quality in pre-training regimes.

Abstract

Multi-token prediction (MTP) is a recently proposed pre-training objective for language models. Rather than predicting only the next token (NTP), MTP predicts the next tokens at each prediction step, using multiple prediction heads. MTP has shown promise in improving downstream performance, inference speed, and training efficiency, particularly for large models. However, prior work has shown that smaller language models (SLMs) struggle with the MTP objective. To address this, we propose a curriculum learning strategy for MTP training, exploring two variants: a forward curriculum, which gradually increases the complexity of the pre-training objective from NTP to MTP, and a reverse curriculum, which does the opposite. Our experiments show that the forward curriculum enables SLMs to better leverage the MTP objective during pre-training, improving downstream NTP performance and generative output quality, while retaining the benefits of self-speculative decoding. The reverse curriculum achieves stronger NTP performance and output quality, but fails to provide any self-speculative decoding benefits.

Paper Structure

This paper contains 30 sections, 4 equations, 3 figures, 7 tables.

Table of Contents

  1. Introduction
  2. Preliminaries
  3. Multi-Token Prediction Objective
  4. Forward and Reverse Multi-Token Curricula
  5. Forward Multi-Token Curriculum
  6. Reverse Multi-Token Curriculum
  7. Experiments
  8. Experimental Setup
  9. Models
  10. Prediction Heads
  11. Dataset
  12. Training setup & Hyperparameters
  13. Experiment 1: Main LM Head Performance
  14. Setup
  15. Results
  16. ...and 15 more sections

Figures (3)

  • Figure 1: Visualization of the forward and reverse MTP curricula. When training a language model on a 3-token-prediction task for $n$ steps, the forward curriculum starts with a vanilla NTP task, adding an additional token to the task every $\frac{n}{3}$ steps. The reverse curriculum does the opposite, starting with a full 3-token-prediction task, and dropping a token from the task every $\frac{n}{3}$ steps.
  • Figure 2: Inference speed-ups achieved by MTP models in relation to the NTP baseline. We measure the speed-ups in terms of the amount of forward passes required to generate a sequence using self-speculative decoding formulation proposed by blockwise_parallel_decoding
  • Figure 3: Breakdown of the per-head acceptance rate of tokens drafted via self-speculative decoding for 3B subword models. The y-axis indicates a head's share of tokens accepted by the model in total.