Robust Length Prediction: A Perspective from Heavy-Tailed Prompt-Conditioned Distributions

Jing Wang; Yu-Yang Qian; Ke Xue; Chao Qian; Peng Zhao; Zhi-Hua Zhou

cs.LGApr 9, 2026

Robust Length Prediction: A Perspective from Heavy-Tailed Prompt-Conditioned Distributions

Authors: Jing Wang, Yu-Yang Qian, Ke Xue, Chao Qian, Peng Zhao, Zhi-Hua Zhou

Organizations: State Key Laboratory of Novel Software Technology, Nanjing University, Nanjing 210023, China · School of Artificial Intelligence, Nanjing University, Nanjing 210023, China

Abstract

Output-length prediction is important for efficient LLM serving, as it directly affects batching, memory reservation, and scheduling. For prompt-only length prediction, most existing methods use a one-shot sampled length as the label, implicitly treating each prompt as if it had one true target length. We show that this is unreliable: even under a fixed model and decoding setup, the same prompt induces a prompt-conditioned output length distribution, not a deterministic scalar, and this distribution is consistent with heavy-tailed behavior. Motivated by this, we cast length prediction as robust estimation from heavy-tailed prompt-conditioned length distributions. We propose prompt-conditioned length distribution (ProD) methods, which construct training targets from multiple independent generations of the same prompt. Two variants are developed to reuse the served LLM's hidden states: \mbox{ProD-M}, which uses a median-based target for robust point prediction, and ProD-D, which uses a distributional target that preserves prompt-conditioned uncertainty. We provide theoretical justifications by analyzing the estimation error under a surrogate model. Experiments across diverse scenarios show consistent gains in prediction quality.

Figures & tables

Panel of Figure 1 — Figure 1: Key observations about prompt-conditioned output length. Figure (a) summarizes prompt-level median-centered noise radius across the Math, Coding, LongSequence, and Chat scenarios via repeated-sampling Median-MAE. Figures (b) and (c) show representative repeated-sampling length distributions for Math, Coding, and LongSequence prompts under Qwen and Llama. Additional per-setting supporting plots are provided in Appendix A.4 .

Method	Qwen					Llama
Method	Math	Coding	LongSeq	Chat	Avg	Math	Coding	LongSeq	Chat	Avg
Constant Median	59.59	56.41	146.11	264.90	131.75	32.07	37.83	95.69	213.31	94.73
$S^{3}$	41.60	57.28	72.21	185.84	89.23	24.41	41.28	50.01	142.28	64.50
TRAIL-mean	44.04	44.70	72.70	143.50	76.24	26.29	37.93	50.99	116.84	58.01
TRAIL-last	35.13	36.63	57.68	127.47	64.23	24.51	35.74	41.04	95.03	49.08
EGTP	49.50	57.14	69.84	276.27	113.19	28.25	37.98	46.73	148.08	65.26

Table 1: Prompt-only length prediction results. Test MAE (tokens) between each method’s prediction and the prompt-level median output length under

16

-sample repeated sampling. Noise Radius is the prompt-level median-centered noise radius. Lower is better; the best trainable result per column is bold . The exact protocol is in Appendix A.2 .

Method	Qwen				Llama
Method	Math	Coding	LongSeq	Chat	Math	Coding	LongSeq	Chat
$S^{3}$	54.79 $\pm$ 0.8	63.38 $\pm$ 1.1	96.00 $\pm$ 6.6	201.33 $\pm$ 4.7	33.31 $\pm$ 1.3	52.57 $\pm$ 2.9	82.97 $\pm$ 3.8	154.78 $\pm$ 4.1
TRAIL-mean	59.74 $\pm$ 3.0	59.21 $\pm$ 2.8	99.58 $\pm$ 6.2	174.14 $\pm$ 5.6	34.59 $\pm$ 1.4	46.71 $\pm$ 3.1	87.13 $\pm$ 4.4	137.98 $\pm$ 6.5
TRAIL-last	54.11 $\pm$ 8.2	56.44 $\pm$ 7.4	88.07 $\pm$ 4.9	154.45 $\pm$ 6.0	32.48 $\pm$ 1.1	45.29 $\pm$ 4.1	76.75 $\pm$ 2.9	112.84 $\pm$ 5.3
EGTP	67.30 $\pm$ 7.5	62.89 $\pm$ 1.0	96.84 $\pm$ 6.1	278.12 $\pm$ 5.3	36.39 $\pm$ 1.6	49.55 $\pm$ 2.9	81.37 $\pm$ 3.6	163.95 $\pm$ 5.0
ProD-M	47.10 $\pm$ 0.8	43.43 $\pm$ 1.1	86.00 $\pm$ 5.2	147.18 $\pm$ 4.6	29.32 $\pm$ 1.0	40.86 $\pm$ 3.8	73.22 $\pm$ 3.2	109.73 $\pm$ 3.9

