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

TL;DR

This paper presents Predictive Scheduling, a framework that uses lightweight predictors to estimate query difficulty and allocate token budgets dynamically, significantly improving inference efficiency and accuracy in large language models.

Contribution

It introduces a novel pre-run prediction method for optimal reasoning length allocation, enhancing LLM inference efficiency and accuracy.

Findings

Up to 7.9% accuracy improvement on GSM8K with predictive scheduling.

Middle transformer layers (12-17) are most informative for size estimation.

Pre-run predictions enable better compute-accuracy trade-offs in LLMs.

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.