Predictive Scheduling for Efficient Inference-Time Reasoning in Large Language Models

TL;DR

This work tackles the inefficiency of a fixed token budget for multi-step reasoning in large language models by introducing Predictive Scheduling, a plug-in framework that pre-runs lightweight predictors to estimate per-query reasoning length or difficulty before generation. It combines hidden-state based MLPs and LoRA-finetuned text models to predict either the optimal per-query trace length or discrete difficulty, and employs a greedy allocation algorithm to distribute a fixed total token budget across a batch to maximize expected accuracy. A systematic analysis reveals middle transformer layers (12–17) carry the strongest signals for size estimation, while linguistic cues better support coarse difficulty classification; difficulty-based allocation, particularly with a LoRA-based classifier, consistently outperforms uniform budgeting and approaches oracle-like performance on GSM8K. The results demonstrate that predictive, pre-run budgeting can substantially reduce compute waste and latency for inference-time reasoning without modifying the underlying LLM, offering a practical pathway toward cost-efficient, latency-sensitive deployments.

Abstract

Large language models (LLMs) achieve state-of-the-art accuracy on complex reasoning tasks by generating multiple chain-of-thought (CoT) traces, but using a fixed token budget per query leads to over-computation on easy inputs and under-computation on hard ones. We introduce Predictive Scheduling, a plug-and-play framework that pre-runs lightweight predictors, an MLP on intermediate transformer hidden states or a LoRA-fine-tuned classifier on raw question text, to estimate each query's optimal reasoning length or difficulty before any full generation. Our greedy batch allocator dynamically distributes a fixed total token budget across queries to maximize expected accuracy. On the GSM8K arithmetic benchmark, predictive scheduling yields up to 7.9 percentage points of absolute accuracy gain over uniform budgeting at identical token cost, closing over 50\% of the gap to an oracle with perfect foresight. A systematic layer-wise study reveals that middle layers (12 - 17) of the transformer carry the richest signals for size estimation. These results demonstrate that pre-run budget prediction enables fine-grained control of the compute-accuracy trade-off, offering a concrete path toward latency-sensitive, cost-efficient LLM deployments.

Figures (11)

• Figure 1: (a) Distribution of correct‐answer probabilities at each token budget; (b) Aggregate KDE over all budgets.
• Figure 2: Visualization of the greedy token allocation algorithm. Left: Early stopping probability vectors showing the likelihood of generating a correct answer given different token budgets for each query. Darker shades indicate higher probability of correctness. Right: Progressive token allocation across queries during algorithm execution. Each cell shows the allocated token budget, and blue outlines highlight the query receiving additional tokens at each step. The algorithm initially allocates a minimum budget to all queries, then iteratively assigns additional tokens to queries with the highest expected marginal accuracy gain, prioritizing easier queries at first but gradually shifting resources to more difficult ones as the budget increases.
• Figure 3: Correlation performance analysis across model layers. The top panel displays the Pearson correlation coefficients for both test and train datasets achieved by MLPs trained on different layer features of the DeepSeek-R1-Distill-Qwen-1.5B model. The middle layers (particularly layer 16) achieve the highest correlation for predicting early stopping performance, suggesting that intermediate representations offer the strongest signal for reasoning difficulty prediction. The bottom left panel shows aggregated performance by layer group. The bottom right panel illustrates the relationship between train and test correlation, with points colored by layer number.
• Figure 4: The top left panel shows test loss (MSE) by layer, demonstrating how prediction error varies across the model's depth. The top right panel illustrates the relationship between MSE and Pearson correlation, revealing that while these metrics are generally inversely related, some layers achieve better correlation despite similar loss values. The bottom panel presents correlation-to-loss ratio for the top performing layers, indicating which layers provide the most predictive value per unit of loss.
• Figure 5: Training DeepSeek-R1-Qwen-Distilled model with LoRA fine-tuning for early stopping prediction. Left: Evaluation MSE Right: Evaluation Pearson correlation coefficient between predicted and ground truth early stopping probabilities.