Robust Length Prediction: A Perspective from Heavy-Tailed Prompt-Conditioned Distributions
This addresses efficiency issues in LLM serving for practitioners by providing more robust length estimation, though it is incremental as it builds on existing prediction methods.
The paper tackles the problem of unreliable output-length prediction for LLM serving by showing that prompts induce heavy-tailed length distributions rather than deterministic lengths, and proposes ProD methods that use multiple generations to improve prediction, achieving consistent gains in experiments.
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 \emph{prompt-conditioned output length distribution}, not a deterministic scalar, and this distribution is consistent with \emph{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.