Fast and Expressive Multi-Token Prediction with Probabilistic Circuits

📄 arXiv: 2511.11346v1 📥 PDF

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

分类: cs.LG

发布日期: 2025-11-14

💡 一句话要点

提出基于概率电路的MTPC框架，加速字节级LLM生成并保持性能。

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 多token预测 概率电路 大型语言模型 字节级LLM 推测解码

📋 核心要点

  1. 现有MTP方法为加速LLM生成，常牺牲表达能力，假设未来token独立性，导致性能下降。
  2. MTPC框架利用概率电路编码未来token的联合分布，在表达能力和延迟之间取得平衡。
  3. 实验表明，MTPC结合推测解码，在加速生成的同时，保证了原始LLM的性能。

📝 摘要（中文）

多token预测(MTP)是加速大型语言模型(LLM)生成的重要策略，尤其是在tokeniser-free但速度极慢的字节级LLM中。然而，现有的MTP方法通常假设未来token之间的独立性，牺牲了表达能力。本文研究了概率电路(PC)框架下MTP中表达能力和延迟之间的权衡。我们提出的框架MTPC允许通过选择不同的电路架构来探索编码未来token联合分布的不同方式，推广了经典模型，如(分层)混合模型、隐马尔可夫模型和张量网络。我们通过改进现有的字节级LLM（如EvaByte）展示了MTPC的有效性。实验表明，当与推测解码结合使用时，MTPC显著加速了生成，同时保证保留原始验证器LLM的性能。我们还严格研究了探索MTPC的可能参数化（如PC架构和验证器与draft LLM之间的部分层共享）时，表达能力和延迟之间的最佳权衡。

🔬 方法详解

问题定义：现有的大型语言模型，特别是字节级别的LLM，在生成文本时速度较慢。多token预测（MTP）是一种加速方法，但现有MTP方法通常假设未来token之间是相互独立的，这牺牲了模型的表达能力，导致生成质量下降。因此，需要在表达能力和生成速度之间找到一个平衡点。

核心思路：论文的核心思路是利用概率电路（Probabilistic Circuits, PCs）来建模未来token的联合概率分布。通过选择不同的PC架构，可以灵活地控制模型的表达能力，从而在表达能力和生成速度之间进行权衡。概率电路能够显式地表示联合概率分布，避免了独立性假设带来的信息损失。

技术框架：MTPC框架包含两个主要部分：一个验证器LLM和一个draft LLM。验证器LLM是原始的LLM，用于评估生成的token序列的质量。Draft LLM基于概率电路MTPC，用于快速生成多个token的草稿。整体流程是：首先，draft LLM使用MTPC生成多个token的草稿；然后，验证器LLM评估这些草稿的质量，并决定是否接受它们。如果草稿被接受，则将其添加到生成的序列中；否则，验证器LLM将生成新的token。

关键创新：该论文的关键创新在于将概率电路引入到多token预测中，提出了MTPC框架。与传统的MTP方法相比，MTPC能够显式地建模未来token之间的依赖关系，从而提高了模型的表达能力。此外，MTPC框架允许通过选择不同的PC架构来灵活地控制模型的复杂度和生成速度。

关键设计：MTPC的关键设计包括：1) 选择合适的PC架构，例如，可以使用混合模型、隐马尔可夫模型或张量网络来建模未来token的联合概率分布。2) 设计合适的损失函数来训练MTPC，例如，可以使用最大似然估计或变分推断来训练模型。3) 考虑验证器LLM和draft LLM之间的层共享，以减少模型的参数量和计算复杂度。实验中探索了不同PC架构和层共享策略对性能的影响。

🖼️ 关键图片

fig_0
fig_1

📊 实验亮点

实验结果表明，MTPC在加速字节级LLM生成方面表现出色。与基于独立性假设的MTP方法相比，MTPC在结合推测解码时，能够显著提高生成速度，同时保证原始验证器LLM的性能。具体而言，在EvaByte模型上，MTPC实现了显著的加速效果，同时保持了与原始模型相当的生成质量。

🎯 应用场景

该研究成果可应用于各种需要快速文本生成的场景，例如机器翻译、文本摘要、对话系统等。通过MTPC框架，可以显著提高生成速度，同时保持生成质量。尤其适用于资源受限的设备或对延迟敏感的应用，具有重要的实际应用价值和商业前景。

📄 摘要（原文）

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.