Pair-In, Pair-Out: Latent Multi-Token Prediction for Efficient LLMs

📄 arXiv: 2605.27255v1 📥 PDF

作者: Wenhui Tan, Minghao Li, Xiaoqian Ma, Siqi Fan, Xiusheng Huang, Liujie Zhang, Ruihua Song, Weihang Chen

分类: cs.CL, cs.AI

发布日期: 2026-05-26

备注: Project Page: GitHub.com/AlbertTan404/PIPO


💡 一句话要点

提出PIPO,通过联合隐空间压缩和多token预测加速LLM推理。

🎯 匹配领域: 支柱二:RL算法与架构 (RL & Architecture) 支柱九:具身大模型 (Embodied Foundation Models)

关键词: 大型语言模型 推理加速 多token预测 隐空间压缩 置信度估计

📋 核心要点

  1. 现有方法或侧重于输入端的隐空间压缩,或侧重于输出端的多token预测,缺乏统一视角。
  2. PIPO将隐空间压缩和多token预测视为镜像操作,并通过置信度头消除验证器开销。
  3. 实验表明,PIPO在多个数据集上显著提升了推理速度和准确率,尤其是在长文本任务中。

📝 摘要(中文)

本文提出Pair-In, Pair-Out (PIPO)方法,旨在提升大型语言模型(LLM)的推理效率,尤其是在长链推理场景下。PIPO将隐空间压缩器和多token预测(MTP)头视为镜像操作,前者将两个输入token压缩为一个隐表示,后者将一个隐状态展开为一个额外的输出token。为了消除MTP中验证器的开销,PIPO训练了一个轻量级的置信度头,用于判断预测的token是否应该被接受。通过利用On-Policy Distillation (OPD)与推测解码的拒绝采样准则的天然匹配,置信度头可以与OPD一同训练,而无需额外的成本。在AIME 2025、GPQA-Diamond、LiveCodeBench v6和LongBench v2等数据集上,使用Qwen3.5-4B和9B作为backbone的实验表明,PIPO在pass@4指标上相比常规解码提升高达+7.15,同时实现了高达2.64倍的首token延迟加速和2.07倍的每token延迟加速。

🔬 方法详解

问题定义:现有的大型语言模型在进行长链推理时,自回归解码成为主要的推理成本瓶颈。现有的加速方法主要集中在输入端的隐空间压缩或者输出端的多token预测,但是这两者通常是独立研究的。此外,输出端的多token预测方法通常需要一个昂贵的验证器来验证预测token的可靠性,增加了计算负担。

核心思路:PIPO的核心思路是将隐空间压缩和多token预测统一起来,把它们看作是镜像操作。具体来说,压缩器将两个输入token压缩成一个隐状态,而多token预测头则将一个隐状态展开成多个输出token。通过这种方式,可以更好地利用模型内部的表示,提高预测的效率和准确性。同时,通过引入置信度头来判断预测token的可靠性,从而避免了验证器的开销。

技术框架:PIPO的整体框架包括三个主要模块:隐空间压缩器、多token预测头和置信度头。隐空间压缩器负责将输入token序列压缩成隐状态序列;多token预测头负责根据隐状态预测多个输出token;置信度头负责判断预测token的可靠性。在训练过程中,PIPO使用On-Policy Distillation (OPD)来训练置信度头,使其能够准确地判断预测token的可靠性。

关键创新:PIPO的关键创新在于将隐空间压缩和多token预测统一起来,并通过置信度头消除验证器的开销。这种方法不仅提高了推理效率,还提高了预测的准确性。此外,PIPO还利用了On-Policy Distillation (OPD)来训练置信度头,使其能够准确地判断预测token的可靠性,而无需额外的训练成本。

关键设计:PIPO的关键设计包括:1) 将隐空间压缩器和多token预测头设计为镜像操作,以便更好地利用模型内部的表示;2) 引入置信度头来判断预测token的可靠性,从而避免了验证器的开销;3) 使用On-Policy Distillation (OPD)来训练置信度头,使其能够准确地判断预测token的可靠性,而无需额外的训练成本。具体的损失函数和网络结构等技术细节在论文中有详细描述,此处不再赘述。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

PIPO在AIME 2025、GPQA-Diamond、LiveCodeBench v6和LongBench v2等数据集上进行了实验,结果表明,PIPO在pass@4指标上相比常规解码提升高达+7.15,同时实现了高达2.64倍的首token延迟加速和2.07倍的每token延迟加速。这些结果表明,PIPO能够显著提高LLM的推理效率和准确性。

🎯 应用场景

PIPO方法可以应用于各种需要高效推理的大型语言模型应用场景,例如智能客服、机器翻译、代码生成等。通过加速推理过程,PIPO可以降低计算成本,提高用户体验,并促进LLM在资源受限设备上的部署。该方法在长文本处理和复杂推理任务中具有显著优势,有望推动LLM在实际应用中的普及。

📄 摘要(原文)

Long chain-of-thought reasoning has made autoregressive decoding the dominant inference cost of modern large language models. Existing methods target either the input side (latent compression) or the output side (speculative decoding and multi-token prediction, MTP), but the two lines of work have been pursued independently. Moreover, output-side methods must incur an expensive verifier pass to validate the unreliable draft tokens predicted by MTP. To address these issues, we propose \textbf{Pair-In, Pair-Out (PIPO)}, which unifies both sides by viewing a latent compressor and an MTP head as mirror-image operations: the compressor folds two input tokens into one latent representation, while the MTP head unfolds one hidden state into one additional output token. To remove the verifier cost without sacrificing reliability, PIPO trains a lightweight confidence head that decides whether draft tokens should be accepted. We observe that On-Policy Distillation (OPD) naturally matches the rejection-sampling criterion of speculative decoding, so the confidence head can be trained alongside OPD with negligible extra cost. Experiments on AIME 2025, GPQA-Diamond, LiveCodeBench v6, and LongBench v2 with Qwen3.5-4B and 9B backbones show that PIPO improves pass@4 over regular decoding by up to $+7.15$ points, while delivering up to $2.64\times$ first-token-latency and $2.07\times$ per-token-latency speedups.