Predictive Scheduling for Efficient Inference-Time Reasoning in Large Language Models
This work addresses cost and latency inefficiencies in LLM deployments for reasoning tasks, offering a plug-and-play solution for compute-accuracy trade-offs.
The paper tackles the problem of inefficient token budgeting in large language models for reasoning tasks by introducing Predictive Scheduling, which uses lightweight predictors to estimate optimal reasoning lengths per query and dynamically allocates a fixed token budget, achieving up to 7.9 percentage points accuracy gain on GSM8K over uniform budgeting.
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.