TokenCast: Forecasting Token Consumption During LLM Agent Execution

AI summary

Overview

• Research area: Machine learning systems for large language model (LLM) agents, specifically cost and resource forecasting for agent execution.
• Technical level: Intermediate. The core idea is accessible, but the paper assumes familiarity with LLM agent loops, token accounting (input, output, reasoning tokens), and gradient-boosted regression models.
• Scope: The paper proposes and evaluates TokenCast, a lightweight predictor that forecasts token consumption at four points during an LLM agent's execution across four benchmarks and six agent LLMs.

What This Paper Is About

When an LLM agent runs the same task twice, the total tokens it burns can differ by more than an order of magnitude, because the agent picks its next steps based on tool feedback and intermediate results, and each step's output is appended to a context that grows for every subsequent call. This makes total cost unpredictable before a run starts, and any prediction must be revised as the run proceeds. TokenCast attacks both problems by learning a small "cost triple" for each execution segment and composing those triples to forecast consumption, updating the forecast from observed execution evidence without making any extra LLM calls.

Key Contributions

1. A segment-cost factorization with an exact composition identity. Each execution segment is represented as (n_A, g_A, b_A) — call count, net change in input length, and a residual after subtracting the starting-input baseline. Composing adjacent segments yields (n_A + n_B, g_A + g_B, b_A + b_B + n_B * g_A), where the cross-term n_B * g_A captures the extra input cost incurred when context from an earlier segment is re-read by every later call.

2. A staged prefix–suffix predictor. TokenCast combines a direct forecast (remaining total cost as one target) with a compositional forecast (current or next prefix segment, its predicted ending boundary, and the following suffix), refreshing features at each prediction point using only execution evidence, with no additional LLM calls.

3. A training pipeline with cross-fitting, cost-weighted losses, and a correction model. Loss weights n_B * |g_A_hat - g_A| and L_B * |n_B_hat - n_B| are derived from the composition identity; cross-fitting generates out-of-fold prefix boundaries so the suffix model is not trained on its own inputs; a final correction model is fit on out-of-fold pipeline outputs. Calibrated LightGBM quantile models produce 0.05 and 0.95 interval endpoints.

4. Large-scale empirical validation. Results are reported across four prediction settings, four benchmarks, and six agent LLMs, with 11,712 execution traces, an ablation study, generalization tests, and an offline budget-control replay.

Main Findings

• Fourth prediction settings are defined along the execution trajectory: Task Start (total consumption before execution), Call Start (consumption of the current call after its request is assembled), In-call Update (updated call prediction as output is generated), and Task Update (remaining consumption after a call completes).

• Accuracy gains are concentrated where execution evidence exists. Across 96 benchmark–model–prediction-point combinations, the MAE reduction against the strongest comparator averages 14.5%. Broken down by prediction point, the averages are -2.2% at Task Start, 1.9% at Call Start, 30.4% at In-call Update, and 27.8% at Task Update. TokenCast trails the strongest comparator in 24 combinations: 15 at Task Start and nine at Call Start, with no losses at In-call Update or Task Update.

• Concrete SWE-bench Verified example with GPT-5.4. At In-call Update, MAE falls 47.9%, from EGTP's 74.6 to 38.9 tokens. At Task Update, MAE falls 30.4%, from TRAIL's 115k to 80k tokens. At Task Start, the best comparator is Self-Prediction and the reduction is 5.3%, from 152k to 144k tokens.

• Error should be read relative to consumption size. For GPT-5.4 at Task Start, TokenCast's MAE of 32.2k tokens on LongBench-v2 corresponds to a WAPE of 4.2%, while its MAE of 8.0k on MMLU-Pro corresponds to 15.8%. Mean target consumptions are 771.2k and 50.6k tokens respectively.

• Interval reliability improves substantially. On anchor tasks, TokenCast's calibrated intervals reduce mean interval score relative to Self-Prediction's native intervals at all four prediction points, by 32.0% on average. At Task Start it covers 82.0% of outcomes against a nominal 90% level, versus Self-Prediction's 52.7%. Coverage is 89.8% vs 73.0% at Call Start, 92.8% vs 75.2% at In-call Update, and 91.2% vs 77.0% at Task Update.

• Budget control saves tokens. On 288 GPT-5.4 runs from 144 SWE-bench Verified tasks, across seven replay budgets, TokenCast uses 21.3% fewer tokens on average while matching fixed-budget trace completion at every budget. The replay measures reaching the recorded terminal state and does not measure task resolution.

