Parallel Token Prediction for Language Models

TL;DR

The paper tackles the autoregressive latency bottleneck in language models by introducing Parallel Token Prediction (PTP), which jointly predicts multiple tokens in a single transformer call through auxiliary random inputs $\{u_i\}$. It proves that PTP is expressive enough to represent arbitrary autoregressive distributions and provides two practical training paths: distillation from a teacher and inverse autoregressive training without supervision, plus an error-correction mechanism to preserve output fidelity. Empirically, PTP achieves state-of-the-art speculative decoding performance on Vicuna-7B (over $4$ tokens per step) and demonstrates strong results on taxi-location data, illustrating both speedups and coherent multi-token generation. Overall, PTP opens a universal framework for fast, reliable long-sequence generation without sacrificing modeling power, with clear implications for real-time and interactive deployment of large language models.

Abstract

We propose Parallel Token Prediction (PTP), a universal framework for parallel sequence generation in language models. PTP jointly predicts multiple dependent tokens in a single transformer call by incorporating the sampling procedure into the model. This reduces the latency bottleneck of autoregressive decoding, and avoids the restrictive independence assumptions common in existing multi-token prediction methods. We prove that PTP can represent arbitrary autoregressive sequence distributions. PTP is trained either by distilling an existing model or through inverse autoregressive training without a teacher. Experimentally, we achieve state-of-the-art speculative decoding performance on Vicuna-7B by accepting over four tokens per step on Spec-Bench. The universality of our framework indicates that parallel generation of long sequences is feasible without loss of modeling power.

Figures (5)

• Figure 1: Our parallelized model generates the same text as its teacher in a fraction of the steps. By the time our model (bottom) has generated an entire function, an autoregressive model (top) only generates the method's signature. Prompt: Write a Python function that computes the factorial of a number. Green tokens are accepted tokens in that step, red tokens are incorrect. Semitransparent tokens are rejected after the first mistake.
• Figure 2: Parallel Token Prediction predicts several tokens in one model call. (a) An autoregressive model predicts the distribution of the single next token $t_i$. From the histogram, that token is chosen with the help of an auxiliary variable $u_i$ that is chosen uniformly at random (see \ref{['fig:pick-token-function']}). (b) One-Hot Parallel Token Prediction merges the sampling into the model by feeding the auxiliaries directly into the model. This allows joint prediction of several tokens. (c) Categorical Parallel Token Prediction models the distribution of each token, but predicts them in parallel using the auxiliary variables.
• Figure 3: Sampling from a discrete distribution. Given a histogram $P_i$(left), compute the inverse cumulative distribution function (right) and look up the token at a random location $u_i \in \mathcal{U}[0, 1]$. Our framework relies on considering both parts jointly.
• Figure 4: Speculative decoding with Parallel Token Prediction on code. The x-axis shows the draft-model parameter count. The left panel reports wall-clock speedup relative to standard autoregressive decoding, and the right panel shows the average number of tokens accepted per step under teacher verification. Curves compare autoregressive draft models to Parallel Token Prediction (PTP) drafts; ($\bullet$) are models trained from scratch for a fixed number of epochs, and ($\star$) is finetuned from the teacher. Across model sizes, PTP drafts achieve higher speedups by generating more than one correct token per step, whereas autoregressive drafts remain sequential. PTP allows parallelism in draft models, yielding larger speedups at equal model size.
• Figure 5: Parallel Token Prediction generates meaningful pairs of tokens. Each panel plots the first and second auxiliary variables $(u_1,u_2)\in[0,1]^2$ on the axes; regions in this unit square correspond to the resulting two-token outputs (green: compatible, red: incompatible). (Left) In a coding problem, autoregressive sampling first selects one of def, import, or n, and then continues with meaningful predictions: a function name to declare, a package to import, or a variable assignment. (Center) Our code completion model produces similarly sensible token pairs, but in a single model call by coordinating predictions through the auxiliaries; only rarely ($<1\%$) does it yield spurious combinations like def sys. (Right) A model that independently predicts future tokens is bound to fail: in about $60\%$ of cases, it combines incompatible tokens because the second token is not informed about the first.

Theorems & Definitions (4)

• Theorem 1
• Theorem 2
• proof
• proof