2-Step Agent: How a Bayesian Decision Maker Learns from AI-Decision Support

The 2-Step Agent framework introduces a general computational model to understand how rational Bayesian decision makers learn from predictions generated by machine learning decision support (ML-DS) in high-stakes fields like healthcare and the judiciary. The gist is that even under ideal conditions, ML-DS can do more harm than good because a single misaligned prior belief can be sufficient to derail the learning process and lead to suboptimal actions.

The 2-Step Agent Framework

The framework models the process in two main steps: (i) how a prediction for a new observation affects the beliefs of a rational Bayesian agent, and (ii) how this change in beliefs affects the estimation of causal effect, the downstream decision, and the subsequent outcome. The first step involves the agent updating their world model using ML-DS predictions as compressed information about the causal mechanism they want to act on. This update is formalized by a Bayesian network Ahist that includes variables for population-level distributions, historical data represented via a plate, and the prediction P redo.

Agent Beliefs and Prior Misalignment

The agent's prior beliefs are encoded in the joint distribution over all variables in the model, which is defined by Ahist and parameter sets αX, αY, αA, NE, and Nθ. Motivating examples illustrate that different prior assumptions about the training data can lead to opposite actions when faced with the same ML-DS information. Specifically, a single misaligned prior belief can be sufficient for ML-DS to lead to worse downstream outcomes compared to no decision support even when the ML model is well-specified and the agent is perfectly rational.

Decision Making and Causal Effect Estimation

The second step defines how the Bayesian agent uses their updated beliefs about the population to decide on a treatment for new patients. This involves estimating the Conditional Average Treatment Effect (CATE) using an inference model Ainf, which draws from a distribution of possible SCMs based on the posterior learned in the first step. The agent chooses treatment a1 if EP Ainf (CATE) > τ, where CATE is calculated based on these beliefs.

Tractability and Theoretical Results

The framework faces a tractability problem in the first step because approximate Bayesian inference struggles to propagate through a plate of completely interchangeable variables. The authors prove that in simple settings, this problem can be solved by reducing the plate to sufficient statistics, resulting in a tractable solution for the linear Gaussian setting. This reduction allows for the computation of updated beliefs about population-level parameters given a prediction from the ML-DS model.

Experimental Findings

The experiments investigate how and how much decision support influences (i) agent’s beliefs, (ii) treatment decisions, and (iii) downstream patient outcomes. Results show that learning from ML-DS manifests through a correlation in the posterior of the agent, where the orange series shows superior outcomes after learning from the interaction with ML-DS. Furthermore, incorrect beliefs about the distribution of covariates can lead to bad decisions and worse outcomes even in near-ideal conditions. The authors show that the effect of ML-DS is defined as E Hist(Y do(A = ML-DS)) − E Hist(Y do(A = no ML-DS)).

Conclusion and Implications

The 2-Step Agent framework provides the first explicit computational model unpacking the interaction between an agent and MLDS, allowing for simulations comparing outcomes with and without ML-DS in an RCT-like fashion. The findings highlight the sensitivity of ML-DS to prior beliefs when using a treatment-naive prediction model, underscoring the importance of providing full documentation and proper training to users of ML-DS. These results apply to any field deploying such systems and to any agent interacting with ML-DS, be it human or artificial. The theory describes what happens after a single interaction of an agent with ML-DS and does not yet cover repeated interactions. The framework is general enough to allow for the divergence of SCM structures, though simulations are only concerned with learning from non-causal prediction models. The code and experiments are available open-source at https://placeholder link.

Implementation Details

