🤖 machine learning

Fast and Expressive Multi-Byte Prediction with Probabilistic Circuits

This paper introduces MTPC, a probabilistic circuit-based framework for multi-token prediction that optimizes the trade-off between expressiveness and latency by encoding joint distributions over future tokens, thereby significantly accelerating byte-level and subword LLM generation while preserving the original model's performance.

Original authors: Andreas Grivas, Lorenzo Loconte, Emile van Krieken, Piotr Nawrot, Yu Zhao, Euan Wielewski, Pasquale Minervini, Edoardo Ponti, Antonio Vergari

Published 2026-06-03

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Andreas Grivas, Lorenzo Loconte, Emile van Krieken, Piotr Nawrot, Yu Zhao, Euan Wielewski, Pasquale Minervini, Edoardo Ponti, Antonio Vergari

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to write a story, but you have a very strict rule: you can only write one letter at a time. Every time you type a letter, you have to stop, think, and ask your super-smart brain (the AI) what comes next. This is how most current Large Language Models (LLMs) work. It's accurate, but it's incredibly slow, especially if you are writing in "bytes" (the raw building blocks of text) rather than whole words, because you have to type thousands of letters to write a single sentence.

The paper introduces a new method called MTPC (Multi-Token Prediction Circuits) to solve this speed problem without losing the quality of the story.

Here is how it works, using some everyday analogies:

1. The Problem: The "Guessing Game" vs. The "Crystal Ball"

To speed things up, researchers tried a trick called Multi-Token Prediction (MTP). Instead of guessing one letter, the AI tries to guess a whole chunk of letters at once (like guessing the next 8 letters of a word).

The Old Way (Independence Assumption): Imagine you are guessing the next 8 letters of a word, but you treat each letter as if it has no relationship to the others. You guess the first letter, then the second, then the third, completely ignoring that if the first letter is "C", the second is unlikely to be "Z".
- The Result: This is fast, but it leads to nonsense. You might get "Cretoria" instead of "Pretoria" or "Craporia" because the model didn't realize those letters should fit together. It's like trying to build a house by picking bricks randomly without checking if they fit.
The New Way (MTPC): The authors say, "Let's stop guessing letters in isolation. Let's guess the whole chunk as a connected group." They use a mathematical tool called a Probabilistic Circuit.
- The Analogy: Think of the old way as a line of people passing a note, where everyone whispers a random word. The new way is like a conductor leading an orchestra. The conductor (the circuit) knows that if the first instrument plays a C-major chord, the next instruments must play notes that fit that chord. It understands the dependencies between the letters.

2. The Toolkit: The "Circuit Architect"

The paper proposes a flexible framework (MTPC) that lets you choose how "connected" the letters should be. They offer different "architectures" (shapes of the circuit) to balance speed and intelligence:

FF (Fully Factorised): The "Random Guess" mode. Fast, but dumb. (The orchestra members play alone).
CP (Canonical Polyadic): A "Group Guess." They guess a few main themes and build the letters around those. A bit smarter.
HMM (Hidden Markov Model): A "Chain Reaction." The first letter influences the second, which influences the third, and so on. This is very smart but slow because you have to wait for one to finish before starting the next.
BTree (Binary Tree): The "Team Huddle." This is the paper's star. Imagine splitting the 8 letters into two groups of 4. The model guesses the first group and the second group at the same time, but they are linked by a "team leader" (a hidden variable) that ensures they agree on the overall theme.
- Why it's great: It gets the intelligence of the "Chain Reaction" but the speed of the "Random Guess" because it does two things at once.

3. The Safety Net: "Speculative Decoding"

You might worry: "If the AI guesses a whole chunk at once, what if it gets it wrong?"

The paper uses a technique called Speculative Decoding.

The Analogy: Imagine a fast runner (the Draft Model) and a slow, ultra-accurate judge (the Verifier).
1. The fast runner sprints ahead and guesses the next 8 letters.
2. The slow judge checks them one by one.
3. If the judge agrees with the runner's guess, great! We keep those letters.
4. If the judge disagrees, we stop right there, discard the bad guesses, and only keep the ones the judge approved.

Because the Draft Model (MTPC) is so good at understanding how letters connect (thanks to the BTree circuit), the judge agrees with the runner much more often than before. This means we get to keep more of the fast guesses, speeding up the whole process.

4. The Results: Speeding Up Without Breaking Things

The authors tested this on two specific AI models:

EvaByte: A model that already writes in bytes.
Llama 3.2 3B (Byte): A popular model converted to write in bytes.

The Findings:

Massive Speedup: Compared to the old "one letter at a time" method, MTPC made EvaByte 5.15 times faster and Llama 2.24 times faster.
Better than the "Independence" Trick: Even compared to other fast methods that just guess letters independently, MTPC was 1.17 times faster.
No Quality Loss: Crucially, because of the "Safety Net" (Speculative Decoding), the final output is exactly the same quality as if the AI had written it one letter at a time. You don't lose accuracy for speed.

Summary

The paper presents a new way to make AI text generation faster by teaching the AI to guess chunks of text as a connected group rather than isolated letters. By using a smart "Binary Tree" structure (BTree) to organize these guesses and a "Judge" to verify them, they achieved a massive speed boost (up to 5x) while guaranteeing the text remains perfect. It's like teaching a typist to type whole words at once, but with a safety net that catches any typos instantly.

Technical Summary: Fast and Expressive Multi-Byte Prediction with Probabilistic Circuits

Problem Statement

