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

AI summary

Overview

Research area: Inference-time efficiency for large language models (LLMs), specifically adaptive compute allocation and batch scheduling for chain-of-thought reasoning.

Technical level: Intermediate. The paper assumes familiarity with transformer internals (hidden states, attention projections), LoRA fine-tuning, and basic allocation/optimization concepts, but the core idea is intuitive.

Scope (one sentence): The paper proposes "Predictive Scheduling," a framework that pre-runs lightweight predictors on transformer hidden states or raw question text to estimate how many reasoning tokens each query needs, then distributes a fixed total token budget across queries with a greedy allocator to maximize expected accuracy.

What This Paper Is About

LLMs solve hard reasoning problems by generating long chain-of-thought traces, but serving them with one fixed token budget per query wastes compute on easy questions and under-serves hard ones. This paper asks whether it is possible to predict, before any generation happens, how much reasoning each query actually needs—or how difficult it is—and then use those predictions to split a fixed batch-level token budget intelligently. The goal is higher accuracy at the same total token cost, without modifying the underlying language model.

Key Contributions

1. Two lightweight early-stopping predictors. An MLP trained on intermediate transformer hidden states, and a LoRA-fine-tuned model that operates directly on raw question text, both producing 16-dimensional vectors of early-stopping (correct-answer) probabilities across token budgets from 16 to 256 tokens in steps of 16.
2. A difficulty classifier. Models that label each query easy, medium, or hard before generation, using both few-shot prompting with GPT (o4-mini) and LoRA fine-tuning of the DeepSeek-R1-Distill-Qwen-1.5B base model.
3. A systematic layer-wise analysis. Separate MLPs trained on each of the 28 transformer layers of DeepSeek-R1-Distill-Qwen-1.5B, showing that middle layers (12–17) carry the strongest predictive signal for reasoning-length estimation, with layer 16 reaching the highest test correlation of 0.742.
4. A greedy batch allocation algorithm. Starting from a minimum of 16 tokens per query, it iteratively assigns additional 16-token windows to whichever query is predicted to gain the most expected accuracy, until the total budget is exhausted or no positive gains remain.

Main Findings

• Adaptive beating uniform, at equal cost: 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.
• Middle layers win: Middle layers (12–17) significantly outperform early and late layers; layer 16 achieves the highest test Pearson correlation of 0.742. Early layers (1–6) show correlations below 0.6, and late layers (21–28) drop to similar levels, producing an inverted U-shape across depth.
• Efficiency of middle layers: Middle layers (12–17) achieve 15–20% higher correlation-to-loss ratio than early or late layers.
• Hidden states beat raw text: The LoRA fine-tuned predictor on raw question text achieved an evaluation MSE of 0.0795 and a Pearson correlation of 0.444, compared to the best MLP's 0.742—lower performance for continuous probability prediction, though still a moderate positive signal.
• Adaptivity depends on budget regime: For average budgets of 16–96 tokens per query, MLP-driven adaptive allocation outperforms uniform allocation, with the largest improvement in the most constrained regime (16–48 tokens). The advantage reverses around 96–128 tokens per query, where prediction errors outweigh the benefits of adaptivity.
• Difficulty stratification is consistent across splits: Difficulty labels were derived from performance under a 256-token budget using percentile thresholds p₂₀ = 0.18 (hard) and p₈₀ = 0.84 (easy). The training split contained 1,506 easy, 4,437 medium, and 1,507 hard examples (7,450 total); the test split contained 271 easy, 760 medium, and 263 hard examples (1,294 total). Averaged early-stopping probability curves are reported as consistent between train and test sets within each difficulty category.
• Not reported: The truncated paper content does not report the final accuracy numbers comparing the few-shot GPT (o4-mini) classifier against the LoRA fine-tuned classifier, nor the end-to-end accuracy of the difficulty-based allocation algorithm (Algorithm 2) relative to baselines.

