← Home

Water-Filling Token Allocation for Reliable LLM Workflows

By James Trappett · 27 May 2026

4 min read

As LLM-based systems move from single-model inference toward complex multi-agent pipelines, a practical engineering problem has emerged: how do you allocate computational resources across agents in a workflow to maximise reliability, without blowing through latency or cost budgets? Most practitioners currently answer this question with intuition or trial-and-error. This paper, available on arXiv, attempts to put that process on a principled mathematical footing.

The contribution is timely. Agentic workflows are increasingly common in production AI systems, combining LLM components with conventional modules such as database queries, solvers, and verification tools. Yet the theoretical machinery for reasoning about their performance has lagged behind the engineering practice. This work takes a step toward closing that gap.

Key Contributions

The paper makes three substantive contributions:

The focus on sequential workflows is a deliberate simplification, but a reasonable starting point. Sequential pipelines are common and analytically tractable, and the results here establish a foundation that more complex topologies could build on.

Methodology

The core modelling choice is a parametric exponential reliability function for LLM agents: reliability increases with token count at a rate governed by a parameter beta specific to each agent. Agents with larger beta reach high reliability quickly; those with smaller beta require substantially more tokens to achieve comparable performance. This is a sensible abstraction, though it does compress a great deal of empirical complexity into a single scalar.

With this model in place, the design problem becomes: maximise the product of agent reliabilities (since agents are sequential) subject to a total token budget. Because the log of a product is a sum of logs, and the log-reliability of each agent is concave in its token allocation, the problem is convex. The authors apply standard Lagrangian methods and derive a water-filling solution analogous to those familiar from information theory and resource allocation in communications systems.

The water-filling rule allocates more tokens to agents with lower beta values. The intuition is straightforward: high-beta agents saturate quickly, so additional tokens yield diminishing returns; low-beta agents have more room to improve and benefit more from additional allocation. The optimal policy equalises marginal gains in log-reliability across all agents, with a shadow price determining the allocation threshold.

Results

The numerical experiments use a five-agent sequential workflow with heterogeneous beta parameters and a total budget of 12,000 output tokens. The water-filling allocation is compared against three baselines: uniform allocation, proportional allocation (tokens scaled with beta), and inverse-proportional allocation (tokens scaled inversely with beta).

The results are clean and the message is clear: the optimal policy consistently outperforms all baselines across a range of token budgets. The proportional baseline performs worst, which makes sense since it allocates more tokens precisely to the agents that need them least. The uniform baseline is competitive at high budgets but suboptimal when resources are scarce, which is where allocation decisions matter most.

The shadow price interpretation is useful for practitioners. It gives a single number that summarises how much reliability the system could gain from a marginal increase in the token budget, which is directly actionable for capacity planning decisions.

Limitations and Open Questions

The paper is honest about its scope, and the limitations are worth being explicit about. The exponential reliability model is a parametric assumption, not an empirically validated law. Fitting beta for real agents on real tasks would require substantial measurement effort, and it is unclear how stable these parameters are across task types or model versions. The model also treats agent reliability as a function of output tokens only, setting aside reasoning tokens and the interaction between the two, which the authors acknowledge as future work.

The restriction to sequential workflows is significant. Many production agentic systems involve branching, parallel execution, or feedback loops. Extending the framework to DAG-structured workflows would require handling correlated failures and more complex reliability aggregation, neither of which is trivial. The shadow price machinery would need to be reconsidered in that setting.

There is also a question of what reliability means operationally. The paper treats it as a scalar quality measure, but in practice, agent failure modes are heterogeneous: an agent might produce a syntactically valid but semantically incorrect output, or fail in ways that are task-specific and hard to capture with a single number. The abstraction is useful for analysis but would need grounding before deployment.

That said, these are limitations of scope rather than flaws in the analysis. The convex optimisation approach is sound, the water-filling result is elegant, and the framework provides a genuine starting point for more complex models.

For researchers working on agentic system design, inference-time compute allocation, or the broader question of how to make LLM pipelines predictable and cost-efficient, this paper is worth reading. The full text is available at arXiv:2605.23929.

LLMMulti-Agent SystemsOptimisationAI SystemsReliability

Related Articles

LLMs Are Overconfident, Especially on Hard TasksQuantifying Redundancy in LLM Reasoning ChainsLCF: Faster LLM Agent Communication via Latent Cache Transfer