Active Inference as Context Acquisition for AI Agents
Sanchayan Dutta, Sai Niranjan Ramachandran, Suvrit Sra
cs.AI, cs.CL, cs.LG
2026-06-08
Context actions scored by expected free energy. Models shrink OQA sets but still trail a DP oracle; two clarifications lift description compliance from 4.2% to 37.5%.
Users rarely specify every constraint, preference, file, or task variable up front. An agent can proceed under a default assumption, or spend tokens on a clarifying question, a retrieval call, a tool call, or a prompt trial. Both choices are expensive. Over-asking burns tokens and patience. Under-asking commits to the wrong assumption.
This paper writes that tradeoff as active inference. The latent task state x can be an intent, a target, a preference, a missing constraint, or the best prompt. Action a is either a context action or a task action. An inner step updates beliefs over x under possible observations. An outer step picks the next move by minimizing expected free energy under cost. When observations are deterministic, the epistemic term is expected information gain. The rule is strict: buy more context only when the expected drop in decision-relevant uncertainty pays for the cost.
One-step active inference is a bilevel program. The inner loop minimizes variational free energy for a hypothetical observation; in the experiments the belief family contains the exact posterior, so the update is closed-form Bayes. The outer loop scores each candidate action by expected risk minus the expected KL change in belief. With flat preferences and noiseless observations, that choice is equivalent to maximizing mutual information I(x; o | a). An information-per-cost stopping rule says to ask only when I(x; o | a) > λ c(a). When terminal loss is logarithmic, the same gap is the drop in Bayes risk.
OQA turns question asking into a controlled Twenty Questions game. A hidden target is drawn uniformly from a finite attribute table. Each turn may query one listed attribute. Answers come from table lookup and are noiseless. The belief is uniform on the remaining consistent set Ct, so entropy is log2 |Ct|. Dynamic programming gives the optimal expected number of questions under the same menu. Planning gap is how many extra questions the model uses versus that oracle. Binary tables use Places, Cars, and Animals at N=25 and 100. Multiway tables are 5-ary with eight attributes (color, shape, material, size, pattern, origin, use case, energy), at N=100, 200, and 300, with 30 targets per tier. Decoding is temperature 0, tools off, one API call per question, fresh session. Models: GPT-5, GPT-4.1, Gemini 2.5 Pro, Gemini 2.0 Flash, Claude Sonnet 4.5, Claude Haiku 4.5, Grok 4.
The same accounting is applied to two token-budgeted prompting studies. In product-description autocompletion, a hidden style U=(tone, length, format) can be queried with up to three fixed templates, costing 24, 48, then 72 tokens, with symmetric label noise ε=0.12 in the belief update. A deterministic verifier checks that feature phrases appear as substrings and that format and length match; tone is not scored. In automated prompt optimization, six system prompts are Bernoulli arms. Each trial is one ARC-Challenge question, scored correct or incorrect, with independent Beta posteriors. Policies include round robin, Thompson sampling, knowledge gradient, mutual information, and an expected-free-energy mix. Training stops at 400 questions or 120,000 tokens.
OQA results are entropy trajectories against the DP oracle. The main text does not table per-model planning gaps. On both binary and multiway tiers, frontier models shrink the consistent set each turn, and the curves stay above the oracle: they make progress, and they still ask questions an optimal tree would skip. On the Animals binary table, duplicate attribute vectors can force a stop with residual entropy above zero.
The autocompletion numbers are concrete. On 48 synthetic product tasks, a no-question baseline that uses a default style passes the verifier 4.17% of the time at about 112 tokens. An activeweighted policy that ranks unasked attributes by entropy per token, with verifier-aware weights, reaches 37.5% compliance at about 219 tokens after a grid search that picks ε=0.01 and at most two clarifications. Format is weighted 1.0, length 0.6, tone 0.25, because the verifier ignores tone. Asking about a variable the final checker does not score spends tokens on information that cannot change the outcome.
In prompt optimization the 400-question cap binds first, at about 5×10^4 tokens, well under the 120,000-token budget. Knowledge gradient concentrates queries early and leaves rarely tested prompts near the 0.5 prior mean. Mutual information and expected free energy cut entropy over which prompt is best earlier; Thompson sampling and round robin spread trials more evenly. Holdout accuracies per policy are not tabulated in the main text.
Product agents already live this decision: ask another question, retrieve, try another prompt, or commit. The paper gives a stop condition that can be computed, and a toy environment where question efficiency is exact. Two pieces transfer immediately: rank context actions by information per cost, and weight questions by the variables the verifier actually scores.
This is not a claim that active inference beats reinforcement learning. The authors treat the difference as operational. Here an action changes what evidence arrives next. Expected free energy puts probing inside the objective; ordinary RL often adds an exploration bonus after the fact. For work that already frames clarification as Bayesian experimental design, the contribution is a control-theoretic vocabulary plus an oracle-backed benchmark.
OQA is noiseless, fixed-menu, table-lookup answering. Real dialogue has vague replies, new constraints, contradictory preferences, and evidence that is not on any prewritten attribute list. Tools, retrieval, and scratchpads are turned off on purpose, so the measurement is bare-model questioning. In a product those tools would store and filter the candidate set exactly, and the efficiency picture would change. The DP oracle is tractable at the released sizes; the number of remaining subsets explodes as N grows.
The two prompting studies are controlled demonstrations. Clarifications use fixed templates and simulated answers. The prompt library has six handwritten variants, scored as multiple-choice correctness. Hosted APIs can drift and break reproduction. An appendix writes the same information accounting onto adaptive prompt attacks, with no attack or defense experiments.
The main text almost never reports mean planning gaps per model, only entropy curves. A reader cannot say how many extra questions GPT-5 asks versus Haiku.