Autoregressive (AR) Large Language Models (LLMs) generate text one token at a time, incurring high latency and energy costs. This bottleneck is particularly severe in byte-level LLMs, which lack tokenizers but require significantly more decoding steps to generate text of the same length compared to subword models. While Multi-Token Prediction (MTP) offers a solution by predicting a window of future tokens simultaneously, existing approaches face a critical trade-off between expressiveness and efficiency:

Independence Assumptions: Many current MTP methods assume future tokens are independent (Fully Factorised, FF). While efficient, this sacrifices expressiveness, leading to "byte-salad" continuations (e.g., nonsensical combinations of valid bytes) and poor acceptance rates during speculative decoding.
Shallow Autoregressive Heads: Alternative methods (e.g., Hydra, EAGLE) retain autoregressive generation within the prediction window. While more expressive, they are closer to accelerated AR decoding than true joint MTP and do not expose the dependency structure as a tunable modeling choice.
Latency vs. Expressiveness: There is a lack of a systematic framework to navigate the spectrum between fully factorised models and complex joint distributions without incurring prohibitive latency.

Methodology: MTPC Framework

The authors propose MTPC (Multi-Token Prediction Circuits), a framework that utilizes Probabilistic Circuits (PCs) to parameterize the joint distribution over future tokens. MTPC treats the joint distribution as a tractable computational graph encoding hierarchical mixture models, allowing for a systematic exploration of architectural variants.

Core Components

Probabilistic Circuit Architectures: MTPC generalizes existing methods and introduces novel structures to model dependencies between future bytes:
- MTPC-FF: A fully factorised model (equivalent to standard independent MTP) represented as a product of input units.
- MTPC-CP: A Canonical Polyadic (CP) decomposition introducing a discrete latent variable to create a mixture of factorised components, recovering some dependencies.
- MTPC-HMM: A Hidden Markov Model parameterisation using a sequence of latent variables to capture distant dependencies via time-dependent transition matrices.
- MTPC-BTree: A novel Binary Tree factorisation that encodes a hierarchy of latent variables. This structure enables parallel sampling of tokens and latent variables, addressing the asymmetry and latency issues found in HMMs.
Parameterisation with LLMs:
- The framework uses an LLM backbone (with optional LoRA adapters) to extract contextual features from the input history.
- A neural network head maps these features to the parameters of the PC (input unit distributions and sum unit weights), creating a neural conditional circuit.
- This allows the model to abstract away specific circuit structures (FF, CP, HMM, BTree) and focus on learning the tensorised parameters.
Speculative Decoding Integration:
- MTPC employs self-speculative decoding, where the draft model (MTPC head) and the verifier (AR LLM) share the same backbone.
- The draft model samples a window of tokens in parallel from the PC.
- The verifier accepts or rejects these tokens in parallel.
- Guarantees: The framework guarantees that greedy decoding outputs are identical to the target AR LLM, and sampling matches the target in expectation, ensuring no loss in generation quality.
Trade-off Dimensions:
- Circuit Architecture: Choosing between FF, CP, HMM, or BTree to balance dependency modeling and computational cost.
- Layer Sharing: Adjusting the number of LoRA layers in the draft model (0 to 4 layers) to control the divergence between draft and verifier representations, further tuning the expressiveness-latency trade-off.

Key Contributions

MTPC Framework: Introduction of a fast MTP framework based on Probabilistic Circuits that relaxes independence assumptions and generalizes tensor decomposition methods (like CP) into a unified PC framework.
Systematic Trade-off Analysis: Rigorous identification of the trade-offs between accepted tokens (expressiveness) and generation latency. The paper demonstrates that optimal throughput is not always achieved by the most expressive model (e.g., HMM) but by architectures like BTree that offer better parallelism.
Empirical Validation: Retrofitting MTPC onto two byte-level LLMs:
- EvaByte (6.5B): A byte-level model pretrained with MTP capabilities.
- Llama 3.2 3B (Byte): A byte-fied version of a popular subword model, not pretrained for MTP.

Experimental Results

The authors evaluated MTPC on NVIDIA L40S and RTX 3090 GPUs using the Tulu 3 dataset.

Throughput Gains:
- EvaByte: MTPC achieved a 5.15× speedup over standard AR generation and a 1.17× speedup over MTP with independence assumptions (FF).
- Llama 3.2 3B: MTPC achieved a 2.24× speedup over AR generation and a 1.17× speedup over FF MTP.
Architecture Performance:
- BTree vs. HMM: While HMMs achieved high acceptance rates, their autoregressive nature introduced high latency, resulting in lower overall throughput. BTree architectures struck the optimal balance, enabling parallel sampling and achieving the highest throughput, particularly for larger window sizes ( $n=16$ ).
- LoRA Layers: Adding LoRA adapters to the draft model increased expressiveness (accepted tokens) but also latency. The optimal configuration varied by model: 1 LoRA layer was optimal for EvaByte, while 2 LoRA layers benefited Llama.
Window Size: Extending the prediction window to $n=16$ significantly boosted throughput for expressive circuits (HMM, BTree) compared to $n=8$ , especially when combined with LoRA adapters.

Significance and Claims

The paper claims that MTPC provides a systematic way to navigate the spectrum of MTP architectural variants, moving beyond the binary choice between independent prediction and shallow autoregressive heads. By leveraging Probabilistic Circuits, the framework allows researchers to explicitly model inter-token dependencies while maintaining tractability.

The authors emphasize that these throughput gains are achieved without sacrificing generation quality, as the speculative decoding process guarantees alignment with the target AR LLM. The work suggests that byte-level LLMs, which are often hindered by computational overhead, can become more mainstream through such efficient multi-byte prediction strategies. The framework is presented as a flexible tool for optimizing the efficiency of next-generation byte-level models, with the code and models made publicly available.