• Prediction overhead is small. Updating after every call takes 32.8 ms of mean cumulative prediction time and makes 19.7 forecasts per run. Updating every three calls cuts this to 11.9 ms and 6.9 forecasts, and a single Task Start forecast takes 2.1 ms. Every-call updating adds under 0.03% to the 129 s median wall time of a GPT-5.4 run.

• The segment triple matters. Replacing it with total token count, input length, and output length raises normalized MAE from 0.71 to 0.88 at the task level and from 0.68 to 0.82 at the call level. Removing composition gives a normalized average MAE of 0.78 and call-level MAE of 0.74. Within the triple, omitting context growth yields 0.76 and omitting the residual yields 0.73.

• Direct and compositional forecasts are complementary. Before execution, direct regression beats composition (0.82 vs 0.87 normalized MAE). The ranking reverses at Task Update, where composition reaches 0.66 and direct regression 0.72. Averaging them gives 0.64, and the full correction model reaches 0.62 while raising interval coverage from 88.1% under direct regression to 90.6%.

• Cross-fitting and cost weighting are the most load-bearing pipeline components. Dropping both raises normalized average MAE from 0.69 to 0.78 and lowers 90% interval coverage from 90.6% to 87.4%. Dropping cross-fitting alone gives 0.73; dropping cost weighting alone gives 0.72. Separate ablations of the correction model and the predicted boundary variables give 0.74 and 0.72.

• LightGBM is the best base predictor tested. Among six base predictors, LightGBM gives the lowest error and highest coverage with 0.8 ms model time.

• Feature importance shifts during a run. Request length leads at Task Start and Call Start (relative importance 0.24 and 0.30). During generation, committed prefix length dominates at 0.58. After a call completes, last-call input tokens, remaining plan items, and calls without progress carry similar importance at 0.28, 0.25, and 0.23.

• Generalization improves with a small amount of target data. On independently released LiveClawBench trajectories, zero-shot leave-one-domain-out transfer yields MAE ratios of 1.31 at Call Start and 1.47 at Task Update relative to Self-Prediction; with 20 target-domain tasks these fall to 0.82 and 0.85. For unseen agent LLMs at Call Start, transfer yields 1.02 times the target history-median MAE without target-model tasks, improving to 0.95 with 20 tasks, and with 3–10 target tasks transfer outperforms target-only training.

• Run-to-run variation is large. Repeated GPT-5.4 executions of the same task show substantial consumption spread across all four benchmarks, and the paper notes that consumption for the same task can differ by up to 30 times across executions.

Methodology in Plain English

The researchers start from a simple observation about how agent billing works. Every call to a model is charged for its input plus its output, and because the conversation history keeps growing, later calls pay for context that earlier calls created. So instead of predicting "how many tokens will the whole task use," they chop the execution into segments and describe each segment by three numbers: how many calls it contains, how much the input length grew by the end of it, and how many tokens it used beyond the trivial baseline of "call count times starting input length."

The useful part is that these three numbers combine with a fixed rule. When you glue two segments together, you must add a cross-term equal to the second segment's call count times the first segment's context growth. That term is exactly the penalty for re-reading old context. This means a prediction can be assembled from parts: predict the next chunk of execution, predict where it leaves the context, then predict the rest using that new context as the baseline.

They implement this with separate LightGBM models — a direct model that predicts remaining total cost as a single number, and prefix and suffix models for the compositional route — plus a final small correction model trained on out-of-fold pipeline outputs. Because the suffix model's inputs depend on the prefix model's predictions, they use cross-fitting (splitting tasks into folds and generating prefix predictions for each fold from models trained on the others) so the suffix model is trained on realistic, out-of-fold boundary errors rather than clean ground truth. Loss weights scaled by n_B and L_B reflect how strongly boundary mistakes propagate. Quantile models plus a calibration step on held-out tasks produce intervals.

They validate on SWE-bench Verified, Search-R1, MMLU-Pro, and LongBench-v2, using DeepSeek Harness and OpenHands to collect 11,712 traces from 240 benchmark tasks with six agent LLMs: GPT-5.4, Claude Opus 4.6, Gemini 3.1 Pro, DeepSeek-V4-Pro, Qwen3.8-27B, and Llama-3.2-3B-Instruct. All runs of the same task stay in one partition, and training, validation, calibration, and test tasks are separated. The local platform provides eight NVIDIA A100 GPUs. Baselines are TRAIL, EGTP, TIE, and Self-Prediction, each adapted to the four prediction points. Metrics are MAE and WAPE, plus empirical coverage, mean width, and mean interval score for 90% intervals.

