PaperlayerView exact revision on arXiv

Source-linked AI summary

Predicting LLM Output Length via Entropy-Guided Representations

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

arXiv:2602.11812v2cs.AI

TL;DR

Variable output lengths waste computation through padding, and auxiliary prompt-only predictors face overhead, generalization, and stochastic-sampling challenges. The paper reuses LLM activations for static and progressive length prediction and evaluates the approach on ForeLen. EGTP reduces MAE by 29.16% against the strongest baseline, and the methods yield significant inference-throughput improvements with a length-aware scheduler.

  • Problem

    Variable output lengths cause padding waste in batched inference, while auxiliary prompt-based predictors incur overhead and struggle with complex or stochastic generation.

    Source: §1 · p. 2

  • Method

    The framework reuses the LLM’s internal activations for entropy-guided static prediction with EGTP and per-step remaining-length updates with PLP.

    Source: §1 · pp. 2–3

  • Results

    29.16% lower MAE than the strongest baseline, averaged across tested models on ForeLen, with significant end-to-end inference throughput improvements.

    Source: §1 · p. 3

  • Takeaways & Limitations

    Length-aware scheduling paired with the framework yields significant inference-throughput improvements, and ForeLen provides evaluation data spanning long sequences, CoT, and RL sampling.

    Source: §1 · p. 3

Abstract

from arXiv · show

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.

1 INTRODUCTION

Variable output lengths waste computation through padding, while prompt-only auxiliary predictors struggle with stochastic generation, accuracy, generalization, and overhead. The paper reuses LLM activations for static and progressive length prediction, evaluates them on ForeLen, and reports improved accuracy and throughput.

1 INTRODUCTION

  • Length variation within batches forces shorter sequences to be padded to the longest, wasting accelerator computation; ForeLen has a wider, longer-tailed distribution, and RL samples vary greatly even for identical prompts.The figure contrasts ForeLen with LMSYS and illustrates same-prompt variation in online RL.Source: §1 · p. 2, Figure 1 · p. 2
  • Prompt-only auxiliary predictors can be unreliable for stochastic one-to-many generation, less accurate in complex settings, and costly because they run separately from the main LLM.Their training benchmarks often contain few long sequences and little complex reasoning.Source: §1 · p. 2
  • EGTP uses token entropy to pool the main LLM’s on-the-fly activations, capturing informative prompt signals for static length prediction with minimal additional cost.The approach reuses internal states rather than deploying a separate predictor.Source: §1 · pp. 2–3
  • PLP updates its estimate of remaining tokens at each decoding step, enabling length-aware scheduling for unpredictable generation such as RL sampling.Its iterative forecasts use current activations and adapt to the generated sequence.Source: §1 · pp. 2–3
  • 29.16% lower MAE than the strongest baseline and 55.09% lower than SSJF-Reg, averaged across tested models; the method also yields significant end-to-end inference throughput improvements.These results are reported on ForeLen.Source: §1 · p. 3
  • ForeLen is a comprehensive length-prediction benchmark covering long sequences, Chain-of-Thought, and reinforcement-learning sampling data.It is designed to test predictors in challenging settings.Source: §1 · p. 3

2 RELATED WORK

Prior work improves serving throughput but leaves padding waste within active batches unresolved, while auxiliary length predictors add deployment and per-request inference costs. PLP addresses stochastic one-to-many sampling with forecasts updated along each generated sequence.

Efficient LLM Serving and Inference Optimization

  • Continuous batching and PagedAttention reduce inter-request idle time and memory fragmentation, but do not address padding waste within a running batch.Length prediction can complement these serving optimizations by enabling more homogeneous batches.Source: §2 · p. 3

LLM Response Length Prediction

  • Prompt-based auxiliary predictors require separate models to be trained, deployed, and run for every request; EGTP instead reuses the LLM’s on-the-fly activations.The related-work passage characterizes EGTP as adding negligible computation.Source: §2 · p. 3

LLM Inference and Sampling in Reinforcement Learning

  • Stochastic RL sampling produces widely varying lengths for the same prompt, so PLP updates its remaining-token estimate from current activations at each autoregressive step.The passage also notes that a standardized benchmark for this one-to-many setting was absent.Source: §2 · p. 3

3 METHOD

