PIPO: Pair‑In Pair‑Out Architecture Cuts Input Length in Half and Doubles Token Output for LLMs

Background and Motivation

Long‑chain reasoning in large language models (LLMs) often requires tens of thousands of tokens, causing the first‑token latency and per‑token compute cost to become a major bottleneck. The growing context length makes the prefilling step increasingly expensive, and standard autoregressive decoding emits only one token per step.

Limitations of Existing Approaches

Current methods address either the input side (e.g., Coconut, Soft Thinking, CoLaR) by compressing tokens into latent representations, or the output side (e.g., speculative decoding, EAGLE, MTP heads) by predicting multiple tokens but still relying on a verifier that repeatedly runs the full backbone.

PIPO Architecture

PIPO unifies both sides with a pair‑level interface:

Pair‑In compressor : concatenates two consecutive token embeddings and maps them via a non‑linear MLP to a single latent that the backbone can consume.

Pair‑Out MTP head : after the backbone predicts the primary token, the MTP head uses the hidden state and the primary token embedding to predict an additional draft token.

This symmetry reduces the effective input length by roughly 50 % and allows each forward step to produce up to two tokens, effectively doubling throughput.

Training Procedure

SFT stage : a next‑pair prediction objective trains the model to output both primary and draft tokens. Randomly replacing some draft positions with PAD embeddings teaches the model to handle rejected drafts.

OPD stage : On‑Policy Distillation (OPD) rolls out the current policy, removes PAD tokens, feeds the sequence to an uncompressed teacher, and aligns the student distribution to the teacher via reverse‑KL. The teacher‑student probability ratio provides free supervision for the confidence head, eliminating the need for a verifier at inference time.

Experimental Setup

Benchmarks: AIME 2025 (math competition), GPQA‑Diamond (graduate‑level science MCQs), LiveCodeBench v6 (coding problems), LongBench v2 (long‑context reasoning). Backbones: Qwen3.5‑4B and Qwen3.5‑9B, all evaluated under a 32K‑slot response budget. Compared methods: Regular decoding, MTP, EAGLE‑2, PIPO‑SFT, and PIPO + OPD.

Results

Pass@4 scores:

Qwen3.5‑4B: PIPO + OPD reaches 67.31, a +3.83 gain over the best baseline.

Qwen3.5‑9B: PIPO + OPD reaches 75.18, a +7.15 gain over the best baseline.

Efficiency:

TTFT (time‑to‑first‑token) improves up to 2.64× (20.3 s → 7.69 s) on 128K inputs.

TPOT (time‑per‑output‑token) improves up to 2.07× on 2K inputs and remains ~1.98× at 128K inputs.

Ablation Studies

Replacing the non‑linear MLP compressor with a linear layer drops pass@4 from 65.01 to 60.38, showing the need for non‑linearity.

Limiting SFT data to only the shortest correct answer reduces pass@4 by 5.14 points, indicating the importance of diverse correct trajectories.

Initializing the compressor with a simple sum (a + b) is critical; random initialization lowers pass@4 from 65.01 to 57.73.

The confidence head outperforms a random‑acceptance baseline across all PAD ratios, achieving 61.93 pass@4 at a PAD ratio of ~0.6 versus 59.64 for random.

Implications

PIPO demonstrates that next‑generation reasoning LLMs should move beyond single‑token input/output loops and repeated verifier calls. By pairing latent compression with multi‑token prediction and distilling verifier knowledge into a lightweight confidence head, PIPO offers a scalable path for efficient long‑chain reasoning and can be extended to larger compression factors, bigger models, open‑ended generation, and multimodal scenarios.