Table 2: Single-sample supervision ablation (single-label evaluation). All predictors are trained with one sampled output length per prompt and evaluated against the same single-label target. We report mean

\pm

std test MAE over

8

trials. Lower is better; the best per column is bold .

Method	Qwen				Llama
Method	Math	Coding	LongSeq	Chat	Math	Coding	LongSeq	Chat
$S^{3}$	43.91 $\pm$ 0.5	57.56 $\pm$ 0.2	75.65 $\pm$ 2.0	186.75 $\pm$ 3.6	25.76 $\pm$ 0.7	44.69 $\pm$ 2.8	55.57 $\pm$ 1.2	146.21 $\pm$ 3.8
TRAIL-mean	49.84 $\pm$ 2.8	52.82 $\pm$ 2.6	80.64 $\pm$ 3.3	158.33 $\pm$ 5.3	27.16 $\pm$ 0.3	36.61 $\pm$ 1.8	60.07 $\pm$ 3.7	129.79 $\pm$ 4.4
TRAIL-last	42.46 $\pm$ 9.6	49.47 $\pm$ 8.0	67.91 $\pm$ 1.2	138.46 $\pm$ 6.3	24.89 $\pm$ 0.2	35.23 $\pm$ 1.6	48.11 $\pm$ 0.9	102.48 $\pm$ 2.6
EGTP	58.56 $\pm$ 9.2	56.90 $\pm$ 0.3	77.14 $\pm$ 3.0	266.55 $\pm$ 3.1	29.70 $\pm$ 1.1	40.80 $\pm$ 2.1	53.55 $\pm$ 1.0	156.07 $\pm$ 4.5
ProD-M	33.56 $\pm$ 0.4	35.01 $\pm$ 1.2	65.44 $\pm$ 1.3	130.60 $\pm$ 2.6	20.79 $\pm$ 0.2	29.05 $\pm$ 1.0	43.45 $\pm$ 1.4	98.95 $\pm$ 1.3

Table 3: Single-sample supervision ablation (median-target evaluation). All predictors are trained with one sampled output length per prompt, but evaluated against the

16

-sample repeated-sampling median target. We report mean

\pm

std test MAE over

8

trials. Lower is better; the best per column is bold .

Panel of Figure 2 — Figure 2: Budget fairness: test MAE vs. repeat sampling number under a fixed inference budget. As the repeat sampling number $k$ increases, only $\lceil B/k\rceil$ unique training prompts are retained. ProD-M and ProD-D are the repeated-sampling predictors; TRAIL-Last is the full-coverage single-sample baseline. All curves report mean $\pm$ std over $8$ trials. Coding is deferred to the appendix.

Appendix figures & tables4 assets

Supplementary material from the paper’s appendix.

Appendix

Panel of Figure 3 — Figure 3: System prompt reduces output-length randomness and MAE noise radius. Qwen2.5-7B-Instruct on 500 MBPP prompts with 16 trials per prompt (8 with system prompt, 8 without). We compare per-prompt mean length, length variance, and two MAE-style dispersion measures (Mean-MAE / Median-MAE). Paired plots (Figure 3(a) and Figure 3(b) ) and shift summaries (Figure 3(d) , Figure 3(e) , and Figure 3(g) ) show that adding the system prompt typically shortens outputs and often reduces variance. MAE distributions (Figure 3(c) and Figure 3(f) ) shift left and become more concentrated, indicating a lower and tighter per-prompt MAE baseline.

