Fast and Expressive Multi-Token Prediction with Probabilistic Circuits
This work addresses the problem of slow generation in tokeniser-free LLMs, offering a method to improve speed without sacrificing expressiveness, though it is incremental as it builds on existing MTP and speculative decoding techniques.
The paper tackles the trade-off between expressiveness and latency in multi-token prediction for large language models by introducing MTPC, a framework using probabilistic circuits to model joint distributions over future tokens, which speeds up generation compared to independence-based methods while maintaining verifier performance.
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.