Technical Summary: Pair-In, Pair-Out (PIPO) for Efficient LLMs

1. Problem Statement

Modern Large Language Models (LLMs) increasingly rely on long chains-of-thought (CoT) reasoning for complex tasks like mathematics and coding. While this paradigm improves accuracy, it creates a significant inference bottleneck: standard autoregressive decoding emits only one token per forward step, making the per-token cost the dominant latency factor for long reasoning traces.

Existing solutions address this from two independent directions:

• Input-side: Latent compression methods (e.g., CoLaR) merge multiple input tokens into a single continuous representation to reduce sequence length.
• Output-side: Multi-token prediction (MTP) and speculative decoding (e.g., EAGLE, Medusa) attempt to emit multiple tokens per step. However, these methods typically require an expensive "verifier" pass (a full forward pass of the large backbone) to validate draft tokens, which often negates the efficiency gains, especially at long context lengths. Furthermore, unverified MTP methods risk propagating unreliable drafts, degrading accuracy.

There is currently no unified framework that simultaneously exploits input compression and output multi-token prediction without incurring the heavy verifier cost.

2. Methodology: Pair-In, Pair-Out (PIPO)

PIPO proposes a unified framework that treats latent compression and multi-token prediction as mirror-image operations around the LLM backbone.

2.1 Core Architecture

The framework operates on token pairs rather than single tokens:

• Pair-In (Compression): A lightweight MLP compressor folds two consecutive input token embeddings ($x_{2i}, x_{2i+1}$) into a single latent representation ($z_i$). This halves the effective input length.
  $$z_i = f_\theta([x_{2i}; x_{2i+1}])$$
  The compressor is initialized such that $f_\theta([a; b]) \approx a + b$ to maintain the backbone's pretraining distribution.
• Pair-Out (Prediction): The backbone processes the compressed latent state and produces a hidden state ($h^b_i$).
  1. Backbone Token: The standard LM head predicts the next backbone token ($x_{2i+2}$).
  2. Draft Token: A Multi-Token Prediction (MTP) head, conditioned on the backbone hidden state and the predicted backbone token, predicts an additional draft token ($x_{2i+3}$).
  3. Confidence Head: A lightweight MLP predicts a confidence score ($c_i$) for the draft token.
     $$c_i = g_\phi([h^b_i; h^d_i])$$
  • Acceptance Logic: If $c_i \geq \tau_c$, the draft token is accepted, and the next input pair consists of two newly generated tokens. If rejected, the draft is replaced by a padding token ($x_{pad}$), and the next input pair is $(x_{2i+2}, x_{pad})$. This ensures the interface remains consistent regardless of acceptance.

2.2 Training Strategy

PIPO utilizes a two-stage training process to bridge the gap between training and inference:

1. Supervised Fine-Tuning (SFT):

   • The model is trained with a next-pair prediction objective using Cross-Entropy (CE) loss for both the backbone and draft tokens.
   • A confidence head is bootstrapped using Binary Cross-Entropy (BCE) loss, aligning the confidence score with the ground-truth draft token's probability.
   • Random Padding: During training, a fraction of draft positions are randomly replaced with padding tokens to expose the model to rejection scenarios.

2. On-Policy Distillation (OPD):

   • PIPO addresses the train-inference distribution shift by rolling out responses using the student's current policy.
   • Key Insight: The role of the "teacher" in OPD is functionally identical to the "verifier" in speculative decoding. Both judge the consistency of a draft against a stronger reference distribution.
   • Free Supervision: The acceptance probability for the confidence head is derived directly from the OPD loss calculation: $y_c = \min(p_t(x)/p_s(x), 1)$, where $p_t$ is the teacher (backbone) probability and $p_s$ is the student (draft) probability.
   • This allows the confidence head to be trained with the rejection-sampling criterion as a label without requiring any extra forward passes or a separate verifier model.

3. Key Contributions

• Unified Framework: PIPO unifies input-side latent compression and output-side multi-token prediction into a symmetric "Pair-In/Pair-Out" interface, effectively halving input length and doubling per-step output.
• Verifier Elimination: By observing that the OPD teacher and speculative verifier play the same role, PIPO trains a lightweight confidence head using the rejection-sampling ratio as a free label. This replaces the expensive per-step backbone verifier pass with a single MLP forward pass.
• Efficiency and Accuracy: The method achieves significant speedups in Time-to-First-Token (TTFT) and Time-Per-Output-Token (TPOT) while improving pass@4 accuracy on complex reasoning benchmarks.

4. Experimental Results

The authors evaluated PIPO on Qwen3.5-4B and Qwen3.5-9B backbones across four benchmarks: AIME 2025, GPQA-Diamond, LiveCodeBench v6, and LongBench v2.

• Accuracy: PIPO + OPD achieved the highest pass@4 scores among all methods.
  • On Qwen3.5-9B, PIPO improved pass@4 by +7.15 points over regular autoregressive decoding.
  • On Qwen3.5-4B, the improvement was +3.83 points over the best baseline.
  • The advantage widened with longer context budgets (up to 32K slots), as PIPO fits more complete reasoning chains into the same budget.
• Efficiency:
  • TTFT: Up to 2.64× speedup (e.g., at 128K input length) due to the compressed prefill phase.
  • TPOT: Up to 2.07× speedup because two tokens are emitted per backbone forward pass.
• Ablation Studies:
  • Non-linear compressors are essential; linear compressors caused a significant drop in performance.
  • Initializing the compressor as an additive sum ($a+b$) is critical for stability.
  • The confidence head effectively rejects unreliable drafts, outperforming random rejection baselines and finding a "sweet spot" (threshold $\tau_c \approx 0.95$) that maximizes pass@4.

5. Significance and Claims

The paper claims that PIPO represents a shift in how efficient LLM decoding is approached:

1. Symmetry: It demonstrates that input compression and output expansion are not competing strategies but complementary, mirror-image operations that can be combined.
2. Cost Amortization: It successfully amortizes the cost of verification. Instead of paying a heavy inference cost (full backbone pass) for every step, the verification logic is distilled into a lightweight training signal (the confidence head), turning a recurring inference cost into a one-time training signal.
3. Practicality: The framework relies mostly on off-the-shelf components (pretrained MTP heads, standard backbones) and requires only small additional modules (compressor and confidence head), making it feasible to deploy on existing strong LLMs.

The authors note limitations, including the focus on pair-wise processing (rather than larger compression factors), the restriction to text-only models, and the evaluation primarily on tasks with verifiable answers (math/coding) rather than open-ended generation. However, they argue that PIPO offers a robust path to reducing the inference costs of long-chain reasoning without sacrificing reliability.