The implementation uses PyMC for inference, reformulating variables using zero-centered normally distributed RVs where needed. The reduction to sufficient statistics is achieved by defining four auxiliary objects: s, W, g, and h. The total scalar degrees of freedom for the reduced plate is k squared + 5k + 2/2. This reduction makes the method both tractable and numerically more stable. The final formula for the slope estimate is given by βˆ 1 = β + τD−1 X W−1/2 g. The number of samples needed using the reduced plate grows as O(k 2), whereas for the full plate it grows as O(kn). The total scalar degrees of freedom are k squared + 5k + 2/2. The classical sufficient statistics can all be expressed from (s, W, g, h). The total number of sufficient statistics is k squared + 5k + 2/2. The final formulas are βˆ 1 = β + τD−1 X W−1/2 g, (1) and βˆ 0 = c + τ√n h − τ α⊤Xµ D−1 X + 1n s⊤ W−1/2 g, (2). The total scalar degrees of freedom is k squared + 5k + 2/2. The final formulas are βˆ 1 = β + τD−1 X W−1/2 g, (1) and βˆ 0 = c + τ√n h − τ α⊤Xµ D−1 X + 1n s⊤ W−1/2 g, (2). The total scalar degrees of freedom is k squared + 5k + 2/2.

Related Work

The literature focuses on how ML-DS is entangled with changes in the data generating process, or how biases in ML predictions amplify through downstream decisions. The authors contrast their framework by focusing on the decision-maker, showing that an accurate prediction paired with a misaligned prior can already be sufficient to worsen outcomes. They add that in their context the signal is an ML prediction produced from training data drawn from the same environment the agent wants to intervene on, yielding a distinctive inference problem. The paper addresses a different mechanism than previous work by explicitly modeling how the agent learns about the data-generating process from the ML prediction.

Motivating Examples

The framework captures motivating examples where an agent's prior beliefs determine how a prediction is interpreted and what action is taken. In the second example, new information from the prediction forces an agent to consider different options by revising their beliefs, which appears as a correlation in the posterior belief in these two variables <ref:2602.

1. Bayesian Decision Maker Framework: Implement a 2-Step Agent to model how a rational agent updates its beliefs about a population and treatment mechanism after observing an ML prediction from training data. This allows AI systems to explicitly capture how a prediction for a new observation affects the beliefs of a rational Bayesian agent.

2. Causal Effect Estimation Under Uncertainty: The system can estimate causal effect by drawing an inference over multiple possible structural causal models based on the updated beliefs. This enables the agent to calculate Conditional Average Treatment Effect (CATE) = E Ainf Y Xo, do(A = a1) - E Ainf Y Xo, do(A = a2), which is crucial for making informed treatment decisions.

3. Tractable Bayesian Inference: Use the derived reduction of the plate to sufficient statistics to make the agent's belief update computationally feasible. This allows for fully simulate what happens with and without decision support for a given agent A and an ML-DS model f, even in complex settings where approximate Bayesian inference fails.

4. Harm Quantification Under Misaligned Priors: The system can identify that a single misaligned prior belief can be sufficient for ML-DS to lead to worse downstream outcomes compared to no decision support even when the ML model is well-specified and the agent is perfectly rational. This provides a mechanism for auditing and mitigating risks associated with biased or incorrect initial assumptions.

5. Robustness to Prior Misalignment: The AI system can be designed such that it learns from ML-DS to estimate better CATE and take better actions when other beliefs are correct, demonstrating the ability to correct its beliefs based on new information.

Prognostic predictions from ML models support human decision making in several fields, including high-stakes ones such as healthcare and the judiciary. Yet, we still lack a clear understanding of how decision makers learn from ML-based decision support (ML-DS). In this paper, we introduce a general computational framework, the 2-Step Agent, to capture how a prediction from an ML model can also be used as a source of useful information about the training data and the causal mechanism that generated it. Our framework models (i) how a prediction from an ML model affects the beliefs of a rational Bayesian agent, and (ii) how this change in beliefs affects the estimation of the causal effect, the downstream decision for a new observation, and the subsequent outcome. In addition to the framework itself, we make three contributions. First, for the multivariate linear-Gaussian setting, we derive a tractable solution for the Bayesian inference problem introduced, i.e. one in which the agent infers information about a population from an ML prediction. Second, we experimentally identify conditions under which ML-DS is beneficial compared to unassisted decision making. Third, we show that a single misaligned prior belief of the user can be sufficient for ML-DS to lead to worse downstream outcomes compared to no decision support even when the prognostic model is well-specified and the agent is perfectly rational. Hence, even under ideal conditions, ML-DS can do more harm than good.