Track papers across labs, topics, and trends.
Search papers, labs, and topics across Lattice.
The paper addresses the problem of inaccurate output length prediction in LLMs by highlighting that prompts induce heavy-tailed distributions of output lengths, rather than a single deterministic length. To address this, they introduce Prompt-conditioned Distribution (ProD) methods, which generate multiple outputs per prompt to construct more robust training targets. Experiments demonstrate that ProD methods, particularly ProD-D which preserves prompt-conditioned uncertainty, consistently improve length prediction accuracy across various scenarios.
LLMs don't have a single "true" output length for a given prompt; they exhibit heavy-tailed distributions of possible lengths, and modeling this distribution improves length prediction.
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.