Why This Matters

Impact on research. The paper reframes agent cost prediction from a single-number estimation problem into a compositional one where the interaction between context growth and future call counts is explicit. The composition identity is exact and associative, which gives a principled target for learning and a principled loss weighting derived from error propagation, rather than treating context accumulation as an unmodeled nuisance. It also pushes evaluation toward the harder "as execution unfolds" settings (In-call Update, Task Update) where the strongest gains appear, and shows that accurate forecasting is achievable without querying an LLM for a cost estimate.

Real-world applications:

• Budget enforcement in agent platforms: a runtime can compare a live forecast against a remaining token budget and decide whether to continue, stop, or simplify a run. The paper's replay shows 21.3% token savings at matched trace completion.
• Cost estimation for users before committing to a job: Task Start forecasts let a user or service decide whether a task is worth launching, though the paper's own averages show only a 1.9% Call Start gain and a -2.2% Task Start average, signaling that pre-execution prediction remains the weak point.
• Capacity planning and scheduling for serving infrastructure: knowing remaining call counts and input lengths supports resource reservation, batching, and queue prioritization, extending the motivation behind request-level length predictors to whole agent tasks.
• Anomaly detection in running agents: a forecast that jumps when a verification step fails or a compatibility workaround is needed gives an operational signal that a run is getting more expensive than expected.

Industry relevance. Token consumption is provider-accounted money. A prediction layer that costs 32.8 ms mean cumulative time per run and adds under 0.03% to a 129 s median GPT-5.4 run is cheap enough to run inline. Because the predictors are learned from recorded traces and improve with 3–20 target tasks rather than requiring thousands, an organization can adapt the approach to its own agent stack and models without retraining from scratch. The code is available at https://github.com/DEFENSE-SEU/TokenCast.

Future Directions

1. Closing the pre-execution gap. Task Start is where TokenCast is weakest: its average reduction against the strongest comparator is -2.2%, and it trails the comparator in 15 of the 24 losing combinations there. Better use of the task description and environment inspection (the direction Self-Prediction takes) is the obvious next target.

2. Tighter intervals at Task Start. Coverage is 82.0% against a nominal 90% level at Task Start, while width is 752.5k tokens — much wider than the 354.3k of Self-Prediction. Achieving both nominal coverage and narrower pre-execution intervals remains open.

3. Integrating forecasts into runtime decision-making. The authors explicitly leave "the integration of these forecasts into a runtime decision framework that actively manages execution under budget constraints" as future work. The replay is offline and measures trace completion, not task resolution, so online policies that trade completion for correctness are untested.

4. Reducing transfer dependence on target tasks. Zero-shot transfer to unseen task domains and unseen agent LLMs still lags Self-Prediction (ratios of 1.31 and 1.47 for domains, 1.02 for models at Call Start). How much target data is truly needed, and whether better feature design can shrink it, is unresolved.

Target Audience

This paper is most useful to systems and infrastructure engineers who serve or operate LLM agents and need to control cost; to researchers working on LLM inference scheduling, resource allocation, and token-length prediction who want a formulation that handles open-ended agent loops rather than fixed workflows; and to practitioners building agent frameworks who want a drop-in forecasting layer that requires no extra LLM calls. Readers familiar with gradient-boosted trees and agent execution traces will get the most from the methodology appendix, while the main results table and the budget-control figure are accessible to anyone who has paid an LLM API bill.

Authors’ abstract

When a large language model (LLM) agent executes the same task, token consumption can vary by over an order of magnitude across runs. The agent chooses its next steps based on tool feedback and intermediate results, while the growing context steadily inflates the input size of every subsequent call. The total consumption of a task is therefore hard to predict before execution and the prediction must be revised as the run unfolds. In this paper, we propose TokenCast, which learns a composable cost representation for each execution segment, recording its own consumption and the context growth it introduces. Composing adjacent segments yields a cumulative estimate that captures the extra input cost incurred when context from earlier segments is re-read by every later call. As execution unfolds, newly observed evidence refreshes the forecast, requiring no additional LLM calls and incurring a mean cumulative prediction time of 32.8 ms per run on SWE-bench Verified. Across 4 task suites and 6 agent models, TokenCast's mean absolute error reduction against the strongest comparator averages 14.5% over 96 evaluated combinations. In offline budget-control replay, TokenCast uses 21.3% fewer tokens on average than a fixed-budget policy at matched trace completion. The code is available at https://github.com/DEFENSE-SEU/TokenCast.