Fast and Expressive Multi-Token Prediction with Probabilistic Circuits

TL;DR

This paper introduces MTPC, a probabilistic circuit framework for multi-token prediction in large language models, balancing expressiveness and speed, and demonstrating significant acceleration without performance loss.

Contribution

The work presents a novel probabilistic circuit approach for multi-token prediction that generalizes existing models and improves generation speed in byte-level LLMs.

Findings

MTPC accelerates generation when combined with speculative decoding.

Retrofitting MTPC retains original LLM performance.

Exploration of trade-offs guides optimal model configurations.

Abstract

Multi-token prediction (MTP) is a prominent strategy to significantly speed up generation in large language models (LLMs), including byte-level LLMs, which are tokeniser-free but prohibitively slow. However, existing MTP methods often sacrifice expressiveness by assuming independence between future tokens. In this work, we investigate the trade-off between expressiveness and latency in MTP within the framework of probabilistic circuits (PCs). Our framework, named MTPC, allows one to explore different ways to encode the joint distributions over future tokens by selecting different circuit architectures, generalising classical models such as (hierarchical) mixture models, hidden Markov models and tensor networks. We show the efficacy of MTPC by retrofitting existing byte-level LLMs, such as EvaByte. Our experiments show that, when combined with speculative decoding, MTPC significantly speeds up generation compared to MTP with independence assumptions, while guaranteeing to retain the performance of the original verifier LLM. We also rigorously study the optimal trade-off between expressiveness and latency when exploring the possible parameterisations of MTPC, such as PC architectures and partial layer sharing between the verifier and draft LLMs.