Fast and Expressive Multi-Token Prediction with Probabilistic Circuits

TL;DR

This work addresses the speed–expressiveness trade-off in multi-token prediction for large language models, with a focus on byte-level LLMs. It introduces MtPC, a framework that parameterises the joint distribution over a window of $n$ future tokens using Probabilistic Circuits, enabling architectures from fully factorised to HMM and BTREE structures. By integrating MtPC with speculative decoding in a self-speculative setup and fine-tuning draft backbones with LoRA adapters, the approach achieves substantial throughput improvements (e.g., up to $5.47\times$ faster than autoregressive generation and $1.22\times$ faster than independence-based MTP) while preserving the original verifier’s performance. Empirical retrofitting of EvaByte demonstrates how longer prediction windows and more expressive PC architectures yield higher throughput, highlighting device-specific trade-offs between acceptance rate and latency. This framework paves the way for richer dependency modelling in MTP with practical latency, enabling scalable, token-window lookahead in byte-level and other token-granularity settings.

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.

Figures (2)

• Figure 1: MtPC allows for exploring the trade-off between efficiency (latency) and expressiveness (token acceptance) with different MTP designs in terms of 1) choice of architecture (\ref{['eq:n-indep-prob']}, \ref{['eq:r-cp']}, \ref{['eq:r-hmm']}, \ref{['eq:btree']}); 2) choice of layers shared between draft and verifier models in self-speculative decoding. Dotted lines indicate iso-throughput (tokens generated per second) regions, highlighting configurations such as \ref{['eq:btree']} for $n=16$ tokens and $2$ LoRA layers that achieve the best throughput.
• Figure 2: PCs allow for modelling a spectrum of dependency structures over sequences of tokens, as shown for the known \ref{['eq:n-indep-prob']} and \ref{['eq:r-cp']} and the novel \ref{['eq:r-hmm']} and \ref{['eq:btree']} MTP variants. Input units are grouped in coloured layers, one for each token, while sum and product layers encoding (hierarchies of) latent variable distributions are in grey. The output unit of each circuit (in blue) computes $q_{\bm{\mathrm{\theta}}}(x_{t+1},\ldots,x_{t+n}\mid \bm{\mathrm{x}}_{\leq t})$. In the figure we omit the dependency on the context $\mathbf{x}_{\leq t}$ for readability.

Theorems & Definitions (1)

• Definition 1: Smoothness and decomposability darwiche2002knowledgechoi2020pc