Fast and Expressive Multi-Byte Prediction with Probabilistic Circuits

Multi-token prediction (MTP) is a prominent strategy to significantly speed up generation in large language models (LLMs), especially in byte-level LLMs, which are tokeniser-free but prohibitively slow. However, many existing MTP methods either assume independence between future tokens, sacrificing expressiveness, or generate tokens one at a time within the window, increasing latency. In this work, we investigate the trade-off between expressiveness and latency in MTP within the framework of probabilistic circuits (PCs). Our framework, MTPC, allows one to explore different ways to encode the joint distributions over future tokens by selecting 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, and byte-fied subword models, such as Llama3.2 3B. Our experiments show that, when combined with speculative decoding, MTPC substantially 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.

In this paper we speed up text generation for AI models that produce text one letter at a time. This is very slow for modern models because each letter requires another pass through an enormous network, slower still than the chunk-at-a-time approach used by models like ChatGPT or Claude, which generate several letters per step. A promising fix is to have a small, fast model draft several letters ahead and let the slow model judge the text quality and accept or reject that draft in a single parallel pass. The catch is a trade-off: if the fast model captures how upcoming letters depend on one another, more of its guesses get accepted and generation speeds up. However, modelling those dependencies costs time and memory, whereas a drafter that ignores them is quick but has more of its guesses rejected. We introduce a framework called MTPC that lets us navigate this trade-off precisely, by choosing how richly the drafter models the dependencies between future letters.