The method combines entropy-guided static prediction, soft-label length regression, and progressive prediction for stochastic generation. It reuses model activations to form length estimates and updates forecasts as tokens are generated.

3 METHOD

  • The framework pairs EGTP for static prediction with PLP for dynamically forecasting length in stochastic environments.EGTP reuses internal activations and token entropy; PLP progressively updates estimates during generation.Source: §3 · p. 4

3.1 ENTROPY-GUIDED TOKEN POOLING (EGTP)

  • EGTP pools prompt hidden states with entropy-derived weights to emphasize tokens that are more informative for predicting length.It computes token entropies, transforms them into softmax weights, and uses those weights to form a single representation.Source: §3.1 · pp. 4–5
  • Pearson’s r = 0.451 between token entropy and importance supports entropy-guided pooling in a BERT linear-head experiment on 10k LMSYS samples.The figure shows average importance across five equal-width entropy bins and a positive entropy–importance association.Source: §3.1 · p. 4, Figure 2 · p. 5

3.2 LENGTH REGRESSION VIA SOFT LABEL DISTRIBUTION

  • The length-regression head uses soft labels to combine outlier robustness with awareness of distance between true and predicted lengths.It predicts a distribution over bins and derives a continuous estimate from that distribution; training combines cross-entropy and MSE losses.Source: §3.2 · pp. 5–6

3.3 PROGRESSIVE LENGTH PREDICTION (PLP)

  • PLP predicts remaining length at each decoding step using the prompt representation and hidden states of generated tokens.It is trained over all timesteps, iteratively refining forecasts to support scheduling adjustments and resource utilization.Source: §3.3 · p. 6

4 EXPERIMENTS

EGTP outperforms comparison methods on LMSYS and ForeLen, and its integration with an SJF scheduler leads on end-to-end metrics. The experiments also evaluate ForeLen’s data coverage and ablate progressive prediction and pooling choices.

4.1 DATASET CONSTRUCTION

  • ForeLen covers long-sequence and complex-reasoning generation as well as dynamic RL sampling, with source-dataset splits keeping validation and test prompts unseen in training.Its RL data comes from GRPO training and records grouped samples with K=4.Source: §4.1 · pp. 6–7

4.2 EXPERIMENT SETTING

  • Experiments compare baselines on LMSYS and ForeLen, evaluating prediction with MAE and end-to-end serving with throughput, JCT, and padding ratio.The comparison includes SSJF-Reg, SSJF-MC, S3, PiA, TPV, TRAIL, and LTR-C.Source: §4.2 · p. 7

4.3 MAIN RESULTS: PREDICTION ACCURACY

  • EGTP has the lowest LMSYS MAE for GPT-4 (87.32) and Claude-2 (68.33), and reduces MAE versus TRAIL on ForeLen across four models.ForeLen comparisons: Qwen2.5 3B, 146.42 to 110.75 (24.4%); Qwen2.5 7B, 137.96 to 103.47 (25.0%); Llama3.2 1B, 151.80 to 105.08 (30.8%); Llama3.2 3B, 157.88 to 101.53 (35.7%).Source: §4.3 · p. 7, Table 1 · p. 8

4.4 END-TO-END SYSTEM PERFORMANCE

  • In Long Sequence and Reasoning workloads, EGTP leads the SJF-scheduled system on throughput and JCT, with padding ratios of 0.18 and 0.09, respectively.The Long Sequence padding ratio is 0.18 versus TRAIL’s 0.51; in Reasoning, 0.09 is a 36% relative reduction in wasted computation versus LTR-C.Source: §4.4 · p. 8, Table 2 · p. 9

4.5 ABLATION STUDY

  • PLP improves MAE across RL, Reasoning, and Long Seq tasks; RL shows the largest drop, from 95.24 to 80.85.Across tasks, gains are rapid in the first decoding steps and then stabilize.Source: §4.3 · pp. 7–8, Figure 3 · p. 9
  • EGTP achieves average MAE 103.47 versus 119.69 for Max Pooling, and an MAE of 81.60 on Long Sequence.The ablation reports EGTP outperforming pooling baselines across all tasks.Source: Table 3 · p. 9, §4.5 · p. 9

5 CONCLUSION