Methodology in Plain English

The authors started with GSM8K, a benchmark of grade school math word problems with 7,450 training and 1,294 test examples. For every question, they generated 100 independent reasoning traces at temperature 0.7 and top-p 0.95, and inserted a fixed "probe" string that forces the model to emit a final answer at every 16 tokens, up to a maximum of 256 tokens. By checking how many of the 100 traces were correct at each probe point, they built a 16-dimensional "early-stopping probability" vector for each question—essentially a curve showing how likely a correct answer is as the token budget grows.

They then trained two kinds of predictors to guess that curve in advance. The first was a small MLP (two fully connected layers, 256 hidden units, ReLU, sigmoid output) fed the 1536-dimensional [CLS] hidden state extracted from each of the 28 transformer layers of DeepSeek-R1-Distill-Qwen-1.5B—one MLP per layer, so they could see which layers carried the most useful signal. The second was a LoRA adapter (rank 16, scaling factor 32) applied to the attention query and value projections, reading the raw question text and passing the final hidden state through a small regression head. A separate LoRA model performed three-way difficulty classification (easy/medium/hard) using cross-entropy loss.

For allocation, the greedy algorithm starts every query at 16 tokens and repeatedly hands out 16-token increments to the query with the highest predicted marginal accuracy gain. They compared this against a uniform baseline (same budget for everyone) and an oracle that uses ground-truth early-stopping probabilities. A second allocation strategy uses the predicted difficulty classes: it averages accuracy curves per difficulty class on the training set, then searches over per-class budgets to maximize expected accuracy subject to the average budget constraint.

Why This Matters

Impact on research: The layer-wise result is a concrete, reusable empirical finding: middle transformer layers (12–17) encode the most useful signal for reasoning-length estimation, which informs where future probing and prediction work should look. The paper also quantifies when adaptivity helps versus hurts (16–96 tokens versus 96–128+ tokens), which is a more nuanced framing than "adaptive is always better."

Real-world applications (as framed by the paper):

• Real-time code autocomplete, where latency budgets are tight and most queries are trivial.
• Interactive tutoring systems, where some questions need long explanations and others need one line.
• Decision-support tools that must balance answer quality against response time.
• Production LLM services and on-device deployments where per-query token cost maps directly to cloud bills and battery drain.

Industry relevance: The framework is plug-and-play—no changes to the underlying language model are required—so it can sit in front of an existing serving stack. Reducing wasted tokens on easy queries directly lowers serving cost and latency, which matters for multi-tenant inference providers and any deployment with a hard per-batch compute ceiling.

Future Directions

• Close the oracle gap. The gap between adaptive allocation and oracle performance remains substantial, especially at higher token budgets, indicating the current predictors are not accurate enough to identify optimal stopping points for long-reasoning queries.
• Fix the high-budget crossover. Adaptive allocation loses to uniform allocation around 96–128 tokens per query; improving prediction accuracy or adding a confidence-aware fallback to uniform budgeting could extend the useful range.
• Generalize beyond one model and benchmark. All results are on GSM8K with DeepSeek-R1-Distill-Qwen-1.5B; whether the layer-16 finding and the 7.9-point gain transfer to other model families, scales, and reasoning benchmarks is untested here.
• Combine signal sources. Hidden states deliver the strongest correlation (0.742) while raw text delivers only 0.444; a hybrid predictor, or one that uses linguistic features for coarse difficulty and hidden states for fine-grained length, is an unexplored direction.

Target Audience

ML systems and inference-infrastructure engineers who serve LLMs under latency or cost constraints; researchers working on adaptive computation, speculative or budgeted decoding, and chain-of-thought efficiency; and graduate students studying how internal transformer representations relate to task difficulty. Readers wanting exact accuracy tables for the difficulty classifier or the difficulty-based allocator will not find them in this excerpt, since those results are not reported in the available content.

Authors’ 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.