Parallel Token Prediction for Language Models

TL;DR

Parallel Token Prediction (PTP) significantly accelerates autoregressive language models by enabling the prediction of multiple tokens simultaneously, maintaining dependency modeling while reducing inference time.

Contribution

The paper introduces PTP, a novel framework that predicts multiple tokens in one forward pass by transforming randomness sources, with proven arbitrary dependency modeling and efficient training methods.

Findings

Achieves 2.4x speedup on a diverse-task benchmark.

Can represent arbitrary token dependencies in a single call.

Provides open-source code and checkpoints.

Abstract

Autoregressive decoding in language models is inherently slow, generating only one token per forward pass. We propose Parallel Token Prediction (PTP), a general-purpose framework for predicting multiple tokens in a single model call. PTP moves the source of randomness from post-hoc sampling to random input variables, making future tokens deterministic functions of those inputs and thus jointly predictable in a single forward pass. We prove that a single PTP call can represent arbitrary dependencies between tokens. PTP is trained by distilling an existing model or through inverse autoregressive training without a teacher. Experimentally, PTP achieves a 2.4x speedup on a diverse-task speculative decoding benchmark. We provide code and checkpoints at https://github.com/mandt-lab/ptp.