The paper proposes length prediction from an LLM’s internal activations, with EGTP for static estimation and PLP for dynamic environments. Validation on ForeLen reports superior prediction accuracy and supports the conclusion that length signals are encoded in hidden states.

5 CONCLUSION

  • The framework predicts sequence length by reusing the model’s internal activations.The paper presents this as avoiding the overhead and generalization failures of separate auxiliary predictors.Source: §5 · p. 10
  • EGTP handles static estimation, while PLP provides progressive prediction in dynamic environments.Source: §5 · p. 10
  • Validation on ForeLen shows superior prediction accuracy, supporting the conclusion that LLM hidden states encode signals for determining output length.Source: §5 · p. 10

B EVALUATION METRICS

The section defines MAE for predictor accuracy and throughput, Job Completion Time (JCT), and padding ratio for end-to-end system performance.

B.1 PREDICTOR PERFORMANCE METRIC

  • MAE measures the average absolute difference between predicted and actual lengths; lower MAE indicates a more accurate predictor.The metric is defined over n samples.Source: §B.1 · p. 16

B.2 END-TO-END SYSTEM PERFORMANCE METRICS

  • Throughput measures the number of jobs processed per unit time, indicating system efficiency.Source: §B.2 · p. 16
  • JCT is the elapsed time from job submission to completion; average JCT gauges system responsiveness.Source: §B.2 · p. 16
  • Padding ratio quantifies padding overhead as the padded resource portion relative to the actual required resource; lower is better.Source: §B.2 · p. 16

C.1 DATASET STATISTICS

ForeLen covers long-sequence, reasoning, and reinforcement-learning data, with varied prompt and response lengths. The appendices also illustrate example tasks and prompting formats, and examine the loss-weighting choice.

C.1 DATASET STATISTICS

  • Long-Sequence, Reasoning, and RL account for 4.7%, 2.3%, and 93.0% of unique prompts, respectively; source datasets are split 3:1:1.Table 4 reports unique-prompt counts by dataset and split.Source: §C.1 · p. 17, Table 4 · p. 17, Figure 4 · p. 17

C.2 PROMPT AND RESPONSE LENGTH DISTRIBUTIONS

  • Prompt lengths range from a few hundred characters for MATH and GSM8K to nearly 6,000 for LongBench and ZeroSCROLLS, while response lengths also vary widely.The benchmark includes outputs from single-word answers to detailed reasoning and code solutions.Source: §C.2 · p. 18, Table 5 · p. 18, Figure 5 · p. 18

C.3 BENCHMARK DATASET EXAMPLES

  • Representative examples illustrate long-context synthesis, multi-step reasoning, and an RL example drawn from GRPO training data.Source: §C.3 · p. 19

C.4 SYSTEM PROMPTS

  • Prompts are direct for long-context QA and more structured for reasoning and coding, with strictness increasing with task complexity and evaluation requirements.Source: §C.4 · p. 21

D.1 LOSS WEIGHTING PARAMETER

  • λ = 0.99 yields the best LMSYS performance, with MAE approximately 68, versus approximately 70 for λ = 1 and 82–84 for λ = 0.The authors attribute the pure-MSE result to scale disparity causing gradient domination; the small MSE contribution at λ = 0.99 complements cross-entropy.Source: Figure 7 · p. 24, §D.1 · p. 24

D.2 ANALYSIS OF LOSS FUNCTION VARIANTS

The appendices compare EGTP with alternative loss designs and baselines, assess final-layer representations, and test generalization to super-long sequences.

D.2 ANALYSIS OF LOSS FUNCTION VARIANTS

  • EGTP outperforms alternative loss designs across LongSeq, Reasoning, and RL, although log-scale regression is competitive on Reasoning.The alternatives include log-scale regression, label smoothing, and label noise.Source: §D.2 · pp. 24–25, Table 6 · p. 25

D.3 LAYER SELECTION STRATEGY

  • On GSM8K, final-layer Layer 24 has RMSE 92.95 ± 2.80 versus Layer 12’s 93.47 ± 1.50; Layer 12 has lower MAE, but Layer 24 is retained for simplicity.Layer 12 MAE is 70.26 ± 1.20 and Layer 24 MAE is 73.93 ± 2.10.Source: §D.3 · p. 25, Table 7 · p. 25

