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

2026-05-26

AI Summary

This work addresses the high computational overhead of autoregressive decoding in large language models during long-chain reasoning. Existing approaches—input compression and multi-token prediction—are typically treated in isolation, with the latter often relying on costly verification mechanisms. To overcome these limitations, the authors propose PIPO, a unified framework that symmetrically integrates input compression and multi-token output prediction through a mirrored architecture. PIPO introduces a lightweight confidence head to replace conventional verifiers and incorporates on-policy distillation tailored to the rejection-sampling mechanism of speculative decoding. Experiments demonstrate that PIPO achieves up to a 7.15-point improvement in pass@4 accuracy across multiple benchmarks, while accelerating first-token latency by 2.64× and per-token latency by 2.07×.

Abstract

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.

Problem

• autoregressive decoding
• latent compression
• multi-token prediction
• speculative decoding
• verifier cost

Innovation

• Pair-In Pair-Out
• Multi-Token Prediction
• Latent Compression
• Speculative Decoding
• On-Policy Distillation