Panel of Figure 4 — Figure 4: Prompt-percentile noise-floor waterfalls. Each curve sorts prompts by prompt-level Median-MAE within the corresponding model and plots the values on a log-scale y-axis against prompt percentile. The plotting floor is used only to visualize zero-valued prompts on the log scale: $300$ prompts for Qwen and $105$ prompts for Llama.

Panel of Figure 5 — Figure 5: Qwen per-setting light/heavy overlays. For each setting, light5 and heavy5 denote the five prompts with the smallest and largest $\max(\mathrm{length})/\mathrm{median}(\mathrm{length})$ among the ten repeated-sampling prompts.

Panel of Figure 6 — Figure 6: Llama per-setting light/heavy overlays. For each setting, light5 and heavy5 denote the five prompts with the smallest and largest $\max(\mathrm{length})/\mathrm{median}(\mathrm{length})$ among the ten repeated-sampling prompts.

LLM Serving

Large language model serving studies how to deliver model responses with low latency, high throughput, and efficient use of compute and memory.

See all 207 papers

OUTLETS: Output-Length Prediction from Speculative Decoding Backbones
Sep 1, 2026Weihuang Wen, Yingying Liu, Yichuan Liu +5LLM Serving LLM Inference Scheduling
The Chinese University of Hong Kong, Shenzhen · The University of Hong Kong · Huawei Technologies Co., Ltd.
When Entropy Is Not Enough: Reclaiming Lost Semantics in LLM Output Length Prediction
Aug 16, 2026Feiyang Ren, Shengtao Wen, Yu Tian +5LLM Serving LLM Inference Scheduling
MIIT Key Laboratory of Pattern Analysis and Machine Intelligence, College of Computer Science and Technology, Nanjing University of Aeronautics and Astronautics · Tsinghua University · Nanjing University of Information Science and Technology
How Much is Left? LLMs Linearly Encode Their Remaining Output Length
Jul 6, 2026Mohamed Amine Merzouk, Dmitri Carpov, Mirko Bronzi +2LLM Interpretability
Mila, Quebec AI Institute · McGill University · LawZero

Sep 1, 2026J/K move · Enter open · S save

OUTLETS: Output-Length Prediction from Speculative Decoding Backbones

Weihuang Wen, Yingying Liu, Yichuan Liu, Wenqi Zeng+4

The Chinese University of Hong Kong, Shenzhen · The University of Hong Kong · Huawei Technologies Co., Ltd.

LLM Serving LLM Inference Scheduling

Open paper AskPDF arXiv

The heavy-tailed distribution of output lengths in Large Language Model (LLM) serving poses major challenges for resource provisioning and cluster scheduling. Although output-length prediction can mitigate these issues, existing approaches have key drawbacks: external proxy models add substantial latency and often have limited fidelity, whereas internal state-based methods are efficient but rely on shallow probes of current model states. We identify a structural connection between speculative decoding (SD) and length prediction: latent representations produced by the draft decoder in advanced frameworks (e.g., EAGLE-3) encode signals that are predictive of generation length. Building on this insight, we introduce OUTLETS (Output-Length Prediction from Speculative Decoding Backbones), which repurposes the speculative backbone as a trajectory-aware length predictor. When its draft representations are already computed for speculative decoding, OUTLETS adds only a lightweight regression head and achieves lower MAE than the evaluated methods. Under saturated disaggregated serving, OUTLETS predictions enable standard scheduling policies to prioritize shorter requests and distribute requests more evenly across decoding instances, reducing short-request P99 latency by 34.8%.

Abstract

LLM Serving

Explore similar work

OUTLETS: Output-Length Prediction from Speculative Decoding Backbones

When Entropy Is Not Enough: Reclaiming Lost Semantics in LLM Output Length Prediction

How Much is Left? LLMs Linearly Encode Their Remaining Output Length