E.1 BASELINES

  • The baselines include auxiliary-model regressors and classifiers, prompt-based prediction, and methods using internal embeddings or progress estimates.Named methods include SSJF, S3, PiA, TPV, TRAIL, and LTR.Source: §E.1 · p. 26

E.2 GENERALIZATION TO SUPER-LONG SEQUENCES

  • On sequences exceeding 17k tokens, EGTP achieves the lowest MAE (195.89) and RMSE (257.49) on euclaise/writingprompts.The authors describe this as strong generalization to lengths beyond the training distribution.Source: §E.2 · p. 26, Table 8 · p. 26

E.3 COMPUTATIONAL EFFICIENCY ANALYSIS

The efficiency analysis measures predictor overhead in isolation and compares latency and memory use against alternatives.

E.3 COMPUTATIONAL EFFICIENCY ANALYSIS

  • The isolated overhead evaluation used GSM8K on a single NVIDIA RTX 4090 with 24 GB of memory.The analysis complements end-to-end system results.Source: §E.3 · p. 26
  • EGTP inference takes approximately 0.66 ms, compared with 2–12 ms for SSJF and TRAIL and approximately 60 ms for PiA.PiA’s recorded time covers only generation of its length-prediction tokens.Source: §E.3 · p. 27, Table 9 · p. 27, Table 10 · p. 27
  • EGTP uses 5–7 MB of VRAM for its lightweight head, versus 238–270 MB for auxiliary-model baselines.Source: §E.3 · p. 27, Table 9 · p. 27, Table 10 · p. 27

E.5 LENGTH PREDICTION PERFORMANCE ON OTHER DATASETS

EGTP has the lowest average MAE for both Qwen2.5 model sizes across Long Sequence, Reasoning, and RL, with TRAIL second. It therefore outperforms the compared baselines in these evaluations.

E.5 LENGTH PREDICTION PERFORMANCE ON OTHER DATASETS

  • EGTP achieves the lowest average MAE for both Qwen2.5 0.5B and 1.5B across Long Sequence, Reasoning, and RL.TRAIL consistently ranks second, but trails EGTP by a notable margin.Source: §E.5 · p. 27
  • TRAIL secures second-best position across the comparisons, behind EGTP by a notable margin.Source: §E.5 · p. 27
  • The comparison evaluates length prediction methods using MAE for Qwen2.5 0.5B and 1.5B across Long Sequence, Reasoning, and RL scenarios.Source: §E.5 · p. 27

E.6 DETAILED PREDICTION ACCURACY RESULTS

The detailed results cover length prediction across Long-Sequence and Reasoning scenarios, with prompts drawn from several datasets and responses generated by multiple LLM families.

E.6 DETAILED PREDICTION ACCURACY RESULTS

  • The detailed comparisons use prompts from LongBench, ZeroSCROLLS, and IFEval, with responses generated by Qwen2.5, Llama3.2, and DeepSeek-R1-Distill models.Source: §E.6 · p. 27
  • Tables 12 and 13 report Long-Sequence performance, with Table 13 providing MAE only.Source: Table 12 · p. 28, Table 13 · p. 28
  • Tables 14 and 15 report length-prediction performance in the Reasoning scenario.Source: Table 14 · p. 29, Table 15 · p. 31

E.7 VISUALIZATION OF THE GRPO TRAINING PROCESS

Across the GRPO training figures, rewards improve stably and consistently across the reported datasets. The visualizations cover code, mathematical, and scientific reasoning tasks and multiple model sizes.

E.7 VISUALIZATION OF THE GRPO TRAINING PROCESS

  • Figures 8, 9, 12, and 13 show reward curves for CRUXEval, MBPP, LiveCodeBench, and GSM8K, respectively.Source: §E.7 · p. 31, Figure 8 · p. 32, Figure 9 · p. 32
  • Figures 10 and 11 show GRPO training processes on MMLU-STEM and MATH, respectively.Source: §E.7 · p. 31, Figure 10 · p. 33, Figure 11 · p. 33
  • Across all figures, the GRPO training process shows stable and consistent improvement in rewards.The experiments train four Qwen2.5 sizes and two Llama3.2 sizes; smaller models use larger batch sizes, which can change the number of training steps.Source: §E.7 · p. 31

Related papers

Loading 2602.11812v2…