Predicting LLM Output Length via Entropy-Guided Representations

Huanyi Xie, Yubin Chen, Liangyu Wang, Lijie Hu, Di Wang

Keywords: Large Language Models, Length Prediction, Progressive Length Prediction

TL;DR: We teach LLMs to predict their own output length by progressively refining estimates from their internal activations, reducing padding waste for both standard inference and dynamic RL training.

Abstract: The long-tailed distribution of sequence lengths in LLM serving and reinforcement learning (RL) sampling causes significant computational waste due to excessive padding in batched inference. Existing methods rely on auxiliary models for static length prediction, but they incur high overhead, generalize poorly, and fail in stochastic "one-to-many" sampling scenarios. We introduce a lightweight framework that reuses the main model's internal hidden states for efficient length prediction. Our framework features two core components: 1) Entropy-Guided Token Pooling (EGTP), which uses on-the-fly activations and token entropy for highly accurate static prediction with negligible cost, and 2) Progressive Length Prediction (PLP), which dynamically estimates the remaining length at each decoding step to handle stochastic generation. To validate our approach, we build and release ForeLen, a comprehensive benchmark with long-sequence, Chain-of-Thought, and RL data. On ForeLen, EGTP achieves state-of-the-art accuracy, reducing MAE by 29.16\% over the best baseline. Integrating our methods with a length-aware scheduler yields significant end-to-end throughput gains. Our work provides a new technical and evaluation baseline for efficient LLM inference.

Primary Area: infrastructure, software libraries, hardware, systems, etc.

Submission Number: 7421