One email each morning — yesterday's AI, sortedGet it in your inbox

METAL LAB

Active Inference as Context Acquisition for AI Agents

arXiv:2608.192022026-08-21

Teaching AI agents to decide when asking a question is worth the tokens it costs

When a user leaves out a detail, an AI agent must choose between guessing, asking a clarifying question, or making a retrieval or tool call, each of which costs tokens. This paper frames that choice as active inference, scoring each possible next action by how much it is expected to cut uncertainty against how much it costs. The authors tested the idea on seven frontier language models using a Twenty-Questions-style benchmark and two prompting experiments under token budgets.

What they did

  1. Frames clarifying questions, retrieval calls, tool calls, and prompt trials as 'context actions' that an agent picks between, alongside acting now or stopping, to reduce uncertainty about a hidden task variable
  2. Splits the decision into an inner step that updates beliefs assuming a hypothetical answer, and an outer step that picks whichever action minimizes expected free energy, a score combining risk, information value, and cost
  3. In deterministic settings this score simplifies to expected information gain, which can be normalized by token cost
  4. Built Optimal Question Asking (OQA), a Twenty-Questions-style benchmark with 25 to 300 candidates, computed an exact dynamic-programming oracle for optimal question count, and benchmarked GPT-5, GPT-4.1, Gemini 2.5 Pro, Gemini 2.0 Flash, Claude Sonnet 4.5, Claude Haiku 4.5, and Grok 4 against it
  5. In a product-description clarification task, the best tuned policy reached 0.375 verifier compliance at about 219 tokens per task, versus a no-question baseline at 0.0417 compliance and about 112 tokens
Figure 1: Binary OQA with N=25. Curves show mean entropy Ht=log2⁡|Ct| across uniformly sampled targets, with ±1 standard deviation, along with the DP oracle.
Figure 1: Binary OQA with N=25. Curves show mean entropy Ht=log2⁡|Ct| across uniformly sampled targets, with ±1 standard deviation, along with the DP oracle.
Figure 2: Binary OQA with N=100. Curves show mean entropy Ht=log2⁡|Ct| across uniformly sampled targets, with ±1 standard deviation, along with the DP oracle. In Animals, duplicate attribute vectors can force stopping with a small residual entropy.
Figure 2: Binary OQA with N=100. Curves show mean entropy Ht=log2⁡|Ct| across uniformly sampled targets, with ±1 standard deviation, along with the DP oracle. In Animals, duplicate attribute vectors can force stopping with a small residual entropy.
Figure 3: Multiway categorical OQA for |𝒳|∈{100,200,300} (top to bottom). Curves show mean posterior entropy over uniformly sampled targets. The DP oracle uses the same query menu and stopping rule. Error bars show ±1 standard deviation when available.
Figure 3: Multiway categorical OQA for |𝒳|∈{100,200,300} (top to bottom). Curves show mean posterior entropy over uniformly sampled targets. The DP oracle uses the same query menu and stopping rule. Error bars show ±1 standard deviation when available.
Table 1: Operational distinction used in this paper. The contrast is not categorical: RL can include information bonuses, and active inference can use RL solvers.
AspectGeneric RL viewActive-inference view used here
ControlOptimize actions for return under the sampled or observed data stream.Choose actions that also determine what evidence arrives next.
ExplorationUsually added through bonuses, entropy, optimism, or posterior sampling.Appears directly as epistemic value in expected free energy.
Belief and queriesBeliefs may be implicit in a value function, recurrent state, or model posterior.An explicit posterior over intent, target, or best prompt makes queries first-class actions.
EvaluationReturn, regret, accuracy, or reward-model score.Entropy drop, oracle gap, bits per token, and final task success.
Figure 4: Prompt autocompletion on 48 synthetic product tasks with ϵ=0.02 and Kmax=3. Panels show clarification counts by policy, the ECDF of total tokens per task, and verifier pass rate versus average tokens. active_weighted improves compliance at a modest token cost relative to baseline and random, while ask_all spends the most tokens.
Figure 4: Prompt autocompletion on 48 synthetic product tasks with ϵ=0.02 and Kmax=3. Panels show clarification counts by policy, the ECDF of total tokens per task, and verifier pass rate versus average tokens. active_weighted improves compliance at a modest token cost relative to baseline and random, while ask_all spends the most tokens.
(b) Total token ECDF
(b) Total token ECDF
(c) Compliance versus average tokens
(c) Compliance versus average tokens

Why it matters

This gives chatbots and agents a principled way to decide, rather than guess, when spending tokens on a clarifying question or tool call actually pays off in accuracy. For any product where token usage translates directly to cost, this kind of cost-aware information accounting could shape how agents are designed to balance efficiency and correctness.

Figure 5: Automated prompt optimization under a token budget. Panels show final posterior mean accuracies, realized information gain per 1,000 tokens, and posterior entropy over the identity of the best prompt. Policies that reduce uncertainty early can avoid committing too soon.
Figure 5: Automated prompt optimization under a token budget. Panels show final posterior mean accuracies, realized information gain per 1,000 tokens, and posterior entropy over the identity of the best prompt. Policies that reduce uncertainty early can avoid committing too soon.
(b) Information gain per 1k tokens
(b) Information gain per 1k tokens
(c) Entropy of best-prompt identity
(c) Entropy of best-prompt identity

Terms in this paper

  • active inference · a framework that jointly scores information-seeking actions and goal-directed actions to pick the best next move
  • expected free energy · a score combining risk, information value, and cost used to rank candidate actions
  • expected information gain · how much an action is predicted to reduce uncertainty about the hidden variable
  • Optimal Question Asking (OQA) · this paper's benchmark where a model narrows down a hidden target by asking questions, like Twenty Questions
  • DP oracle · a dynamic-programming solution that computes the theoretically optimal minimum number of questions needed under the same rules

Original abstract (English)

Interactive AI agents must acquire the right context as efficiently as possible. When a user omits a constraint, preference, file, or task variable, an agent can proceed with a default assumption or spend tokens on a clarifying question, retrieval call, tool call, or prompt trial. We formulate this tradeoff as active inference for context acquisition. An inner inference step updates beliefs over a latent task state, and an outer decision selects the next context action, task action, or stop action to minimize expected free energy under cost. In deterministic settings, the epistemic term reduces to expected information gain, optionally normalized by token cost. We instantiate the framework in Optimal Question Asking (OQA), with exact posteriors and a dynamic programming oracle, and benchmark frontier language models on binary and multiway categorical tasks from 25 to 300 candidates. We also study clarification before generation and automated prompt optimization under token budgets. The formulation is model-agnostic and views active inference as a design principle for the context-acquisition layer of AI agents.

Authors · Sanchayan Dutta, Sai Niranjan Ramachandran, Suvrit Sra

Read on arXiv

Latest papers

All papers →

Latest from METAL LAB

Figures: Sanchayan Dutta et al., arXiv:2608.19202, arxiv-nonexclusive